From 882f9a9e8a6fb741eb5cb5ec07f1cd0de9bb32c2 Mon Sep 17 00:00:00 2001 From: isvoid Date: Fri, 7 Aug 2026 13:47:59 -0700 Subject: [PATCH] Use CUDA driver for CUDA device enumeration --- cuda_core/cuda/core/_device.pyx | 9 ++++++--- cuda_core/docs/source/release/1.2.0-notes.rst | 5 +++++ cuda_core/examples/show_device_properties.py | 9 +++++---- cuda_core/examples/simple_multi_gpu_example.py | 4 ++-- cuda_core/tests/example_tests/test_basic_examples.py | 4 ++-- cuda_core/tests/system/test_system_device.py | 10 +++++++--- cuda_core/tests/system/test_system_system.py | 6 ++++-- cuda_core/tests/test_green_context.py | 8 ++++---- cuda_core/tests/test_memory.py | 11 ++++------- cuda_core/tests/test_memory_peer_access.py | 4 ++-- cuda_core/tests/test_object_protocols.py | 3 +-- cuda_core/tests/test_stream.py | 4 +--- cuda_core/tests/test_tensor_map.py | 9 ++++----- 13 files changed, 47 insertions(+), 39 deletions(-) diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx index a0a0f472f2b..c148ebb78f8 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/cuda_core/cuda/core/_device.pyx @@ -1033,9 +1033,12 @@ class Device: tuple of Device A tuple containing instances of available devices. """ - from cuda.core import system - total = system.get_num_devices() - return tuple(cls(device_id) for device_id in range(total)) + return cls._get_all_devices_from_cuda_driver() + + @classmethod + def _get_all_devices_from_cuda_driver(cls): + Device_ensure_cuda_initialized() + return tuple(Device_ensure_tls_devices(cls)) def to_system_device(self) -> 'cuda.core.system.Device': """ diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 120d2c2a253..215e90ec0eb 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -73,6 +73,11 @@ Fixes and enhancements Windows, both ``ctypes.CFUNCTYPE`` and ``ctypes.WINFUNCTYPE`` are accepted. (`#2439 `__) +- CUDA device enumeration now queries the CUDA driver rather than using the + NVML system-device count. This prevents non-CUDA accelerators, such as an NPU, + from being treated as CUDA devices by :meth:`Device.get_all_devices`, examples, + and tests. + Deprecation Notices ------------------- diff --git a/cuda_core/examples/show_device_properties.py b/cuda_core/examples/show_device_properties.py index 566f6890948..dcfb716af1d 100644 --- a/cuda_core/examples/show_device_properties.py +++ b/cuda_core/examples/show_device_properties.py @@ -15,7 +15,7 @@ import sys -from cuda.core import Device, system +from cuda.core import Device # Convert boolean to YES or NO string @@ -219,11 +219,12 @@ def print_device_properties(properties): # Print info about all CUDA devices in the system def show_device_properties(): - ndev = system.get_num_devices() + devices = Device.get_all_devices() + ndev = len(devices) print(f"Number of GPUs: {ndev}") - for device_id in range(ndev): - device = Device(device_id) + for device in devices: + device_id = device.device_id print(f"DEVICE {device.name} (id={device_id})") device.set_current() diff --git a/cuda_core/examples/simple_multi_gpu_example.py b/cuda_core/examples/simple_multi_gpu_example.py index e4d7a1ccfb5..840051eb746 100644 --- a/cuda_core/examples/simple_multi_gpu_example.py +++ b/cuda_core/examples/simple_multi_gpu_example.py @@ -17,7 +17,7 @@ import cupy as cp -from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch, system +from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch dtype = cp.float32 size = 50000 @@ -35,7 +35,7 @@ def __cuda_stream__(self): def main(): - if system.get_num_devices() < 2: + if len(Device.get_all_devices()) < 2: print("this example requires at least 2 GPUs", file=sys.stderr) sys.exit(1) diff --git a/cuda_core/tests/example_tests/test_basic_examples.py b/cuda_core/tests/example_tests/test_basic_examples.py index bf423758366..c57207a51a6 100644 --- a/cuda_core/tests/example_tests/test_basic_examples.py +++ b/cuda_core/tests/example_tests/test_basic_examples.py @@ -13,7 +13,7 @@ import pytest from cuda_python_test_helpers.pep723 import has_package_requirements_or_skip -from cuda.core import Device, ManagedMemoryResource, system +from cuda.core import Device, ManagedMemoryResource from cuda.core._program import _can_load_generated_ptx @@ -22,7 +22,7 @@ def has_compute_capability_9_or_higher() -> bool: def has_multiple_devices() -> bool: - return system.get_num_devices() >= 2 + return len(Device.get_all_devices()) >= 2 def has_display() -> bool: diff --git a/cuda_core/tests/system/test_system_device.py b/cuda_core/tests/system/test_system_device.py index 3b2131548b9..6f09bac587b 100644 --- a/cuda_core/tests/system/test_system_device.py +++ b/cuda_core/tests/system/test_system_device.py @@ -16,6 +16,7 @@ import helpers import pytest +from cuda.core import Device as CudaDevice from cuda.core import system from cuda.core.system import typing @@ -132,7 +133,8 @@ def test_numa_node_id(): def test_device_cuda_compute_capability(): - for device in system.Device.get_all_devices(): + for cuda_device in CudaDevice.get_all_devices(): + device = cuda_device.to_system_device() cuda_compute_capability = device.cuda_compute_capability assert isinstance(cuda_compute_capability, tuple) assert len(cuda_compute_capability) == 2 @@ -310,7 +312,8 @@ def test_device_brand(): def test_device_pci_bus_id(): - for device in system.Device.get_all_devices(): + for cuda_device in CudaDevice.get_all_devices(): + device = cuda_device.to_system_device() pci_bus_id = device.pci_info.bus_id assert isinstance(pci_bus_id, str) @@ -796,7 +799,8 @@ def test_pstates(): def test_compute_running_processes(): - for device in system.Device.get_all_devices(): + for cuda_device in CudaDevice.get_all_devices(): + device = cuda_device.to_system_device() with unsupported_before(device, "FERMI"): processes = device.compute_running_processes assert isinstance(processes, list) diff --git a/cuda_core/tests/system/test_system_system.py b/cuda_core/tests/system/test_system_system.py index 28173836ff4..9fc600b05da 100644 --- a/cuda_core/tests/system/test_system_system.py +++ b/cuda_core/tests/system/test_system_system.py @@ -9,6 +9,7 @@ from cuda_python_test_helpers.arch_check import skip_if_nvml_unsupported from cuda.bindings import driver +from cuda.core import Device as CudaDevice from cuda.core import system from cuda.core._utils.cuda_utils import handle_return @@ -57,8 +58,9 @@ def test_nvml_version(): @skip_if_nvml_unsupported def test_get_process_name(): - for device in system.Device.get_all_devices(): - x = device.compute_running_processes + for cuda_device in CudaDevice.get_all_devices(): + device = cuda_device.to_system_device() + _ = device.compute_running_processes try: process_name = system.get_process_name(os.getpid()) diff --git a/cuda_core/tests/test_green_context.py b/cuda_core/tests/test_green_context.py index 24ad6db482b..5f97959f48b 100644 --- a/cuda_core/tests/test_green_context.py +++ b/cuda_core/tests/test_green_context.py @@ -291,12 +291,12 @@ def test_configure_scope_with_enum(self, wq_resource, scope): assert wq_resource.sharing_scope is scope def test_device_id_matches_source_multi_gpu(self): - from cuda.core import Device, system + from cuda.core import Device - if system.get_num_devices() < 2: + devices = Device.get_all_devices() + if len(devices) < 2: pytest.skip("requires 2+ GPUs") - dev0 = Device(0) - dev1 = Device(1) + dev0, dev1 = devices[:2] try: wq0 = dev0.resources.workqueue wq1 = dev1.resources.workqueue diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 8d86fa32432..dff42239bd4 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -39,9 +39,6 @@ VirtualMemoryResource, VirtualMemoryResourceOptions, ) -from cuda.core import ( - system as ccx_system, -) from cuda.core._dlpack import DLDeviceType from cuda.core._memory._ipc import IPCBufferDescriptor from cuda.core._utils.cuda_utils import CUDAError, handle_return @@ -365,7 +362,7 @@ def test_buffer_external_host(): @pytest.mark.parametrize("change_device", [True, False]) def test_buffer_external_device(change_device): - n = ccx_system.get_num_devices() + n = len(Device.get_all_devices()) if n < 1: pytest.skip("No devices found") dev_id = n - 1 @@ -389,7 +386,7 @@ def test_buffer_external_device(change_device): @pytest.mark.parametrize("change_device", [True, False]) def test_buffer_external_pinned_alloc(change_device): - n = ccx_system.get_num_devices() + n = len(Device.get_all_devices()) if n < 1: pytest.skip("No devices found") dev_id = n - 1 @@ -414,7 +411,7 @@ def test_buffer_external_pinned_alloc(change_device): @pytest.mark.parametrize("change_device", [True, False]) def test_buffer_external_pinned_registered(change_device): - n = ccx_system.get_num_devices() + n = len(Device.get_all_devices()) if n < 1: pytest.skip("No devices found") dev_id = n - 1 @@ -447,7 +444,7 @@ def test_buffer_external_pinned_registered(change_device): @pytest.mark.parametrize("change_device", [True, False]) def test_buffer_external_managed(change_device): - n = ccx_system.get_num_devices() + n = len(Device.get_all_devices()) if n < 1: pytest.skip("No devices found") dev_id = n - 1 diff --git a/cuda_core/tests/test_memory_peer_access.py b/cuda_core/tests/test_memory_peer_access.py index 4763b761ad4..c32c21adbc7 100644 --- a/cuda_core/tests/test_memory_peer_access.py +++ b/cuda_core/tests/test_memory_peer_access.py @@ -6,7 +6,7 @@ from helpers.collection_interface_testers import assert_single_member_mutable_set_interface from helpers.constants import POOL_SIZE -from cuda.core import Device, DeviceMemoryResource, DeviceMemoryResourceOptions, system +from cuda.core import Device, DeviceMemoryResource, DeviceMemoryResourceOptions from cuda.core._memory import _peer_access_utils from cuda.core._memory._peer_access_utils import PeerAccessibleBySetProxy from cuda.core._utils.cuda_utils import CUDAError @@ -231,7 +231,7 @@ def test_peer_accessible_by_silently_ignores_owner(isolated_dmr_x2): def test_peer_accessible_by_rejects_invalid_inputs(isolated_dmr_x2): """``add`` raises on out-of-range/unsupported inputs; lenient methods do not.""" dmr, dev0, dev1 = isolated_dmr_x2 - bad_id = system.get_num_devices() # one past the last valid device ordinal + bad_id = len(Device.get_all_devices()) # one past the last valid CUDA device ordinal # add: validates strictly, propagates errors from Device(bad_id) with pytest.raises((ValueError, CUDAError)): diff --git a/cuda_core/tests/test_object_protocols.py b/cuda_core/tests/test_object_protocols.py index e01f7a86941..840cdba1647 100644 --- a/cuda_core/tests/test_object_protocols.py +++ b/cuda_core/tests/test_object_protocols.py @@ -27,7 +27,6 @@ LaunchConfig, Program, Stream, - system, ) from cuda.core._program import _can_load_generated_ptx from cuda.core.graph import GraphDefinition @@ -170,7 +169,7 @@ def sample_program_nvvm(init_cuda): @pytest.fixture def sample_device_alt(init_cuda): """An alternate Device object (requires multi-GPU).""" - if system.get_num_devices() < 2: + if len(Device.get_all_devices()) < 2: pytest.skip("requires multi-GPU") device_alt = Device(1) device_alt.set_current() diff --git a/cuda_core/tests/test_stream.py b/cuda_core/tests/test_stream.py index 55f34bbc9ec..759be542483 100644 --- a/cuda_core/tests/test_stream.py +++ b/cuda_core/tests/test_stream.py @@ -408,9 +408,7 @@ def test_default_stream_per_thread_when_env_set(monkeypatch): def _skip_unless_multi_gpu(): - from cuda.core import system - - if system.get_num_devices() < 2: + if len(Device.get_all_devices()) < 2: pytest.skip("requires 2+ GPUs") diff --git a/cuda_core/tests/test_tensor_map.py b/cuda_core/tests/test_tensor_map.py index 6f63938710f..a9e6432b33e 100644 --- a/cuda_core/tests/test_tensor_map.py +++ b/cuda_core/tests/test_tensor_map.py @@ -9,7 +9,6 @@ Device, ManagedMemoryResourceOptions, TensorMapDescriptor, - system, ) from cuda.core._dlpack import DLDeviceType from cuda.core._tensor_map import ( @@ -386,7 +385,7 @@ def test_replace_address_requires_device_accessible(self, dev, skip_if_no_tma): desc.replace_address(host_arr) def test_replace_address_rejects_tensor_from_other_device(self, dev, skip_if_no_tma): - if system.get_num_devices() < 2: + if len(Device.get_all_devices()) < 2: pytest.skip("requires multi-GPU") dev0 = dev @@ -407,7 +406,7 @@ def test_replace_address_rejects_tensor_from_other_device(self, dev, skip_if_no_ desc.replace_address(buf1) def test_replace_address_accepts_managed_buffer_on_nonzero_device(self, init_cuda): - if system.get_num_devices() < 2: + if len(Device.get_all_devices()) < 2: pytest.skip("requires multi-GPU") dev1 = Device(1) @@ -431,7 +430,7 @@ class TestTensorMapMultiDeviceValidation: """Test multi-device validation for descriptor creation.""" def test_from_tiled_rejects_tensor_from_other_device(self, init_cuda): - if system.get_num_devices() < 2: + if len(Device.get_all_devices()) < 2: pytest.skip("requires multi-GPU") dev0 = Device(0) @@ -451,7 +450,7 @@ def test_from_tiled_rejects_tensor_from_other_device(self, init_cuda): ) def test_from_tiled_accepts_managed_buffer_on_nonzero_device(self, init_cuda): - if system.get_num_devices() < 2: + if len(Device.get_all_devices()) < 2: pytest.skip("requires multi-GPU") dev1 = Device(1)