diff --git a/README.md b/README.md index 4af1bd5..3874950 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,19 @@ stats = dev.metric(fast=True) print(stats.memory_used, stats.utilization) ``` +### NVIDIA NVML and `nvidia-smi` caching + +For NVIDIA devices, Device-SMI first tries to call NVML directly via `libnvidia-ml` so it does not spawn `nvidia-smi` repeatedly. When NVML is not available, it falls back to `nvidia-smi`. + +`nvidia-smi` caching is disabled by default (`nvidia_smi_cache_ttl = 0`). Set the TTL to a positive value to enable a thread-safe LRU cache that returns stale data within the TTL window: + +```py +from device_smi import Device + +Device.config.nvidia_smi_cache_ttl = 2.0 # seconds; 0 disables caching +Device.config.nvidia_smi_cache_maxsize = 32 +``` + ## Roadmap - Support Intel/Gaudi diff --git a/device_smi/__init__.py b/device_smi/__init__.py index 549158f..b5406c7 100644 --- a/device_smi/__init__.py +++ b/device_smi/__init__.py @@ -1 +1,2 @@ +from .config import config as config from .device import Device as Device diff --git a/device_smi/config.py b/device_smi/config.py new file mode 100644 index 0000000..29953d2 --- /dev/null +++ b/device_smi/config.py @@ -0,0 +1,38 @@ +class Config: + """Module-level configuration for ``device-smi``.""" + + def __init__(self): + # A TTL of 0 disables the nvidia-smi cache entirely so every call + # runs the subprocess. Set it >0 to enable caching. + self._nvidia_smi_cache_ttl = 0.0 + self._nvidia_smi_cache_maxsize = 16 + + @property + def nvidia_smi_cache_ttl(self) -> float: + """How long ``nvidia-smi`` results are cached, in seconds. + + ``0.0`` disables caching and every call spawns a fresh subprocess. + """ + return self._nvidia_smi_cache_ttl + + @nvidia_smi_cache_ttl.setter + def nvidia_smi_cache_ttl(self, value: float) -> None: + value = float(value) + if value < 0: + raise ValueError("nvidia_smi_cache_ttl must be greater than or equal to 0") + self._nvidia_smi_cache_ttl = value + + @property + def nvidia_smi_cache_maxsize(self) -> int: + """Maximum number of distinct ``nvidia-smi`` query results kept in the LRU cache.""" + return self._nvidia_smi_cache_maxsize + + @nvidia_smi_cache_maxsize.setter + def nvidia_smi_cache_maxsize(self, value: int) -> None: + value = int(value) + if value <= 0: + raise ValueError("nvidia_smi_cache_maxsize must be greater than 0") + self._nvidia_smi_cache_maxsize = value + + +config = Config() diff --git a/device_smi/device.py b/device_smi/device.py index ac2968e..b47c264 100644 --- a/device_smi/device.py +++ b/device_smi/device.py @@ -7,6 +7,7 @@ from .amd import AMDDevice from .apple import AppleDevice from .base import _run +from .config import config from .cpu import CPUDevice from .intel import IntelDevice from .nvidia import NvidiaDevice @@ -25,6 +26,8 @@ def _get_torch_runtime(): class Device: + config = config + def __init__(self, device, *, fast_metrics_interval: float = 0.200): # init attribute first to avoid IDE not attr warning # CPU/GPU Device diff --git a/device_smi/nvidia.py b/device_smi/nvidia.py index 04be0fc..425ce49 100644 --- a/device_smi/nvidia.py +++ b/device_smi/nvidia.py @@ -1,13 +1,399 @@ +import ctypes +import ctypes.util +import logging import os +import platform +import re +import threading +import time import warnings +from collections import OrderedDict from .base import GPU, BaseMetrics, GPUDevice, Pcie, _run +from .config import config + +logger = logging.getLogger(__name__) class NvidiaGPUMetrics(BaseMetrics): pass +class _nvmlMemory_t(ctypes.Structure): + _fields_ = [ + ("total", ctypes.c_ulonglong), + ("free", ctypes.c_ulonglong), + ("used", ctypes.c_ulonglong), + ] + + +class _nvmlUtilization_t(ctypes.Structure): + _fields_ = [ + ("gpu", ctypes.c_uint), + ("memory", ctypes.c_uint), + ] + + +class _nvmlPciInfo_t(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_library_candidates(): + system = platform.system() + if system == "Linux": + return [ + "/usr/lib/wsl/lib/libnvidia-ml.so.1", + "/usr/lib/x86_64-linux-gnu/libnvidia-ml.so.1", + "/usr/lib/x86_64-linux-gnu/libnvidia-ml.so", + "/usr/lib64/libnvidia-ml.so.1", + "/usr/lib64/libnvidia-ml.so", + "libnvidia-ml.so.1", + "libnvidia-ml.so", + ] + if system == "Windows": + return [ + os.path.join( + os.environ.get("SYSTEMROOT", r"C:\Windows"), "System32", "nvml.dll" + ), + os.path.join( + os.environ.get("ProgramW6432", r"C:\Program Files"), + "NVIDIA Corporation", + "NVSMI", + "nvml.dll", + ), + "nvml.dll", + ] + return ["libnvidia-ml.dylib"] + + +def _load_nvml_library(): + lib_path = ctypes.util.find_library("nvidia-ml") + candidates = [lib_path] if lib_path else [] + candidates.extend(_nvml_library_candidates()) + for path in candidates: + if not path: + continue + try: + return ctypes.CDLL(path) + except OSError: + continue + return None + + +class _NVML: + """Lightweight ctypes binding to the NVIDIA Management Library (NVML).""" + + _instance = None + _lock = threading.Lock() + + def __new__(cls): + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + cls._instance._available = False + cls._instance._lib = None + return cls._instance + + def __init__(self): + with self._lock: + if self._initialized: + return + lib = _load_nvml_library() + if lib is None: + self._initialized = True + return + self._lib = lib + try: + self._bind_functions() + code = self._init() + except (OSError, AttributeError, RuntimeError) as exc: + logger.debug("NVML initialization failed: %s", exc) + code = -1 + self._initialized = True + self._available = code == 0 + if self._available: + import atexit + + atexit.register(self._shutdown) + + @property + def available(self) -> bool: + return self._available + + def _bind_functions(self): + spec = { + "nvmlInit_v2": ([], ctypes.c_int), + "nvmlInit": ([], ctypes.c_int), + "nvmlShutdown": ([], ctypes.c_int), + "nvmlErrorString": ([ctypes.c_int], ctypes.c_char_p), + "nvmlSystemGetDriverVersion": ([ctypes.c_char_p, ctypes.c_uint], ctypes.c_int), + "nvmlDeviceGetCount_v2": ([ctypes.POINTER(ctypes.c_uint)], ctypes.c_int), + "nvmlDeviceGetCount": ([ctypes.POINTER(ctypes.c_uint)], ctypes.c_int), + "nvmlDeviceGetHandleByIndex_v2": ( + [ctypes.c_uint, ctypes.POINTER(ctypes.c_void_p)], + ctypes.c_int, + ), + "nvmlDeviceGetHandleByIndex": ( + [ctypes.c_uint, ctypes.POINTER(ctypes.c_void_p)], + ctypes.c_int, + ), + "nvmlDeviceGetHandleByUUID": ( + [ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)], + ctypes.c_int, + ), + "nvmlDeviceGetHandleByPciBusId_v2": ( + [ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)], + ctypes.c_int, + ), + "nvmlDeviceGetHandleByPciBusId": ( + [ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)], + ctypes.c_int, + ), + "nvmlDeviceGetName": ( + [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint], + ctypes.c_int, + ), + "nvmlDeviceGetPciInfo_v3": ( + [ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_t)], + ctypes.c_int, + ), + "nvmlDeviceGetMaxPcieLinkGeneration": ( + [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint)], + ctypes.c_int, + ), + "nvmlDeviceGetCurrPcieLinkGeneration": ( + [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint)], + ctypes.c_int, + ), + "nvmlDeviceGetMemoryInfo": ( + [ctypes.c_void_p, ctypes.POINTER(_nvmlMemory_t)], + ctypes.c_int, + ), + "nvmlDeviceGetUtilizationRates": ( + [ctypes.c_void_p, ctypes.POINTER(_nvmlUtilization_t)], + ctypes.c_int, + ), + "nvmlDeviceGetVbiosVersion": ( + [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint], + ctypes.c_int, + ), + "nvmlDeviceGetCudaComputeCapability": ( + [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ], + ctypes.c_int, + ), + } + for name, (argtypes, restype) in spec.items(): + fn = getattr(self._lib, name, None) + if fn is not None: + try: + fn.argtypes = argtypes + fn.restype = restype + except AttributeError: + fn = None + setattr(self, f"_{name}", fn) + + 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) + + def _check(self, code): + if code == 0: + return + msg = "Unknown NVML error" + if self._nvmlErrorString is not None: + try: + raw = self._nvmlErrorString(code) + if raw: + msg = raw.decode("utf-8") + except (AttributeError, UnicodeDecodeError) as exc: + logger.debug("Failed to decode NVML error string: %s", exc) + raise RuntimeError(f"NVML error {code}: {msg}") + + def _call(self, names, *args): + for name in names: + fn = getattr(self, f"_{name}", None) + if fn is not None: + return fn(*args) + raise RuntimeError(f"NVML functions {names} not available") + + def device_count(self): + if not self._available: + raise RuntimeError("NVML is not available") + count = ctypes.c_uint() + self._check(self._call(["nvmlDeviceGetCount_v2", "nvmlDeviceGetCount"], ctypes.byref(count))) + return count.value + + def device_handle(self, gpu_id: str): + if not self._available: + raise RuntimeError("NVML is not available") + handle = ctypes.c_void_p() + if gpu_id.isdigit(): + self._check( + self._call( + ["nvmlDeviceGetHandleByIndex_v2", "nvmlDeviceGetHandleByIndex"], + int(gpu_id), + ctypes.byref(handle), + ) + ) + elif gpu_id.startswith("GPU-"): + self._check( + self._call( + ["nvmlDeviceGetHandleByUUID"], + gpu_id.encode("utf-8"), + ctypes.byref(handle), + ) + ) + elif ":" in gpu_id: + self._check( + self._call( + [ + "nvmlDeviceGetHandleByPciBusId_v2", + "nvmlDeviceGetHandleByPciBusId", + ], + gpu_id.encode("utf-8"), + ctypes.byref(handle), + ) + ) + else: + self._check( + self._call( + ["nvmlDeviceGetHandleByIndex_v2", "nvmlDeviceGetHandleByIndex"], + int(gpu_id), + ctypes.byref(handle), + ) + ) + return handle + + def device_name(self, handle): + buf = ctypes.create_string_buffer(96) + self._check(self._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 + + def max_pcie_generation(self, handle): + gen = ctypes.c_uint() + self._check(self._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))) + return gen.value + + def memory_total(self, handle): + mem = _nvmlMemory_t() + self._check(self._nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem))) + return mem.total + + def memory_used(self, handle): + mem = _nvmlMemory_t() + self._check(self._nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem))) + return mem.used + + def utilization_gpu(self, handle): + util = _nvmlUtilization_t() + self._check(self._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)) + 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)) + 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) + ) + ) + return f"{major.value}.{minor.value}" + + +_nvml = _NVML() + + +class _NvidiaSmiCache: + """Thread-safe LRU cache for ``nvidia-smi`` output with a per-entry TTL. + + All reads and writes to the shared ``OrderedDict`` are serialized by a + ``threading.Lock``, so the cache is safe to use in free-threaded (GIL=0) + Python builds. + """ + + def __init__(self): + self._cache = OrderedDict() + self._lock = threading.Lock() + + def __call__(self, key, fetch): + with self._lock: + ttl = config.nvidia_smi_cache_ttl + if ttl == 0: + return fetch() + + with self._lock: + now = time.monotonic() + if key in self._cache: + ts, value = self._cache[key] + if now - ts < ttl: + self._cache.move_to_end(key) + return value + del self._cache[key] + value = fetch() + self._cache[key] = (now, value) + self._cache.move_to_end(key) + maxsize = config.nvidia_smi_cache_maxsize + while len(self._cache) > maxsize: + self._cache.popitem(last=False) + return value + + +_nvidia_smi_cache = _NvidiaSmiCache() + + +def _run_nvidia_smi(args, line_start=None, seperator=None): + """Run ``nvidia-smi`` at most once per TTL; cache keyed by the command.""" + + key = (tuple(args), line_start, seperator) + return _nvidia_smi_cache(key, lambda: _run(args, line_start=line_start, seperator=seperator)) + + class NvidiaDevice(GPUDevice): fast_metrics_same_as_slow = False @@ -15,71 +401,127 @@ def __init__(self, cls, index): super().__init__(cls, index) self.gpu_id = self._get_gpu_id() - try: - args = [ - "nvidia-smi", - f"--id={self.gpu_id}", - "--query-gpu=" - "name," - "memory.total," - "pci.bus_id," - "pcie.link.gen.max," - "pcie.link.gen.current," - "driver_version", - "--format=csv,noheader,nounits", - ] - - result = _run(args=args, seperator="\n") - - model, total_memory, pci_bus_id, pcie_gen, pcie_width, driver = (result[0].split(", ")) - - result = _run(args=["nvidia-smi", "-q", "-i", f"{self.gpu_id}"], seperator="\n") - firmware = " ".join([line.split(":", 1)[1].strip() for line in result if "VBIOS" in line]) - - if model.lower().startswith("nvidia"): - model = model[len("nvidia"):] - - compute_cap = ( - _run(["nvidia-smi", "--format=csv", "--query-gpu=compute_cap", "-i", f"{self.gpu_id}"]) - .removeprefix("compute_cap\n") - ) + if _nvml.available: + try: + self._init_from_nvml(cls) + return + except (RuntimeError, ValueError, IndexError, AttributeError) as exc: + logger.debug("NVML device initialization failed, falling back to nvidia-smi: %s", exc) - cls.model = model.strip().lower() - cls.memory_total = int(total_memory) * 1024 * 1024 # bytes - cls.vendor = "nvidia" - cls.features = [compute_cap] - cls.pcie = Pcie(gen=int(pcie_gen), speed=int(pcie_width), id=pci_bus_id) - cls.gpu = GPU(driver=driver, firmware=firmware) - except FileNotFoundError: - raise FileNotFoundError() - except Exception as e: - raise e + self._init_from_nvidia_smi(cls) def _get_gpu_id(self): - gpu_count = len(_run(["nvidia-smi", "--list-gpus"]).splitlines()) + if _nvml.available: + try: + gpu_count = _nvml.device_count() + except (RuntimeError, OSError, ValueError) as exc: + logger.debug("Failed to get NVML device count: %s", exc) + gpu_count = 0 + else: + try: + gpu_count = len(_run_nvidia_smi(["nvidia-smi", "--list-gpus"]).splitlines()) + except (RuntimeError, OSError, ValueError) as exc: + logger.debug("Failed to list NVIDIA GPUs: %s", exc) + gpu_count = 0 + cudas = os.environ.get("CUDA_VISIBLE_DEVICES", "") cuda_list = 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.", RuntimeWarning, 2) + warnings.warn( + "Detected different devices in the system. Please make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to avoid unexpected behavior.", + RuntimeWarning, + 2, + ) if cuda_list and len(cuda_list) > self.index: return cuda_list[self.index] else: return str(self.index) + def _init_from_nvml(self, cls): + handle = _nvml.device_handle(self.gpu_id) + + 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) + + if model.lower().startswith("nvidia"): + model = model[len("nvidia"):] + + 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.gpu = GPU(driver=driver, firmware=firmware) + + def _init_from_nvidia_smi(self, cls): + args = [ + "nvidia-smi", + f"--id={self.gpu_id}", + "--query-gpu=name,memory.total,pci.bus_id,pcie.link.gen.max,pcie.link.gen.current,driver_version", + "--format=csv,noheader,nounits", + ] + + result = _run_nvidia_smi(args=args, seperator="\n") + + model, total_memory, pci_bus_id, pcie_gen, pcie_width, driver = (result[0].split(", ")) + + result = _run_nvidia_smi(args=["nvidia-smi", "-q", "-i", f"{self.gpu_id}"], seperator="\n") + firmware = " ".join([line.split(":", 1)[1].strip() for line in result if "VBIOS" in line]) + + if model.lower().startswith("nvidia"): + model = model[len("nvidia"):] + + compute_cap = ( + _run_nvidia_smi(["nvidia-smi", "--format=csv", "--query-gpu=compute_cap", "-i", f"{self.gpu_id}"]) + .removeprefix("compute_cap\n") + ) + + cls.model = model.strip().lower() + cls.memory_total = int(total_memory) * 1024 * 1024 # bytes + cls.vendor = "nvidia" + cls.features = [compute_cap] + cls.pcie = Pcie(gen=int(pcie_gen), speed=int(pcie_width), id=pci_bus_id) + cls.gpu = GPU(driver=driver, firmware=firmware) + + def _metrics_from_nvml(self): + handle = _nvml.device_handle(self.gpu_id) + return NvidiaGPUMetrics( + memory_used=_nvml.memory_used(handle), + memory_process=0, + utilization=_nvml.utilization_gpu(handle), + ) + + def _metrics_from_nvidia_smi(self): + args = [ + "nvidia-smi", + f"--id={self.gpu_id}", + "--query-gpu=memory.used,utilization.gpu", + "--format=csv,noheader,nounits", + ] + used_memory, utilization = _run_nvidia_smi(args=args, seperator="\n")[0].split(", ") + + return NvidiaGPUMetrics( + memory_used=int(used_memory) * 1024 * 1024, # bytes + memory_process=0, # Bytes, TODO, get this + utilization=float(utilization), + ) + def metrics(self): try: - args = ["nvidia-smi", f"--id={self.gpu_id}", "--query-gpu=memory.used,utilization.gpu", "--format=csv,noheader,nounits"] - used_memory, utilization = _run(args=args, seperator="\n")[0].split(", ") - - return NvidiaGPUMetrics( - memory_used=int(used_memory) * 1024 * 1024, # bytes - memory_process=0, # Bytes, TODO, get this - utilization=float(utilization), - ) + if _nvml.available: + try: + return self._metrics_from_nvml() + except (RuntimeError, ValueError, IndexError, AttributeError) as exc: + logger.debug("NVML metrics failed, falling back to nvidia-smi: %s", exc) + return self._metrics_from_nvidia_smi() except FileNotFoundError: raise FileNotFoundError( "The 'nvidia-smi' command was not found. Please ensure that the 'nvidia-utils' package is installed." ) - except Exception as e: - raise e diff --git a/pyproject.toml b/pyproject.toml index d2635f0..9bffbee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" [project] name = "Device-SMI" -version = "0.5.6" +version = "0.5.7" description = "Retrieve gpu, cpu, and npu device info and properties from Linux/MacOS with zero package dependency." readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3" diff --git a/tests/test_nvidia_cache.py b/tests/test_nvidia_cache.py new file mode 100644 index 0000000..58495b5 --- /dev/null +++ b/tests/test_nvidia_cache.py @@ -0,0 +1,167 @@ +import threading +import time + +import pytest + +from device_smi import Device, config +from device_smi.nvidia import _NvidiaSmiCache, _nvml, _run_nvidia_smi + + +def test_config_ttl_and_maxsize(): + original_ttl = config.nvidia_smi_cache_ttl + original_maxsize = config.nvidia_smi_cache_maxsize + try: + assert config.nvidia_smi_cache_ttl == 0.0 + + config.nvidia_smi_cache_ttl = 2.0 + config.nvidia_smi_cache_maxsize = 8 + assert config.nvidia_smi_cache_ttl == 2.0 + assert config.nvidia_smi_cache_maxsize == 8 + + # TTL may be 0 (disabled); maxsize must remain positive. + config.nvidia_smi_cache_ttl = 0 + with pytest.raises(ValueError): + config.nvidia_smi_cache_maxsize = 0 + with pytest.raises(ValueError): + config.nvidia_smi_cache_maxsize = -1 + with pytest.raises(ValueError): + config.nvidia_smi_cache_ttl = -1 + finally: + config.nvidia_smi_cache_ttl = original_ttl + config.nvidia_smi_cache_maxsize = original_maxsize + + +def test_device_config_exposed(): + assert Device.config is config + assert hasattr(config, "nvidia_smi_cache_ttl") + assert hasattr(config, "nvidia_smi_cache_maxsize") + + +def test_nvidia_smi_cache_disabled_when_ttl_zero(): + original_ttl = config.nvidia_smi_cache_ttl + config.nvidia_smi_cache_ttl = 0 + try: + cache = _NvidiaSmiCache() + call_count = 0 + + def fetch(): + nonlocal call_count + call_count += 1 + return f"result-{call_count}" + + v1 = cache("key", fetch) + v2 = cache("key", fetch) + assert v1 == "result-1" + assert v2 == "result-2" + assert call_count == 2 + finally: + config.nvidia_smi_cache_ttl = original_ttl + + +def test_nvidia_smi_cache_returns_cached_value_within_ttl(): + original_ttl = config.nvidia_smi_cache_ttl + config.nvidia_smi_cache_ttl = 0.2 + try: + cache = _NvidiaSmiCache() + call_count = 0 + + def fetch(): + nonlocal call_count + call_count += 1 + return f"result-{call_count}" + + v1 = cache("key", fetch) + v2 = cache("key", fetch) + assert v1 == v2 == "result-1" + assert call_count == 1 + + time.sleep(0.25) + v3 = cache("key", fetch) + assert v3 == "result-2" + assert call_count == 2 + finally: + config.nvidia_smi_cache_ttl = original_ttl + + +def test_nvidia_smi_cache_lru_eviction(): + original_ttl = config.nvidia_smi_cache_ttl + original_maxsize = config.nvidia_smi_cache_maxsize + config.nvidia_smi_cache_ttl = 2.0 + config.nvidia_smi_cache_maxsize = 2 + try: + cache = _NvidiaSmiCache() + cache("a", lambda: 1) + cache("b", lambda: 2) + cache("a", lambda: 1) # refresh a + cache("c", lambda: 3) # evicts b + assert "b" not in cache._cache + assert "a" in cache._cache + assert "c" in cache._cache + finally: + config.nvidia_smi_cache_ttl = original_ttl + config.nvidia_smi_cache_maxsize = original_maxsize + + +def test_run_nvidia_smi_uses_cache(monkeypatch): + original_ttl = config.nvidia_smi_cache_ttl + config.nvidia_smi_cache_ttl = 0.5 + try: + call_count = 0 + + def fake_run(args, line_start=None, seperator=None): + nonlocal call_count + call_count += 1 + return ["line"] if seperator else "str" + + monkeypatch.setattr("device_smi.nvidia._run", fake_run) + + from device_smi import nvidia as nvidia_module + + old_cache = nvidia_module._nvidia_smi_cache + nvidia_module._nvidia_smi_cache = _NvidiaSmiCache() + try: + r1 = _run_nvidia_smi(["nvidia-smi", "--query-gpu=name"], seperator="\n") + r2 = _run_nvidia_smi(["nvidia-smi", "--query-gpu=name"], seperator="\n") + assert r1 == r2 == ["line"] + assert call_count == 1 + finally: + nvidia_module._nvidia_smi_cache = old_cache + finally: + config.nvidia_smi_cache_ttl = original_ttl + + +def test_nvidia_smi_cache_is_thread_safe(): + original_ttl = config.nvidia_smi_cache_ttl + config.nvidia_smi_cache_ttl = 2.0 + try: + cache = _NvidiaSmiCache() + call_count = 0 + results = [] + lock = threading.Lock() + + def fetch(): + nonlocal call_count + call_count += 1 + time.sleep(0.001) + return call_count + + def worker(): + value = cache("key", fetch) + with lock: + results.append(value) + + threads = [threading.Thread(target=worker) for _ in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + + # All threads should observe the same cached value. + assert len(set(results)) == 1 + finally: + config.nvidia_smi_cache_ttl = original_ttl + + +def test_nvml_singleton_does_not_crash_without_gpu(): + assert _nvml._initialized is True + assert isinstance(_nvml.available, bool)