cleanup chromium event log on dynamo exit rather than on entry (pytor…

…ch#143175)

clearing at dynamo start is an issue because it throws away events from compiled autograd

Pull Request resolved: pytorch#143175
Approved by: https://github.com/Skylion007, https://github.com/jamesjwu
ghstack dependencies: pytorch#141907

torch/_dynamo/convert_frame.py

@@ -894,7 +894,7 @@ def count_args(code: CodeType) -> int:
     with _use_lazy_graph_module(config.use_lazy_graph_module), compile_context(
         CompileContext(compile_id)
     ), chromium_event_timed(
-        "dynamo", reset_event_log=True, log_pt2_compile_event=True
+        "dynamo", reset_event_log_on_exit=True, log_pt2_compile_event=True
     ), metrics_context:
         restart_reasons: set[str] = set()
         # This is shared across restarts

torch/_dynamo/utils.py

@@ -1255,7 +1255,7 @@ def log_event_start(
         :param str event_name Name of event to appear in trace
         :param time_ns Timestamp in nanoseconds
         :param metadata: Any extra metadata associated with this event
-        :param log_pt_compile_event: If True, log to pt2_compile_events
+        :param log_pt2_compile_event: If True, log to pt2_compile_events
         :param compile_id: Explicit compile_id (rather than using the current context)
         """
         compile_id = compile_id or torch._guards.CompileContext.current_compile_id()
@@ -1440,7 +1440,7 @@ def get_chromium_event_logger() -> ChromiumEventLogger:
 @contextmanager
 def chromium_event_timed(
     event_name: str,
-    reset_event_log: bool = False,
+    reset_event_log_on_exit: bool = False,
     log_pt2_compile_event: bool = False,
 ) -> Generator[Any, None, None]:
     """
@@ -1449,8 +1449,6 @@ def chromium_event_timed(
     instead. Use this context manager only if you want to avoid dynamo_timed.
     """
     chromium_event_log = get_chromium_event_logger()
-    if reset_event_log:
-        chromium_event_log.reset()
     chromium_start_time = time.time_ns()
     chromium_event_log.log_event_start(
         event_name,
@@ -1468,6 +1466,8 @@ def chromium_event_timed(
             chromium_start_time,
             log_pt2_compile_event,
         )
+        if reset_event_log_on_exit:
+            chromium_event_log.reset()
 
 
 @dataclasses.dataclass