Fix memory leak in optional_import traceback handling#8782
Fix memory leak in optional_import traceback handling#8782ericspod merged 6 commits intoProject-MONAI:devfrom
Conversation
…I#7480, Project-MONAI#7727) When an import fails, `optional_import` captured the live traceback object and stored it in a `_LazyRaise` closure. The traceback held references to stack frames containing CUDA tensors, preventing garbage collection and causing GPU memory leaks. Replace the live traceback with a string-formatted copy via `traceback.format_exception()` and clear the original with `import_exception.__traceback__ = None`. The formatted traceback is appended to the error message so debugging information is preserved. Also move three hot-path `optional_import` calls in `monai/transforms/utils.py` (cucim.skimage, cucim morphology EDT) from function bodies to module level, eliminating repeated leaked tracebacks on every invocation. Fixes Project-MONAI#7480, fixes Project-MONAI#7727 Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a module-level note in monai/transforms/utils.py clarifying CuCIM-related imports must remain inside functions (no module-level CuCIM import). Updates monai/utils/module.py to capture formatted tracebacks as a string (introducing Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/utils/test_optional_import.py (1)
80-97: Prefix unused variable with underscore.Line 91:
modis never used. Rename to_modto indicate intentional discard.🔧 Suggested fix
- mod, flag = optional_import("nonexistent_module_for_leak_test") + _mod, flag = optional_import("nonexistent_module_for_leak_test")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/utils/test_optional_import.py` around lines 80 - 97, In test_no_traceback_leak, the unused variable `mod` returned from optional_import should be renamed to `_mod` to signal intentional discard; update the tuple unpacking in the _do_import function from `mod, flag = optional_import("nonexistent_module_for_leak_test")` to `_mod, flag = optional_import("nonexistent_module_for_leak_test")` (keeping the rest of test_no_traceback_leak, _Marker, and gc check unchanged).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/utils/test_optional_import.py`:
- Around line 80-97: In test_no_traceback_leak, the unused variable `mod`
returned from optional_import should be renamed to `_mod` to signal intentional
discard; update the tuple unpacking in the _do_import function from `mod, flag =
optional_import("nonexistent_module_for_leak_test")` to `_mod, flag =
optional_import("nonexistent_module_for_leak_test")` (keeping the rest of
test_no_traceback_leak, _Marker, and gc check unchanged).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 7b9386de-fac6-49d8-a44b-50a0d17f850d
📒 Files selected for processing (3)
monai/transforms/utils.pymonai/utils/module.pytests/utils/test_optional_import.py
|
Thank you very much for the review. Will get to this soon! |
Restore cucim imports inside functions to avoid slow import times and buggy behaviour. Add comment explaining why cucim is not at module level. Also rename unused `mod` to `_mod` in test per CodeRabbit nitpick. Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/utils/test_optional_import.py (1)
83-93: Add docstrings for new local definitions.
_Markerand_do_importare new definitions but have no docstrings.Suggested patch
class _Marker: + """Marker object used to validate that frame references are released.""" pass def _do_import(): + """Run a failing optional import and return a weak reference. + + Returns: + weakref.ReferenceType: Weak reference to a local marker instance. + """ marker = _Marker() ref = weakref.ref(marker)As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/utils/test_optional_import.py` around lines 83 - 93, Add Google-style docstrings for the new local definitions _Marker (briefly stating it’s a sentinel class used to detect GC/frame leaks) and for _do_import (documenting its purpose, that it creates a _Marker instance, returns a weakref.ref to that marker, and that it calls optional_import for a nonexistent module expecting flag==False). Ensure the _do_import docstring includes a "Returns" section describing the weakref.ref return value and any expected behavior (no exceptions expected), and keep the wording concise and consistent with other test docstrings.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/utils/test_optional_import.py`:
- Around line 103-104: The test currently uses a bare attribute access
`mod.something` inside the with self.assertRaises(OptionalImportError) block
which triggers Ruff B018; replace that expression with an explicit getattr call
(e.g., getattr(mod, "something")) so __getattr__ is invoked without a useless
expression statement; update the test in tests/utils/test_optional_import.py
where the context manager asserts OptionalImportError to use getattr(mod,
"something") instead of the bare `mod.something`.
---
Nitpick comments:
In `@tests/utils/test_optional_import.py`:
- Around line 83-93: Add Google-style docstrings for the new local definitions
_Marker (briefly stating it’s a sentinel class used to detect GC/frame leaks)
and for _do_import (documenting its purpose, that it creates a _Marker instance,
returns a weakref.ref to that marker, and that it calls optional_import for a
nonexistent module expecting flag==False). Ensure the _do_import docstring
includes a "Returns" section describing the weakref.ref return value and any
expected behavior (no exceptions expected), and keep the wording concise and
consistent with other test docstrings.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 7c66820b-ac88-48da-9477-d9421a7fc627
📒 Files selected for processing (2)
monai/transforms/utils.pytests/utils/test_optional_import.py
✅ Files skipped from review due to trivial changes (1)
- monai/transforms/utils.py
Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com> Signed-off-by: Soumya Snigdha Kundu <soumyawork15@gmail.com>
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/utils/test_optional_import.py (1)
83-87: Add docstrings to new nested definitions.Lines 83-87 add
_Markerand_do_importwithout docstrings. Please document them (Google-style), including return/raises for_do_import.As per coding guidelines, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/utils/test_optional_import.py` around lines 83 - 87, Add Google-style docstrings for the new nested class and function: document the _Marker class with a brief one-line description of its purpose, and add a Google-style docstring to _do_import that describes what it does, its return value, and any exceptions it may raise (include an Args section only if there are parameters). Ensure the docstrings follow the Google format with sections like Returns and Raises for _do_import and a short summary for _Marker.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/utils/test_optional_import.py`:
- Around line 83-87: Add Google-style docstrings for the new nested class and
function: document the _Marker class with a brief one-line description of its
purpose, and add a Google-style docstring to _do_import that describes what it
does, its return value, and any exceptions it may raise (include an Args section
only if there are parameters). Ensure the docstrings follow the Google format
with sections like Returns and Raises for _do_import and a short summary for
_Marker.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 49aa01a5-de1b-4471-b7df-ed7377162130
📒 Files selected for processing (1)
tests/utils/test_optional_import.py
ericspod
left a comment
ericspod
left a comment
There was a problem hiding this comment.
The pre-commit undid the getattr change in the test but I think it's fine now anyway, thanks.
2 participants
Summary
optional_importwhere stored traceback objects retain references to entire call stacks, preventing garbage collectionimport_exception.__traceback__ = Noneoptional_importcalls for cucim to module level inmonai/transforms/utils.pyFixes #7480, #7727
Details
monai/utils/module.py:__traceback__object withtraceback.format_exception()stringimport_exception.__traceback__ = Noneto break reference chainmonai/transforms/utils.py:optional_importcalls to module level (consistent with existing skimage/scipy/cupy pattern)distance_transform_edtTest plan
test_no_traceback_leak— weakref regression test for the leaktest_failed_import_shows_traceback_string— verifies traceback in error messagetest_optional_import.pytests pass