From 9d72fe73850a08ec0a9ec0d3f57fe54ecee69db7 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 10 Dec 2024 16:20:28 -0800 Subject: [PATCH 1/3] move decode into try block --- distributed/diagnostics/nvml.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/distributed/diagnostics/nvml.py b/distributed/diagnostics/nvml.py index f621a7eb766..360a984f9c0 100644 --- a/distributed/diagnostics/nvml.py +++ b/distributed/diagnostics/nvml.py @@ -104,9 +104,11 @@ def init_once(): NVML_STATE = NVMLState.DISABLED_LIBRARY_NOT_FOUND return - if _in_wsl() and parse_version( - pynvml.nvmlSystemGetDriverVersion().decode() - ) < parse_version(MINIMUM_WSL_VERSION): + try: + driver_vsn = pynvml.nvmlSystemGetDriverVersion().decode() + except AttributeError: + driver_vsn = pynvml.nvmlSystemGetDriverVersion() + if _in_wsl() and parse_version(driver_vsn) < parse_version(MINIMUM_WSL_VERSION): NVML_STATE = NVMLState.DISABLED_WSL_INSUFFICIENT_DRIVER return else: From b004b1f88fc7c4203b54f871c2cf92e92b21ba99 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 10 Dec 2024 18:27:08 -0800 Subject: [PATCH 2/3] use _maybe_decode --- distributed/diagnostics/nvml.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/distributed/diagnostics/nvml.py b/distributed/diagnostics/nvml.py index 360a984f9c0..fee6ef46e48 100644 --- a/distributed/diagnostics/nvml.py +++ b/distributed/diagnostics/nvml.py @@ -68,6 +68,11 @@ def _in_wsl(): return "microsoft-standard" in uname().release +def _maybe_decode(value): + """Decode if bytes instance""" + return value.decode() if isinstance(value, bytes) else value + + def init_once(): """Idempotent (per-process) initialization of PyNVML @@ -104,11 +109,9 @@ def init_once(): NVML_STATE = NVMLState.DISABLED_LIBRARY_NOT_FOUND return - try: - driver_vsn = pynvml.nvmlSystemGetDriverVersion().decode() - except AttributeError: - driver_vsn = pynvml.nvmlSystemGetDriverVersion() - if _in_wsl() and parse_version(driver_vsn) < parse_version(MINIMUM_WSL_VERSION): + if _in_wsl() and parse_version( + _maybe_decode(pynvml.nvmlSystemGetDriverVersion()) + ) < parse_version(MINIMUM_WSL_VERSION): NVML_STATE = NVMLState.DISABLED_WSL_INSUFFICIENT_DRIVER return else: From 0bdd63d8c7689a597f527e460d04d1eddf3d23d9 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 11 Dec 2024 06:18:10 -0800 Subject: [PATCH 3/3] address code review --- distributed/diagnostics/nvml.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/distributed/diagnostics/nvml.py b/distributed/diagnostics/nvml.py index fee6ef46e48..092848c6877 100644 --- a/distributed/diagnostics/nvml.py +++ b/distributed/diagnostics/nvml.py @@ -8,6 +8,7 @@ from packaging.version import parse as parse_version import dask +from dask.utils import ensure_unicode try: import pynvml @@ -68,11 +69,6 @@ def _in_wsl(): return "microsoft-standard" in uname().release -def _maybe_decode(value): - """Decode if bytes instance""" - return value.decode() if isinstance(value, bytes) else value - - def init_once(): """Idempotent (per-process) initialization of PyNVML @@ -110,7 +106,7 @@ def init_once(): return if _in_wsl() and parse_version( - _maybe_decode(pynvml.nvmlSystemGetDriverVersion()) + ensure_unicode(pynvml.nvmlSystemGetDriverVersion()) ) < parse_version(MINIMUM_WSL_VERSION): NVML_STATE = NVMLState.DISABLED_WSL_INSUFFICIENT_DRIVER return