"""GPU load readout via nvidia-smi (present in the container thanks to the NVIDIA Container Toolkit's `utility` capability). Returns None if unavailable — the UI just shows n/a (e.g. CPU-fallback run). Reads are CACHED and de-duplicated: the UI meter polls fast, /status reads it, and the autoscaler samples it — if each spawned its own `nvidia-smi` (slow on a busy GPU) those blocking subprocesses would pile up in the server's thread pool and make the Start/Stop buttons feel dead. So a short TTL serves recent callers from cache, and only ONE probe runs at a time (others get the last value).""" import subprocess import threading import time _TTL = 1.0 # seconds a sample is reused before re-probing _lock = threading.Lock() _cache: dict | None = None _cache_t = 0.0 _probing = False def _probe() -> dict | None: try: out = subprocess.run( [ "nvidia-smi", "--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu", "--format=csv,noheader,nounits", ], capture_output=True, text=True, timeout=5, check=True, ).stdout.strip().splitlines() except (OSError, subprocess.SubprocessError): return None if not out: return None parts = [p.strip() for p in out[0].split(",")] try: return { "util_pct": int(float(parts[0])), "mem_used_mb": int(float(parts[1])), "mem_total_mb": int(float(parts[2])), "temp_c": int(float(parts[3])), } except (ValueError, IndexError): return None def read_gpu(max_age: float = _TTL) -> dict | None: """Latest GPU reading, cached. Serves from cache when fresh; when stale, exactly one caller re-probes while the rest get the last value — so request threads never block behind more than one `nvidia-smi`.""" global _cache, _cache_t, _probing now = time.monotonic() with _lock: fresh = _cache is not None and (now - _cache_t) < max_age if fresh or _probing: # fresh, or a probe is already running return _cache _probing = True try: val = _probe() finally: with _lock: _cache = val _cache_t = time.monotonic() _probing = False return val