From abe2e3151e77fdbcad68268111aaabf1e73c3dbe Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Mon, 22 Jan 2024 23:58:29 +0800
Subject: [PATCH] Fix CUDA_VISIBLE_DEVICES setting ignored (#7408)

Fixes #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>
---
 monai/metrics/utils.py            | 6 ++----
 tests/test_set_visible_devices.py | 7 +++++++
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/monai/metrics/utils.py b/monai/metrics/utils.py
index d4b8f6e9b6..e7057256fb 100644
--- a/monai/metrics/utils.py
+++ b/monai/metrics/utils.py
@@ -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",
@@ -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
diff --git a/tests/test_set_visible_devices.py b/tests/test_set_visible_devices.py
index 53703e107a..993e8a4ac2 100644
--- a/tests/test_set_visible_devices.py
+++ b/tests/test_set_visible_devices.py
@@ -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()