Fixes #8697 GPU memory leak by checking both image and label tensors for CUDA device #8698
Conversation
Fixes Project-MONAI#8697 GPU memory leak by checking both image and label tensors for CUDA device Signed-off-by: benediktjohannes <benedikt.johannes.hofer@gmail.com>
benediktjohannes
requested review from
KumoLiu,
Nic-Ma and
ericspod
as code owners
📝 WalkthroughWalkthroughFgImageStats.call and LabelStats.call now create local Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
📜 Recent review detailsConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Cache: Disabled due to data retention organization setting Knowledge base: Disabled due to 📒 Files selected for processing (1)
🧰 Additional context used📓 Path-based instructions (1)**/*.py⚙️ CodeRabbit configuration file
Files:
🪛 Ruff (0.14.11)monai/auto3dseg/analyzer.py485-485: Avoid specifying long messages outside the exception class (TRY003) 🔇 Additional comments (2)
✏️ Tip: You can disable this entire section by setting 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.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In @monai/auto3dseg/analyzer.py:
- Around line 471-477: Refactor the later references that still use
d[self.image_key] and d[self.label_key] to reuse the local variables
image_tensor and label_tensor (so subsequent logic around self.do_ccp and CUDA
detection uses those locals consistently), and add a GPU unit test in
tests/apps/test_auto3dseg.py (for LabelStats) decorated with @skip_if_no_cuda
that constructs a mixed-device input (one tensor on CPU, the other on CUDA) to
exercise the new mixed-device detection and ensure behavior is correct under
self.do_ccp.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (1)
monai/auto3dseg/analyzer.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
⚙️ CodeRabbit configuration file
Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. 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. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.
Files:
monai/auto3dseg/analyzer.py
🧬 Code graph analysis (1)
monai/auto3dseg/analyzer.py (1)
monai/data/meta_tensor.py (1)
MetaTensor(52-609)
Small refactoring (usage of local variable instead for performance) Signed-off-by: benediktjohannes <benedikt.johannes.hofer@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
monai/auto3dseg/analyzer.py (1)
487-488: Type inconsistency in empty foreground fallback.Line 488 uses
torch.Tensor([0])while line 354 (FgImageStats) usesMetaTensor([0.0]). Sinceget_foreground_labelreturnsMetaTensor, the fallback should also useMetaTensorfor consistency.Suggested fix
- nda_foregrounds = [nda if nda.numel() > 0 else torch.Tensor([0]) for nda in nda_foregrounds] + nda_foregrounds = [nda if nda.numel() > 0 else MetaTensor([0.0]) for nda in nda_foregrounds]
🧹 Nitpick comments (1)
monai/auto3dseg/analyzer.py (1)
481-481: Inconsistent variable usage.Use
image_tensor.shape[0]instead ofd[self.image_key].shape[0]for consistency with the refactoring.Suggested fix
- ndas: list[MetaTensor] = [image_tensor[i] for i in range(d[self.image_key].shape[0])] # type: ignore + ndas: list[MetaTensor] = [image_tensor[i] for i in range(image_tensor.shape[0])] # type: ignore
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (1)
monai/auto3dseg/analyzer.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
⚙️ CodeRabbit configuration file
Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. 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. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.
Files:
monai/auto3dseg/analyzer.py
🧬 Code graph analysis (1)
monai/auto3dseg/analyzer.py (1)
monai/data/meta_tensor.py (2)
MetaTensor(52-609)astype(434-461)
🔇 Additional comments (1)
monai/auto3dseg/analyzer.py (1)
471-477: CUDA detection logic looks correct.Properly checks both tensors for CUDA presence, addressing the mixed-device memory leak scenario.
Another small refactoring Signed-off-by: benediktjohannes <benedikt.johannes.hofer@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
monai/auto3dseg/analyzer.py (1)
488-488: Type inconsistency:torch.Tensor([0])vsMetaTensor([0.0]).Line 488 differs from line 354 which uses
MetaTensor([0.0]), creating both type and dtype inconsistency. Whennda_foregroundsis empty, it becomestorch.Tensor([0])(int), whereas elsewhere it'sMetaTensor(float). This inconsistency should be resolved for type safety:- nda_foregrounds = [nda if nda.numel() > 0 else torch.Tensor([0]) for nda in nda_foregrounds] + nda_foregrounds = [nda if nda.numel() > 0 else MetaTensor([0.0]) for nda in nda_foregrounds]
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (1)
monai/auto3dseg/analyzer.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
⚙️ CodeRabbit configuration file
Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. 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. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.
Files:
monai/auto3dseg/analyzer.py
🧬 Code graph analysis (1)
monai/auto3dseg/analyzer.py (1)
monai/data/meta_tensor.py (2)
MetaTensor(52-609)astype(434-461)
🔇 Additional comments (2)
monai/auto3dseg/analyzer.py (2)
471-477: CUDA detection logic looks correct.The fix properly addresses the GPU memory leak by checking both
image_tensorandlabel_tensorfor CUDA device presence before callingtorch.cuda.empty_cache()at line 517. This handles mixed-device scenarios as intended.
481-482: LGTM on tensor extraction refactor.Using the local
image_tensorandlabel_tensorvariables improves readability and avoids repeated dictionary lookups.
Corrected type inconsistency Signed-off-by: benediktjohannes <benedikt.johannes.hofer@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In @monai/auto3dseg/analyzer.py:
- Around line 481-482: The code calls label_tensor.astype(...) which only exists
on MetaTensor but earlier isinstance checks allow plain torch.Tensor; update the
ndas_label assignment to handle both: if isinstance(label_tensor, MetaTensor)
use label_tensor.astype(torch.int16), else if isinstance(label_tensor,
torch.Tensor) use label_tensor.to(torch.int16) (or wrap/convert to MetaTensor as
appropriate), ensuring ndas_label is a MetaTensor-compatible tensor; adjust the
isinstance logic around image_tensor/label_tensor if needed to keep types
consistent.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (1)
monai/auto3dseg/analyzer.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
⚙️ CodeRabbit configuration file
Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. 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. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.
Files:
monai/auto3dseg/analyzer.py
🧬 Code graph analysis (1)
monai/auto3dseg/analyzer.py (2)
monai/data/meta_tensor.py (2)
MetaTensor(52-609)astype(434-461)monai/auto3dseg/utils.py (1)
get_foreground_label(70-83)
🪛 Ruff (0.14.11)
monai/auto3dseg/analyzer.py
485-485: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (2)
monai/auto3dseg/analyzer.py (2)
471-477: CUDA detection logic looks correct.Checking both tensors ensures
torch.cuda.empty_cache()is called whenever CuPy-backed CCP might encounter PyTorch-held GPU memory. This should fix the memory leak in mixed-device scenarios.
488-488: MetaTensor fallback is appropriate.
get_foreground_labelreturnsMetaTensor, so usingMetaTensor([0.0])for empty foregrounds maintains type consistency for downstream operations.
Fixed type check and astype() usage Signed-off-by: benediktjohannes <benedikt.johannes.hofer@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
monai/auto3dseg/analyzer.py (2)
471-477: Consider usingtorch.Tensorinstead ofMetaTensorfor broader CUDA detection.The current check only detects CUDA tensors if they are specifically
MetaTensorinstances. SinceMetaTensorinherits fromtorch.Tensor, and all tensors have the.deviceattribute, usingtorch.Tensorwould be more robust:using_cuda = ( - isinstance(image_tensor, MetaTensor) and image_tensor.device.type == "cuda" + isinstance(image_tensor, torch.Tensor) and image_tensor.device.type == "cuda" ) or ( - isinstance(label_tensor, MetaTensor) and label_tensor.device.type == "cuda" + isinstance(label_tensor, torch.Tensor) and label_tensor.device.type == "cuda" )This would catch CUDA tensors even if someone passes a plain
torch.Tensor(though line 482 would fail later due to.astype()). The isinstance guard still protects against numpy arrays.
484-485: Minor: Consider extracting the error message.Static analysis flags the long inline message. Acceptable as-is, but could be a constant if reused elsewhere.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (1)
monai/auto3dseg/analyzer.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
⚙️ CodeRabbit configuration file
Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. 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. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.
Files:
monai/auto3dseg/analyzer.py
🧬 Code graph analysis (1)
monai/auto3dseg/analyzer.py (2)
monai/data/meta_tensor.py (1)
MetaTensor(52-609)monai/auto3dseg/utils.py (1)
get_foreground_label(70-83)
🪛 Ruff (0.14.11)
monai/auto3dseg/analyzer.py
485-485: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (2)
monai/auto3dseg/analyzer.py (2)
481-482: LGTM.Using local variables
image_tensorandlabel_tensorimproves readability and avoids repeated dictionary lookups.
488-488: Good fix for type consistency.Using
MetaTensor([0.0])instead oftorch.Tensor([0])maintains consistent typing throughout the pipeline when handling empty foregrounds. This matches the same pattern at line 354 inFgImageStats.
Using torch.Tensor instead of MetaTensor for broader CUDA detection Signed-off-by: benediktjohannes <benedikt.johannes.hofer@gmail.com>
|
I've fixed all issues mentioned by CodeRabbit 👍 |
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
| using_cuda = ( | ||
| isinstance(image_tensor, torch.Tensor) and image_tensor.device.type == "cuda" | ||
| ) or ( | ||
| isinstance(label_tensor, torch.Tensor) and label_tensor.device.type == "cuda" | ||
| ) |
There was a problem hiding this comment.
| using_cuda = ( | |
| isinstance(image_tensor, torch.Tensor) and image_tensor.device.type == "cuda" | |
| ) or ( | |
| isinstance(label_tensor, torch.Tensor) and label_tensor.device.type == "cuda" | |
| ) | |
| using_cuda = any(isinstance(t, torch.Tensor) and t.device.type == "cuda" for t in (image_tensor, label_tensor)) |
|
Hi @benediktjohannes thanks for the contribution. I think it's fine just from looking at the changes but we'll run the tests now and see. I would also like @mingxin-zheng to have a look as the file originates with work I think. |
2 participants
Modified device detection to check BOTH image and label tensors
torch.cuda.empty_cache() now called if EITHER tensor is on GPU
Prevents GPU memory leaks in mixed device scenarios