Coverage for src/lilbee/providers/fleet/gpu_backends/nvidia.py: 100%

1"""CUDA GPU utilization via nvidia-smi.""" 

2 

3from __future__ import annotations 

4 

5import threading 

6import time 

7 

8from lilbee.providers.fleet.devices import MIB 

9from lilbee.providers.fleet.gpu_backends.base import UtilSample, run_smi 

10 

11_TOOL = "nvidia-smi" 

12_QUERY = "index,utilization.gpu,memory.used,memory.total" 

13_FIELDS = 4 

14_TIMEOUT_S = 5.0 

15# Coalesce concurrent streams: cache one probe just under the SSE tick interval. 

16_CACHE_TTL_S = 0.9 

17 

18 

19class SmiCache: 

20 """Shares one nvidia-smi probe across concurrent callers within the TTL. 

21 

22 Not a ``cachetools.TTLCache``: this holds the lock *across* the probe so 

23 concurrent callers share one result. ``@cached(cache, lock=...)`` releases 

24 the lock around the call, so three concurrent misses spawn three 

25 subprocesses (measured). ``test_concurrent_threads_probe_once`` pins it. 

26 """ 

27 

28 def __init__(self) -> None: 

29 self._lock = threading.Lock() 

30 self._at = 0.0 

31 self._value: dict[int, UtilSample] = {} 

32 self._primed = False 

33 

34 def stats(self) -> dict[int, UtilSample]: 

35 with self._lock: 

36 now = time.monotonic() 

37 if not self._primed or now - self._at >= _CACHE_TTL_S: 

38 self._value = _parse_smi_output(_smi_output()) 

39 self._at = now 

40 self._primed = True 

41 return self._value 

42 

43 def reset(self) -> None: 

44 with self._lock: 

45 self._primed = False 

46 

47 

48_cache = SmiCache() 

49 

50 

51class NvidiaBackend: 

52 """CUDA util backend: nvidia-smi with TTL caching.""" 

53 

54 def sample(self, indices: frozenset[int]) -> dict[int, UtilSample]: 

55 return {i: s for i, s in _cache.stats().items() if i in indices} 

56 

57 

58def _smi_output() -> str: 

59 """nvidia-smi CSV stdout, or "" when it can't run.""" 

60 return run_smi(_TOOL, [f"--query-gpu={_QUERY}", "--format=csv,noheader,nounits"], _TIMEOUT_S) 

61 

62 

63def _as_int(field: str) -> int | None: 

64 """A CSV field as an int, or ``None`` for nvidia-smi's [N/A].""" 

65 try: 

66 return int(field) 

67 except ValueError: 

68 return None 

69 

70 

71def _parse_smi_output(out: str) -> dict[int, UtilSample]: 

72 """Parse four-column nvidia-smi CSV into UtilSample per index.""" 

73 samples: dict[int, UtilSample] = {} 

74 for line in out.splitlines(): 

75 parts = [p.strip() for p in line.split(",")] 

76 if len(parts) != _FIELDS: 

77 continue 

78 try: 

79 index = int(parts[0]) 

80 except ValueError: 

81 # No index, nothing to attribute the reading to. 

82 continue 

83 # Each column degrades on its own. nvidia-smi prints [N/A] for 

84 # utilization on some cards and inside some VMs, and dropping the row for 

85 # it threw away that card's memory reading too, so a GPU reporting its 

86 # VRAM perfectly well disappeared from the panel entirely. 

87 util = _as_int(parts[1]) 

88 used_mib = _as_int(parts[2]) 

89 total_mib = _as_int(parts[3]) 

90 if used_mib is None or total_mib is None: 

91 continue 

92 total = total_mib * MIB 

93 samples[index] = UtilSample( 

94 index=index, 

95 utilization_pct=util, 

96 temperature_c=None, # four-column query omits temperature 

97 free_bytes=max(total - used_mib * MIB, 0), 

98 total_bytes=total, 

99 ) 

100 return samples