Skip to content

Commit

Permalink
Fix CUDA_VISIBLE_DEVICES setting ignored (Project-MONAI#7408)
Browse files Browse the repository at this point in the history
Fixes Project-MONAI#7407

### Description

Move `optional import cucim` inside the function to avoid using all
GPUs.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [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).
- [ ] 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`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Mark Graham <markgraham539@gmail.com>
  • Loading branch information
2 people authored and marksgraham committed Jan 30, 2024
1 parent ef912c8 commit 21f9ae1
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 4 deletions.
6 changes: 2 additions & 4 deletions monai/metrics/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,10 +38,6 @@
binary_erosion, _ = optional_import("scipy.ndimage.morphology", name="binary_erosion")
distance_transform_edt, _ = optional_import("scipy.ndimage.morphology", name="distance_transform_edt")
distance_transform_cdt, _ = optional_import("scipy.ndimage.morphology", name="distance_transform_cdt")
cucim_binary_erosion, has_cucim_binary_erosion = optional_import("cucim.skimage.morphology", name="binary_erosion")
cucim_distance_transform_edt, has_cucim_distance_transform_edt = optional_import(
"cucim.core.operations.morphology", name="distance_transform_edt"
)

__all__ = [
"ignore_background",
Expand Down Expand Up @@ -179,6 +175,8 @@ def get_mask_edges(
always_return_as_numpy: whether to a numpy array regardless of the input type.
If False, return the same type as inputs.
"""
# move in the funciton to avoid using all the GPUs
cucim_binary_erosion, has_cucim_binary_erosion = optional_import("cucim.skimage.morphology", name="binary_erosion")
if seg_pred.shape != seg_gt.shape:
raise ValueError(f"seg_pred and seg_gt should have same shapes, got {seg_pred.shape} and {seg_gt.shape}.")
converter: Any
Expand Down
7 changes: 7 additions & 0 deletions tests/test_set_visible_devices.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,13 @@ def test_visible_devices(self):
)
self.assertEqual(num_gpus_before, num_gpus_after)

# test import monai won't affect setting CUDA_VISIBLE_DEVICES
num_gpus_after_monai = self.run_process_and_get_exit_code(
'python -c "import os; import torch; import monai; '
+ "os.environ['CUDA_VISIBLE_DEVICES'] = '0'; exit(torch.cuda.device_count())\""
)
self.assertEqual(num_gpus_after_monai, 1)


if __name__ == "__main__":
unittest.main()

0 comments on commit 21f9ae1

Please sign in to comment.