From 50324fcf3ba97610ff93ab8765814a2349edd42a Mon Sep 17 00:00:00 2001 From: Xuan Dai Date: Sun, 4 Oct 2026 20:07:11 +0900 Subject: [PATCH] fix(profiler): gracefully handle stopping unstarted actions in AdvancedProfiler (#9136) `AdvancedProfiler.teardown()` cleared `profiled_actions` unconditionally, so a profiling context that was still active at teardown had its action dropped from the dict. When that context exited, its `finally` block called `stop()` on an action that no longer existed and raised: ValueError: Attempting to stop recording an action (run_test_evaluation) which was never started. Changes: - `teardown()` now disables any still-running profiles before clearing the dict, so no profile is left enabled after the profiler goes away. - `stop()` logs at debug level and returns instead of raising when the action was never started. This path is reachable from `teardown()` by construction, so raising there was never actionable for the user. - Removes the now-inaccurate `Raises:` section from the `AdvancedProfiler` docstring. - Adds a regression test for stopping an action that was never started, and for teardown while an action is still active. Fixes #9136 --- src/lightning/pytorch/CHANGELOG.md | 2 + src/lightning/pytorch/profilers/advanced.py | 14 ++++-- .../tests_pytorch/profilers/test_profiler.py | 43 +++++++++++++++++-- 3 files changed, 51 insertions(+), 8 deletions(-) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 1c73ea347751a..57f4dc3ce07a0 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/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)) diff --git a/src/lightning/pytorch/profilers/advanced.py b/src/lightning/pytorch/profilers/advanced.py index 6fd16ab991736..2c628e4fdaf01 100644 --- a/src/lightning/pytorch/profilers/advanced.py +++ b/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: diff --git a/tests/tests_pytorch/profilers/test_profiler.py b/tests/tests_pytorch/profilers/test_profiler.py index 67db953facbbb..195852f9161dd 100644 --- a/tests/tests_pytorch/profilers/test_profiler.py +++ b/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)