Support pynvml>=11.5 in WSL (#8962)

* move decode into try block * use _maybe_decode * address code review
dask · Dec 11, 2024 · 3ef88b3 · 3ef88b3
1 parent 6eabc90
commit 3ef88b3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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
@@ -105,7 +106,7 @@ def init_once():
             return
 
         if _in_wsl() and parse_version(
-            pynvml.nvmlSystemGetDriverVersion().decode()
+            ensure_unicode(pynvml.nvmlSystemGetDriverVersion())
         ) < parse_version(MINIMUM_WSL_VERSION):
             NVML_STATE = NVMLState.DISABLED_WSL_INSUFFICIENT_DRIVER
             return