pynvml compatibility

Seems like some version pynvml switched from returning bytes to str (I
can pin down the exact version if needed). This updates our test and
type annotations to accomodate either.

Author: TomAugspurger

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:

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)