Migrate from pynvml to cuda.core.system

Author: mdboom

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -6,7 +6,8 @@ channels:
 - conda-forge
 dependencies:
 - click >=8.1
-- cuda-core>=0.3.2
+- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
+- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=12.9
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.6.*,>=0.0.0a0
 - numpy>=1.23,<3.0
 - numpydoc>=1.1.0
-- nvidia-ml-py>=12
 - pandas>=1.3
 - pre-commit
 - pytest

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -6,7 +6,8 @@ channels:
 - conda-forge
 dependencies:
 - click >=8.1
-- cuda-core>=0.3.2
+- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
+- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=12.9
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.6.*,>=0.0.0a0
 - numpy>=1.23,<3.0
 - numpydoc>=1.1.0
-- nvidia-ml-py>=12
 - pandas>=1.3
 - pre-commit
 - pytest

conda/environments/all_cuda-131_arch-aarch64.yaml

@@ -6,7 +6,8 @@ channels:
 - conda-forge
 dependencies:
 - click >=8.1
-- cuda-core>=0.3.2
+- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
+- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=13.1
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.6.*,>=0.0.0a0
 - numpy>=1.23,<3.0
 - numpydoc>=1.1.0
-- nvidia-ml-py>=12
 - pandas>=1.3
 - pre-commit
 - pytest

conda/environments/all_cuda-131_arch-x86_64.yaml

@@ -6,7 +6,8 @@ channels:
 - conda-forge
 dependencies:
 - click >=8.1
-- cuda-core>=0.3.2
+- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
+- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=13.1
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.6.*,>=0.0.0a0
 - numpy>=1.23,<3.0
 - numpydoc>=1.1.0
-- nvidia-ml-py>=12
 - pandas>=1.3
 - pre-commit
 - pytest

conda/recipes/dask-cuda/recipe.yaml

@@ -37,10 +37,8 @@ requirements:
   run:
     - python
     - click >=8.1
-    - cuda-core >=0.3
+    - cuda-core >=0.7.1
     - numpy >=1.23,<3.0
-    # 'nvidia-ml-py' provides the 'pynvml' module
-    - nvidia-ml-py>=12
     - pandas >=1.3
     - rapids-dask-dependency =${{ minor_version }}
     - zict >=2.0.0

dask_cuda/tests/test_initialize.py

@@ -20,7 +20,7 @@
 from distributed.deploy.local import LocalCluster
 
 from dask_cuda.initialize import initialize
-from dask_cuda.utils import get_ucx_config, get_gpu_handle
+from dask_cuda.utils import get_ucx_config, get_gpu
 from dask_cuda.utils_test import IncreasedCloseTimeoutNanny
 
 if CUDA_CORE_0_5_0():
@@ -39,13 +39,8 @@
 
 def _has_v100_gpu():
     """Return True if the first GPU (index 0) is a V100."""
-    import pynvml
-
-    handle = get_gpu_handle(0)
-    name = pynvml.nvmlDeviceGetName(handle)
-    if isinstance(name, bytes):
-        name = name.decode("utf-8", errors="ignore")
-    return "V100" in name
+    device = get_gpu(0)
+    return "V100" in device.name
 
 
 def _test_initialize_ucx_tcp():

dask_cuda/tests/test_utils.py

@@ -4,7 +4,6 @@
 import os
 from unittest.mock import patch
 
-import pynvml
 import pytest
 
 try:
@@ -16,6 +15,9 @@
     Device = cuda.core.experimental.Device
 
 
+from cuda.core import system
+
+
 from dask.config import canonical_name
 
 from dask_cuda.utils import (
@@ -30,7 +32,6 @@
     nvml_device_index,
     parse_cuda_visible_device,
     parse_device_memory_limit,
-    unpack_bitmask,
 )
 
 
@@ -62,30 +63,6 @@ def test_get_n_gpus():
     assert get_n_gpus() == 3
 
 
-@pytest.mark.parametrize(
-    "params",
-    [
-        {
-            "input": [1152920405096267775, 0],
-            "output": [i for i in range(20)] + [i + 40 for i in range(20)],
-        },
-        {
-            "input": [17293823668613283840, 65535],
-            "output": [i + 20 for i in range(20)] + [i + 60 for i in range(20)],
-        },
-        {"input": [18446744073709551615, 0], "output": [i for i in range(64)]},
-        {"input": [0, 18446744073709551615], "output": [i + 64 for i in range(64)]},
-    ],
-)
-def test_unpack_bitmask(params):
-    assert unpack_bitmask(params["input"]) == params["output"]
-
-
-def test_unpack_bitmask_single_value():
-    with pytest.raises(TypeError):
-        unpack_bitmask(1)
-
-
 def test_cpu_affinity():
     for i in range(get_n_gpus()):
         affinity = get_cpu_affinity(i)
@@ -220,15 +197,10 @@ def test_get_ucx_config(enable_tcp_over_ucx, enable_infiniband, enable_nvlink):
 
 
 def test_parse_visible_devices():
-    pynvml.nvmlInit()
     indices = []
     uuids = []
-    for index in range(get_gpu_count()):
-        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
-        try:
-            uuid = pynvml.nvmlDeviceGetUUID(handle).decode("utf-8")
-        except AttributeError:
-            uuid = pynvml.nvmlDeviceGetUUID(handle)
+    for index, device in enumerate(system.Device.get_all_devices()):
+        uuid = device.uuid
 
         assert parse_cuda_visible_device(index) == index
         assert parse_cuda_visible_device(uuid) == uuid
@@ -350,12 +322,10 @@ def test_has_device_memory_resoure():
 
 
 def test_parse_visible_mig_devices():
-    pynvml.nvmlInit()
-    for index in range(get_gpu_count()):
-        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
+    for device in system.Device.get_all_devices():
         try:
-            mode = pynvml.nvmlDeviceGetMigMode(handle)[0]
-        except pynvml.NVMLError:
+            mode = device.mig.mode
+        except system.NvmlError:
             # if not a MIG device, i.e. a normal GPU, skip
             continue
         if mode:
@@ -364,14 +334,6 @@ def test_parse_visible_mig_devices():
             # in that GPU is <= to count, where count gives us the
             # maximum number of MIG devices/instances that can exist
             # under a given parent NVML device.
-            count = pynvml.nvmlDeviceGetMaxMigDeviceCount(handle)
-            miguuids = []
-            for i in range(count):
-                try:
-                    mighandle = pynvml.nvmlDeviceGetMigDeviceHandleByIndex(
-                        device=handle, index=i
-                    )
-                    miguuids.append(mighandle)
-                except pynvml.NVMLError:
-                    pass
-            assert len(miguuids) <= count
+            mig_devices = list(device.mig.get_all_devices())
+            count = device.mig.get_device_count()
+            assert len(mig_devices) <= count