diff --git a/device_smi/nvidia.py b/device_smi/nvidia.py index 425ce49..a6d6170 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 @@ -34,7 +41,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 +53,19 @@ 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), + ] + + def _nvml_library_candidates(): system = platform.system() if system == "Linux": @@ -101,6 +121,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 +144,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 +186,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 +240,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 +280,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 +289,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,60 +320,91 @@ 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))) - bus_id = info.busId.decode("utf-8").strip() - if not bus_id: - 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() - 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 +489,9 @@ 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.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.", @@ -442,12 +508,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 +523,22 @@ 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 [] + 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 + ) 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..51ccd80 --- /dev/null +++ b/tests/test_nvidia_nvml.py @@ -0,0 +1,403 @@ +import ctypes +from types import SimpleNamespace + +import pytest + +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 + nvml._lib = object() + nvml._shutdown_registered = False + for name in ( + "nvmlDeviceGetHandleByIndex_v2", + "nvmlDeviceGetHandleByIndex", + "nvmlDeviceGetHandleByUUID", + "nvmlDeviceGetHandleByPciBusId_v2", + "nvmlDeviceGetHandleByPciBusId", + "nvmlDeviceGetPciInfo_v3", + "nvmlDeviceGetPciInfo_v2", + "nvmlDeviceGetPciInfo", + "nvmlDeviceGetVbiosVersion", + "nvmlDeviceGetCudaComputeCapability", + "nvmlDeviceGetMaxPcieLinkGeneration", + "nvmlDeviceGetCurrPcieLinkGeneration", + "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") + + +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 = [] + + 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( + ("gpu_id", "symbol", "expected_argument"), + [ + ("GPU-deadbeef", "nvmlDeviceGetHandleByUUID", b"GPU-deadbeef"), + (" 7 ", "nvmlDeviceGetHandleByIndex_v2", 7), + ( + " 0000:65:00.0 ", + "nvmlDeviceGetHandleByPciBusId_v2", + b"0000:65:00.0", + ), + ], +) +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(c_layout)).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(PciInfoV3CLayout)).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" + + +@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"] + + +@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 + 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" + + +@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_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" + 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_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) + 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 + + +@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