Coverage for src/lilbee/app/placement.py: 100%
139 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-12 00:44 +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, 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 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
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, ...] = ()
96def _active_spec() -> PlacementSpec | None:
97 raw = cfg.placement
98 return PlacementSpec.from_json(raw) if raw else None
101def _view(
102 resolved: ResolvedPlacement,
103 *,
104 manual: bool,
105 spec_json: str | None,
106 rejected_spec_json: str | None = None,
107) -> PlacementView:
108 gpus = tuple(
109 GpuInfo(
110 index=d.index,
111 backend=d.backend,
112 label=f"{d.backend}{d.index}",
113 name=d.name,
114 total_bytes=d.total_bytes,
115 free_bytes=d.free_bytes,
116 )
117 for d in resolved.devices
118 )
119 by_role: dict[WorkerRole, RolePlacementView] = {}
120 for plan in resolved.instances:
121 existing = by_role.get(plan.role)
122 if existing is not None:
123 devices = tuple(sorted(set(existing.devices) | set(plan.devices)))
124 by_role[plan.role] = replace(existing, devices=devices, replicas=existing.replicas + 1)
125 else:
126 by_role[plan.role] = RolePlacementView(
127 role=plan.role,
128 model=resolved.model_refs.get(plan.role, ""),
129 devices=plan.devices,
130 tensor_split=plan.tensor_split or None,
131 replicas=1,
132 )
133 return PlacementView(
134 gpus=gpus,
135 roles=tuple(by_role.values()),
136 unplaceable=resolved.unplaceable_roles,
137 manual=manual,
138 spec_json=spec_json,
139 tight=tuple(
140 TightRole(role=role, shortfall_bytes=shortfall)
141 for role, shortfall in sorted(resolved.tight_roles.items(), key=lambda kv: kv[0].value)
142 ),
143 skipped_not_installed=tuple(
144 SkippedRole(role=role, model=ref)
145 for role, ref in resolved.skipped_not_installed.items()
146 ),
147 co_tenants=tuple(sorted(resolved.co_tenants, key=lambda role: role.value)),
148 rejected_spec_json=rejected_spec_json,
149 )
152def get_placement() -> PlacementView:
153 """The current effective placement (manual if a spec is set, else auto).
155 A saved spec that no longer fits the hardware is not the effective placement:
156 the fleet runs the auto plan, and this reports that rather than a manual layout
157 nothing is using.
158 """
159 spec = _active_spec()
160 resolved = resolve_placement_plan(spec, fall_back_to_auto=True)
161 if spec is None:
162 return _view(resolved, manual=False, spec_json=None)
163 if not resolved.spec_applied:
164 return _view(resolved, manual=False, spec_json=None, rejected_spec_json=spec.to_json())
165 return _view(resolved, manual=True, spec_json=spec.to_json())
168def preview_placement(spec: PlacementSpec | None = None) -> PlacementView:
169 """Dry-run: what spec (or auto, when None) would place. No persistence or reload."""
170 resolved = resolve_placement_plan(spec)
171 return _view(resolved, manual=spec is not None, spec_json=spec.to_json() if spec else None)
174def placement_refused_message() -> str:
175 """Shared refusal for placement changes on the shared HTTP server.
177 Kept in one place so the REST routes and the HTTP-mounted MCP tools
178 cannot drift apart.
179 """
180 return (
181 "Changing placement on the HTTP server is unavailable: it rebuilds the shared "
182 "fleet for every connected client. Enable allow_http_placement "
183 "(LILBEE_ALLOW_HTTP_PLACEMENT) on a single-client deployment, or change it "
184 "from the CLI or TUI."
185 )
188def set_placement(spec: PlacementSpec | None) -> PlacementView:
189 """Validate, persist to config.toml, apply to the live fleet, and return the new view.
191 Raises PlacementError before any write when the spec does not fit the hardware.
192 The live fleet applies the change surgically (``reload_placement`` restarts
193 only the roles whose placement moved), so an untouched role's loaded model
194 stays resident; with no services built there is nothing running and the next
195 use plans fresh. On the live path the planner re-plans against its clean-box
196 plan snapshot (see ``planning.capture_plan_probe``): probing under a loaded
197 fleet would report our own residency as unavailable and poison the chat
198 context sizing, while charging stays against total capacity (bb-a8f).
199 """
200 resolved = resolve_placement_plan(spec)
201 if spec is None:
202 settings.delete_values(cfg.data_root, [_PLACEMENT_KEY])
203 cfg.placement = None
204 else:
205 spec_json = spec.to_json()
206 settings.update_values(cfg.data_root, {_PLACEMENT_KEY: spec_json})
207 cfg.placement = spec_json
208 services = peek_services()
209 if services is None:
210 clear_read_device_cache() # nothing running; let the next boot probe fresh
211 else:
212 services.provider.reload_placement(wait=True)
213 return _view(resolved, manual=spec is not None, spec_json=spec.to_json() if spec else None)
216def wait_chat_ready(
217 timeout_s: float = _CHAT_READY_TIMEOUT_S,
218 *,
219 on_progress: Callable[[WarmProgress], None] | None = None,
220 should_abort: Callable[[], bool] | None = None,
221) -> bool:
222 """Block while a chat warm is in flight; True once a prompt can be served.
224 ``reload_placement(wait=True)`` returns once the proxies are healthy while the
225 restarted model still warms off-thread, so a chat request sent right after an
226 apply hits the busy 429 path. Callers that gate user input on the reload call
227 this to hold until the model actually serves. Waits only while a warm is
228 actively in flight: with no fleet, no warm, or a failed/finished warm it
229 returns at once, so a change that never restarts chat cannot stall the caller.
230 The brief grace covers the reload kicking its warm on a separate thread.
232 ``on_progress`` receives each actively-reporting warm snapshot so the caller
233 can render the load. ``should_abort`` is polled every cycle; True ends the
234 wait at once, so a cancelled prompt never pins its worker thread.
235 """
236 services = peek_services()
237 if services is None:
238 return False
239 provider = services.provider
240 started = time.monotonic()
241 deadline = started + timeout_s
242 grace_deadline = started + _CHAT_READY_GRACE_S
243 while time.monotonic() < deadline:
244 if provider.role_ready(WorkerRole.CHAT):
245 return True
246 if should_abort is not None and should_abort():
247 return False
248 snapshot = provider.warm_progress()
249 # A requested warm counts as in flight before it stamps a phase: the fleet
250 # spawns and health-checks llama-swap first, which takes seconds.
251 if is_active_warm(snapshot):
252 if on_progress is not None and snapshot is not None:
253 on_progress(snapshot)
254 grace_deadline = time.monotonic() + _CHAT_READY_GRACE_S
255 elif provider.warm_pending():
256 grace_deadline = time.monotonic() + _CHAT_READY_GRACE_S
257 elif time.monotonic() > grace_deadline:
258 return False
259 time.sleep(_CHAT_READY_POLL_S)
260 return False
263def request_engine_warm() -> None:
264 """Kick the provider's warm-up when nothing is loaded or loading.
266 ``warm_up_pool`` is idempotent (a no-op while a warm is in flight or the
267 fleet is up), so a prompt sent after a failed boot warm drives a fresh
268 engine start instead of bouncing for the rest of the session.
269 """
270 services = peek_services()
271 if services is None:
272 return
273 services.provider.warm_up_pool()
276def chat_engine_ready() -> bool:
277 """Whether a chat prompt can be served right now.
279 Positive readiness, not absence-of-warm: before the services container is
280 built nothing is loading yet and nothing can answer, which a warm snapshot
281 cannot distinguish from a finished load.
282 """
283 services = peek_services()
284 if services is None:
285 return False
286 return services.provider.role_ready(WorkerRole.CHAT)
289def active_chat_warm_progress() -> WarmProgress | None:
290 """The chat warm snapshot while a cold load is genuinely in flight, else None.
292 A surface gates interactive input on this: ``None`` covers ready, no fleet, a
293 missing model, and a finished or failed warm, so nothing traps the input in a
294 locked state. Non-``None`` carries the phase and byte progress to render.
296 The in-process warm snapshot is checked before ``role_ready`` because the task
297 bar polls this on every tick: the snapshot is a free attribute read, while
298 ``role_ready`` is an HTTP probe of the engine. Ordering it this way keeps the
299 probe to the seconds a load is actually in flight instead of firing forever on
300 an idle TUI. The two orders return the same answer.
301 """
302 services = peek_services()
303 if services is None:
304 return None
305 provider = services.provider
306 snapshot = provider.warm_progress()
307 if not is_active_warm(snapshot):
308 return None
309 return None if provider.role_ready(WorkerRole.CHAT) else snapshot
312def chat_warm_error() -> str | None:
313 """The failed chat warm's error text, or None when no failure is on record."""
314 services = peek_services()
315 if services is None:
316 return None
317 snapshot = services.provider.warm_progress()
318 if snapshot is not None and snapshot.phase is WarmPhase.ERROR:
319 return snapshot.error or ""
320 return None