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

1"""Shared types and helpers for per-vendor GPU utilization backends.""" 

2 

3from __future__ import annotations 

4 

5import re 

6import shutil 

7import subprocess 

8from collections.abc import Sequence 

9from dataclasses import dataclass 

10from typing import Protocol 

11 

12from lilbee.providers.fleet.proc import run_bounded 

13 

14_SMI_TIMEOUT_S = 5.0 

15# Bounded wait for a timed-out smi tool to die before abandoning it: a driver- 

16# wedged sampler must not hang the util-sampling thread (it holds a lock). 

17_SMI_KILL_WAIT_S = 5.0 

18 

19 

20@dataclass(frozen=True) 

21class UtilSample: 

22 """Live utilization + temperature reading for one GPU index.""" 

23 

24 index: int 

25 utilization_pct: int | None 

26 temperature_c: int | None 

27 # VRAM: optional; populated when the backend's tool also reports memory. 

28 # 0 is used as sentinel (no data), not "0 bytes used". 

29 free_bytes: int 

30 total_bytes: int 

31 

32 

33class UtilBackend(Protocol): 

34 """One vendor's live-util probe: given indices, returns per-index samples.""" 

35 

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

37 

38 

39def run_smi(tool: str, args: Sequence[str], timeout: float = _SMI_TIMEOUT_S) -> str: 

40 """Resolve tool via shutil.which, run it, return stdout on rc==0 else "". 

41 

42 Returns "" when the tool is not found, exits non-zero, or raises. 

43 """ 

44 binary = shutil.which(tool) or tool 

45 try: 

46 stdout, returncode = run_bounded( 

47 [binary, *args], timeout_s=timeout, kill_wait_s=_SMI_KILL_WAIT_S, label=tool 

48 ) 

49 except (OSError, subprocess.SubprocessError): 

50 return "" 

51 return stdout if returncode == 0 else "" 

52 

53 

54def extract_int(obj: object, keys: tuple[str, ...]) -> int | None: 

55 """Return the first key's value coerced to int, or None. 

56 

57 Accepts integer, float, and decimal-string values (e.g. "35.0" from rocm-smi). 

58 """ 

59 if not isinstance(obj, dict): 

60 return None 

61 for key in keys: 

62 val = obj.get(key) 

63 if val is not None: 

64 try: 

65 return int(float(val)) 

66 except (ValueError, TypeError): 

67 pass 

68 return None 

69 

70 

71def parse_device_index(key: str) -> int | None: 

72 """Extract a GPU index from keys like 'card0', 'GPU[0]', '0'.""" 

73 m = re.search(r"\d+", key) 

74 return int(m.group()) if m else None 

75 

76 

77def _coerce_metric(val: object) -> int | None: 

78 """Coerce a metric value to int: a bare number, a decimal string, or {"value": N}.""" 

79 if isinstance(val, dict): 

80 val = val.get("value") 

81 if isinstance(val, bool): 

82 return None 

83 if isinstance(val, (int, float)): 

84 return int(val) 

85 if isinstance(val, str): 

86 try: 

87 return int(float(val)) 

88 except ValueError: 

89 return None 

90 return None 

91 

92 

93def find_metric(obj: object, keys: tuple[str, ...]) -> int | None: 

94 """Find the first of *keys* anywhere in a nested dict and coerce it to int. 

95 

96 SMI tools nest the same reading differently across versions -- flat 

97 (``{"gfx_activity": 45}``), value-wrapped (``{"gfx_activity": {"value": 45}}``), 

98 or under a block (``{"usage": {"gfx_activity": {"value": 45}}}``). Searching by 

99 key at any depth reads all three without hard-coding one layout. 

100 """ 

101 if not isinstance(obj, dict): 

102 return None 

103 for key in keys: 

104 if key in obj: 

105 found = _coerce_metric(obj[key]) 

106 if found is not None: 

107 return found 

108 for val in obj.values(): 

109 if isinstance(val, dict): 

110 found = find_metric(val, keys) 

111 if found is not None: 

112 return found 

113 return None