Coverage for src/lilbee/cli/tui/widgets/gpu_fleet_panel.py: 100%

1"""Live GPU fleet panel widget for the Fleet view. 

2 

3Renders per-card utilization and VRAM bars styled via rose-pine theme tokens and 

4refreshes on a ~1 s interval. The probe runs off the UI thread so the 

5event loop is never blocked by the nvidia-smi subprocess. 

6""" 

7 

8from __future__ import annotations 

9 

10import logging 

11from collections.abc import Sequence 

12from pathlib import Path 

13from typing import TYPE_CHECKING, ClassVar 

14 

15from textual import work 

16from textual.timer import Timer 

17from textual.widgets import Static 

18 

19from lilbee.cli.tui import messages as msg 

20from lilbee.cli.tui.thread_safe import call_from_thread 

21from lilbee.providers.fleet.gpu_stats import GpuStat, intel_util_hint, probe_gpu_stats 

22 

23if TYPE_CHECKING: 

24 from lilbee.providers.fleet.gpu_stats import DeviceLike 

25 

26log = logging.getLogger(__name__) 

27 

28_CSS_FILE = Path(__file__).parent / "gpu_fleet_panel.tcss" 

29 

30# Bar render constants. Fill/track glyphs come from msg.progress_bar_glyphs() 

31# at render time: block elements where the terminal tiles them, box-drawing 

32# elsewhere. 

33_BAR_WIDTH = 16 # cells per bar 

34_BULLET = "●" # colored dot before card label 

35 

36# Utilization thresholds (%) 

37_UTIL_WARM = 40 

38_UTIL_HOT = 80 

39 

40# VRAM saturation threshold (fraction) 

41_VRAM_HIGH = 0.85 

42 

43# Temperature threshold (°C) above which temp is shown in error color 

44_TEMP_HOT = 75 

45 

46# Refresh interval 

47_TICK_INTERVAL_S = 1.0 

48 

49_GIB = 1024**3 

50 

51 

52# Shown as the util reading when utilization_pct is None 

53_UTIL_DASH = " -- " 

54 

55# Rich color used when a theme token is absent (renders in terminal default) 

56_COLOR_FALLBACK = "default" 

57 

58# Badge role/model separator 

59_BADGE_SEP = " - " 

60 

61 

62def _theme_color(theme: dict[str, str], token: str) -> str: 

63 """Resolve a theme token to a Rich color, falling back to terminal default.""" 

64 return theme.get(token, _COLOR_FALLBACK) 

65 

66 

67def _heat_color(utilization_pct: int, theme: dict[str, str]) -> str: 

68 """Theme-token heat color keyed on utilization percentage.""" 

69 if utilization_pct < _UTIL_WARM: 

70 return _theme_color(theme, "success") 

71 if utilization_pct < _UTIL_HOT: 

72 return _theme_color(theme, "warning") 

73 return _theme_color(theme, "error") 

74 

75 

76def _bar(fraction: float, width: int, fill_color: str, track_color: str) -> str: 

77 """Render a filled bar with a track for the remainder.""" 

78 fill, track = msg.progress_bar_glyphs() 

79 clamped = max(0.0, min(1.0, fraction)) 

80 filled = round(clamped * width) 

81 return f"[{fill_color}]{fill * filled}[/][{track_color}]{track * (width - filled)}[/]" 

82 

83 

84def _badge_role_markup(badge: str, secondary: str, muted: str) -> str: 

85 """Return Rich markup for the role portion of a badge. 

86 

87 Splits "role - model" so the role renders in secondary and the model in muted. 

88 """ 

89 if _BADGE_SEP in badge: 

90 role_part, model_part = badge.split(_BADGE_SEP, 1) 

91 return f"[{secondary}]{role_part}[/][{muted}]{_BADGE_SEP}{model_part}[/]" 

92 return f"[{secondary}]{badge}[/]" 

93 

94 

95def _render_row( 

96 s: GpuStat, 

97 label: str, 

98 badge: str, 

99 theme: dict[str, str], 

100) -> str: 

101 """Render one GPU as a single table row: bullet, label, util bar, vram bar, badge.""" 

102 muted = _theme_color(theme, "text-muted") 

103 secondary = _theme_color(theme, "secondary") 

104 foreground = _theme_color(theme, "foreground") 

105 panel_color = _theme_color(theme, "panel") 

106 primary = _theme_color(theme, "primary") 

107 error = _theme_color(theme, "error") 

108 

109 used_bytes = s.total_bytes - s.free_bytes 

110 vram_frac = used_bytes / s.total_bytes if s.total_bytes else 0.0 

111 used_gib = used_bytes / _GIB 

112 total_gib = s.total_bytes / _GIB 

113 vram_color = error if vram_frac >= _VRAM_HIGH else primary 

114 

115 if s.utilization_pct is not None: 

116 heat = _heat_color(s.utilization_pct, theme) 

117 dot_color = heat 

118 util_bar = _bar(s.utilization_pct / 100, _BAR_WIDTH, heat, panel_color) 

119 util_pct = f"[{heat}]{s.utilization_pct:>3}%[/]" 

120 else: 

121 dot_color = muted 

122 util_bar = f"[{panel_color}]{msg.progress_bar_glyphs()[1] * _BAR_WIDTH}[/]" 

123 util_pct = f"[{muted}]{_UTIL_DASH}[/]" 

124 

125 vram_bar = _bar(vram_frac, _BAR_WIDTH, vram_color, panel_color) 

126 vram_txt = f"[{muted}]{used_gib:>5.1f}/{total_gib:>2.0f}G[/]" 

127 badge_txt = f" {_badge_role_markup(badge, secondary, muted)}" if badge else "" 

128 return ( 

129 f" [{dot_color}]{_BULLET}[/] [bold {foreground}]{label:<6}[/]" 

130 f" {util_bar} {util_pct} {vram_bar} {vram_txt}{badge_txt}" 

131 ) 

132 

133 

134def _render_stats( 

135 stats: dict[int, GpuStat], 

136 labels: dict[int, str], 

137 roles: dict[int, str], 

138 theme: dict[str, str], 

139 *, 

140 probed: bool, 

141) -> str: 

142 """Build the unified GPU table (one row per GPU) from a stat snapshot. 

143 

144 With no stats, ``probed`` picks the placeholder: the empty-GPUs text only 

145 once the device probe has completed, else the probing placeholder so a cold 

146 start never flashes '(no GPUs detected)' before the first device list lands. 

147 """ 

148 if not stats: 

149 muted = _theme_color(theme, "text-muted") 

150 text = msg.FLEET_NO_GPUS if probed else msg.FLEET_GPU_PROBING 

151 return f"[{muted}] {text}[/]" 

152 return "\n".join( 

153 _render_row(stats[idx], labels.get(idx, f"GPU{idx}"), roles.get(idx, ""), theme) 

154 for idx in sorted(stats) 

155 ) 

156 

157 

158class GpuFleetPanel(Static): 

159 """Live GPU utilization and VRAM panel, refreshed every ~1 s. 

160 

161 Mount on any Screen or Widget; devices must be set via `set_devices` 

162 before the first tick fires meaningful data. 

163 """ 

164 

165 DEFAULT_CSS: ClassVar[str] = _CSS_FILE.read_text(encoding="utf-8") 

166 

167 def __init__(self) -> None: 

168 # Initial content is the probing state, not FLEET_NO_GPUS, so a multi-GPU 

169 # box doesn't flash the empty state while the first sample is in flight. 

170 super().__init__(f" {msg.FLEET_GPU_PROBING}", id="gpu-fleet-panel") 

171 self._devices: Sequence[DeviceLike] = [] 

172 self._labels: dict[int, str] = {} 

173 self._roles: dict[int, str] = {} 

174 self._timer: Timer | None = None 

175 # Single-flight: True while a probe worker is running, so a probe 

176 # slower than the tick interval skips ticks instead of stacking 

177 # threads (and their nvidia-smi subprocesses) without bound. 

178 self._probing = False 

179 # False until the parent device probe reports its result via set_devices; 

180 # gates the empty-GPUs text so a cold start holds the probing placeholder. 

181 self._probed = False 

182 # The device probe's failure reason; while set, ticks pause and the panel 

183 # shows the failure instead of stats. 

184 self._probe_error: str | None = None 

185 

186 def set_devices( 

187 self, 

188 devices: Sequence[DeviceLike], 

189 *, 

190 labels: dict[int, str], 

191 roles: dict[int, str] | None = None, 

192 ) -> None: 

193 """Register the GPU devices the panel probes on each tick. 

194 

195 `labels` maps device index to the short display label (e.g. "CUDA0"). 

196 `roles` maps device index to a badge string (e.g. "chat - Qwen3-235B"), "" when idle. 

197 """ 

198 self._devices = list(devices) 

199 self._labels = dict(labels) 

200 self._roles = dict(roles) if roles is not None else {} 

201 self._probed = True 

202 self._probe_error = None 

203 

204 def set_probe_failed(self, reason: str) -> None: 

205 """Render the device probe's failure instead of the probing placeholder.""" 

206 self._probed = True 

207 self._probe_error = reason 

208 theme = self._resolve_theme() 

209 error = _theme_color(theme, "error") 

210 self.update(f"[{error}] {msg.FLEET_GPU_PROBE_FAILED.format(reason=reason)}[/]") 

211 

212 def on_mount(self) -> None: 

213 self._timer = self.set_interval(_TICK_INTERVAL_S, self._request_stats) 

214 # Populate immediately instead of waiting one full interval 

215 self._request_stats() 

216 

217 def on_unmount(self) -> None: 

218 if self._timer is not None: 

219 self._timer.stop() 

220 self._timer = None 

221 

222 def _request_stats(self) -> None: 

223 """Kick off an off-thread stats probe unless one is already in flight.""" 

224 if self._probing or self._probe_error is not None: 

225 return 

226 self._probing = True 

227 self._probe_worker( 

228 list(self._devices), 

229 self._labels.copy(), 

230 self._roles.copy(), 

231 ) 

232 

233 def _resolve_theme(self) -> dict[str, str]: 

234 """Return the current app theme variables with `$` prefix stripped.""" 

235 try: 

236 raw: dict[str, str] = self.app.theme_variables 

237 return {k.lstrip("$"): v for k, v in raw.items()} 

238 except AttributeError: 

239 return {} 

240 

241 @work(thread=True, exit_on_error=False) 

242 def _probe_worker( 

243 self, 

244 devices: list[DeviceLike], 

245 labels: dict[int, str], 

246 roles: dict[int, str], 

247 ) -> None: 

248 """Probe GPU stats off the UI thread and push the result back.""" 

249 try: 

250 stats = probe_gpu_stats(devices) 

251 except Exception: 

252 log.debug("gpu_fleet_panel: probe failed", exc_info=True) 

253 return 

254 finally: 

255 call_from_thread(self, setattr, self, "_probing", False) 

256 call_from_thread(self, self._apply_stats, stats, labels, roles) 

257 

258 def _apply_stats( 

259 self, 

260 stats: dict[int, GpuStat], 

261 labels: dict[int, str], 

262 roles: dict[int, str], 

263 ) -> None: 

264 """Update the rendered content with fresh stat data (main thread).""" 

265 if self._probe_error is not None: 

266 # A stats worker already in flight when the probe failure landed 

267 # must not repaint the empty state over the failure message. 

268 return 

269 theme = self._resolve_theme() 

270 markup = _render_stats(stats, labels, roles, theme, probed=self._probed) 

271 hint = intel_util_hint(self._devices, stats) 

272 if hint: 

273 muted = _theme_color(theme, "text-muted") 

274 markup += f"\n[{muted}] {msg.intel_util_hint_text(hint)}[/]" 

275 self.update(markup)