Coverage for src/lilbee/cli/tui/widgets/fleet_body.py: 100%
318 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""FleetBody: the GPU table, live panel, and interactive placement editor.
3Reusable widget that can be mounted both as the top-level Fleet view and as
4a modal overlay.
5"""
7from __future__ import annotations
9import contextlib
10import logging
11import re
12from collections.abc import Callable
13from dataclasses import dataclass
14from pathlib import Path
15from typing import TYPE_CHECKING, ClassVar
17from textual import events, work
18from textual.app import ComposeResult
19from textual.binding import Binding, BindingType
20from textual.containers import Horizontal, Vertical
21from textual.css.query import NoMatches
22from textual.message import Message
23from textual.reactive import reactive
24from textual.widget import Widget
25from textual.widgets import Label, Static
27from lilbee.app.placement import (
28 get_placement,
29 preview_placement,
30 set_placement,
31 wait_chat_ready,
32)
33from lilbee.cli.tui import messages as msg
34from lilbee.cli.tui.thread_safe import call_from_thread
35from lilbee.cli.tui.widgets.gpu_fleet_panel import GpuFleetPanel
36from lilbee.providers.fleet.placement_spec import PlacementError, PlacementSpec, RolePlacement
37from lilbee.providers.roles import REPLICATED_ROLES, WorkerRole
39if TYPE_CHECKING:
40 from lilbee.app.placement import PlacementView
42log = logging.getLogger(__name__)
44_GGUF_SHARD_RE = re.compile(r"-\d{5}-of-\d{5}$")
45_GGUF_QUANT_RE = re.compile(r"-(?:Q\d[\w.]*|IQ\d[\w.]*|F16|BF16|FP16|F32)$", re.IGNORECASE)
46_GGUF_REPO_SUFFIXES = ("-GGUF", "-gguf")
49def _clean_model_name(ref: str) -> str:
50 """Short, human-friendly model name from a GGUF reference.
52 "Qwen/Qwen3-235B-A22B-GGUF/Q4_K_M/...-00001-of-00005.gguf" -> "Qwen3-235B-A22B".
53 """
54 parts = ref.split("/")
55 for component in parts[1:]: # skip the org; the repo name is the friendly one
56 if component.endswith(_GGUF_REPO_SUFFIXES):
57 return component.rsplit("-", 1)[0]
58 stem = parts[-1].removesuffix(".gguf")
59 stem = _GGUF_SHARD_RE.sub("", stem)
60 return _GGUF_QUANT_RE.sub("", stem)
63_CSS_FILE = Path(__file__).parent / "fleet_body.tcss"
65_EDITOR_ID = "#placement-editor"
66_SKIPPED_ID = "#placement-skipped"
67_TITLE_ID = "#placement-title"
68_STATE_ID = "#placement-state"
69_COMMANDS_ID = "#placement-commands"
70_HINT_WIDGET_ID = "#placement-hint"
71_FLEET_PANEL_ID = "#gpu-fleet-panel"
73# Only the replicated roles show a replica stepper; the others always serve one.
74_REPLICA_ROLES = REPLICATED_ROLES
75_HINT = (
76 "Toggle a GPU for each role; -/+ sets replicas. ctrl+r preview · ctrl+s apply · ctrl+x auto"
77)
78_CMD_PREVIEW = "cmd-preview"
79_CMD_APPLY = "cmd-apply"
80_CMD_AUTO = "cmd-auto"
81# GPUs shown per page in the placement grid; more than this paginate.
82_PLACEMENT_PAGE_SIZE = 8
83# rerank is a single pinned instance on one card (a small cross-encoder that never
84# tensor-splits), so its GPU choice is single-select, unlike the multi-GPU roles.
85_SINGLE_ROLES = (WorkerRole.RERANK,)
86# Editor row order: the multi-GPU roles first, rerank last -- the single-card odd
87# one out sits at the bottom instead of between the replicated roles.
88_EDITOR_ROLE_ORDER = (WorkerRole.CHAT, WorkerRole.EMBED, WorkerRole.VISION, WorkerRole.RERANK)
91def _role_kind(role: WorkerRole) -> str:
92 """How a role occupies GPUs: 'mirror' (a copy per card), 'single' (one pinned
93 card), or 'split' (one model tensor-split across cards)."""
94 if role in _REPLICA_ROLES:
95 return "mirror"
96 if role in _SINGLE_ROLES:
97 return "single"
98 return "split"
101class FleetPill(Static, can_focus=True):
102 """Focusable one-line pill; Enter / Space / click presses it.
104 The editor's toggles, steppers, pager, and command controls are all pills
105 (the ``Static, can_focus=True`` + bindings pattern from
106 ``widgets/confirm_dialog.py`` / ``model_bar.py::ChatModePill``): state and
107 focus ride the fill and text style, so a row costs one line instead of the
108 three rows of Button chrome.
109 """
111 class Pressed(Message):
112 """Posted on activation; carries the pill for id-based dispatch."""
114 def __init__(self, pill: FleetPill) -> None:
115 super().__init__()
116 self.pill = pill
118 BINDINGS: ClassVar[list[BindingType]] = [
119 Binding("enter", "press", "Press", show=False),
120 Binding("space", "press", "Press", show=False),
121 ]
123 def press(self) -> None:
124 """Activate this pill (shared by the mouse and keyboard paths)."""
125 self.post_message(self.Pressed(self))
127 def action_press(self) -> None:
128 self.press()
130 def on_click(self, event: events.Click) -> None:
131 event.stop()
132 self.press()
135@dataclass
136class _RoleEdit:
137 """Mutable editor state for one role.
139 ``tensor_split`` carries a manual split from the loaded view so re-applying an
140 unedited placement preserves it (an even split would OOM the smaller of two
141 unequal cards). It is cleared the moment the user toggles this role's devices,
142 since a split sized for the old card set no longer applies to the new one.
143 """
145 role: WorkerRole
146 model: str
147 devices: set[int]
148 replicas: int
149 tensor_split: tuple[int, ...] | None = None
152class FleetBody(Widget):
153 """Live fleet GPU table and interactive placement editor."""
155 DEFAULT_CSS: ClassVar[str] = _CSS_FILE.read_text(encoding="utf-8")
157 applying: reactive[bool] = reactive(False)
159 class PlacementReloading(Message):
160 """Posted while an apply/clear reloads the fleet, so the chat screen can
161 hold submissions until the reload finishes instead of hitting a 429."""
163 def __init__(self, active: bool) -> None:
164 self.active = active
165 super().__init__()
167 def __init__(self) -> None:
168 super().__init__(id="fleet-body")
169 self._edits: dict[WorkerRole, _RoleEdit] = {}
170 self._device_indices: tuple[int, ...] = ()
171 self._view_manual = False
172 self._page = 0
173 self._command_actions: dict[str, Callable[[], None]] = {
174 _CMD_PREVIEW: self.action_preview,
175 _CMD_APPLY: self.action_apply,
176 _CMD_AUTO: self.action_clear,
177 }
179 def watch_applying(self, applying: bool) -> None:
180 """Disable the editor controls while an apply/clear is in flight, surface a
181 'Rebuilding fleet…' status so the reload isn't a silent idle screen, and tell
182 the chat screen to hold submissions until the fleet finishes reloading."""
183 self.post_message(self.PlacementReloading(applying))
184 with contextlib.suppress(NoMatches):
185 self.query_one(_EDITOR_ID, Vertical).disabled = applying
186 if applying:
187 self._show_rebuilding()
189 def _show_rebuilding(self) -> None:
190 """Hold a 'Rebuilding fleet…' status until the reloaded fleet reports ready.
192 The reload restarts a role and warms its model off-thread (tens of seconds
193 for a tensor-split giant); without a positive status the editor is just a
194 greyed, silent screen and the user cannot tell the apply took, or whether a
195 stale draft still stands. The re-render that ``_change_placement`` runs once
196 the fleet is ready (or the change failed) clears this back to the live state.
197 """
198 with contextlib.suppress(NoMatches):
199 state = self.query_one(_STATE_ID, Static)
200 state.update(msg.FLEET_STATE_REBUILDING)
201 state.set_class(False, "-edited")
202 state.set_class(False, "-manual")
203 state.set_class(True, "-rebuilding")
205 def compose(self) -> ComposeResult:
206 with Vertical(id="placement-layout"):
207 with Horizontal(id="placement-titlebar"):
208 yield Static(msg.FLEET_TITLE, id="placement-title")
209 yield Static("", id="placement-state")
210 help_icon = Static(msg.FLEET_HELP_ICON, id="placement-help")
211 help_icon.tooltip = msg.FLEET_HELP_TOOLTIP
212 yield help_icon
213 yield GpuFleetPanel()
214 yield Static("", id="placement-skipped")
215 yield Vertical(id="placement-editor")
216 with Horizontal(id="placement-commands"):
217 yield FleetPill(msg.FLEET_CMD_PREVIEW, id=_CMD_PREVIEW, classes="cmd-pill")
218 yield FleetPill(msg.FLEET_CMD_APPLY, id=_CMD_APPLY, classes="cmd-pill")
219 yield FleetPill(msg.FLEET_CMD_AUTO, id=_CMD_AUTO, classes="cmd-pill")
220 yield Static(_HINT, id="placement-hint")
222 def on_mount(self) -> None:
223 self._load_worker()
225 @work(thread=True, exit_on_error=False)
226 def _load_worker(self) -> None:
227 """Fetch placement off the UI thread and populate the widget.
229 get_placement resolves the plan, which can spawn a device-probe
230 subprocess on a cold cache -- too slow for the event loop.
231 """
232 try:
233 view = get_placement()
234 except Exception as exc:
235 log.debug("Failed to load placement", exc_info=True)
236 call_from_thread(self, self._render_load_failure, str(exc))
237 return
238 call_from_thread(self, self._render_view, view)
240 def _render_load_failure(self, reason: str) -> None:
241 """Name the placement-load failure in the panel instead of probing forever."""
242 self.query_one(_FLEET_PANEL_ID, GpuFleetPanel).set_probe_failed(reason)
243 self.notify(reason, severity="error")
245 # -- rendering -------------------------------------------------------
247 def _render_view(self, view: PlacementView) -> None:
248 """Reset the widget (title, live table, editor) from a resolved placement view."""
249 self._view_manual = view.manual
250 self._device_indices = tuple(g.index for g in view.gpus)
251 edits = {
252 r.role: _RoleEdit(r.role, r.model, set(r.devices), r.replicas, r.tensor_split)
253 for r in view.roles
254 }
255 # Rows render in _EDITOR_ROLE_ORDER; any role beyond it keeps plan order.
256 self._edits = {role: edits.pop(role) for role in _EDITOR_ROLE_ORDER if role in edits}
257 self._edits.update(edits)
258 self._page = 0
259 self._build_editor()
260 self._refresh_title(dirty=False)
261 self._update_fleet_panel(view)
262 self._render_skipped(view)
263 if view.co_tenants:
264 names = ", ".join(role.value for role in view.co_tenants)
265 self.notify(f"Sharing memory, one loaded at a time: {names}", severity="information")
266 if view.unplaceable:
267 names = ", ".join(role.value for role in view.unplaceable)
268 self.notify(f"Does not fit: {names}", severity="warning")
269 if view.rejected_spec_json:
270 self.notify(msg.FLEET_SAVED_PLACEMENT_IGNORED, severity="warning")
272 def _render_skipped(self, view: PlacementView) -> None:
273 """Show a 'not downloaded' line per role skipped for a missing model."""
274 widget = self.query_one(_SKIPPED_ID, Static)
275 widget.display = bool(view.skipped_not_installed)
276 widget.update(
277 "\n".join(
278 msg.FLEET_MODEL_NOT_DOWNLOADED.format(
279 role=skipped.role.value, model=_clean_model_name(skipped.model)
280 )
281 for skipped in view.skipped_not_installed
282 )
283 )
285 def _page_devices(self) -> tuple[int, ...]:
286 """The GPU indices visible on the current page."""
287 start = self._page * _PLACEMENT_PAGE_SIZE
288 return self._device_indices[start : start + _PLACEMENT_PAGE_SIZE]
290 def _page_count(self) -> int:
291 """Number of GPU pages at the current fleet size."""
292 n = len(self._device_indices)
293 return max(1, (n + _PLACEMENT_PAGE_SIZE - 1) // _PLACEMENT_PAGE_SIZE)
295 def _build_editor(self) -> None:
296 """Rebuild the placement grid: GPU header, one row per role, optional pager.
298 A single-GPU fleet has nothing to arrange -- every role can only live on
299 that card -- so the grid, command pills, and hint collapse to a one-line
300 note and the drawer stays a pure live monitor.
301 """
302 container = self.query_one(_EDITOR_ID, Vertical)
303 container.remove_children()
304 single = len(self._device_indices) <= 1
305 for selector in (_COMMANDS_ID, _HINT_WIDGET_ID):
306 self.query_one(selector).display = not single
307 if single:
308 container.mount(Label(msg.FLEET_SINGLE_GPU_NOTE, classes="single-gpu-note"))
309 return
310 devices = self._page_devices()
311 widgets: list[Horizontal] = [self._gpu_header_row(devices)]
312 for role, edit in self._edits.items():
313 kind = _role_kind(role)
314 children: list[FleetPill | Label] = [Label(f"{role.value:<7}", classes="role-name")]
315 for idx in devices:
316 on = " on" if idx in edit.devices else ""
317 children.append(
318 FleetPill(
319 f" {idx} ", id=f"dev-{role.value}-{idx}", classes=f"dev-toggle {kind}{on}"
320 )
321 )
322 if role in _REPLICA_ROLES:
323 children.append(FleetPill(" - ", id=f"rep-{role.value}-dec", classes="rep-pill"))
324 children.append(
325 Label(f"x{edit.replicas}", id=f"repn-{role.value}", classes="rep-count")
326 )
327 children.append(FleetPill(" + ", id=f"rep-{role.value}-inc", classes="rep-pill"))
328 elif role in _SINGLE_ROLES:
329 children.append(Label(msg.FLEET_TAG_SINGLE, classes="role-tag"))
330 elif len(edit.devices) > 1:
331 children.append(Label(msg.FLEET_TAG_SPLIT, classes="role-tag"))
332 widgets.append(Horizontal(*children, classes="role-row"))
333 container.mount(*widgets)
334 if self._page_count() > 1:
335 container.mount(self._pager_row())
337 def _gpu_header_row(self, devices: tuple[int, ...]) -> Horizontal:
338 """A header labelling the visible GPU columns."""
339 cells: list[Label] = [Label("GPU", classes="role-name gpu-hdr-lead")]
340 cells += [Label(str(idx), classes="gpu-hdr") for idx in devices]
341 return Horizontal(*cells, classes="gpu-header-row")
343 def _pager_row(self) -> Horizontal:
344 """Prev/next controls and a page indicator for fleets past one page."""
345 first = self._page * _PLACEMENT_PAGE_SIZE
346 last = first + len(self._page_devices()) - 1
347 info = f"GPUs {first}-{last} · page {self._page + 1}/{self._page_count()}"
348 return Horizontal(
349 FleetPill(" ◄ ", id="pg-prev", classes="pg-pill"),
350 Label(info, classes="pg-info"),
351 FleetPill(" ► ", id="pg-next", classes="pg-pill"),
352 classes="pager-row",
353 )
355 def _refresh_title(self, *, dirty: bool) -> None:
356 """Reflect the placement mode in the state segment beside the title pill."""
357 state = self.query_one(_STATE_ID, Static)
358 if dirty:
359 state.update(msg.FLEET_STATE_EDITED)
360 else:
361 state.update(msg.FLEET_STATE_MANUAL if self._view_manual else msg.FLEET_STATE_AUTO)
362 state.set_class(dirty, "-edited")
363 state.set_class(not dirty and self._view_manual, "-manual")
364 state.set_class(False, "-rebuilding")
366 def _update_fleet_panel(self, view: PlacementView) -> None:
367 """Push the current device list and roles into the fleet panel."""
368 try:
369 panel = self.query_one(_FLEET_PANEL_ID, GpuFleetPanel)
370 except NoMatches:
371 return
372 labels = {g.index: g.label for g in view.gpus}
373 roles: dict[int, str] = {}
374 for r in view.roles:
375 short_model = _clean_model_name(r.model) if r.model else ""
376 badge = f"{r.role.value} - {short_model}" if short_model else r.role.value
377 for idx in r.devices:
378 roles[idx] = badge
379 panel.set_devices(view.gpus, labels=labels, roles=roles)
381 # -- editor state ----------------------------------------------------
383 def _spec_from_editor(self) -> PlacementSpec | None:
384 """Build a PlacementSpec from the editor; None when nothing is configured."""
385 if not self._edits:
386 return None
387 roles: dict[WorkerRole, RolePlacement] = {}
388 for edit in self._edits.values():
389 if not edit.devices:
390 raise PlacementError(f"{edit.role.value} needs at least one GPU")
391 devices = tuple(sorted(edit.devices))
392 # Keep a loaded manual split only while it still matches the device set;
393 # a stale-length split would be rejected by the spec validator.
394 split = (
395 edit.tensor_split
396 if edit.tensor_split and len(edit.tensor_split) == len(devices)
397 else None
398 )
399 roles[edit.role] = RolePlacement(
400 devices=devices, replicas=edit.replicas, tensor_split=split
401 )
402 return PlacementSpec(roles=roles)
404 def on_fleet_pill_pressed(self, event: FleetPill.Pressed) -> None:
405 """Handle a GPU toggle, a replica -/+ press, or a command pill."""
406 bid = event.pill.id or ""
407 command = self._command_actions.get(bid)
408 if command is not None:
409 command()
410 return
411 if bid in ("pg-prev", "pg-next"):
412 step = -1 if bid == "pg-prev" else 1
413 self._page = min(max(0, self._page + step), self._page_count() - 1)
414 self._build_editor()
415 return
416 if bid.startswith("dev-"):
417 role_value, idx_str = bid.removeprefix("dev-").rsplit("-", 1)
418 role = WorkerRole(role_value)
419 edit = self._edits[role]
420 idx = int(idx_str)
421 # Editing the device set invalidates any loaded manual split (its
422 # length was sized for the old cards); fall back to a capacity split.
423 edit.tensor_split = None
424 if role in _SINGLE_ROLES:
425 # Single pinned instance: the picked card becomes the only one.
426 edit.devices = {idx}
427 for other in self._page_devices():
428 self.query_one(f"#dev-{role.value}-{other}", FleetPill).set_class(
429 other == idx, "on"
430 )
431 elif idx in edit.devices:
432 if len(edit.devices) > 1: # keep at least one GPU per role
433 edit.devices.discard(idx)
434 event.pill.remove_class("on")
435 else:
436 edit.devices.add(idx)
437 event.pill.add_class("on")
438 elif bid.startswith("rep-"):
439 role_value, op = bid.removeprefix("rep-").rsplit("-", 1)
440 role = WorkerRole(role_value)
441 edit = self._edits[role]
442 edit.replicas = max(1, edit.replicas + (1 if op == "inc" else -1))
443 self.query_one(f"#repn-{role.value}", Label).update(f"x{edit.replicas}")
444 else:
445 return
446 self._refresh_title(dirty=True)
448 # -- actions ---------------------------------------------------------
450 def action_preview(self) -> None:
451 """Resolve the edited placement against the hardware without applying it."""
452 try:
453 spec = self._spec_from_editor()
454 except PlacementError as exc:
455 self.notify(str(exc), severity="error")
456 return
457 self._preview_worker(spec)
459 @work(thread=True, exit_on_error=False)
460 def _preview_worker(self, spec: PlacementSpec | None) -> None:
461 try:
462 view = preview_placement(spec)
463 call_from_thread(self, self._render_view, view)
464 except Exception as exc:
465 call_from_thread(self, self.notify, str(exc), severity="error")
467 def action_apply(self) -> None:
468 """Apply the edited placement (persists and reloads the fleet)."""
469 if self.applying:
470 return
471 try:
472 spec = self._spec_from_editor()
473 except PlacementError as exc:
474 self.notify(str(exc), severity="error")
475 return
476 self.applying = True
477 self._apply_worker(spec)
479 @work(thread=True, exit_on_error=False)
480 def _apply_worker(self, spec: PlacementSpec | None) -> None:
481 self._change_placement(lambda: set_placement(spec))
483 def action_clear(self) -> None:
484 """Restore automatic placement."""
485 if self.applying:
486 return
487 self.applying = True
488 self._clear_worker()
490 @work(thread=True, exit_on_error=False)
491 def _clear_worker(self) -> None:
492 self._change_placement(lambda: set_placement(None))
494 def _change_placement(self, change: Callable[[], PlacementView]) -> None:
495 """Apply a placement change off the UI thread, then re-render the live view.
497 ``applying`` (and with it the chat screen's submit hold) stays up until the
498 restarted chat role actually serves: releasing when the reload returns
499 still leaves the model warming, and a prompt sent then errors with the
500 busy 429 instead of an answer. The view is re-rendered from the live
501 placement whether the change applied or failed, so a rejected or failed
502 draft returns to the current read-only state instead of stranding the
503 stale 'edited'/won't-fit editor.
504 """
505 error: str | None = None
506 try:
507 change()
508 wait_chat_ready()
509 except Exception as exc:
510 error = str(exc)
511 try:
512 view = get_placement()
513 call_from_thread(self, self._render_view, view)
514 except Exception as exc:
515 error = error or str(exc)
516 if error is not None:
517 call_from_thread(self, self.notify, error, severity="error")
518 call_from_thread(self, setattr, self, "applying", False)