From 3b9b74006655baaeba04827c39e9208c5ab7795c Mon Sep 17 00:00:00 2001 From: Michael Droettboom Date: Tue, 28 Apr 2026 12:38:37 -0400 Subject: [PATCH 1/3] Migrate pynvml to cuda.core.system --- .../all_cuda-129_arch-aarch64.yaml | 3 +- .../all_cuda-129_arch-x86_64.yaml | 3 +- .../all_cuda-131_arch-aarch64.yaml | 3 +- .../all_cuda-131_arch-x86_64.yaml | 3 +- conda/recipes/ucxx/recipe.yaml | 11 +++--- dependencies.yaml | 6 ++-- python/ucxx/pyproject.toml | 4 +-- python/ucxx/ucxx/__init__.py | 36 ++++++------------- python/ucxx/ucxx/_lib_async/utils_test.py | 10 +++--- 9 files changed, 34 insertions(+), 45 deletions(-) diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 3047261a4..d72a9d669 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -10,6 +10,8 @@ dependencies: - c-compiler - cloudpickle - cmake>=3.26.4,!=3.30.0 +- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* +- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 - cuda-cudart-dev - cuda-nvcc @@ -28,7 +30,6 @@ dependencies: - ninja - numba-cuda>=0.22.1,<0.29.0 - numpy>=1.23,<3.0 -- nvidia-ml-py>=12 - pip - pkg-config - pre-commit diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index ab416cffc..d2d08b7b1 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -10,6 +10,8 @@ dependencies: - c-compiler - cloudpickle - cmake>=3.26.4,!=3.30.0 +- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* +- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 - cuda-cudart-dev - cuda-nvcc @@ -28,7 +30,6 @@ dependencies: - ninja - numba-cuda>=0.22.1,<0.29.0 - numpy>=1.23,<3.0 -- nvidia-ml-py>=12 - pip - pkg-config - pre-commit diff --git a/conda/environments/all_cuda-131_arch-aarch64.yaml b/conda/environments/all_cuda-131_arch-aarch64.yaml index 83114c6bc..f601417dc 100644 --- a/conda/environments/all_cuda-131_arch-aarch64.yaml +++ b/conda/environments/all_cuda-131_arch-aarch64.yaml @@ -10,6 +10,8 @@ dependencies: - c-compiler - cloudpickle - cmake>=3.26.4,!=3.30.0 +- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* +- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 - cuda-cudart-dev - cuda-nvcc @@ -28,7 +30,6 @@ dependencies: - ninja - numba-cuda>=0.22.1,<0.29.0 - numpy>=1.23,<3.0 -- nvidia-ml-py>=12 - pip - pkg-config - pre-commit diff --git a/conda/environments/all_cuda-131_arch-x86_64.yaml b/conda/environments/all_cuda-131_arch-x86_64.yaml index b44847323..ab6b38088 100644 --- a/conda/environments/all_cuda-131_arch-x86_64.yaml +++ b/conda/environments/all_cuda-131_arch-x86_64.yaml @@ -10,6 +10,8 @@ dependencies: - c-compiler - cloudpickle - cmake>=3.26.4,!=3.30.0 +- cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* +- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 - cuda-cudart-dev - cuda-nvcc @@ -28,7 +30,6 @@ dependencies: - ninja - numba-cuda>=0.22.1,<0.29.0 - numpy>=1.23,<3.0 -- nvidia-ml-py>=12 - pip - pkg-config - pre-commit diff --git a/conda/recipes/ucxx/recipe.yaml b/conda/recipes/ucxx/recipe.yaml index 02ddcce1b..9e473762d 100644 --- a/conda/recipes/ucxx/recipe.yaml +++ b/conda/recipes/ucxx/recipe.yaml @@ -87,8 +87,6 @@ outputs: host: - cuda-version =${{ cuda_version }} - cython >=3.2.2 - # 'nvidia-ml-py' provides the 'pynvml' module - - nvidia-ml-py>=12 - pip - python =${{ py_abi_min }} - python-abi3 ${{ py_abi_min }}.* @@ -98,16 +96,19 @@ outputs: - ucx - libucxx =${{ version }} - cuda-cudart-dev + - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* + # TODO: Change to cuda-core >= 1.0.0 once that's released + - cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core run: - - cuda-core >=0.3.2 - numpy >=1.23,<3.0 - # 'nvidia-ml-py' provides the 'pynvml' module - - nvidia-ml-py>=12 - python - ucx >=1.18.0,<1.21.0 - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - ${{ pin_compatible("rmm", upper_bound="x.x") }} - libucxx =${{ version }} + - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* + # TODO: Change to cuda-core >= 1.0.0 once that's released + - cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core run_constraints: - cupy >=13.6.0 ignore_run_exports: diff --git a/dependencies.yaml b/dependencies.yaml index 5f1b91d71..fe0ee3121 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -346,9 +346,9 @@ dependencies: - output_types: [conda, requirements, pyproject] packages: - &numpy numpy>=1.23,<3.0 - # 'nvidia-ml-py' provides the 'pynvml' module - - nvidia-ml-py>=12 - - cuda-core>=0.3.2 + - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* + # TODO: Change to cuda-core >= 1.0.0 once that's released + - cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core run_python_distributed_ucxx: common: - output_types: [conda, requirements, pyproject] diff --git a/python/ucxx/pyproject.toml b/python/ucxx/pyproject.toml index 15de3d59c..6fe6993d0 100644 --- a/python/ucxx/pyproject.toml +++ b/python/ucxx/pyproject.toml @@ -19,10 +19,10 @@ authors = [ license = "BSD-3-Clause" requires-python = ">=3.11" dependencies = [ - "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", "libucxx==0.50.*,>=0.0.0a0", "numpy>=1.23,<3.0", - "nvidia-ml-py>=12", "rmm==26.6.*,>=0.0.0a0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. classifiers = [ diff --git a/python/ucxx/ucxx/__init__.py b/python/ucxx/ucxx/__init__.py index a727cbe41..5e07990ed 100644 --- a/python/ucxx/ucxx/__init__.py +++ b/python/ucxx/ucxx/__init__.py @@ -33,10 +33,7 @@ from .core import * # noqa: E402, F403 from .utils import get_address, get_ucxpy_logger # noqa: E402 -try: - import pynvml -except ImportError: - pynvml = None +from cuda.core import system _ucx_version = get_ucx_version() # noqa: F405 __ucx_min_version__ = "1.18.0" @@ -62,28 +59,17 @@ os.environ["UCX_RNDV_FRAG_MEM_TYPE"] = "cuda" if ( - pynvml is not None and "UCX_CUDA_COPY_MAX_REG_RATIO" not in os.environ and _ucx_version >= (1, 12, 0) ): try: - pynvml.nvmlInit() - device_count = pynvml.nvmlDeviceGetCount() + device_count = system.Device.get_device_count() large_bar1 = [False] * device_count - def _is_mig_device(handle): - try: - pynvml.nvmlDeviceGetMigMode(handle)[0] - except pynvml.NVMLError: - return False - return True - - for dev_idx in range(device_count): - handle = pynvml.nvmlDeviceGetHandleByIndex(dev_idx) - + for dev_idx, device in enumerate(system.Device.get_all_devices()): try: - total_memory = pynvml.nvmlDeviceGetMemoryInfo(handle).total - except pynvml.NVMLError_NotSupported: + total_memory = device.memory_info.total + except system.NotSupportedError: total_memory = None # Ignore MIG devices and devices with no memory resource (i.e., only @@ -91,12 +77,12 @@ def _is_mig_device(handle): # now. Increasing `UCX_CUDA_COPY_MAX_REG_RATIO` should be thoroughly # tested, as it's not yet clear whether it would be safe to set `1.0` # for those instances too. - if _is_mig_device(handle) or total_memory is None: + if device.mig.is_mig_device or total_memory is None: continue try: - bar1_total = pynvml.nvmlDeviceGetBAR1MemoryInfo(handle).bar1Total - except pynvml.NVMLError_NotSupported: + bar1_total = device.bar1_memory_info.total + except system.NotSupportedError: # Bar1 access not supported on this device, set it to # zero (always lower than device memory). bar1_total = 0 @@ -108,9 +94,9 @@ def _is_mig_device(handle): logger.info("Setting UCX_CUDA_COPY_MAX_REG_RATIO=1.0") os.environ["UCX_CUDA_COPY_MAX_REG_RATIO"] = "1.0" except ( - pynvml.NVMLError_LibraryNotFound, - pynvml.NVMLError_DriverNotLoaded, - pynvml.NVMLError_Unknown, + system.NotFoundError, + system.DriverNotLoadedError, + system.UnknownError, ): pass diff --git a/python/ucxx/ucxx/_lib_async/utils_test.py b/python/ucxx/ucxx/_lib_async/utils_test.py index e6b6359e5..da959961d 100644 --- a/python/ucxx/ucxx/_lib_async/utils_test.py +++ b/python/ucxx/ucxx/_lib_async/utils_test.py @@ -14,6 +14,9 @@ from ucxx._lib_async.pytest_stash_keys import ASYNCIO_PLUGIN_TIMEOUT_STASH_KEY +from cuda.core import system + + normal_env = { "UCX_RNDV_SCHEME": "put_zcopy", "UCX_MEMTYPE_CACHE": "n", @@ -27,12 +30,7 @@ def set_env(): def get_num_gpus(): - import pynvml - - pynvml.nvmlInit() - ngpus = pynvml.nvmlDeviceGetCount() - pynvml.nvmlShutdown() - return ngpus + return system.Device.get_device_count() def get_cuda_devices(): From 7197e0944aa1f4858d48a4814ec0adc55438fd43 Mon Sep 17 00:00:00 2001 From: Michael Droettboom Date: Tue, 19 May 2026 09:40:13 -0400 Subject: [PATCH 2/3] Update to cuda_core 1.0.0 --- conda/environments/all_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-129_arch-x86_64.yaml | 2 +- conda/environments/all_cuda-131_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-131_arch-x86_64.yaml | 2 +- conda/recipes/ucxx/recipe.yaml | 8 +++----- dependencies.yaml | 3 +-- python/ucxx/pyproject.toml | 2 +- python/ucxx/ucxx/_cuda_context.py | 2 +- 8 files changed, 10 insertions(+), 13 deletions(-) diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index d72a9d669..99b321e15 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -11,8 +11,8 @@ dependencies: - cloudpickle - cmake>=3.26.4,!=3.30.0 - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* -- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 +- cuda-core>=1.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-version=12.9 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index d2d08b7b1..137fa7a35 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -11,8 +11,8 @@ dependencies: - cloudpickle - cmake>=3.26.4,!=3.30.0 - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* -- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 +- cuda-core>=1.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-version=12.9 diff --git a/conda/environments/all_cuda-131_arch-aarch64.yaml b/conda/environments/all_cuda-131_arch-aarch64.yaml index f601417dc..5973e014c 100644 --- a/conda/environments/all_cuda-131_arch-aarch64.yaml +++ b/conda/environments/all_cuda-131_arch-aarch64.yaml @@ -11,8 +11,8 @@ dependencies: - cloudpickle - cmake>=3.26.4,!=3.30.0 - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* -- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 +- cuda-core>=1.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-version=13.1 diff --git a/conda/environments/all_cuda-131_arch-x86_64.yaml b/conda/environments/all_cuda-131_arch-x86_64.yaml index ab6b38088..b78ecc973 100644 --- a/conda/environments/all_cuda-131_arch-x86_64.yaml +++ b/conda/environments/all_cuda-131_arch-x86_64.yaml @@ -11,8 +11,8 @@ dependencies: - cloudpickle - cmake>=3.26.4,!=3.30.0 - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* -- cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core - cuda-core>=0.3.2 +- cuda-core>=1.0.0 - cuda-cudart-dev - cuda-nvcc - cuda-version=13.1 diff --git a/conda/recipes/ucxx/recipe.yaml b/conda/recipes/ucxx/recipe.yaml index 9e473762d..13b4a0573 100644 --- a/conda/recipes/ucxx/recipe.yaml +++ b/conda/recipes/ucxx/recipe.yaml @@ -97,8 +97,7 @@ outputs: - libucxx =${{ version }} - cuda-cudart-dev - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* - # TODO: Change to cuda-core >= 1.0.0 once that's released - - cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core + - cuda-core>=1.0.0 run: - numpy >=1.23,<3.0 - python @@ -107,8 +106,7 @@ outputs: - ${{ pin_compatible("rmm", upper_bound="x.x") }} - libucxx =${{ version }} - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* - # TODO: Change to cuda-core >= 1.0.0 once that's released - - cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core + - cuda-core>=1.0.0 run_constraints: - cupy >=13.6.0 ignore_run_exports: @@ -242,7 +240,7 @@ outputs: - setuptools>=77.0.0 - wheel run: - - cuda-core >=0.3.2 + - cuda-core>=1.0.0 - python - pyyaml >=6 - rapids-dask-dependency ${{ rapids_version }} diff --git a/dependencies.yaml b/dependencies.yaml index fe0ee3121..b392c1021 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -347,8 +347,7 @@ dependencies: packages: - &numpy numpy>=1.23,<3.0 - cuda-bindings>=12.9.6,!=13.0.*,!=13.1.* - # TODO: Change to cuda-core >= 1.0.0 once that's released - - cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core + - cuda-core>=1.0.0 run_python_distributed_ucxx: common: - output_types: [conda, requirements, pyproject] diff --git a/python/ucxx/pyproject.toml b/python/ucxx/pyproject.toml index 6fe6993d0..6e45916ff 100644 --- a/python/ucxx/pyproject.toml +++ b/python/ucxx/pyproject.toml @@ -20,7 +20,7 @@ license = "BSD-3-Clause" requires-python = ">=3.11" dependencies = [ "cuda-bindings>=12.9.6,!=13.0.*,!=13.1.*", - "cuda-core @ git+https://github.com/nvidia/cuda-python@main#subdirectory=cuda_core", + "cuda-core>=1.0.0", "libucxx==0.50.*,>=0.0.0a0", "numpy>=1.23,<3.0", "rmm==26.6.*,>=0.0.0a0", diff --git a/python/ucxx/ucxx/_cuda_context.py b/python/ucxx/ucxx/_cuda_context.py index a4d32f518..abb54aeac 100644 --- a/python/ucxx/ucxx/_cuda_context.py +++ b/python/ucxx/ucxx/_cuda_context.py @@ -21,7 +21,7 @@ def _get_device_class(): return Device except ImportError as e: raise ImportError( - "CUDA context management requires cuda-core (cuda-core>=0.3.2)." + "CUDA context management requires cuda-core (cuda-core>=1.0.0)." ) from e From 4e248fc5185f8d6000ea41a207ca8cc35e6b138b Mon Sep 17 00:00:00 2001 From: Michael Droettboom Date: Tue, 19 May 2026 11:05:10 -0400 Subject: [PATCH 3/3] Fix style issues --- python/ucxx/ucxx/__init__.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/python/ucxx/ucxx/__init__.py b/python/ucxx/ucxx/__init__.py index 5e07990ed..5f7fd3e61 100644 --- a/python/ucxx/ucxx/__init__.py +++ b/python/ucxx/ucxx/__init__.py @@ -33,7 +33,7 @@ from .core import * # noqa: E402, F403 from .utils import get_address, get_ucxpy_logger # noqa: E402 -from cuda.core import system +from cuda.core import system # noqa: E402 _ucx_version = get_ucx_version() # noqa: F405 __ucx_min_version__ = "1.18.0" @@ -58,10 +58,7 @@ logger.info("Setting UCX_RNDV_FRAG_MEM_TYPE=cuda") os.environ["UCX_RNDV_FRAG_MEM_TYPE"] = "cuda" -if ( - and "UCX_CUDA_COPY_MAX_REG_RATIO" not in os.environ - and _ucx_version >= (1, 12, 0) -): +if "UCX_CUDA_COPY_MAX_REG_RATIO" not in os.environ and _ucx_version >= (1, 12, 0): try: device_count = system.Device.get_device_count() large_bar1 = [False] * device_count