Skip to content

Commit c4412eb

Browse files
yinkevericspod
andauthored
Fix tracing-state leak after transform exceptions (#9019)
Fixes #7701. ### Description `TraceableTransform.trace_transform()` now restores the previous tracing state even when code inside the context manager raises. The original exception still propagates unchanged. The previous implementation restored `self.tracing` only after a normal `yield`, so an exception permanently leaked the temporary tracing value into the transform. This caused later inverse transforms to fail with `Transform Tracing must be enabled` after a caught CUDA OOM or another runtime exception. ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality). - [ ] Breaking change (fix or new feature that would cause existing functionality to change). - [x] New tests added to cover the changes. - [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [ ] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [x] In-line docstrings updated. - [ ] Documentation updated, tested `make html` command in the `docs/` folder. ### Validation - The regression fails on `upstream/dev` because the tracing state remains changed after an exception. - `python -m pytest tests/transforms/inverse/test_traceable_transform.py -q`: 2 passed. - Black, isort, Ruff, mypy, documentation build, packaging, DCO, and pre-commit checks pass in upstream CI. - `git diff --check` passes. Signed-off-by: kyinhub <kevinpyin@gmail.com> Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
1 parent 21cbeda commit c4412eb

2 files changed

Lines changed: 24 additions & 3 deletions

File tree

monai/transforms/inverse.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -402,11 +402,20 @@ def pop_transform(self, data, key: Hashable = None, check: bool = True):
402402

403403
@contextmanager
404404
def trace_transform(self, to_trace: bool):
405-
"""Temporarily set the tracing status of a transform with a context manager."""
405+
"""Temporarily set the tracing status of a transform.
406+
407+
The previous tracing state is restored when the context exits normally
408+
or because of an exception.
409+
410+
Args:
411+
to_trace: tracing state to use within the context.
412+
"""
406413
prev = self.tracing
407414
self.tracing = to_trace
408-
yield
409-
self.tracing = prev
415+
try:
416+
yield
417+
finally:
418+
self.tracing = prev
410419

411420

412421
class InvertibleTransform(TraceableTransform, InvertibleTrait):

tests/transforms/inverse/test_traceable_transform.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,18 @@ def pop(self, data):
2929

3030
class TestTraceable(unittest.TestCase):
3131

32+
def test_trace_transform_restores_state_after_exception(self):
33+
"""Verify tracing state is restored after an exception."""
34+
transform = _TraceTest()
35+
transform.tracing = True
36+
37+
with self.assertRaisesRegex(RuntimeError, "expected failure"):
38+
with transform.trace_transform(False):
39+
self.assertFalse(transform.tracing)
40+
raise RuntimeError("expected failure")
41+
42+
self.assertTrue(transform.tracing)
43+
3244
def test_default(self):
3345
expected_key = "_transforms"
3446
a = _TraceTest()

0 commit comments

Comments
 (0)