diff --git a/cuda_bindings/tests/nvml/conftest.py b/cuda_bindings/tests/nvml/conftest.py index 7fb1aed4be4..ca61416c137 100644 --- a/cuda_bindings/tests/nvml/conftest.py +++ b/cuda_bindings/tests/nvml/conftest.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from collections import namedtuple import pytest from cuda_python_test_helpers.arch_check import unsupported_before # noqa: F401 @@ -26,56 +25,19 @@ def nvml_init(): yield -@pytest.fixture(scope="session", autouse=True) -def device_info(): - dev_count = None - bus_id_to_board_details = {} - - with NVMLInitializer(): - dev_count = nvml.device_get_count_v2() - - # Store some details for each device now when we know NVML is in known state - for i in range(dev_count): - try: - dev = nvml.device_get_handle_by_index_v2(i) - except nvml.NoPermissionError: - continue - pci_info = nvml.device_get_pci_info_v3(dev) - - name = nvml.device_get_name(dev) - # Get architecture name ex: Ampere, Kepler - arch_id = nvml.device_get_architecture(dev) - - BoardCfg = namedtuple("BoardCfg", "name, ids_arr") - board = BoardCfg(name, ids_arr=[(pci_info.pci_device_id, pci_info.pci_sub_system_id)]) - - try: - serial = nvml.device_get_serial(dev) - except nvml.NvmlError: - serial = None - - bus_id = pci_info.bus_id - device_id = pci_info.device_ - uuid = nvml.device_get_uuid(dev) - - BoardDetails = namedtuple("BoardDetails", "name, board, arch_id, bus_id, device_id, serial") - bus_id_to_board_details[uuid] = BoardDetails(name, board, arch_id, bus_id, device_id, serial) - - return bus_id_to_board_details - - -def get_devices(device_info): - for uuid in list(device_info.keys()): +def get_devices(): + dev_count = nvml.device_get_count_v2() + for i in range(dev_count): try: - yield nvml.device_get_handle_by_uuid(uuid) + yield nvml.device_get_handle_by_index_v2(i) except nvml.NoPermissionError: continue # ignore devices that can't be accessed @pytest.fixture -def all_devices(device_info): +def all_devices(): with NVMLInitializer(): - yield sorted(set(get_devices(device_info))) + yield sorted(set(get_devices())) @pytest.fixture diff --git a/cuda_bindings/tests/nvml/test_cuda.py b/cuda_bindings/tests/nvml/test_cuda.py index 7a782e7403c..c9997af45fb 100644 --- a/cuda_bindings/tests/nvml/test_cuda.py +++ b/cuda_bindings/tests/nvml/test_cuda.py @@ -58,8 +58,9 @@ def test_cuda_device_order(): cuda_devices = get_cuda_device_names() nvml_devices = get_nvml_device_names() - if any("Thor" in device["name"] for device in nvml_devices): - pytest.skip("Skipping test on Thor, which has non-standard device naming") + for kind in ("Orin", "Thor"): + if any(kind in device["name"] for device in nvml_devices): + pytest.skip(f"Skipping test on {kind}, which has non-standard device naming") return if "CUDA_VISIBLE_DEVICES" not in os.environ: diff --git a/cuda_bindings/tests/nvml/test_device.py b/cuda_bindings/tests/nvml/test_device.py index 225e1b13c5e..24f3ffd7d77 100644 --- a/cuda_bindings/tests/nvml/test_device.py +++ b/cuda_bindings/tests/nvml/test_device.py @@ -63,11 +63,15 @@ def test_grid_licensable_features(all_devices): nvml.GridLicenseExpiry(feature.license_expiry) -def test_get_handle_by_uuidv(all_devices): +def test_get_handle_by_uuidv(all_devices, subtests): for device in all_devices: - uuid = nvml.device_get_uuid(device) - new_handle = nvml.device_get_handle_by_uuidv(nvml.UUIDType.ASCII, uuid.encode("ascii")) - assert new_handle == device + with subtests.test(device_index=nvml.device_get_index(device)): + uuid = nvml.device_get_uuid(device) + if "Orin" in nvml.device_get_name(device) and len(uuid) == 36: + pytest.skip("UUID lookup is unsupported on Orin, which reports a UUID without a GPU- prefix") + with unsupported_before(device, None): + new_handle = nvml.device_get_handle_by_uuidv(nvml.UUIDType.ASCII, uuid.encode("ascii")) + assert new_handle == device def test_get_nv_link_supported_bw_modes(all_devices, subtests): diff --git a/cuda_bindings/tests/nvml/test_pynvml.py b/cuda_bindings/tests/nvml/test_pynvml.py index 5f946bab7b0..bb62b716fcd 100644 --- a/cuda_bindings/tests/nvml/test_pynvml.py +++ b/cuda_bindings/tests/nvml/test_pynvml.py @@ -52,9 +52,13 @@ def test_device_get_attributes(mig_handles): pytest.skip("No MIG devices found") -def test_device_get_handle_by_uuid(ngpus, uuids): - handles = [nvml.device_get_handle_by_uuid(uuids[i]) for i in range(ngpus)] - assert len(handles) == ngpus +def test_device_get_handle_by_uuid(ngpus, handles, uuids, subtests): + for i in range(ngpus): + with subtests.test(device_index=i): + uuid = uuids[i] + if "Orin" in nvml.device_get_name(handles[i]) and len(uuid) == 36: + pytest.skip("UUID lookup is unsupported on Orin, which reports a UUID without a GPU- prefix") + assert nvml.device_get_handle_by_uuid(uuid) == handles[i] def test_device_get_handle_by_pci_bus_id(ngpus, pci_info): @@ -68,7 +72,7 @@ def test_device_get_memory_affinity(handles, scope, subtests): size = 1024 for device_index, handle in enumerate(handles): with subtests.test(device_index=device_index): - with unsupported_before(handle, nvml.DeviceArch.KEPLER): + with unsupported_before(handle, None): node_set = nvml.device_get_memory_affinity(handle, size, scope) assert node_set is not None assert len(node_set) == size @@ -80,7 +84,7 @@ def test_device_get_cpu_affinity_within_scope(handles, scope, subtests): size = 1024 for device_index, handle in enumerate(handles): with subtests.test(device_index=device_index): - with unsupported_before(handle, nvml.DeviceArch.KEPLER): + with unsupported_before(handle, None): cpu_set = nvml.device_get_cpu_affinity_within_scope(handle, size, scope) assert cpu_set is not None assert len(cpu_set) == size