Coverage for src/lilbee/app/placement.py: 100%
142 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Surface-agnostic placement use-cases: inspect, preview, and set GPU placement."""
3from __future__ import annotations
5import time
6from collections.abc import Callable
7from dataclasses import dataclass, replace
9from lilbee.app.services import peek_services
10from lilbee.core import settings
11from lilbee.core.config import cfg
12from lilbee.providers.fleet.placement_spec import PlacementSpec
13from lilbee.providers.fleet.planning import (
14 ResolvedPlacement,
15 clear_read_device_cache,
16 resolve_placement_plan,
17)
18from lilbee.providers.roles import WorkerRole
19from lilbee.providers.warm_progress import WarmPhase, WarmProgress
21_PLACEMENT_KEY = "placement"
23# Ceiling and cadence for waiting out the post-reload chat warm: a cold
24# tensor-split giant off a slow filesystem takes minutes, and the wait stops
25# early when nothing is warming or the warm failed. The grace covers the gap
26# between reload_placement returning and its off-thread warm stamping a phase.
27_CHAT_READY_TIMEOUT_S = 1800.0
28_CHAT_READY_POLL_S = 0.5
29_CHAT_READY_GRACE_S = 3.0
30_ACTIVE_WARM_PHASES = frozenset(
31 {WarmPhase.STARTING, WarmPhase.READING_WEIGHTS, WarmPhase.LOADING_ENGINE}
32)
35@dataclass(frozen=True)
36class GpuInfo:
37 """One detected GPU as a surface can render it."""
39 index: int
40 backend: str
41 label: str
42 name: str
43 total_bytes: int
44 free_bytes: int
47@dataclass(frozen=True)
48class RolePlacementView:
49 """Where one role's model is placed in the resolved plan."""
51 role: WorkerRole
52 model: str
53 devices: tuple[int, ...]
54 tensor_split: tuple[int, ...] | None
55 replicas: int
58@dataclass(frozen=True)
59class SkippedRole:
60 """A configured role left unplaced because its model isn't downloaded."""
62 role: WorkerRole
63 model: str
66@dataclass(frozen=True)
67class TightRole:
68 """A placed role whose estimate exceeds the memory on the card it landed on."""
70 role: WorkerRole
71 shortfall_bytes: int
74@dataclass(frozen=True)
75class PlacementView:
76 """The full placement picture: GPUs, per-role placement, and whether manual."""
78 gpus: tuple[GpuInfo, ...]
79 roles: tuple[RolePlacementView, ...]
80 unplaceable: tuple[WorkerRole, ...]
81 manual: bool
82 spec_json: str | None
83 # Configured roles absent from the plan because their model isn't installed,
84 # so a surface can show "not downloaded" instead of an unexplained empty table.
85 skipped_not_installed: tuple[SkippedRole, ...] = ()
86 # Roles sharing one swap group: each is placed, but only one is resident at a
87 # time, so their footprints do not sum against the card they name.
88 co_tenants: tuple[WorkerRole, ...] = ()
89 # A saved spec this hardware no longer satisfies. The auto plan is what runs,
90 # but the spec stays in config.toml and reapplies once it fits again, so a
91 # surface has to say it is there rather than report placement as plain auto.
92 rejected_spec_json: str | None = None
93 # Roles placed on a card that cannot hold them, with the shortfall in bytes.
94 # They load on demand and may fail; a view that omits this shows them as
95 # comfortably placed right up until they do.
96 tight: tuple[TightRole, ...] = ()
99def _active_spec() -> PlacementSpec | None:
100 raw = cfg.placement
101 return PlacementSpec.from_json(raw) if raw else None
104def _view(
105 resolved: ResolvedPlacement,
106 *,
107 manual: bool,
108 spec_json: str | None,
109 rejected_spec_json: str | None = None,
110) -> PlacementView:
111 gpus = tuple(
112 GpuInfo(
113 index=d.index,
114 backend=d.backend,
115 label=f"{d.backend}{d.index}",
116 name=d.name,
117 total_bytes=d.total_bytes,
118 free_bytes=d.free_bytes,
119 )
120 for d in resolved.devices
121 )
122 by_role: dict[WorkerRole, RolePlacementView] = {}
123 for plan in resolved.instances:
124 existing = by_role.get(plan.role)
125 if existing is not None:
126 devices = tuple(sorted(set(existing.devices) | set(plan.devices)))
127 by_role[plan.role] = replace(existing, devices=devices, replicas=existing.replicas + 1)
128 else:
129 by_role[plan.role] = RolePlacementView(
130 role=plan.role,
131 model=resolved.model_refs.get(plan.role, ""),
132 devices=plan.devices,
133 tensor_split=plan.tensor_split or None,
134 replicas=1,
135 )
136 return PlacementView(
137 gpus=gpus,
138 roles=tuple(by_role.values()),
139 unplaceable=resolved.unplaceable_roles,
140 manual=manual,
141 spec_json=spec_json,
142 tight=tuple(
143 TightRole(role=role, shortfall_bytes=shortfall)
144 for role, shortfall in sorted(resolved.tight_roles.items(), key=lambda kv: kv[0].value)
145 ),
146 skipped_not_installed=tuple(
147 SkippedRole(role=role, model=ref)
148 for role, ref in resolved.skipped_not_installed.items()
149 ),
150 co_tenants=tuple(sorted(resolved.co_tenants, key=lambda role: role.value)),
151 rejected_spec_json=rejected_spec_json,
152 )
155def get_placement() -> PlacementView:
156 """The current effective placement (manual if a spec is set, else auto).
158 A saved spec that no longer fits the hardware is not the effective placement:
159 the fleet runs the auto plan, and this reports that rather than a manual layout
160 nothing is using.
161 """
162 spec = _active_spec()
163 resolved = resolve_placement_plan(spec, fall_back_to_auto=True)
164 if spec is None:
165 return _view(resolved, manual=False, spec_json=None)
166 if not resolved.spec_applied:
167 return _view(resolved, manual=False, spec_json=None, rejected_spec_json=spec.to_json())
168 return _view(resolved, manual=True, spec_json=spec.to_json())
171def preview_placement(spec: PlacementSpec | None = None) -> PlacementView:
172 """Dry-run: what spec (or auto, when None) would place. No persistence or reload."""
173 resolved = resolve_placement_plan(spec)
174 return _view(resolved, manual=spec is not None, spec_json=spec.to_json() if spec else None)
177def placement_refused_message() -> str:
178 """Shared refusal for placement changes on the shared HTTP server.
180 Kept in one place so the REST routes and the HTTP-mounted MCP tools
181 cannot drift apart.
182 """
183 return (
184 "Changing placement on the HTTP server is unavailable: it rebuilds the shared "
185 "fleet for every connected client. Enable allow_http_placement "
186 "(LILBEE_ALLOW_HTTP_PLACEMENT) on a single-client deployment, or change it "
187 "from the CLI or TUI."
188 )
191def set_placement(spec: PlacementSpec | None) -> PlacementView:
192 """Validate, persist to config.toml, apply to the live fleet, and return the new view.
194 Raises PlacementError before any write when the spec does not fit the hardware.
195 The live fleet applies the change surgically (``reload_placement`` restarts
196 only the roles whose placement moved), so an untouched role's loaded model
197 stays resident; with no services built there is nothing running and the next
198 use plans fresh. On the live path the planner re-plans against its clean-box
199 plan snapshot (see ``planning.capture_plan_probe``): probing under a loaded
200 fleet would report our own residency as unavailable and poison the chat
201 context sizing, while charging stays against total capacity (bb-a8f).
202 """
203 resolved = resolve_placement_plan(spec)
204 if spec is None:
205 settings.delete_values(cfg.data_root, [_PLACEMENT_KEY])
206 cfg.placement = None
207 else:
208 spec_json = spec.to_json()
209 settings.update_values(cfg.data_root, {_PLACEMENT_KEY: spec_json})
210 cfg.placement = spec_json
211 services = peek_services()
212 if services is None:
213 clear_read_device_cache() # nothing running; let the next boot probe fresh
214 else:
215 services.provider.reload_placement(wait=True)
216 return _view(resolved, manual=spec is not None, spec_json=spec.to_json() if spec else None)
219def warm_is_reporting(snapshot: WarmProgress | None) -> bool:
220 """Whether *snapshot* is a warm that is actively loading, rather than idle or done."""
221 return snapshot is not None and snapshot.phase in _ACTIVE_WARM_PHASES
224def wait_chat_ready(
225 timeout_s: float = _CHAT_READY_TIMEOUT_S,
226 *,
227 on_progress: Callable[[WarmProgress], None] | None = None,
228 should_abort: Callable[[], bool] | None = None,
229) -> bool:
230 """Block while a chat warm is in flight; True once a prompt can be served.
232 ``reload_placement(wait=True)`` returns once the proxies are healthy while the
233 restarted model still warms off-thread, so a chat request sent right after an
234 apply hits the busy 429 path. Callers that gate user input on the reload call
235 this to hold until the model actually serves. Waits only while a warm is
236 actively in flight: with no fleet, no warm, or a failed/finished warm it
237 returns at once, so a change that never restarts chat cannot stall the caller.
238 The brief grace covers the reload kicking its warm on a separate thread.
240 ``on_progress`` receives each actively-reporting warm snapshot so the caller
241 can render the load. ``should_abort`` is polled every cycle; True ends the
242 wait at once, so a cancelled prompt never pins its worker thread.
243 """
244 services = peek_services()
245 if services is None:
246 return False
247 provider = services.provider
248 started = time.monotonic()
249 deadline = started + timeout_s
250 grace_deadline = started + _CHAT_READY_GRACE_S
251 while time.monotonic() < deadline:
252 if provider.role_ready(WorkerRole.CHAT):
253 return True
254 if should_abort is not None and should_abort():
255 return False
256 snapshot = provider.warm_progress()
257 # A requested warm counts as in flight before it stamps a phase: the fleet
258 # spawns and health-checks llama-swap first, which takes seconds.
259 if warm_is_reporting(snapshot):
260 if on_progress is not None and snapshot is not None:
261 on_progress(snapshot)
262 grace_deadline = time.monotonic() + _CHAT_READY_GRACE_S
263 elif provider.warm_pending():
264 grace_deadline = time.monotonic() + _CHAT_READY_GRACE_S
265 elif time.monotonic() > grace_deadline:
266 return False
267 time.sleep(_CHAT_READY_POLL_S)
268 return False
271def request_engine_warm() -> None:
272 """Kick the provider's warm-up when nothing is loaded or loading.
274 ``warm_up_pool`` is idempotent (a no-op while a warm is in flight or the
275 fleet is up), so a prompt sent after a failed boot warm drives a fresh
276 engine start instead of bouncing for the rest of the session.
277 """
278 services = peek_services()
279 if services is None:
280 return
281 services.provider.warm_up_pool()
284def chat_engine_ready() -> bool:
285 """Whether a chat prompt can be served right now.
287 Positive readiness, not absence-of-warm: before the services container is
288 built nothing is loading yet and nothing can answer, which a warm snapshot
289 cannot distinguish from a finished load.
290 """
291 services = peek_services()
292 if services is None:
293 return False
294 return services.provider.role_ready(WorkerRole.CHAT)
297def active_chat_warm_progress() -> WarmProgress | None:
298 """The chat warm snapshot while a cold load is genuinely in flight, else None.
300 A surface gates interactive input on this: ``None`` covers ready, no fleet, a
301 missing model, and a finished or failed warm, so nothing traps the input in a
302 locked state. Non-``None`` carries the phase and byte progress to render.
304 The in-process warm snapshot is checked before ``role_ready`` because the task
305 bar polls this on every tick: the snapshot is a free attribute read, while
306 ``role_ready`` is an HTTP probe of the engine. Ordering it this way keeps the
307 probe to the seconds a load is actually in flight instead of firing forever on
308 an idle TUI. The two orders return the same answer.
309 """
310 services = peek_services()
311 if services is None:
312 return None
313 provider = services.provider
314 snapshot = provider.warm_progress()
315 if not warm_is_reporting(snapshot):
316 return None
317 return None if provider.role_ready(WorkerRole.CHAT) else snapshot
320def chat_warm_error() -> str | None:
321 """The failed chat warm's error text, or None when no failure is on record."""
322 services = peek_services()
323 if services is None:
324 return None
325 snapshot = services.provider.warm_progress()
326 if snapshot is not None and snapshot.phase is WarmPhase.ERROR:
327 return snapshot.error or ""
328 return None