Raise error when CUDA_VISIBLE_DEVICES is set but empty

ml-energy · Aug 8, 2024 · 9415761 · 9415761
1 parent c563e1d
commit 9415761
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/zeus/device/gpu/amd.py b/zeus/device/gpu/amd.py
@@ -301,6 +301,12 @@ def _init_gpus(self) -> None:
         if (visible_device := os.environ.get("HIP_VISIBLE_DEVICES")) is not None or (
             visible_device := os.environ.get("CUDA_VISIBLE_DEVICES")
         ) is not None:
+            if not visible_device:
+                raise gpu_common.ZeusGPUInitError(
+                    "HIP_VISIBLE_DEVICES or CUDA_VISIBLE_DEVICES is set but empty. "
+                    "You can use either one for AMD GPUs, but it should either be unset "
+                    "or a comma-separated list of GPU indices."
+                )
             visible_indices = [int(idx) for idx in visible_device.split(",")]
         else:
             visible_indices = list(range(len(amdsmi.amdsmi_get_processor_handles())))

diff --git a/zeus/device/gpu/nvidia.py b/zeus/device/gpu/nvidia.py
@@ -368,6 +368,11 @@ def gpus(self) -> Sequence[gpu_common.GPU]:
     def _init_gpus(self) -> None:
         # Must respect `CUDA_VISIBLE_DEVICES` if set
         if (visible_device := os.environ.get("CUDA_VISIBLE_DEVICES")) is not None:
+            if not visible_device:
+                raise gpu_common.ZeusGPUInitError(
+                    "CUDA_VISIBLE_DEVICES is set to an empty string. "
+                    "It should either be unset or a comma-separated list of GPU indices."
+                )
             visible_indices = [int(idx) for idx in visible_device.split(",")]
         else:
             visible_indices = list(range(pynvml.nvmlDeviceGetCount()))