Migrate pynvml to cuda.core.system by mdboom · Pull Request #640 · rapidsai/ucxx

mdboom · 2026-04-28T16:41:12Z

This migrates from pynvml.py to the new Cython/cybind-based cuda.core.system.

I was unable to get my local development set up for testing this one, so it's untested.

copy-pr-bot · 2026-04-28T16:41:16Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

pentschev · 2026-04-30T07:57:04Z

/ok to test 3b9b740

….system

pentschev · 2026-05-19T14:12:05Z

/ok to test 408fe66

pentschev · 2026-05-19T14:19:28Z

@mdboom there seems to be some style issues, could you try to automatically fix them with pre-commit install / pre-commit run -a?

mdboom · 2026-05-19T15:05:56Z

/ok to test 4e248fc

pentschev · 2026-05-19T15:43:31Z

/ok to test 4e248fc

pentschev · 2026-05-19T15:44:38Z

Unfortunately you're not part of the rapidsai org, so I think you have no power for triggering testing.

pentschev · 2026-05-19T20:29:36Z

/ok to test 9dfabb5

pentschev · 2026-05-20T07:29:02Z

@mdboom there are errors occurring that I think are legit issues, see here for example:

 │ │ Traceback (most recent call last):
 │ │   File "$SRC_DIR/conda_build_script.py", line 1, in <module>
 │ │     import ucxx
 │ │   File "$PREFIX/lib/python3.14/site-packages/ucxx/__init__.py", line 63, in <module>
 │ │     device_count = system.Device.get_device_count()
 │ │   File "cuda/core/system/_device.pyx", line 410, in cuda.core.system._device.Device.get_device_count
 │ │   File "cuda/core/system/_nvml_context.pxd", line 43, in cuda.core.system._nvml_context.initialize
 │ │   File "cuda/core/system/_nvml_context.pyx", line 32, in cuda.core.system._nvml_context._initialize
 │ │   File "cuda/core/system/_nvml_context.pyx", line 42, in cuda.core.system._nvml_context._initialize
 │ │   File "cuda/bindings/nvml.pyx", line 20860, in cuda.bindings.nvml.init_v2
 │ │   File "cuda/bindings/nvml.pyx", line 20866, in cuda.bindings.nvml.init_v2
 │ │     __status__ = nvmlInit_v2()
 │ │   File "cuda/bindings/cynvml.pyx", line 15, in cuda.bindings.cynvml.nvmlInit_v2
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 3892, in cuda.bindings._internal.nvml._nvmlInit_v2
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 2832, in cuda.bindings._internal.nvml._check_or_init_nvml
 │ │     return _init_nvml()
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 417, in cuda.bindings._internal.nvml._init_nvml
 │ │     with gil, __symbol_lock:
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 427, in cuda.bindings._internal.nvml._init_nvml
 │ │     handle = load_library()
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 409, in cuda.bindings._internal.nvml.load_library
 │ │     cdef uintptr_t handle = load_nvidia_dynamic_lib("nvml")._handle_uint
 │ │   File "$PREFIX/lib/python3.14/site-packages/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py", line 303, in load_nvidia_dynamic_lib
 │ │     return _load_lib_no_cache(libname)
 │ │   File "$PREFIX/lib/python3.14/site-packages/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py", line 158, in _load_lib_no_cache
 │ │     return _load_driver_lib_no_cache(desc)
 │ │   File "$PREFIX/lib/python3.14/site-packages/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py", line 71, in _load_driver_lib_no_cache
 │ │     raise DynamicLibNotFoundError(
 │ │     ...<2 lines>...
 │ │     )
 │ │ cuda.pathfinder._dynamic_libs.load_dl_common.DynamicLibNotFoundError: "nvml" is an NVIDIA driver library and can only be found via system search. Ensure the NVIDIA display driver is installed.

Are we supposed to bring NVML from conda to satisfy this?

mdboom · 2026-06-10T12:37:44Z

@mdboom there are errors occurring that I think are legit issues, see here for example:

 │ │ Traceback (most recent call last):
 │ │   File "$SRC_DIR/conda_build_script.py", line 1, in <module>
 │ │     import ucxx
 │ │   File "$PREFIX/lib/python3.14/site-packages/ucxx/__init__.py", line 63, in <module>
 │ │     device_count = system.Device.get_device_count()
 │ │   File "cuda/core/system/_device.pyx", line 410, in cuda.core.system._device.Device.get_device_count
 │ │   File "cuda/core/system/_nvml_context.pxd", line 43, in cuda.core.system._nvml_context.initialize
 │ │   File "cuda/core/system/_nvml_context.pyx", line 32, in cuda.core.system._nvml_context._initialize
 │ │   File "cuda/core/system/_nvml_context.pyx", line 42, in cuda.core.system._nvml_context._initialize
 │ │   File "cuda/bindings/nvml.pyx", line 20860, in cuda.bindings.nvml.init_v2
 │ │   File "cuda/bindings/nvml.pyx", line 20866, in cuda.bindings.nvml.init_v2
 │ │     __status__ = nvmlInit_v2()
 │ │   File "cuda/bindings/cynvml.pyx", line 15, in cuda.bindings.cynvml.nvmlInit_v2
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 3892, in cuda.bindings._internal.nvml._nvmlInit_v2
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 2832, in cuda.bindings._internal.nvml._check_or_init_nvml
 │ │     return _init_nvml()
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 417, in cuda.bindings._internal.nvml._init_nvml
 │ │     with gil, __symbol_lock:
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 427, in cuda.bindings._internal.nvml._init_nvml
 │ │     handle = load_library()
 │ │   File "cuda/bindings/_internal/nvml.pyx", line 409, in cuda.bindings._internal.nvml.load_library
 │ │     cdef uintptr_t handle = load_nvidia_dynamic_lib("nvml")._handle_uint
 │ │   File "$PREFIX/lib/python3.14/site-packages/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py", line 303, in load_nvidia_dynamic_lib
 │ │     return _load_lib_no_cache(libname)
 │ │   File "$PREFIX/lib/python3.14/site-packages/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py", line 158, in _load_lib_no_cache
 │ │     return _load_driver_lib_no_cache(desc)
 │ │   File "$PREFIX/lib/python3.14/site-packages/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py", line 71, in _load_driver_lib_no_cache
 │ │     raise DynamicLibNotFoundError(
 │ │     ...<2 lines>...
 │ │     )
 │ │ cuda.pathfinder._dynamic_libs.load_dl_common.DynamicLibNotFoundError: "nvml" is an NVIDIA driver library and can only be found via system search. Ensure the NVIDIA display driver is installed.

Are we supposed to bring NVML from conda to satisfy this?

Interesting. NVML should be installed alongside the driver, so should basically "always" be there (as long as the system can do anything with the GPU). I thought the search here was equivalent to what pynvml was doing, but perhaps not. I'll poke around and see what I find.

Created this bug to track the issue: NVIDIA/cuda-python#2189

rwgk · 2026-06-10T16:25:03Z

Hi Mike, I think the failing CI case is a subtle difference in where the missing-NVML-library error is surfaced after the migration from pynvml to cuda.core.system.

In the old code, UCXX did:

try:
    pynvml.nvmlInit()
    device_count = pynvml.nvmlDeviceGetCount()
    ...
except (
    pynvml.NVMLError_LibraryNotFound,
    pynvml.NVMLError_DriverNotLoaded,
    pynvml.NVMLError_Unknown,
):
    pass

For nvidia-ml-py, nvmlDeviceGetCount() itself does not load libnvidia-ml.so.1; nvmlInit() does. On Linux, nvmlInit() eventually calls:

CDLL("libnvidia-ml.so.1")

If that CDLL call fails, pynvml converts the OSError into pynvml.NVMLError_LibraryNotFound, which UCXX already catches and treats as "NVML is unavailable, skip this heuristic".

With the current new code:

try:
    device_count = system.Device.get_device_count()
    ...
except (
    system.NotFoundError,
    system.DriverNotLoadedError,
    system.UnknownError,
):
    pass

the missing-library error is currently escaping as:

cuda.pathfinder.DynamicLibNotFoundError

from load_nvidia_dynamic_lib("nvml"), before it is translated into a cuda.core.system / cuda.bindings.nvml LibraryNotFoundError. That makes this PR less forgiving than the old pynvml code in environments where NVML is unavailable.

My recommendation would be to add DynamicLibNotFoundError to the try/except around the cuda.core.system call:

from cuda.pathfinder import DynamicLibNotFoundError

try:
    device_count = system.Device.get_device_count()
    ...
except (
    DynamicLibNotFoundError,
    system.LibraryNotFoundError,
    system.DriverNotLoadedError,
    system.UnknownError,
):
    pass

The fact that DynamicLibNotFoundError currently leaks may also be worth fixing in cuda-python so cuda.core.system consistently translates pathfinder's missing-library error into its public NVML exception type. But for this UCXX PR, catching DynamicLibNotFoundError directly seems like a reasonable compatibility workaround.

mdboom · 2026-06-11T13:12:59Z

Thanks for the analysis, @rwgk. We can discuss in NVIDIA/cuda-python#2189 whether cuda.core.system can or should remap the exception. But in the meantime, I think fixing this here as you suggest makes sense.

Migrate pynvml to cuda.core.system

3b9b740

mdboom mentioned this pull request Apr 28, 2026

EPIC: Port Rapids projects from pynvml to cuda.core.system NVIDIA/cuda-python#1915

Open

wence- approved these changes Apr 30, 2026

View reviewed changes

mdboom added 2 commits May 19, 2026 09:40

Update to cuda_core 1.0.0

7197e09

Merge remote-tracking branch 'upstream/main' into pynvml-to-cuda.core…

408fe66

….system

mdboom marked this pull request as ready for review May 19, 2026 13:43

mdboom requested review from a team as code owners May 19, 2026 13:43

mdboom requested a review from msarahan May 19, 2026 13:43

pentschev added improvement Improves an existing functionality non-breaking Introduces a non-breaking change ucxx labels May 19, 2026

Fix style issues

4e248fc

Merge branch 'main' into pynvml-to-cuda.core.system

9dfabb5

mdboom mentioned this pull request Jun 10, 2026

[BUG]: cuda.bindings.nvml doesn't find libnvidia-ml.so.1 in situations where pynvml (ctypes) does NVIDIA/cuda-python#2189

Closed

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Migrate pynvml to cuda.core.system#640

Migrate pynvml to cuda.core.system#640
mdboom wants to merge 5 commits into
rapidsai:mainfrom
mdboom:pynvml-to-cuda.core.system

mdboom commented Apr 28, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Apr 28, 2026

Uh oh!

pentschev commented Apr 30, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

mdboom commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 20, 2026

Uh oh!

mdboom commented Jun 10, 2026 •

edited

Loading

Uh oh!

rwgk commented Jun 10, 2026

Uh oh!

mdboom commented Jun 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

mdboom commented Apr 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Apr 28, 2026

Uh oh!

pentschev commented Apr 30, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

mdboom commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 19, 2026

Uh oh!

pentschev commented May 20, 2026

Uh oh!

mdboom commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rwgk commented Jun 10, 2026

Uh oh!

mdboom commented Jun 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mdboom commented Apr 28, 2026 •

edited

Loading

mdboom commented Jun 10, 2026 •

edited

Loading