Skip to content

Migrate pynvml.py to cuda.core.system #8022

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mdboom wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Migrate pynvml.py to cuda.core.system #8022

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ee4b32c
Migrate pynvml.py to cuda.core.system
mdboom Apr 28, 2026
9e27372
Update to cuda_core 1.0.0
mdboom May 19, 2026
ebc64c6
Merge remote-tracking branch 'upstream/main' into pynvml-to-cuda.core…
mdboom May 19, 2026
2f3b843
Merge remote-tracking branch 'upstream/main' into pynvml-to-cuda.core…
mdboom Jun 10, 2026
0cb7bb0
Merge remote-tracking branch 'origin/main' into pr-8022
csadorf Jun 11, 2026
1f7cf38
Fix import order
csadorf Jun 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion conda/environments/all_cuda-129_arch-aarch64.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ dependencies:
- ccache
- certifi
- cmake>=4.0
- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
- cuda-core>=1.0.0
- cuda-cudart-dev
- cuda-nvcc
- cuda-profiler-api
Expand Down Expand Up @@ -53,7 +55,6 @@ dependencies:
- numpydoc
- numpydoc<1.9
- nvforest==26.8.*,>=0.0.0a0
- nvidia-ml-py>=12
- onnxruntime
- packaging
- pre-commit
Expand Down
3 changes: 2 additions & 1 deletion conda/environments/all_cuda-129_arch-x86_64.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ dependencies:
- ccache
- certifi
- cmake>=4.0
- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
- cuda-core>=1.0.0
- cuda-cudart-dev
- cuda-nvcc
- cuda-profiler-api
Expand Down Expand Up @@ -52,7 +54,6 @@ dependencies:
- numpydoc
- numpydoc<1.9
- nvforest==26.8.*,>=0.0.0a0
- nvidia-ml-py>=12
- onnxruntime
- packaging
- pre-commit
Expand Down
3 changes: 2 additions & 1 deletion conda/environments/all_cuda-132_arch-aarch64.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ dependencies:
- ccache
- certifi
- cmake>=4.0
- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
- cuda-core>=1.0.0
- cuda-cudart-dev
- cuda-nvcc
- cuda-profiler-api
Expand Down Expand Up @@ -53,7 +55,6 @@ dependencies:
- numpydoc
- numpydoc<1.9
- nvforest==26.8.*,>=0.0.0a0
- nvidia-ml-py>=12
- onnxruntime
- packaging
- pre-commit
Expand Down
3 changes: 2 additions & 1 deletion conda/environments/all_cuda-132_arch-x86_64.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ dependencies:
- ccache
- certifi
- cmake>=4.0
- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*
- cuda-core>=1.0.0
- cuda-cudart-dev
- cuda-nvcc
- cuda-profiler-api
Expand Down Expand Up @@ -52,7 +54,6 @@ dependencies:
- numpydoc
- numpydoc<1.9
- nvforest==26.8.*,>=0.0.0a0
- nvidia-ml-py>=12
- onnxruntime
- packaging
- pre-commit
Expand Down
4 changes: 2 additions & 2 deletions dependencies.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -621,8 +621,8 @@ dependencies:
- nltk
# upstream sklearn docstring tests require numpydoc<1.9
- numpydoc<1.9
# 'nvidia-ml-py' provides the 'pynvml' module
- nvidia-ml-py>=12
- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*

@betatim betatim May 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do the versions here correspond to the version of CUDA that the user has installed?

Will this also work for the -cu13 package that we upload to pypi?

- cuda-core>=1.0.0
# TODO: uncomment these two lines once `onnxruntime` has Python 3.14 support (and conda packages built)
# - *onnxruntime
# - *skl2onnx
Expand Down
3 changes: 2 additions & 1 deletion python/cuml/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,8 @@ classifiers = [
[project.optional-dependencies]
test = [
"certifi",
"cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*",
"cuda-core>=1.0.0",
"cython>=3.2.2",
"hdbscan>=0.8.39",
"hypothesis>=6.0,<7",
Expand All @@ -118,7 +120,6 @@ test = [
"msgspec",
"nltk",
"numpydoc<1.9",
"nvidia-ml-py>=12",
"pynndescent",
"pytest",
"pytest-benchmark",
Expand Down
34 changes: 15 additions & 19 deletions python/cuml/tests/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@
import cupy as cp
import hypothesis
import numpy as np
import pynvml
import pytest
from cuda.core import system
from sklearn import datasets

from cuml.testing.datasets import make_text_classification_dataset
Expand Down Expand Up @@ -277,7 +277,7 @@ def pytest_pyfunc_call(pyfuncitem):
pytest.skip("Test requires cudf.pandas accelerator")


def _get_pynvml_device_handle(device_id=0):
def _get_nvml_device_handle(device_id=0):
"""Get GPU handle from device index or UUID.

Parameters
Expand All @@ -289,36 +289,32 @@ def _get_pynvml_device_handle(device_id=0):
------
ValueError
If acquiring the device handle for the device specified failed.
pynvml.NVMLError
cuda.core.system.NvmlError
If any NVML error occurred while initializing.

Returns
-------
A pynvml handle to the device.
A cuda.core.system.Device.

Examples
--------
>>> _get_pynvml_device_handle(device_id=0)
>>> _get_nvml_device_handle(device_id=0)

>>> _get_pynvml_device_handle(device_id="GPU-9fb42d6f-7d6b-368f-f79c-3c3e784c93f6")
>>> _get_nvml_device_handle(device_id="GPU-9fb42d6f-7d6b-368f-f79c-3c3e784c93f6")
"""
pynvml.nvmlInit()

try:
if device_id and not str(device_id).isnumeric():
# This means device_id is UUID.
# This works for both MIG and non-MIG device UUIDs.
handle = pynvml.nvmlDeviceGetHandleByUUID(str.encode(device_id))
if pynvml.nvmlDeviceIsMigDeviceHandle(handle):
device = system.Device(uuid=device_id)
if device.mig.is_mig_device:
# Additionally get parent device handle
# if the device itself is a MIG instance
handle = pynvml.nvmlDeviceGetDeviceHandleFromMigDeviceHandle(
handle
)
device = device.mig.parent
else:
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
return handle
except pynvml.NVMLError:
device = system.Device(index=device_id)
return device
except system.NvmlError:
raise ValueError(f"Invalid device index or UUID: {device_id}")


Expand All @@ -336,12 +332,12 @@ def _get_gpu_memory(device_index=0):
have a dedicated memory resource, as is usually the case for system on a chip (SoC)
devices.
"""
handle = _get_pynvml_device_handle(device_index)
device = _get_nvml_device_handle(device_index)

try:
# Return total memory in GB
return ceil(pynvml.nvmlDeviceGetMemoryInfo(handle).total / 2**30)
except pynvml.NVMLError_NotSupported:
return ceil(device.memory_info.total / 2**30)
except system.NotSupportedError:
return None


Expand Down
Loading