Coverage for src/lilbee/providers/fleet/gpu_backends/amd.py: 100%
81 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""ROCm/HIP GPU utilization via rocm-smi (preferred) or amd-smi (fallback).
3rocm-smi leads because it reports power draw, the board power cap, and VRAM in
4one call; amd-smi (checked on AMDSMI 26.0.2 / MI300X) exposes no power cap in
5metric or static output, so its path cannot derive power-based activity and
6reports the raw busy flag instead."""
8from __future__ import annotations
10import json
12from lilbee.providers.fleet.gpu_backends.base import (
13 UtilSample,
14 extract_int,
15 find_metric,
16 parse_device_index,
17 run_smi,
18)
20_TOOL_AMD_SMI = "amd-smi"
21_TOOL_ROCM_SMI = "rocm-smi"
23_AMD_SMI_ARGS = ("metric", "--usage", "--temperature", "--json")
24_ROCM_SMI_ARGS = (
25 "--showuse",
26 "--showmeminfo",
27 "vram",
28 "--showtemp",
29 "--showpower",
30 "--showmaxpower",
31 "--json",
32)
34# rocm-smi key names for VRAM (byte values).
35_ROCM_VRAM_TOTAL_KEY = "VRAM Total Memory (B)"
36_ROCM_VRAM_USED_KEY = "VRAM Total Used Memory (B)"
38# amdgpu reports 100% busy whenever a compute context is resident, idle or not,
39# so power draw as a fraction of the board cap is the activity signal on AMD.
40# Consumer cards report "Average Graphics Package Power"; datacenter cards
41# (MI300X) report "Current Socket Graphics Package Power".
42_ROCM_POWER_KEYS = (
43 "Average Graphics Package Power (W)",
44 "Current Socket Graphics Package Power (W)",
45)
46_ROCM_POWER_CAP_KEYS = ("Max Graphics Package Power (W)",)
48_TIMEOUT_S = 5.0
51class AmdBackend:
52 """ROCm/HIP util via rocm-smi, falling back to amd-smi."""
54 def sample(self, indices: frozenset[int]) -> dict[int, UtilSample]:
55 result = _rocm_smi_samples(indices)
56 if not result:
57 result = _amd_smi_samples(indices)
58 return result
61def _amd_smi_output() -> str:
62 """amd-smi JSON stdout, or "" when it can't run."""
63 return run_smi(_TOOL_AMD_SMI, list(_AMD_SMI_ARGS), _TIMEOUT_S)
66def _rocm_smi_output() -> str:
67 """rocm-smi JSON stdout, or "" when it can't run."""
68 return run_smi(_TOOL_ROCM_SMI, list(_ROCM_SMI_ARGS), _TIMEOUT_S)
71def _parse_amd_smi(raw: str, indices: frozenset[int]) -> dict[int, UtilSample]:
72 """Parse amd-smi JSON into UtilSample per index, or {} on failure.
74 amd-smi nests the same readings differently across versions: flat
75 ({"gfx_activity": 72}), value-wrapped ({"gfx_activity": {"value": 72}}), or under
76 a block ({"usage": {"gfx_activity": {"value": 72}}, "temperature": {"edge":
77 {"value": 61}}}). ``find_metric`` reads any of these by key at any depth. The
78 index-carrying key ("gpu") is read at the top level so a nested "gpu" block
79 can't be mistaken for it.
80 """
81 if not raw:
82 return {}
83 try:
84 data = json.loads(raw)
85 except json.JSONDecodeError:
86 return {}
87 # amd-smi metric --json emits a list of GPU objects, {"gpu": [...]}, or
88 # {"gpu_data": [...]} (AMDSMI 26.x, seen on MI300X).
89 if isinstance(data, list):
90 items: list[object] = data
91 else:
92 items = data.get("gpu_data") or data.get("gpu", [])
93 samples: dict[int, UtilSample] = {}
94 for item in items:
95 if not isinstance(item, dict):
96 continue
97 raw_index = item.get("gpu") if "gpu" in item else item.get("id", -1)
98 index = extract_int({"i": raw_index}, ("i",))
99 if index is None or index not in indices:
100 continue
101 util = find_metric(item, ("gfx_activity", "gfx_busy_percent", "gpu_activity"))
102 temp_block = item.get("temperature")
103 if isinstance(temp_block, dict):
104 temp = find_metric(temp_block, ("edge", "junction", "hotspot"))
105 else:
106 temp = find_metric(item, ("temperature_c", "temp_edge"))
107 # VRAM not reliably present in metric mode; leave 0 so the orchestrator
108 # keeps structural VRAM.
109 samples[index] = UtilSample(
110 index=index,
111 utilization_pct=util,
112 temperature_c=temp,
113 free_bytes=0,
114 total_bytes=0,
115 )
116 return samples
119def _parse_rocm_smi(raw: str, indices: frozenset[int]) -> dict[int, UtilSample]:
120 """Parse rocm-smi JSON into UtilSample per index, or {} on failure.
122 utilization_pct is power draw as a percent of the board power cap when both
123 keys are present; the busy flag is the fallback for older rocm-smi output.
124 VRAM is parsed when the rocm-smi VRAM keys are present (byte values);
125 absent keys leave free_bytes/total_bytes as 0 (structural-fallback sentinel).
126 """
127 if not raw:
128 return {}
129 try:
130 data = json.loads(raw)
131 except json.JSONDecodeError:
132 return {}
133 # rocm-smi --json emits {"card0": {...}, "card1": {...}} or {"GPU[0]": {...}}.
134 if not isinstance(data, dict):
135 return {}
136 samples: dict[int, UtilSample] = {}
137 for key, val in data.items():
138 if not isinstance(val, dict):
139 continue
140 index = parse_device_index(key)
141 if index is None or index not in indices:
142 continue
143 util = extract_int(val, ("GPU use (%)", "GPU_UTIL", "gfx_activity"))
144 power = extract_int(val, _ROCM_POWER_KEYS)
145 cap = extract_int(val, _ROCM_POWER_CAP_KEYS)
146 if power is not None and cap:
147 util = min(100, round(power * 100 / cap))
148 # Datacenter cards (MI300X) have no edge sensor; junction carries it.
149 temp = extract_int(
150 val,
151 (
152 "Temperature (Sensor edge) (C)",
153 "Temperature (Sensor junction) (C)",
154 "temp_edge",
155 "temp",
156 ),
157 )
158 # VRAM keys carry byte strings; parse when present.
159 total_b = extract_int(val, (_ROCM_VRAM_TOTAL_KEY,))
160 used_b = extract_int(val, (_ROCM_VRAM_USED_KEY,))
161 if total_b is not None:
162 free_b = max(total_b - (used_b or 0), 0)
163 else:
164 total_b = 0
165 free_b = 0
166 samples[index] = UtilSample(
167 index=index,
168 utilization_pct=util,
169 temperature_c=temp,
170 free_bytes=free_b,
171 total_bytes=total_b,
172 )
173 return samples
176def _amd_smi_samples(indices: frozenset[int]) -> dict[int, UtilSample]:
177 return _parse_amd_smi(_amd_smi_output(), indices)
180def _rocm_smi_samples(indices: frozenset[int]) -> dict[int, UtilSample]:
181 return _parse_rocm_smi(_rocm_smi_output(), indices)