Coverage for src/lilbee/providers/fleet/gpu_stats.py: 100%
81 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Live per-GPU activity stats for the placement view.
3devices.py enumerates GPUs once (structural: index, name, total VRAM). This reads
4the moving numbers, compute utilization and current free memory, so a client can
5animate a per-card load bar.
7Vendor-specific probing lives in gpu_backends/; this module groups devices by
8backend, dispatches to resolve_backend(), and merges live samples into GpuStat
9entries. Adding a new vendor means one file in gpu_backends/ and one registry
10line; this file is not touched.
11"""
13from __future__ import annotations
15import os
16import threading
17import time
18from collections.abc import Sequence
19from dataclasses import dataclass
20from typing import Protocol
22from lilbee.providers.fleet.devices import _CUDA_ORDER_VAR, _PCI_BUS_ID_ORDER
23from lilbee.providers.fleet.gpu_backends import (
24 IntelUtilHint,
25 UtilSample,
26 resolve_backend,
27 util_backend_name,
28)
29from lilbee.providers.fleet.gpu_backends import (
30 intel_util_hint as _detect_intel_util_hint,
31)
34class DeviceLike(Protocol):
35 """Structural view of a probed GPU (FleetDevice or app-layer GpuInfo)."""
37 @property
38 def index(self) -> int: ...
39 @property
40 def backend(self) -> str: ...
41 @property
42 def name(self) -> str: ...
43 @property
44 def total_bytes(self) -> int: ...
45 @property
46 def free_bytes(self) -> int: ...
49@dataclass(frozen=True)
50class GpuStat:
51 """A live snapshot of one GPU, keyed to its structural index."""
53 index: int
54 utilization_pct: int | None
55 free_bytes: int
56 total_bytes: int
57 temperature_c: int | None = None
60def _safe_sample(
61 backend_name: str,
62 indices: frozenset[int],
63) -> dict[int, UtilSample]:
64 """Dispatch to the vendor backend; return {} on any failure."""
65 backend = resolve_backend(backend_name)
66 if backend is None:
67 return {}
68 try:
69 return backend.sample(indices)
70 except Exception: # backends must never crash the sampler
71 return {}
74# One in-flight probe per device set, shared by every concurrent caller. The
75# probe shells out to a vendor SMI tool (five-second timeout) and the Intel paths
76# sleep and scan /proc on top, so N open placement views would otherwise mean N
77# concurrent subprocesses every tick against the same hardware. A sample is
78# reused for slightly under one tick, which is fresh enough for a live view and
79# turns the per-client cost back into a per-machine one.
80_SHARED_SAMPLE_TTL_S = 0.9
81_shared_lock = threading.Lock()
82_shared_sample: dict[tuple[int, ...], tuple[float, dict[int, GpuStat]]] = {}
85def probe_gpu_stats_shared(devices: Sequence[DeviceLike]) -> dict[int, GpuStat]:
86 """``probe_gpu_stats``, coalesced across concurrent callers.
88 Callers arriving while a sample is fresh reuse it; the rest serialise on the
89 lock so exactly one probe runs per device set per interval, rather than one
90 per subscriber. Kept separate from ``probe_gpu_stats`` so one-shot callers
91 still get an uncached reading.
92 """
93 key = tuple(sorted(d.index for d in devices))
94 with _shared_lock:
95 cached = _shared_sample.get(key)
96 if cached is not None and time.monotonic() - cached[0] < _SHARED_SAMPLE_TTL_S:
97 return cached[1]
98 stats = probe_gpu_stats(devices)
99 _shared_sample[key] = (time.monotonic(), stats)
100 return stats
103# Which visible-devices variable masks each util backend's index space. The
104# engine numbers its devices after the mask, the SMI tools before it, so under
105# CUDA_VISIBLE_DEVICES=2,3 the fleet holds devices 0 and 1 while nvidia-smi keeps
106# reporting 0..3. Merging those two spaces by ordinal attributes another
107# tenant's cards' utilization, temperature and free memory to this fleet.
108_BACKEND_VISIBLE_VARS: dict[str, tuple[str, ...]] = {
109 "CUDA": ("CUDA_VISIBLE_DEVICES",),
110 # Innermost mask first. ROCr filters, then HIP re-indexes within the
111 # survivors, so walking outward from the fleet's index means undoing HIP
112 # before ROCr.
113 "ROCm": ("HIP_VISIBLE_DEVICES", "GPU_DEVICE_ORDINAL", "ROCR_VISIBLE_DEVICES"),
114 "HIP": ("HIP_VISIBLE_DEVICES", "GPU_DEVICE_ORDINAL", "ROCR_VISIBLE_DEVICES"),
115}
118_CUDA_BACKEND = "CUDA"
121def _cuda_order_is_reordered() -> bool:
122 """Whether CUDA_DEVICE_ORDER puts the runtime out of step with nvidia-smi.
124 Reads the same variable and default the probe writes, so the two cannot
125 drift apart: unset means the probe supplies bus order and the two index
126 spaces agree, and a preset value is respected there and honoured here.
127 """
128 order = os.environ.get(_CUDA_ORDER_VAR, "").strip().upper()
129 return bool(order) and order != _PCI_BUS_ID_ORDER
132def _physical_index(backend_name: str, fleet_index: int) -> int | None:
133 """The index the vendor's SMI tool uses for the fleet's device *fleet_index*.
135 Resolved by walking the visible-devices masks outward from the fleet's own
136 index, undoing each one in the reverse of the order the runtime applied it.
137 ``None`` when a mask cannot be undone, which is the whole point of the
138 return type.
140 A mask naming devices by UUID or MIG instance cannot be inverted by
141 arithmetic, and neither can an index the mask does not reach. Returning the
142 fleet index there was a guess, and on a masked host it is reliably the wrong
143 card: another tenant's utilization, temperature and free VRAM then drive both
144 the GPU panel and the adaptive ingest throttle. No sample is the honest
145 answer, and the caller already has a device-reported baseline to fall back on.
147 An unset mask, or a backend with no mask, leaves the index alone: that is
148 every unmasked host. Intel is absent because ONEAPI_DEVICE_SELECTOR is a
149 selector grammar rather than an index list, and xpu-smi is called with the
150 fleet index directly.
152 A reordered index space is the same problem without a mask. CUDA_DEVICE_ORDER
153 set to FASTEST_FIRST makes the runtime enumerate by speed while nvidia-smi
154 keeps enumerating by bus, so the two index spaces name different cards and no
155 arithmetic relates them. Correlating on a stable identity (UUID or PCI bus
156 id) is what would make this structural; until the probe carries one, an
157 honest refusal beats a confident mismatch.
158 """
159 if backend_name == _CUDA_BACKEND and _cuda_order_is_reordered():
160 return None
161 index = fleet_index
162 for var in _BACKEND_VISIBLE_VARS.get(backend_name, ()):
163 raw = os.environ.get(var)
164 if not raw:
165 continue
166 entries = [e.strip() for e in raw.split(",")]
167 if not all(e.isdigit() for e in entries) or index >= len(entries):
168 return None
169 index = int(entries[index])
170 return index
173def probe_gpu_stats(devices: Sequence[DeviceLike]) -> dict[int, GpuStat]:
174 """Live stats keyed by device index. Empty when no devices are given.
176 Groups devices by vendor backend, dispatches once per group, and merges
177 live util/temp samples into GpuStat entries. Any index the backend can't
178 cover falls back to structural VRAM with util=None and temperature_c=None.
179 """
180 # Structural fallbacks: util=None, temp=None, VRAM from the probe.
181 stats: dict[int, GpuStat] = {
182 d.index: GpuStat(d.index, None, d.free_bytes, d.total_bytes) for d in devices
183 }
185 # Group by the util backend, not the raw inference backend: a Vulkan-exposed
186 # consumer GPU routes to its vendor's util source by name.
187 by_backend: dict[str, list[DeviceLike]] = {}
188 for d in devices:
189 by_backend.setdefault(util_backend_name(d.backend, d.name), []).append(d)
191 for backend_name, group in by_backend.items():
192 # Ask the SMI tool about its own indices, and merge its answers back into
193 # the engine's; the two spaces differ whenever a visible-devices mask is set.
194 # A device whose mask cannot be inverted is left out entirely rather than
195 # merged under a guessed index: it keeps the structural reading the probe
196 # already gave it, which is right, instead of another card's live one.
197 fleet_by_physical = {
198 physical: d.index
199 for d in group
200 if (physical := _physical_index(backend_name, d.index)) is not None
201 }
202 for physical, sample in _safe_sample(backend_name, frozenset(fleet_by_physical)).items():
203 index = fleet_by_physical.get(physical)
204 if index is None or index not in stats:
205 continue
206 base = stats[index]
207 # Keep structural VRAM when the backend returns the 0/0 sentinel
208 # (amd-smi metric mode and xpu-smi stats both omit total VRAM); a
209 # backend that does report memory takes precedence.
210 free = sample.free_bytes if sample.free_bytes or sample.total_bytes else base.free_bytes
211 total = sample.total_bytes if sample.total_bytes else base.total_bytes
212 stats[index] = GpuStat(
213 index=index,
214 utilization_pct=sample.utilization_pct,
215 free_bytes=free,
216 total_bytes=total,
217 temperature_c=sample.temperature_c,
218 )
220 return {i: stats[i] for i in sorted(stats)}
223def intel_util_hint(
224 devices: Sequence[DeviceLike], stats: dict[int, GpuStat]
225) -> IntelUtilHint | None:
226 """The fix that would unblock an Intel GPU's missing util reading, else None.
228 Fires only when an Intel device's util is actually missing; a surface turns
229 the hint into a localized message (grant when intel_gpu_top is installed but
230 blocked, install when it is absent, e.g. kernels too old for fdinfo).
231 """
232 for d in devices:
233 if util_backend_name(d.backend, d.name) != "SYCL":
234 continue
235 stat = stats.get(d.index)
236 if stat is not None and stat.utilization_pct is None:
237 return _detect_intel_util_hint()
238 return None
241def probe_intel_util_hint(devices: Sequence[DeviceLike]) -> IntelUtilHint | None:
242 """Probe live stats and evaluate the Intel util hint in one call."""
243 return intel_util_hint(devices, probe_gpu_stats(devices))