diff --git a/dask_cuda/tests/test_utils.py b/dask_cuda/tests/test_utils.py index ca17c097c..34e63f1b4 100644 --- a/dask_cuda/tests/test_utils.py +++ b/dask_cuda/tests/test_utils.py @@ -189,13 +189,16 @@ def test_parse_visible_devices(): uuids = [] for index in range(get_gpu_count()): handle = pynvml.nvmlDeviceGetHandleByIndex(index) - uuid = pynvml.nvmlDeviceGetUUID(handle).decode("utf-8") + try: + uuid = pynvml.nvmlDeviceGetUUID(handle).decode("utf-8") + except AttributeError: + uuid = pynvml.nvmlDeviceGetUUID(handle) assert parse_cuda_visible_device(index) == index assert parse_cuda_visible_device(uuid) == uuid indices.append(str(index)) - uuids.append(pynvml.nvmlDeviceGetUUID(handle).decode("utf-8")) + uuids.append(uuid) index_devices = ",".join(indices) with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": index_devices}): diff --git a/dask_cuda/utils.py b/dask_cuda/utils.py index 1a24d80b0..5e558fbc5 100644 --- a/dask_cuda/utils.py +++ b/dask_cuda/utils.py @@ -676,7 +676,10 @@ def get_gpu_uuid_from_index(device_index=0): pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(device_index) - return pynvml.nvmlDeviceGetUUID(handle).decode("utf-8") + try: + return pynvml.nvmlDeviceGetUUID(handle).decode("utf-8") + except AttributeError: + return pynvml.nvmlDeviceGetUUID(handle) def get_worker_config(dask_worker):