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