fix(profiler): gracefully handle stopping unstarted actions in Advanc…

src/lightning/pytorch/CHANGELOG.md

@@ -36,6 +36,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `AdvancedProfiler` leaving cProfile enabled if teardown occurs before a profiling context exits ([#9136](https://github.com/Lightning-AI/pytorch-lightning/issues/9136))
+
 - Fixed `LightningCLI` emitting `jsonargparse` deprecation warnings ([#21900](https://github.com/Lightning-AI/pytorch-lightning/issues/21900))
 
 - Fixed `RichProgressBar` showing a nonsensical negative epoch total (e.g. `Epoch 5/-2`) when `Trainer(max_epochs=-1)` (unlimited epochs) is used and training stops via another condition ([#21925](https://github.com/Lightning-AI/pytorch-lightning/issues/21925))

src/lightning/pytorch/profilers/advanced.py

@@ -61,9 +61,6 @@ def __init__(
 
             dump_stats: Whether to save raw profiler results. When ``True`` then ``dirpath`` must be provided.
 
-        Raises:
-            ValueError:
-                If you attempt to stop recording an action which was never started.
         """
         super().__init__(dirpath=dirpath, filename=filename)
         self.profiled_actions: dict[str, cProfile.Profile] = defaultdict(cProfile.Profile)
@@ -81,7 +78,14 @@ def start(self, action_name: str) -> None:
     def stop(self, action_name: str) -> None:
         pr = self.profiled_actions.get(action_name)
         if pr is None:
-            raise ValueError(f"Attempting to stop recording an action ({action_name}) which was never started.")
+            # The profiler can be torn down while a profiling context is active. In that case,
+            # the context's finally block calls stop() after teardown has cleared this action.
+            # See: https://github.com/Lightning-AI/pytorch-lightning/issues/9136
+            log.debug(
+                "Attempting to stop recording an action (%s) which was never started. Skipping.",
+                action_name,
+            )
+            return
         pr.disable()
 
     def _dump_stats(self, action_name: str, profile: cProfile.Profile) -> None:
@@ -115,6 +119,8 @@ def summary(self) -> str:
     @override
     def teardown(self, stage: Optional[str]) -> None:
         super().teardown(stage=stage)
+        for pr in self.profiled_actions.values():
+            pr.disable()
         self.profiled_actions.clear()
 
     def __reduce__(self) -> tuple:

tests/tests_pytorch/profilers/test_profiler.py

@@ -396,16 +396,51 @@ def test_advanced_profiler_dump_states_sanitizes_filename(tmp_path, char):
     assert (tmp_path / prof_name).read_bytes()
 
 
-def test_advanced_profiler_value_errors(advanced_profiler):
-    """Ensure errors are raised where expected."""
+def test_advanced_profiler_stop_unstarted_action(advanced_profiler):
+    """Ensure stopping an unstarted action does not raise an error.
+
+    See: https://github.com/Lightning-AI/pytorch-lightning/issues/9136
+
+    Teardown can clear an active action before its profiling context calls stop()
+    in the finally block. This should be handled gracefully.
+
+    """
     action = "test"
-    with pytest.raises(ValueError, match="Attempting to stop recording*"):
-        advanced_profiler.stop(action)
+    # Should NOT raise, just log and return
+    advanced_profiler.stop(action)
 
+    # Starting and stopping should still work normally
     advanced_profiler.start(action)
     advanced_profiler.stop(action)
 
 
+def test_advanced_profiler_teardown_disables_active_action(advanced_profiler):
+    """Ensure teardown stops recording and an active context exits safely.
+
+    Regression test for https://github.com/Lightning-AI/pytorch-lightning/issues/9136
+
+    """
+
+    def tracked() -> None:
+        pass
+
+    with advanced_profiler.profile("active_action"):
+        active_profile = advanced_profiler.profiled_actions["active_action"]
+        advanced_profiler.teardown(stage=None)
+
+        def tracked_call_count() -> int:
+            return next(
+                (stat.callcount for stat in active_profile.getstats() if stat.code is tracked.__code__),
+                0,
+            )
+
+        before = tracked_call_count()
+        tracked()
+        assert tracked_call_count() == before
+
+    assert not advanced_profiler.profiled_actions
+
+
 def test_advanced_profiler_deepcopy(advanced_profiler):
     advanced_profiler.describe()
     assert deepcopy(advanced_profiler)