pynvml string/bytes compatibility (#8981)

dask · Jan 10, 2025 · 0657de2 · 0657de2
1 parent c2ba834
commit 0657de2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/distributed/diagnostics/nvml.py b/distributed/diagnostics/nvml.py
@@ -32,7 +32,8 @@ class NVMLState(IntEnum):
 
 
 class CudaDeviceInfo(NamedTuple):
-    uuid: bytes | None = None
+    # Older versions of pynvml returned bytes, newer versions return str.
+    uuid: str | bytes | None = None
     device_index: int | None = None
     mig_index: int | None = None
 
@@ -278,13 +279,13 @@ def get_device_index_and_uuid(device):
     Examples
     --------
     >>> get_device_index_and_uuid(0)  # doctest: +SKIP
-    {'device-index': 0, 'uuid': b'GPU-e1006a74-5836-264f-5c26-53d19d212dfe'}
+    {'device-index': 0, 'uuid': 'GPU-e1006a74-5836-264f-5c26-53d19d212dfe'}
 
     >>> get_device_index_and_uuid('GPU-e1006a74-5836-264f-5c26-53d19d212dfe')  # doctest: +SKIP
-    {'device-index': 0, 'uuid': b'GPU-e1006a74-5836-264f-5c26-53d19d212dfe'}
+    {'device-index': 0, 'uuid': 'GPU-e1006a74-5836-264f-5c26-53d19d212dfe'}
 
     >>> get_device_index_and_uuid('MIG-7feb6df5-eccf-5faa-ab00-9a441867e237')  # doctest: +SKIP
-    {'device-index': 0, 'uuid': b'MIG-7feb6df5-eccf-5faa-ab00-9a441867e237'}
+    {'device-index': 0, 'uuid': 'MIG-7feb6df5-eccf-5faa-ab00-9a441867e237'}
     """
     init_once()
     try:

diff --git a/distributed/diagnostics/tests/test_nvml.py b/distributed/diagnostics/tests/test_nvml.py
@@ -11,6 +11,7 @@
 pynvml = pytest.importorskip("pynvml")
 
 import dask
+from dask.utils import ensure_unicode
 
 from distributed.diagnostics import nvml
 from distributed.utils_test import gen_cluster
@@ -66,7 +67,7 @@ def run_has_cuda_context(queue):
         assert (
             ctx.has_context
             and ctx.device_info.device_index == 0
-            and isinstance(ctx.device_info.uuid, bytes)
+            and isinstance(ctx.device_info.uuid, str)
         )
 
         queue.put(None)
@@ -127,7 +128,7 @@ def test_visible_devices_uuid():
     assert info.uuid
 
     with mock.patch.dict(
-        os.environ, {"CUDA_VISIBLE_DEVICES": info.uuid.decode("utf-8")}
+        os.environ, {"CUDA_VISIBLE_DEVICES": ensure_unicode(info.uuid)}
     ):
         h = nvml._pynvml_handles()
         h_expected = pynvml.nvmlDeviceGetHandleByIndex(0)
@@ -147,7 +148,7 @@ def test_visible_devices_uuid_2(index):
     assert info.uuid
 
     with mock.patch.dict(
-        os.environ, {"CUDA_VISIBLE_DEVICES": info.uuid.decode("utf-8")}
+        os.environ, {"CUDA_VISIBLE_DEVICES": ensure_unicode(info.uuid)}
     ):
         h = nvml._pynvml_handles()
         h_expected = pynvml.nvmlDeviceGetHandleByIndex(index)