From 459a44e1c3272d0e0edc249efade35f1fe2a3355 Mon Sep 17 00:00:00 2001 From: Qubitium-ModelCloud Date: Sun, 23 Aug 2026 00:44:18 +0000 Subject: [PATCH 1/3] Fix direct NVML compatibility paths --- device_smi/nvidia.py | 127 +++++++++++++++++++++-------- tests/test_nvidia_nvml.py | 165 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 259 insertions(+), 33 deletions(-) create mode 100644 tests/test_nvidia_nvml.py diff --git a/device_smi/nvidia.py b/device_smi/nvidia.py index 425ce49..32e8970 100644 --- a/device_smi/nvidia.py +++ b/device_smi/nvidia.py @@ -34,7 +34,7 @@ class _nvmlUtilization_t(ctypes.Structure): ] -class _nvmlPciInfo_t(ctypes.Structure): +class _nvmlPciInfo_v3_t(ctypes.Structure): _fields_ = [ ("busIdLegacy", ctypes.c_char * 16), ("domain", ctypes.c_uint), @@ -46,6 +46,23 @@ class _nvmlPciInfo_t(ctypes.Structure): ] +class _nvmlPciInfo_v2_t(ctypes.Structure): + """Legacy PCI info ABI used by both the v2 and unversioned symbols.""" + + _fields_ = [ + ("busId", ctypes.c_char * 16), + ("domain", ctypes.c_uint), + ("bus", ctypes.c_uint), + ("device", ctypes.c_uint), + ("pciDeviceId", ctypes.c_uint), + ("pciSubSystemId", ctypes.c_uint), + ("reserved0", ctypes.c_uint), + ("reserved1", ctypes.c_uint), + ("reserved2", ctypes.c_uint), + ("reserved3", ctypes.c_uint), + ] + + def _nvml_library_candidates(): system = platform.system() if system == "Linux": @@ -101,6 +118,7 @@ def __new__(cls): cls._instance._initialized = False cls._instance._available = False cls._instance._lib = None + cls._instance._shutdown_registered = False return cls._instance def __init__(self): @@ -123,7 +141,9 @@ def __init__(self): if self._available: import atexit - atexit.register(self._shutdown) + if not self._shutdown_registered: + atexit.register(self._shutdown) + self._shutdown_registered = True @property def available(self) -> bool: @@ -163,7 +183,15 @@ def _bind_functions(self): ctypes.c_int, ), "nvmlDeviceGetPciInfo_v3": ( - [ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_t)], + [ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_v3_t)], + ctypes.c_int, + ), + "nvmlDeviceGetPciInfo_v2": ( + [ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_v2_t)], + ctypes.c_int, + ), + "nvmlDeviceGetPciInfo": ( + [ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_v2_t)], ctypes.c_int, ), "nvmlDeviceGetMaxPcieLinkGeneration": ( @@ -209,13 +237,14 @@ def _init(self): return self._call(["nvmlInit_v2", "nvmlInit"]) def _shutdown(self): - if not self._available or self._lib is None: - return - try: - if self._nvmlShutdown is not None: - self._nvmlShutdown() - except (RuntimeError, AttributeError) as exc: - logger.debug("NVML shutdown failed: %s", exc) + with self._lock: + if not self._available or self._lib is None: + return + self._available = False + try: + self._check(self._call(["nvmlShutdown"])) + except (OSError, RuntimeError, AttributeError) as exc: + logger.debug("NVML shutdown failed: %s", exc) def _check(self, code): if code == 0: @@ -248,6 +277,7 @@ def device_handle(self, gpu_id: str): if not self._available: raise RuntimeError("NVML is not available") handle = ctypes.c_void_p() + gpu_id = gpu_id.strip() if gpu_id.isdigit(): self._check( self._call( @@ -256,7 +286,7 @@ def device_handle(self, gpu_id: str): ctypes.byref(handle), ) ) - elif gpu_id.startswith("GPU-"): + elif gpu_id.startswith(("GPU-", "MIG-")): self._check( self._call( ["nvmlDeviceGetHandleByUUID"], @@ -287,14 +317,19 @@ def device_handle(self, gpu_id: str): def device_name(self, handle): buf = ctypes.create_string_buffer(96) - self._check(self._nvmlDeviceGetName(handle, buf, 96)) + self._check(self._call(["nvmlDeviceGetName"], handle, buf, 96)) return buf.value.decode("utf-8").strip() def pci_info(self, handle): - info = _nvmlPciInfo_t() - self._check(self._nvmlDeviceGetPciInfo_v3(handle, ctypes.byref(info))) + if self._nvmlDeviceGetPciInfo_v3 is not None: + info = _nvmlPciInfo_v3_t() + names = ["nvmlDeviceGetPciInfo_v3"] + else: + info = _nvmlPciInfo_v2_t() + names = ["nvmlDeviceGetPciInfo_v2", "nvmlDeviceGetPciInfo"] + self._check(self._call(names, handle, ctypes.byref(info))) bus_id = info.busId.decode("utf-8").strip() - if not bus_id: + if not bus_id and isinstance(info, _nvmlPciInfo_v3_t): bus_id = info.busIdLegacy.decode("utf-8").strip() if not re.match(r"^[0-9a-fA-F]{4,8}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$", bus_id): raise RuntimeError(f"Invalid NVML PCI bus id: {bus_id!r}") @@ -302,45 +337,58 @@ def pci_info(self, handle): def max_pcie_generation(self, handle): gen = ctypes.c_uint() - self._check(self._nvmlDeviceGetMaxPcieLinkGeneration(handle, ctypes.byref(gen))) + self._check( + self._call( + ["nvmlDeviceGetMaxPcieLinkGeneration"], handle, ctypes.byref(gen) + ) + ) return gen.value def current_pcie_generation(self, handle): gen = ctypes.c_uint() - self._check(self._nvmlDeviceGetCurrPcieLinkGeneration(handle, ctypes.byref(gen))) + self._check( + self._call( + ["nvmlDeviceGetCurrPcieLinkGeneration"], handle, ctypes.byref(gen) + ) + ) return gen.value def memory_total(self, handle): mem = _nvmlMemory_t() - self._check(self._nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem))) + self._check(self._call(["nvmlDeviceGetMemoryInfo"], handle, ctypes.byref(mem))) return mem.total def memory_used(self, handle): mem = _nvmlMemory_t() - self._check(self._nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem))) + self._check(self._call(["nvmlDeviceGetMemoryInfo"], handle, ctypes.byref(mem))) return mem.used def utilization_gpu(self, handle): util = _nvmlUtilization_t() - self._check(self._nvmlDeviceGetUtilizationRates(handle, ctypes.byref(util))) + self._check( + self._call(["nvmlDeviceGetUtilizationRates"], handle, ctypes.byref(util)) + ) return float(util.gpu) def driver_version(self): buf = ctypes.create_string_buffer(80) - self._check(self._nvmlSystemGetDriverVersion(buf, 80)) + self._check(self._call(["nvmlSystemGetDriverVersion"], buf, 80)) return buf.value.decode("utf-8").strip() def vbios_version(self, handle): buf = ctypes.create_string_buffer(32) - self._check(self._nvmlDeviceGetVbiosVersion(handle, buf, 32)) + self._check(self._call(["nvmlDeviceGetVbiosVersion"], handle, buf, 32)) return buf.value.decode("utf-8").strip() def compute_capability(self, handle): major = ctypes.c_int() minor = ctypes.c_int() self._check( - self._nvmlDeviceGetCudaComputeCapability( - handle, ctypes.byref(major), ctypes.byref(minor) + self._call( + ["nvmlDeviceGetCudaComputeCapability"], + handle, + ctypes.byref(major), + ctypes.byref(minor), ) ) return f"{major.value}.{minor.value}" @@ -425,7 +473,7 @@ def _get_gpu_id(self): gpu_count = 0 cudas = os.environ.get("CUDA_VISIBLE_DEVICES", "") - cuda_list = cudas.split(",") if cudas else [] + cuda_list = [cuda.strip() for cuda in cudas.split(",")] if cudas else [] if gpu_count > 0 and os.environ.get("CUDA_DEVICE_ORDER", "") != "PCI_BUS_ID": warnings.warn( "Detected different devices in the system. Please make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to avoid unexpected behavior.", @@ -442,12 +490,14 @@ def _init_from_nvml(self, cls): model = _nvml.device_name(handle) total_memory = _nvml.memory_total(handle) - pci_bus_id = _nvml.pci_info(handle) - pcie_gen_max = _nvml.max_pcie_generation(handle) - pcie_gen_current = _nvml.current_pcie_generation(handle) - driver = _nvml.driver_version() - firmware = _nvml.vbios_version(handle) - compute_cap = _nvml.compute_capability(handle) + driver = self._optional_nvml(_nvml.driver_version, "") + firmware = self._optional_nvml(_nvml.vbios_version, "", handle) + compute_cap = self._optional_nvml(_nvml.compute_capability, None, handle) + pci_bus_id = self._optional_nvml(_nvml.pci_info, None, handle) + pcie_gen_max = self._optional_nvml(_nvml.max_pcie_generation, None, handle) + pcie_gen_current = self._optional_nvml( + _nvml.current_pcie_generation, None, handle + ) if model.lower().startswith("nvidia"): model = model[len("nvidia"):] @@ -455,10 +505,21 @@ def _init_from_nvml(self, cls): cls.model = model.strip().lower() cls.memory_total = int(total_memory) cls.vendor = "nvidia" - cls.features = [compute_cap] - cls.pcie = Pcie(gen=int(pcie_gen_max), speed=int(pcie_gen_current), id=pci_bus_id) + cls.features = [compute_cap] if compute_cap else [] + if None not in (pci_bus_id, pcie_gen_max, pcie_gen_current): + cls.pcie = Pcie( + gen=int(pcie_gen_max), speed=int(pcie_gen_current), id=pci_bus_id + ) cls.gpu = GPU(driver=driver, firmware=firmware) + @staticmethod + def _optional_nvml(fetch, default, *args): + try: + return fetch(*args) + except (OSError, RuntimeError, ValueError, AttributeError) as exc: + logger.debug("Optional NVML metadata unavailable: %s", exc) + return default + def _init_from_nvidia_smi(self, cls): args = [ "nvidia-smi", diff --git a/tests/test_nvidia_nvml.py b/tests/test_nvidia_nvml.py new file mode 100644 index 0000000..cc1f8bd --- /dev/null +++ b/tests/test_nvidia_nvml.py @@ -0,0 +1,165 @@ +import ctypes +from types import SimpleNamespace + +import pytest + +from device_smi import nvidia + + +def nvml_stub(**functions): + nvml = object.__new__(nvidia._NVML) + nvml._available = True + nvml._lib = object() + nvml._shutdown_registered = False + for name in ( + "nvmlDeviceGetHandleByIndex_v2", + "nvmlDeviceGetHandleByIndex", + "nvmlDeviceGetHandleByUUID", + "nvmlDeviceGetHandleByPciBusId_v2", + "nvmlDeviceGetHandleByPciBusId", + "nvmlDeviceGetPciInfo_v3", + "nvmlDeviceGetPciInfo_v2", + "nvmlDeviceGetPciInfo", + "nvmlErrorString", + "nvmlShutdown", + ): + setattr(nvml, f"_{name}", functions.get(name)) + return nvml + + +def test_missing_symbol_raises_controlled_error(): + nvml = nvml_stub() + + with pytest.raises(RuntimeError, match="nvmlDeviceGetHandleByUUID"): + nvml.device_handle("GPU-deadbeef") + + +@pytest.mark.parametrize("gpu_id", ["MIG-deadbeef", "MIG-GPU-deadbeef/1/2"]) +def test_mig_identifiers_use_uuid_lookup(gpu_id): + calls = [] + + def by_uuid(uuid, handle): + calls.append(uuid) + ctypes.cast(handle, ctypes.POINTER(ctypes.c_void_p)).contents.value = 123 + return 0 + + nvml = nvml_stub(nvmlDeviceGetHandleByUUID=by_uuid) + + assert nvml.device_handle(f" {gpu_id} ").value == 123 + assert calls == [gpu_id.encode()] + + +@pytest.mark.parametrize( + ("symbol", "expected_type"), + [ + ("nvmlDeviceGetPciInfo_v3", nvidia._nvmlPciInfo_v3_t), + ("nvmlDeviceGetPciInfo_v2", nvidia._nvmlPciInfo_v2_t), + ("nvmlDeviceGetPciInfo", nvidia._nvmlPciInfo_v2_t), + ], +) +def test_pci_info_selects_newest_available_abi(symbol, expected_type): + calls = [] + + def get_pci_info(handle, info_pointer): + info = ctypes.cast(info_pointer, ctypes.POINTER(expected_type)).contents + info.busId = b"0000:01:00.0" + calls.append(handle) + return 0 + + nvml = nvml_stub(**{symbol: get_pci_info}) + + assert nvml.pci_info(ctypes.c_void_p(7)) == "0000:01:00.0" + assert len(calls) == 1 + + +def test_pci_info_prefers_v3(): + def unexpected(*args): + raise AssertionError("older PCI symbol called") + + def v3(handle, info_pointer): + info = ctypes.cast( + info_pointer, ctypes.POINTER(nvidia._nvmlPciInfo_v3_t) + ).contents + info.busId = b"0000:02:00.0" + return 0 + + nvml = nvml_stub( + nvmlDeviceGetPciInfo_v3=v3, + nvmlDeviceGetPciInfo_v2=unexpected, + nvmlDeviceGetPciInfo=unexpected, + ) + + assert nvml.pci_info(None) == "0000:02:00.0" + + +def test_cuda_visible_devices_is_trimmed(monkeypatch): + device = object.__new__(nvidia.NvidiaDevice) + device.index = 1 + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", " GPU-first , MIG-GPU-second/1/2 ") + monkeypatch.setattr(nvidia._nvml, "_available", False) + monkeypatch.setattr(nvidia, "_run_nvidia_smi", lambda *args, **kwargs: "") + + assert device._get_gpu_id() == "MIG-GPU-second/1/2" + + +def test_optional_metadata_failure_keeps_nvml_device(monkeypatch): + device = object.__new__(nvidia.NvidiaDevice) + device.gpu_id = "0" + target = SimpleNamespace(pcie=None) + monkeypatch.setattr(nvidia._nvml, "device_handle", lambda gpu_id: object()) + monkeypatch.setattr(nvidia._nvml, "device_name", lambda handle: "NVIDIA Test") + monkeypatch.setattr(nvidia._nvml, "memory_total", lambda handle: 1024) + monkeypatch.setattr(nvidia._nvml, "driver_version", lambda: "555.1") + for name in ( + "vbios_version", + "compute_capability", + "pci_info", + "max_pcie_generation", + "current_pcie_generation", + ): + monkeypatch.setattr( + nvidia._nvml, + name, + lambda *args, _name=name: (_ for _ in ()).throw(RuntimeError(_name)), + ) + + device._init_from_nvml(target) + + assert target.model == "test" + assert target.memory_total == 1024 + assert target.features == [] + assert target.pcie is None + assert target.gpu.driver == "555.1" + assert target.gpu.firmware == "" + + +def test_core_nvml_failure_still_falls_back_to_nvidia_smi(monkeypatch): + calls = [] + monkeypatch.setattr(nvidia._nvml, "_available", True) + monkeypatch.setattr(nvidia.NvidiaDevice, "_get_gpu_id", lambda self: "0") + monkeypatch.setattr( + nvidia.NvidiaDevice, + "_init_from_nvml", + lambda self, cls: (_ for _ in ()).throw(RuntimeError("core failure")), + ) + monkeypatch.setattr( + nvidia.NvidiaDevice, + "_init_from_nvidia_smi", + lambda self, cls: calls.append(cls), + ) + target = SimpleNamespace() + + nvidia.NvidiaDevice(target, 0) + + assert calls == [target] + + +def test_shutdown_is_idempotent_and_calls_nvml_once(): + calls = [] + nvml = nvml_stub(nvmlShutdown=lambda: calls.append(True) or 0) + + nvml._shutdown() + nvml._shutdown() + + assert calls == [True] + assert nvml.available is False From da1cba161af33b7c996ed8a3e1f3ab6208054cdd Mon Sep 17 00:00:00 2001 From: Qubitium-ModelCloud Date: Sun, 23 Aug 2026 00:51:29 +0000 Subject: [PATCH 2/3] Correct legacy NVML PCI ABI --- device_smi/nvidia.py | 53 +++++++---- tests/test_nvidia_nvml.py | 188 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 214 insertions(+), 27 deletions(-) diff --git a/device_smi/nvidia.py b/device_smi/nvidia.py index 32e8970..f36efee 100644 --- a/device_smi/nvidia.py +++ b/device_smi/nvidia.py @@ -14,6 +14,13 @@ logger = logging.getLogger(__name__) +_NVML_ERROR_NOT_SUPPORTED = 3 +_NVML_ERROR_FUNCTION_NOT_FOUND = 13 +_NVML_COMPATIBILITY_ERRORS = { + _NVML_ERROR_NOT_SUPPORTED, + _NVML_ERROR_FUNCTION_NOT_FOUND, +} + class NvidiaGPUMetrics(BaseMetrics): pass @@ -56,10 +63,6 @@ class _nvmlPciInfo_v2_t(ctypes.Structure): ("device", ctypes.c_uint), ("pciDeviceId", ctypes.c_uint), ("pciSubSystemId", ctypes.c_uint), - ("reserved0", ctypes.c_uint), - ("reserved1", ctypes.c_uint), - ("reserved2", ctypes.c_uint), - ("reserved3", ctypes.c_uint), ] @@ -321,19 +324,32 @@ def device_name(self, handle): return buf.value.decode("utf-8").strip() def pci_info(self, handle): - if self._nvmlDeviceGetPciInfo_v3 is not None: - info = _nvmlPciInfo_v3_t() - names = ["nvmlDeviceGetPciInfo_v3"] - else: - info = _nvmlPciInfo_v2_t() - names = ["nvmlDeviceGetPciInfo_v2", "nvmlDeviceGetPciInfo"] - self._check(self._call(names, handle, ctypes.byref(info))) - bus_id = info.busId.decode("utf-8").strip() - if not bus_id and isinstance(info, _nvmlPciInfo_v3_t): - bus_id = info.busIdLegacy.decode("utf-8").strip() - if not re.match(r"^[0-9a-fA-F]{4,8}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$", bus_id): - raise RuntimeError(f"Invalid NVML PCI bus id: {bus_id!r}") - return bus_id + attempts = ( + ("nvmlDeviceGetPciInfo_v3", _nvmlPciInfo_v3_t), + ("nvmlDeviceGetPciInfo_v2", _nvmlPciInfo_v2_t), + ("nvmlDeviceGetPciInfo", _nvmlPciInfo_v2_t), + ) + attempted = [] + for name, info_type in attempts: + fn = getattr(self, f"_{name}", None) + if fn is None: + continue + attempted.append(name) + info = info_type() + code = self._call([name], handle, ctypes.byref(info)) + if code in _NVML_COMPATIBILITY_ERRORS: + continue + self._check(code) + bus_id = info.busId.decode("utf-8").strip() + if not bus_id and isinstance(info, _nvmlPciInfo_v3_t): + bus_id = info.busIdLegacy.decode("utf-8").strip() + if not re.match( + r"^[0-9a-fA-F]{4,8}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$", + bus_id, + ): + raise RuntimeError(f"Invalid NVML PCI bus id: {bus_id!r}") + return bus_id + raise RuntimeError(f"NVML PCI info functions unavailable: {attempted}") def max_pcie_generation(self, handle): gen = ctypes.c_uint() @@ -473,7 +489,7 @@ def _get_gpu_id(self): gpu_count = 0 cudas = os.environ.get("CUDA_VISIBLE_DEVICES", "") - cuda_list = [cuda.strip() for cuda in cudas.split(",")] if cudas else [] + cuda_list = [cuda.strip() for cuda in cudas.split(",")] if cudas.strip() else [] if gpu_count > 0 and os.environ.get("CUDA_DEVICE_ORDER", "") != "PCI_BUS_ID": warnings.warn( "Detected different devices in the system. Please make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to avoid unexpected behavior.", @@ -506,6 +522,7 @@ def _init_from_nvml(self, cls): cls.memory_total = int(total_memory) cls.vendor = "nvidia" cls.features = [compute_cap] if compute_cap else [] + cls.pcie = None if None not in (pci_bus_id, pcie_gen_max, pcie_gen_current): cls.pcie = Pcie( gen=int(pcie_gen_max), speed=int(pcie_gen_current), id=pci_bus_id diff --git a/tests/test_nvidia_nvml.py b/tests/test_nvidia_nvml.py index cc1f8bd..818e02e 100644 --- a/tests/test_nvidia_nvml.py +++ b/tests/test_nvidia_nvml.py @@ -6,6 +6,29 @@ from device_smi import nvidia +class LegacyPciInfoCLayout(ctypes.Structure): + _fields_ = [ + ("busId", ctypes.c_char * 16), + ("domain", ctypes.c_uint), + ("bus", ctypes.c_uint), + ("device", ctypes.c_uint), + ("pciDeviceId", ctypes.c_uint), + ("pciSubSystemId", ctypes.c_uint), + ] + + +class PciInfoV3CLayout(ctypes.Structure): + _fields_ = [ + ("busIdLegacy", ctypes.c_char * 16), + ("domain", ctypes.c_uint), + ("bus", ctypes.c_uint), + ("device", ctypes.c_uint), + ("pciDeviceId", ctypes.c_uint), + ("pciSubSystemId", ctypes.c_uint), + ("busId", ctypes.c_char * 32), + ] + + def nvml_stub(**functions): nvml = object.__new__(nvidia._NVML) nvml._available = True @@ -34,6 +57,23 @@ def test_missing_symbol_raises_controlled_error(): nvml.device_handle("GPU-deadbeef") +def test_legacy_pci_info_matches_c_abi_layout(): + expected_offsets = { + "busId": 0, + "domain": 16, + "bus": 20, + "device": 24, + "pciDeviceId": 28, + "pciSubSystemId": 32, + } + + assert ctypes.sizeof(LegacyPciInfoCLayout) == 36 + assert ctypes.sizeof(nvidia._nvmlPciInfo_v2_t) == 36 + for field, expected_offset in expected_offsets.items(): + assert getattr(LegacyPciInfoCLayout, field).offset == expected_offset + assert getattr(nvidia._nvmlPciInfo_v2_t, field).offset == expected_offset + + @pytest.mark.parametrize("gpu_id", ["MIG-deadbeef", "MIG-GPU-deadbeef/1/2"]) def test_mig_identifiers_use_uuid_lookup(gpu_id): calls = [] @@ -50,18 +90,44 @@ def by_uuid(uuid, handle): @pytest.mark.parametrize( - ("symbol", "expected_type"), + ("gpu_id", "symbol", "expected_argument"), [ - ("nvmlDeviceGetPciInfo_v3", nvidia._nvmlPciInfo_v3_t), - ("nvmlDeviceGetPciInfo_v2", nvidia._nvmlPciInfo_v2_t), - ("nvmlDeviceGetPciInfo", nvidia._nvmlPciInfo_v2_t), + ("GPU-deadbeef", "nvmlDeviceGetHandleByUUID", b"GPU-deadbeef"), + (" 7 ", "nvmlDeviceGetHandleByIndex_v2", 7), + ( + " 0000:65:00.0 ", + "nvmlDeviceGetHandleByPciBusId_v2", + b"0000:65:00.0", + ), ], ) -def test_pci_info_selects_newest_available_abi(symbol, expected_type): +def test_device_handle_routing(gpu_id, symbol, expected_argument): + calls = [] + + def lookup(identifier, handle): + calls.append(identifier) + ctypes.cast(handle, ctypes.POINTER(ctypes.c_void_p)).contents.value = 42 + return 0 + + nvml = nvml_stub(**{symbol: lookup}) + + assert nvml.device_handle(gpu_id).value == 42 + assert calls == [expected_argument] + + +@pytest.mark.parametrize( + ("symbol", "c_layout"), + [ + ("nvmlDeviceGetPciInfo_v3", PciInfoV3CLayout), + ("nvmlDeviceGetPciInfo_v2", LegacyPciInfoCLayout), + ("nvmlDeviceGetPciInfo", LegacyPciInfoCLayout), + ], +) +def test_pci_info_selects_newest_available_abi(symbol, c_layout): calls = [] def get_pci_info(handle, info_pointer): - info = ctypes.cast(info_pointer, ctypes.POINTER(expected_type)).contents + info = ctypes.cast(info_pointer, ctypes.POINTER(c_layout)).contents info.busId = b"0000:01:00.0" calls.append(handle) return 0 @@ -77,9 +143,7 @@ def unexpected(*args): raise AssertionError("older PCI symbol called") def v3(handle, info_pointer): - info = ctypes.cast( - info_pointer, ctypes.POINTER(nvidia._nvmlPciInfo_v3_t) - ).contents + info = ctypes.cast(info_pointer, ctypes.POINTER(PciInfoV3CLayout)).contents info.busId = b"0000:02:00.0" return 0 @@ -92,6 +156,57 @@ def v3(handle, info_pointer): assert nvml.pci_info(None) == "0000:02:00.0" +@pytest.mark.parametrize("compatibility_error", [3, 13]) +def test_pci_info_falls_back_after_v3_compatibility_error(compatibility_error): + calls = [] + + def v3(handle, info_pointer): + calls.append("v3") + return compatibility_error + + def v2(handle, info_pointer): + calls.append("v2") + info = ctypes.cast( + info_pointer, ctypes.POINTER(LegacyPciInfoCLayout) + ).contents + info.busId = b"0000:03:00.0" + return 0 + + nvml = nvml_stub( + nvmlDeviceGetPciInfo_v3=v3, + nvmlDeviceGetPciInfo_v2=v2, + nvmlDeviceGetPciInfo=lambda *args: pytest.fail("legacy symbol called"), + ) + + assert nvml.pci_info(None) == "0000:03:00.0" + assert calls == ["v3", "v2"] + + +@pytest.mark.parametrize("compatibility_error", [3, 13]) +def test_pci_info_falls_back_from_v2_to_legacy(compatibility_error): + calls = [] + + def v2(handle, info_pointer): + calls.append("v2") + return compatibility_error + + def legacy(handle, info_pointer): + calls.append("legacy") + info = ctypes.cast( + info_pointer, ctypes.POINTER(LegacyPciInfoCLayout) + ).contents + info.busId = b"0000:04:00.0" + return 0 + + nvml = nvml_stub( + nvmlDeviceGetPciInfo_v2=v2, + nvmlDeviceGetPciInfo=legacy, + ) + + assert nvml.pci_info(None) == "0000:04:00.0" + assert calls == ["v2", "legacy"] + + def test_cuda_visible_devices_is_trimmed(monkeypatch): device = object.__new__(nvidia.NvidiaDevice) device.index = 1 @@ -102,6 +217,17 @@ def test_cuda_visible_devices_is_trimmed(monkeypatch): assert device._get_gpu_id() == "MIG-GPU-second/1/2" +@pytest.mark.parametrize("visible_devices", ["", " "]) +def test_empty_cuda_visible_devices_uses_device_index(monkeypatch, visible_devices): + device = object.__new__(nvidia.NvidiaDevice) + device.index = 2 + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", visible_devices) + monkeypatch.setattr(nvidia._nvml, "_available", False) + monkeypatch.setattr(nvidia, "_run_nvidia_smi", lambda *args, **kwargs: "") + + assert device._get_gpu_id() == "2" + + def test_optional_metadata_failure_keeps_nvml_device(monkeypatch): device = object.__new__(nvidia.NvidiaDevice) device.gpu_id = "0" @@ -133,6 +259,37 @@ def test_optional_metadata_failure_keeps_nvml_device(monkeypatch): assert target.gpu.firmware == "" +def test_one_pcie_metadata_failure_clears_stale_pcie_only(monkeypatch): + device = object.__new__(nvidia.NvidiaDevice) + device.gpu_id = "0" + stale_pcie = object() + target = SimpleNamespace(pcie=stale_pcie) + monkeypatch.setattr(nvidia._nvml, "device_handle", lambda gpu_id: object()) + monkeypatch.setattr(nvidia._nvml, "device_name", lambda handle: "NVIDIA Test") + monkeypatch.setattr(nvidia._nvml, "memory_total", lambda handle: 2048) + monkeypatch.setattr(nvidia._nvml, "driver_version", lambda: "555.2") + monkeypatch.setattr(nvidia._nvml, "vbios_version", lambda handle: "VBIOS-1") + monkeypatch.setattr( + nvidia._nvml, "compute_capability", lambda handle: "9.0" + ) + monkeypatch.setattr(nvidia._nvml, "pci_info", lambda handle: "0000:05:00.0") + monkeypatch.setattr( + nvidia._nvml, + "max_pcie_generation", + lambda handle: (_ for _ in ()).throw(RuntimeError("not supported")), + ) + monkeypatch.setattr(nvidia._nvml, "current_pcie_generation", lambda handle: 4) + + device._init_from_nvml(target) + + assert target.model == "test" + assert target.memory_total == 2048 + assert target.features == ["9.0"] + assert target.gpu.driver == "555.2" + assert target.gpu.firmware == "VBIOS-1" + assert target.pcie is None + + def test_core_nvml_failure_still_falls_back_to_nvidia_smi(monkeypatch): calls = [] monkeypatch.setattr(nvidia._nvml, "_available", True) @@ -163,3 +320,16 @@ def test_shutdown_is_idempotent_and_calls_nvml_once(): assert calls == [True] assert nvml.available is False + + +@pytest.mark.parametrize( + "shutdown", + [None, lambda: 3, lambda: (_ for _ in ()).throw(OSError("shutdown failed"))], +) +def test_missing_or_erroring_shutdown_is_safe_and_idempotent(shutdown): + nvml = nvml_stub(nvmlShutdown=shutdown) + + nvml._shutdown() + nvml._shutdown() + + assert nvml.available is False From d0187ec231f73ddad4ea9b36ffd71592c87dbd0f Mon Sep 17 00:00:00 2001 From: Qubitium-ModelCloud Date: Sun, 23 Aug 2026 00:54:36 +0000 Subject: [PATCH 3/3] Harden NVML compatibility regressions --- device_smi/nvidia.py | 2 ++ tests/test_nvidia_nvml.py | 68 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/device_smi/nvidia.py b/device_smi/nvidia.py index f36efee..a6d6170 100644 --- a/device_smi/nvidia.py +++ b/device_smi/nvidia.py @@ -490,6 +490,8 @@ def _get_gpu_id(self): cudas = os.environ.get("CUDA_VISIBLE_DEVICES", "") cuda_list = [cuda.strip() for cuda in cudas.split(",")] if cudas.strip() else [] + if any(not cuda for cuda in cuda_list): + raise ValueError("CUDA_VISIBLE_DEVICES contains an empty device entry") if gpu_count > 0 and os.environ.get("CUDA_DEVICE_ORDER", "") != "PCI_BUS_ID": warnings.warn( "Detected different devices in the system. Please make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to avoid unexpected behavior.", diff --git a/tests/test_nvidia_nvml.py b/tests/test_nvidia_nvml.py index 818e02e..51ccd80 100644 --- a/tests/test_nvidia_nvml.py +++ b/tests/test_nvidia_nvml.py @@ -43,6 +43,10 @@ def nvml_stub(**functions): "nvmlDeviceGetPciInfo_v3", "nvmlDeviceGetPciInfo_v2", "nvmlDeviceGetPciInfo", + "nvmlDeviceGetVbiosVersion", + "nvmlDeviceGetCudaComputeCapability", + "nvmlDeviceGetMaxPcieLinkGeneration", + "nvmlDeviceGetCurrPcieLinkGeneration", "nvmlErrorString", "nvmlShutdown", ): @@ -207,6 +211,59 @@ def legacy(handle, info_pointer): assert calls == ["v2", "legacy"] +@pytest.mark.parametrize("nvml_error", [1, 6]) +def test_pci_info_does_not_fallback_for_noncompatibility_error(nvml_error): + calls = [] + + def v3(handle, info_pointer): + calls.append("v3") + return nvml_error + + def unexpected(*args): + calls.append("older") + return 0 + + nvml = nvml_stub( + nvmlDeviceGetPciInfo_v3=v3, + nvmlDeviceGetPciInfo_v2=unexpected, + nvmlDeviceGetPciInfo=unexpected, + ) + + with pytest.raises(RuntimeError, match=rf"NVML error {nvml_error}:"): + nvml.pci_info(None) + assert calls == ["v3"] + + +@pytest.mark.parametrize( + ("method_name", "args", "missing_symbol"), + [ + ("vbios_version", (ctypes.c_void_p(1),), "nvmlDeviceGetVbiosVersion"), + ( + "compute_capability", + (ctypes.c_void_p(1),), + "nvmlDeviceGetCudaComputeCapability", + ), + ( + "max_pcie_generation", + (ctypes.c_void_p(1),), + "nvmlDeviceGetMaxPcieLinkGeneration", + ), + ( + "current_pcie_generation", + (ctypes.c_void_p(1),), + "nvmlDeviceGetCurrPcieLinkGeneration", + ), + ], +) +def test_missing_optional_bound_symbol_raises_controlled_error( + method_name, args, missing_symbol +): + nvml = nvml_stub() + + with pytest.raises(RuntimeError, match=missing_symbol): + getattr(nvml, method_name)(*args) + + def test_cuda_visible_devices_is_trimmed(monkeypatch): device = object.__new__(nvidia.NvidiaDevice) device.index = 1 @@ -228,6 +285,17 @@ def test_empty_cuda_visible_devices_uses_device_index(monkeypatch, visible_devic assert device._get_gpu_id() == "2" +def test_cuda_visible_devices_rejects_interior_empty_entry(monkeypatch): + device = object.__new__(nvidia.NvidiaDevice) + device.index = 1 + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0, ,2") + monkeypatch.setattr(nvidia._nvml, "_available", False) + monkeypatch.setattr(nvidia, "_run_nvidia_smi", lambda *args, **kwargs: "") + + with pytest.raises(ValueError, match="empty device entry"): + device._get_gpu_id() + + def test_optional_metadata_failure_keeps_nvml_device(monkeypatch): device = object.__new__(nvidia.NvidiaDevice) device.gpu_id = "0"