From 8134e6bebf50e4b4b428e07eebc98c3b2e851774 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Wed, 15 Feb 2023 14:11:28 +0100 Subject: [PATCH] Fix for bytes/str discrepancy after PyNVML update (#1118) A PyNVML update has changed how some objects were previously returned bytes but now return str, which is now handled appropriately with this change. Authors: - Peter Andreas Entschev (https://github.com/pentschev) Approvers: - Lawrence Mitchell (https://github.com/wence-) URL: https://github.com/rapidsai/dask-cuda/pull/1118 --- dask_cuda/tests/test_utils.py | 7 +++++-- dask_cuda/utils.py | 5 ++++- 2 files changed, 9 insertions(+), 3 deletions(-) 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):