Fix memory leak in optional_import traceback handling

[aymuos15]

Summary

• Fix memory leak in optional_import where stored traceback objects retain references to entire call stacks, preventing garbage collection
• Format traceback to string immediately and clear import_exception.__traceback__ = None
• Hoist 3 optional_import calls for cucim to module level in monai/transforms/utils.py

Fixes #7480, #7727

Details

monai/utils/module.py:

• Replace storing live __traceback__ object with traceback.format_exception() string
• Clear import_exception.__traceback__ = None to break reference chain
• Embed formatted traceback in error message under "Original traceback:" section

monai/transforms/utils.py:

• Move cucim optional_import calls to module level (consistent with existing skimage/scipy/cupy pattern)
• Fixes incidental name shadowing in distance_transform_edt

Test plan

• test_no_traceback_leak — weakref regression test for the leak
• test_failed_import_shows_traceback_string — verifies traceback in error message
• All 12 test_optional_import.py tests pass
• pre-commit, black, ruff, pytype, mypy — all clean

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes refactor optional CucIM/CuPy dependency handling to prevent GPU memory leaks. Private import flags and guards are introduced in monai/transforms/utils.py to conditionally branch on CuCIM availability. Exception handling in monai/utils/module.py is updated to capture tracebacks as strings instead of retaining traceback objects, reducing stack frame references. New tests verify no stack frame leakage occurs and traceback strings are preserved in error messages.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 70.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title clearly and concisely summarizes the main change: fixing a memory leak in optional_import traceback handling.
Description check	✅ Passed	Description covers objectives, technical details, and test plan, matching the template's essential sections adequately.
Linked Issues check	✅ Passed	Changes directly address #7480 (GPU memory leak via traceback retention) and #7727 (name shadowing) by formatting tracebacks to strings, clearing references, and hoisting imports.
Out of Scope Changes check	✅ Passed	All changes align with stated objectives: memory leak fix in module.py, cucim import hoisting in transforms/utils.py, and test coverage additions.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

📝 Coding Plan

• Generate coding plan for human review comments

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/utils/test_optional_import.py (1)

80-97: Prefix unused variable with underscore.

Line 91: mod is never used. Rename to _mod to 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

📥 Commits

Reviewing files that changed from the base of the PR and between daaedaa and 5395848.

📒 Files selected for processing (3)

• monai/transforms/utils.py
• monai/utils/module.py
• tests/utils/test_optional_import.py

[ericspod]

Hi @aymuos15 thanks for these fixes. I commented on the cucim change which mostly should be reverted, the changes for the traceback leaking look fine to me so I hope this solves the problems the issues have raised.

[ericspod]

I believe we had strange import issues if we had cucim at the module level, such as very slow import or some other buggy behaviour. Our solution was to put importation into functions so that this happened only when the library was needed, which was very rarely. Please undo this change but replace it with a note that cucim is deliberately not imported at the module level, there should have been one already we neglected to add in the past.