Migrate from pynvml to cuda.core.system

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>=1.0.0
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=12.9
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.8.*,>=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>=1.0.0
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=12.9
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.8.*,>=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-132_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>=1.0.0
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=13.2
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.8.*,>=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-132_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>=1.0.0
 - cuda-nvcc-impl
 - cuda-nvrtc
 - cuda-version=13.2
@@ -16,7 +17,6 @@ dependencies:
 - kvikio==26.8.*,>=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,9 @@ requirements:
   run:
     - python
     - click >=8.1
-    - cuda-core >=0.3
+    - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
+    - cuda-core >=1.0.0
     - 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/_compat.py

@@ -17,13 +17,3 @@ def DASK_2025_4_0():
     # dask 2025.4.0 isn't currently released, so we're relying
     # on strictly greater than here.
     return get_dask_version() > packaging.version.parse("2025.3.0")
-
-
-@functools.lru_cache(maxsize=None)
-def get_cuda_core_version() -> packaging.version.Version:
-    return packaging.version.parse(importlib.metadata.version("cuda-core"))
-
-
-@functools.lru_cache(maxsize=None)
-def CUDA_CORE_0_5_0():
-    return get_cuda_core_version() >= packaging.version.parse("0.5.0")

dask_cuda/initialize.py

@@ -6,13 +6,7 @@
 
 import click
 
-try:
-    from cuda.core import Device
-except ImportError:
-    # Remove when cuda-core>=0.5
-    import cuda.core.experimental
-
-    Device = cuda.core.experimental.Device
+from cuda.core import Device
 
 import dask
 from distributed.diagnostics.nvml import (

dask_cuda/tests/test_initialize.py

@@ -9,8 +9,6 @@
 import tempfile
 import textwrap
 
-from dask_cuda._compat import CUDA_CORE_0_5_0
-
 import numpy
 import psutil
 import pytest
@@ -20,13 +18,10 @@
 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():
-    from cuda.core import system
-else:
-    from cuda.core.experimental import system
+from cuda.core import system
 
 
 mp = mp.get_context("spawn")  # type: ignore
@@ -39,13 +34,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():
@@ -268,12 +258,7 @@ def _test_cuda_context_warning_with_subprocess_warnings(protocol):
         # Problematic library that creates CUDA context at import time
         import os
 
-        try:
-            from cuda.core import Device
-        except ImportError:
-            # cuda-core < 0.5
-            import cuda.core.experimental
-            Device = cuda.core.experimental.Device
+        from cuda.core import Device
 
         try:
             # Create CUDA context at import time, this will be inherited by subprocesses
@@ -378,10 +363,7 @@ def _test_cuda_context_warning_with_subprocess_warnings(protocol):
             ):
                 warnings_assigned_device_found.append(line)
 
-        if CUDA_CORE_0_5_0():
-            num_devices = system.get_num_devices()
-        else:
-            num_devices = system.num_devices
+        num_devices = system.get_num_devices()
 
         # Every worker raises the warning once. With protocol="ucx" the warning is
         # raised once more by the parent process.

dask_cuda/tests/test_utils.py

@@ -4,16 +4,12 @@
 import os
 from unittest.mock import patch
 
-import pynvml
 import pytest
 
-try:
-    from cuda.core import Device
-except ImportError:
-    # Remove when cuda-core>=0.5
-    import cuda.core.experimental
+from cuda.core import Device
 
-    Device = cuda.core.experimental.Device
+
+from cuda.core import system
 
 
 from dask.config import canonical_name
@@ -30,7 +26,6 @@
     nvml_device_index,
     parse_cuda_visible_device,
     parse_device_memory_limit,
-    unpack_bitmask,
 )
 
 
@@ -62,30 +57,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)
@@ -232,15 +203,10 @@ def test_get_ucx_config_automatic_transport_selection():
 
 
 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
@@ -362,12 +328,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:
@@ -376,14 +340,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