Coverage for src/lilbee/providers/fleet/provider.py: 100%
995 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""FleetProvider: the local llama-server engine for every role.
3On first use it plans GPU placement and starts one llama-swap process per swap
4group, each fronting that group's llama-server(s); each call routes to its role's
5proxy by replica model id. Per-group processes let a reload restart only the
6groups whose launches changed, so a placement or model change never unloads an
7untouched group's model. There is no in-process fallback, so a missing role
8surfaces a user-facing ``ProviderError``. Model management
9(list/show/capabilities) reads the registry and GGUF headers directly and needs no
10running server. See docs/architecture.md for swap tenancy.
11"""
13from __future__ import annotations
15import functools
16import logging
17import re
18import sys
19import threading
20import time
21from concurrent.futures import ThreadPoolExecutor
22from contextlib import contextmanager
23from pathlib import Path
24from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypeVar, overload
26import httpx
28from lilbee.catalog import clean_display_name
29from lilbee.core.config import cfg
30from lilbee.core.vectors import Vector
31from lilbee.modelhub.registry import ModelRegistry
32from lilbee.providers.base import (
33 GENERATION_RESERVE_TOKENS,
34 ProviderError,
35 ProviderErrorKind,
36 prompt_token_budget,
37)
38from lilbee.providers.fleet import planning
39from lilbee.providers.fleet.binary import engine_pin, resolve_llama_server
40from lilbee.providers.fleet.client import (
41 ChatDeadlineError,
42 LlamaServerClient,
43 is_connection_failure,
44 is_load_capacity_failure,
45 is_rebuildable_failure,
46 retry_on_busy,
47)
48from lilbee.providers.fleet.contract import (
49 chat_ctx_covers,
50 contract_matches,
51 decoded_launches,
52 served_pairs,
53)
54from lilbee.providers.fleet.groups import SwapGroup, group_for
55from lilbee.providers.fleet.ingest_warmth import ingest_keep_warm
56from lilbee.providers.fleet.launch import InstanceLaunch
57from lilbee.providers.fleet.swap_config import cold_load_timeout_s
58from lilbee.providers.fleet.swap_manager import (
59 SwapManager,
60 SwapState,
61 engine_record_exists,
62 find_live_state,
63 reap_stale,
64 state_is_healthy,
65 stop_engine,
66)
67from lilbee.providers.fleet.windowing import window_messages
68from lilbee.providers.model_ref import parse_model_ref
69from lilbee.providers.roles import MODEL_FIELD_TO_ROLE, WorkerRole, configured_model_message
70from lilbee.providers.warm_progress import WarmPhase, WarmProgress, WarmProgressTracker
71from lilbee.runtime.engine_lock import (
72 ENGINE_DIR_ENV,
73 UserLockHold,
74 build_lock,
75 hold_user_lock,
76 keep_warm_requested,
77 kernel_arbitrates_locks,
78 live_users_exist,
79 machine_engine_dir,
80 private_engine_dir,
81 request_keep_warm,
82 withdraw_keep_warm,
83)
85log = logging.getLogger(__name__)
87# How long a shutdown waits for an in-flight build to finish before tearing down
88# regardless. Generous against a legitimate llama-swap spawn (a 30 s boot budget)
89# and far short of any supervisor's patience for a process that will not exit.
90_SHUTDOWN_BUILD_LOCK_WAIT_S = 45.0
92if TYPE_CHECKING:
93 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
95 from lilbee.providers.base import (
96 ChatMessage,
97 ChatResult,
98 ChatStreamItem,
99 ChatToolResult,
100 ClosableIterator,
101 )
103# User-facing name for this engine in error messages.
104_PROVIDER_NAME = "llama-server"
105# Tokens held back from the served context for the model's own generation when the
106# request does not cap it, plus a margin for chat-template overhead and estimate drift.
107# Minimal input used to pre-load a role's upstream during warm-up (llama-swap
108# starts an upstream on its first request, so warming issues one cheap call).
109_WARM_PROMPT = "warm"
110_WARM_MAX_TOKENS = 1
111# Read size for paging chat shards into the page cache during warm; large enough
112# to keep sequential reads efficient without holding much resident at once.
113_PREWARM_CHUNK_BYTES = 8 * 1024 * 1024
114# Shards fully paged in this boot, keyed on (path, size, mtime_ns); a fleet
115# rebuild (e.g. a placement change) skips re-reading a hot cache. Module-level so
116# it survives reset_services() replacing the provider instance.
117_PREWARMED_SHARDS: set[tuple[str, int, int]] = set()
118# Per-role client request budget: the first request covers the lazy cold load plus
119# generation, so the weights-scaled cold-load budget plus the margin raises this floor.
120_REQUEST_TIMEOUT_FLOOR_S = 900.0
121_REQUEST_TIMEOUT_GENERATION_MARGIN_S = 120.0
122# Jinja chat templates flag tool support by referencing one of these names as an
123# identifier inside a ``{% ... %}`` / ``{{ ... }}`` block (not free-text prose).
124# The server parses tool calls natively via ``--jinja``; this probe only decides
125# whether to offer tools to a given model at all.
126_TOOL_TEMPLATE_PATTERN = re.compile(r"\{[%{][^}]*\b(?:tools|tool_calls|functions|function_calls)\b")
127# Attempt cap for the busy-retry only when a page has no deadline (ocr_timeout=0,
128# "no limit"): it backstops the retry so a persistently busy fleet can't spin
129# forever. A page with a deadline retries until that deadline instead (see
130# _ocr_dispatch), so the count doesn't bound the common case.
131_VISION_BUSY_RETRIES = 18
132# How often a waiter blocked on full replicas re-polls their health: an
133# unhealthy replica re-admits itself by cool-down expiry, which notifies nobody.
134_DISPATCH_HEALTH_RECHECK_S = 0.5
135_T = TypeVar("_T")
138def _prewarm_key(shard: Path) -> tuple[str, int, int]:
139 """The prewarm identity of *shard*: same path, size, and mtime -> same pages."""
140 stat = shard.stat()
141 return (str(shard), stat.st_size, stat.st_mtime_ns)
144def _request_timeout_s(weights_bytes: int) -> float:
145 """Per-client request budget: the floor, or the cold-load budget plus margin."""
146 return max(
147 _REQUEST_TIMEOUT_FLOOR_S,
148 cold_load_timeout_s(weights_bytes) + _REQUEST_TIMEOUT_GENERATION_MARGIN_S,
149 )
152def _launches_by_group(
153 plan: planning.FleetPlan,
154) -> dict[SwapGroup, tuple[InstanceLaunch, ...]]:
155 """Group a plan's launches by swap group, replica order preserved within each group.
157 Co-tenant roles land in one group, so llama-swap evicts between them rather
158 than holding both resident.
159 """
160 grouped: dict[SwapGroup, list[InstanceLaunch]] = {}
161 for launch in plan.launches:
162 grouped.setdefault(group_for(launch.role, plan.co_tenants), []).append(launch)
163 return {group: tuple(group_launches) for group, group_launches in grouped.items()}
166def _by_role(launches: list[InstanceLaunch]) -> dict[WorkerRole, list[InstanceLaunch]]:
167 """Split one group's launches per role, replica order preserved."""
168 grouped: dict[WorkerRole, list[InstanceLaunch]] = {}
169 for launch in launches:
170 grouped.setdefault(launch.role, []).append(launch)
171 return grouped
174def _least_in_flight(clients: list[LlamaServerClient]) -> LlamaServerClient:
175 """Pick the healthy client with the fewest in-flight requests.
177 Falls back to the full pool when every client is marked unhealthy, so a
178 fully-dead pool still gets a call (which surfaces the error and lets a
179 recovered replica mark itself healthy again).
180 """
181 healthy = [client for client in clients if client.healthy]
182 return min(healthy or clients, key=lambda c: c.in_flight)
185# Serializes pick-and-reserve so concurrent routers see each other's assignment.
186# Held only for the O(replicas) selection, never across the request itself.
187_ROUTE_LOCK = threading.Lock()
190def _reserve_least_in_flight(clients: list[LlamaServerClient]) -> LlamaServerClient:
191 """Atomically pick the least-loaded healthy client and reserve a slot on it.
193 Selection and reservation are one critical section: without it, concurrent
194 callers all read the same idlest replica before any of them increments its
195 counter and route there together (a thundering herd that starves the rest of
196 the fleet). The caller must :meth:`~LlamaServerClient.release` the slot.
197 """
198 with _ROUTE_LOCK:
199 client = _least_in_flight(clients)
200 client.reserve()
201 return client
204def _healthy_groups_ours(
205 states: dict[SwapGroup, SwapState], pin: str, wanted: set[tuple[WorkerRole, str]]
206) -> bool:
207 """Whether every healthy group in *states* is pin-equal and serves only wanted pairs.
209 True marks the incumbent as this contract's own engine that a full bind
210 could not cover (a dead group, or config grew a role): the ladder rebuilds
211 it in place even with live users, since those users need the rebuild too.
212 False (a foreign pin or a model outside the contract) keeps the incumbent
213 protected while in use. Vacuously False with no healthy group.
214 """
215 if not states:
216 return False
217 for state in states.values():
218 if not contract_matches(state, (), pin):
219 return False
220 pairs = served_pairs(state)
221 if pairs is None or not pairs <= wanted:
222 return False
223 return True
226def _healthy_states(engine_dir: Path) -> dict[SwapGroup, SwapState]:
227 """One probe pass over *engine_dir*: the recorded, answering group states.
229 The ladder's single view of a dir. Bind eligibility and the replaceability
230 check both read this snapshot, so they cannot disagree about an engine that
231 died between them, and one wedged proxy port is paid for once per ladder
232 pass rather than once per decision -- all of it under the build lock, which
233 every other lilbee start is waiting on.
234 """
235 found: dict[SwapGroup, SwapState] = {}
236 for group in SwapGroup:
237 state = find_live_state(engine_dir, group)
238 if state is not None and state_is_healthy(state):
239 found[group] = state
240 return found
243def _bindable_group(
244 state: SwapState, pin: str, wanted: set[tuple[WorkerRole, str]]
245) -> tuple[SwapState, list[InstanceLaunch], set[tuple[WorkerRole, str]]] | None:
246 """*state*'s launches and the wanted pairs it covers, or ``None``.
248 ``None`` for every reason an already-healthy group is not bindable by us:
249 a foreign pin, an undecodable contract, or serving nothing we want.
250 """
251 if not contract_matches(state, (), pin):
252 # Pin mismatch or undecodable contract: not bindable by us.
253 return None
254 launches = decoded_launches(state)
255 if launches is None:
256 return None
257 pairs = {(launch.role, launch.model) for launch in launches} & wanted
258 return (state, launches, pairs) if pairs else None
261class _PrimedStream:
262 """A stream re-fronted with its eagerly-pulled first frame.
264 close() always reaches the source stream, even before any iteration, so a
265 caller that truncates immediately still releases the fleet's in-flight
266 request slot (an unstarted chaining generator would silently drop it).
267 """
269 def __init__(self, first: ChatStreamItem, source: ClosableIterator[ChatStreamItem]) -> None:
270 self._first: list[ChatStreamItem] = [first]
271 self._source = source
273 def __iter__(self) -> _PrimedStream:
274 return self
276 def __next__(self) -> ChatStreamItem:
277 if self._first:
278 return self._first.pop()
279 return next(self._source)
281 def close(self) -> None:
282 self._source.close()
285def _primed_stream(items: ClosableIterator[ChatStreamItem]) -> ClosableIterator[ChatStreamItem]:
286 """Pull the first frame of *items* now, so a dead engine raises to the caller.
288 The stream connects lazily on first iteration; without priming, a proxy
289 that died raises only inside the consumer's loop, past any rediscovery.
290 """
291 try:
292 first = next(items)
293 except StopIteration:
294 return items # already exhausted; still closable
295 return _PrimedStream(first, items)
298def _call_with_failover(
299 clients: list[LlamaServerClient],
300 call: Callable[[LlamaServerClient], _T],
301) -> _T:
302 """Run *call* on the least-busy healthy client, retrying once on another replica.
304 The client is reserved at selection so concurrent ingest threads spread
305 across replicas. A connection-level failure marks the client unhealthy and
306 retries once on a different replica; with no other replica the failure
307 surfaces. The reservation is released once the call resolves.
308 """
309 client = _reserve_least_in_flight(clients)
310 try:
311 result = call(client)
312 except Exception as exc:
313 if not is_connection_failure(exc):
314 raise
315 client.mark_unhealthy()
316 return _retry_on_other_replica(clients, client, call, exc)
317 else:
318 client.mark_healthy()
319 return result
320 finally:
321 client.release()
324def _retry_on_other_replica(
325 clients: list[LlamaServerClient],
326 failed: LlamaServerClient,
327 call: Callable[[LlamaServerClient], _T],
328 cause: Exception,
329) -> _T:
330 """Retry *call* once on a replica other than *failed*, marking its health."""
331 others = [c for c in clients if c is not failed]
332 if not others:
333 raise _no_healthy_replica_error() from cause
334 retry = _reserve_least_in_flight(others)
335 try:
336 retry_result = call(retry)
337 except Exception as retry_exc:
338 if is_connection_failure(retry_exc):
339 retry.mark_unhealthy()
340 raise
341 else:
342 retry.mark_healthy()
343 return retry_result
344 finally:
345 retry.release()
348def _no_healthy_replica_error() -> ProviderError:
349 """User-facing error for a call with no healthy replica left to retry on."""
350 return ProviderError(
351 "The model server is not responding and no healthy replica is available. "
352 "It may be restarting; try again in a moment.",
353 provider=_PROVIDER_NAME,
354 kind=ProviderErrorKind.CONNECTION,
355 )
358# Env vars a launch pins its devices with, one per backend (Metal has none).
359_VISIBLE_DEVICE_ENV_VARS = (
360 "CUDA_VISIBLE_DEVICES",
361 "ROCR_VISIBLE_DEVICES",
362 "HIP_VISIBLE_DEVICES",
363 "GGML_VK_VISIBLE_DEVICES",
364 "ONEAPI_DEVICE_SELECTOR",
365)
368def _role_device_sets(
369 launches: Iterable[InstanceLaunch],
370) -> dict[WorkerRole, frozenset[str]]:
371 """Backend-qualified device tokens each role's launches pin, by role.
373 A role's set is the union across its replicas. Roles whose launches carry
374 no visibility env (Metal, or an unpinned backend) are absent: without
375 pinning there is no proof of sharing, so they keep the concurrent warm.
376 """
377 sets: dict[WorkerRole, set[str]] = {}
378 for launch in launches:
379 for var in _VISIBLE_DEVICE_ENV_VARS:
380 value = launch.env_overrides.get(var)
381 if value:
382 sets.setdefault(launch.role, set()).update(
383 f"{var}={part.strip()}" for part in value.split(",")
384 )
385 return {role: frozenset(tokens) for role, tokens in sets.items()}
388def _warm_chains(
389 warm_roles: list[WorkerRole], device_sets: dict[WorkerRole, frozenset[str]]
390) -> list[list[WorkerRole]]:
391 """Group *warm_roles* into chains warmed sequentially; chains run in parallel.
393 Roles with overlapping device sets land in one chain, merged transitively.
394 Within a chain chat goes last: it sizes its KV against the headroom the
395 settled residents leave, so it must not race their loads. A role with no
396 device set shares nothing provable and gets its own chain.
397 """
398 chains: list[tuple[set[str], list[WorkerRole]]] = []
399 ordered = sorted(warm_roles, key=lambda r: (r is WorkerRole.CHAT, list(WorkerRole).index(r)))
400 for role in ordered:
401 tokens = device_sets.get(role)
402 if not tokens:
403 chains.append((set(), [role]))
404 continue
405 merged_tokens, merged_roles = set(tokens), [role]
406 kept: list[tuple[set[str], list[WorkerRole]]] = []
407 for chain_tokens, chain_roles in chains:
408 if chain_tokens & merged_tokens:
409 merged_tokens |= chain_tokens
410 merged_roles = chain_roles + merged_roles
411 else:
412 kept.append((chain_tokens, chain_roles))
413 kept.append((merged_tokens, merged_roles))
414 chains = kept
415 return [roles for _tokens, roles in chains]
418def _warm_role(role: WorkerRole, client: LlamaServerClient) -> None:
419 """Send the cheapest request that loads *role*'s upstream behind llama-swap.
421 Vision is skipped (its load is heavy and it warms on the first OCR); chat,
422 embed, and rerank each issue a minimal call to trigger the upstream start.
423 """
424 if role is WorkerRole.CHAT:
425 client.chat(
426 [{"role": "user", "content": _WARM_PROMPT}],
427 options={"max_tokens": _WARM_MAX_TOKENS},
428 stream=False,
429 )
430 elif role is WorkerRole.EMBED:
431 client.embed([_WARM_PROMPT])
432 elif role is WorkerRole.RERANK:
433 client.rerank(_WARM_PROMPT, [_WARM_PROMPT])
436@functools.lru_cache(maxsize=32)
437def _supports_tools_cached(path_str: str, _mtime_ns: int) -> bool:
438 """Memoised tool-template probe keyed on the GGUF's path + mtime.
440 The mtime arg participates in the cache key only; a re-quantised file at the
441 same path invalidates automatically because its mtime changes.
442 """
443 from lilbee.providers.gguf_meta import read_gguf_metadata
445 meta = read_gguf_metadata(Path(path_str))
446 if not isinstance(meta, dict):
447 return False
448 template = meta.get("chat_template")
449 if not isinstance(template, str):
450 return False
451 return _TOOL_TEMPLATE_PATTERN.search(template) is not None
454class _VisionReplica(NamedTuple):
455 """One vision server paired with its fitted ``--parallel`` slot count."""
457 client: LlamaServerClient
458 slots: int
461class _PageBudgetExhausted(Exception): # noqa: N818 - internal control flow, not an error API
462 """A page's document-wide OCR budget ran out before its slot came up."""
465class _VisionDispatcher:
466 """Process-wide per-replica slot assignment for vision requests.
468 The ingest file fan-out runs many OCR requests at once; each request is
469 assigned one specific replica and only while that replica has a free
470 continuous-batching slot, so lilbee's own traffic can never oversubscribe a
471 vision server into a 429 (an aggregate cap plus racy least-busy routing
472 can). Requests past capacity wait in-process until any usable replica
473 frees a slot; unhealthy replicas take no new work until their half-open
474 cool-down re-admits them.
475 """
477 def __init__(self) -> None:
478 self._cond = threading.Condition()
479 self._assigned: dict[LlamaServerClient, int] = {}
481 @contextmanager
482 def slot(self, pool: Sequence[_VisionReplica]) -> Iterator[LlamaServerClient]:
483 """Hold one batching slot on the pool's best replica; yields that client."""
484 client = self._acquire(pool)
485 try:
486 yield client
487 finally:
488 self._release(client)
490 def _acquire(self, pool: Sequence[_VisionReplica]) -> LlamaServerClient:
491 with self._cond:
492 while True:
493 client = self._pick(pool)
494 if client is not None:
495 self._assigned[client] = self._assigned.get(client, 0) + 1
496 return client
497 # The timed wait re-polls health: a replica can become routable
498 # again by cool-down expiry alone, which notifies no waiter.
499 self._cond.wait(timeout=_DISPATCH_HEALTH_RECHECK_S)
501 def _pick(self, pool: Sequence[_VisionReplica]) -> LlamaServerClient | None:
502 """The usable replica with the most free slots, or None while all are full.
504 Falls back to the full pool when every replica is unhealthy (mirrors
505 ``_least_in_flight``), so a dead pool surfaces the error instead of
506 queueing forever.
507 """
508 usable = [replica for replica in pool if replica.client.healthy] or list(pool)
509 best = max(usable, key=self._free_slots)
510 return best.client if self._free_slots(best) > 0 else None
512 def _free_slots(self, replica: _VisionReplica) -> int:
513 return replica.slots - self._assigned.get(replica.client, 0)
515 def _release(self, client: LlamaServerClient) -> None:
516 with self._cond:
517 remaining = self._assigned.get(client, 0) - 1
518 if remaining <= 0:
519 self._assigned.pop(client, None)
520 else:
521 self._assigned[client] = remaining
522 self._cond.notify_all()
525_VISION_DISPATCHER = _VisionDispatcher()
528def _dispatch_vision(pool: Sequence[_VisionReplica], call: Callable[[LlamaServerClient], _T]) -> _T:
529 """Run *call* on a replica with a free batching slot, failing over once.
531 Blocks until a slot frees rather than racing requests at a full server. A
532 connection-level failure marks the replica unhealthy and retries once on
533 another replica's slot; with no other replica the failure surfaces.
534 """
535 with _VISION_DISPATCHER.slot(pool) as client:
536 try:
537 result = call(client)
538 except Exception as exc:
539 if not is_connection_failure(exc):
540 raise
541 client.mark_unhealthy()
542 failed, cause = client, exc
543 else:
544 client.mark_healthy()
545 return result
546 others = [replica for replica in pool if replica.client is not failed]
547 if not others:
548 raise _no_healthy_replica_error() from cause
549 with _VISION_DISPATCHER.slot(others) as retry_client:
550 try:
551 retry_result = call(retry_client)
552 except Exception as retry_exc:
553 if is_connection_failure(retry_exc):
554 retry_client.mark_unhealthy()
555 raise
556 retry_client.mark_healthy()
557 return retry_result
560def _vision_call(
561 client: LlamaServerClient, messages: Sequence[Mapping[str, Any]], timeout: float | None
562) -> str:
563 """Run a vision chat on *client*, enforcing *timeout* like the in-process OCR.
565 Caps generation at ``cfg.vision_ocr_max_tokens`` so a runaway repetition loop
566 on one page (seen looping to tens of thousands of chars) can't dominate a
567 scan's OCR time; a real page stays well under the cap. A timeout surfaces as
568 a ``ProviderError`` so the page-level OCR caller can fail just that page.
569 Callers hold a dispatcher slot, so queue time isn't billed against the timeout.
570 """
572 options = {"max_tokens": cfg.vision_ocr_max_tokens}
573 if timeout and timeout > 0:
574 return _bounded_vision_chat(client, messages, options, timeout)
575 return client.chat(messages, options=options, stream=False)
578def _bounded_vision_chat(
579 client: LlamaServerClient,
580 messages: Sequence[Mapping[str, Any]],
581 options: dict[str, Any],
582 timeout: float,
583) -> str:
584 """One vision chat streamed under a total *timeout*, released promptly on expiry.
586 ``chat_bounded`` streams the response in this thread and closes it (freeing the
587 in-flight slot) once the deadline passes, so a trickling upstream can't outlive
588 the caller. Its deadline signal is re-worded as the vision OCR timeout.
589 """
590 try:
591 return client.chat_bounded(messages, options=options, deadline_s=timeout)
592 except ChatDeadlineError:
593 raise ProviderError(
594 f"Vision OCR timed out after {timeout:.0f}s.",
595 provider=_PROVIDER_NAME,
596 ) from None
599def _ocr_dispatch(
600 pool: Sequence[_VisionReplica],
601 messages: Sequence[Mapping[str, Any]],
602 deadline: float | None,
603) -> str:
604 """OCR *messages* on a free replica slot, retrying transient failures until *deadline*.
606 Backpressure (the dispatcher blocking until a slot frees) makes a
607 self-inflicted 429 unreachable; a residual busy response is a still-warming
608 server or foreign traffic, and a gateway error is a replica restarting
609 mid-run. The retry is deadline-bound rather than
610 attempt-bound so a page on a deep queue waits for a genuinely free slot until
611 its own budget passes instead of dropping after a fixed count. Each attempt
612 is bounded by the budget remaining before *deadline*; an exhausted budget
613 raises :class:`_PageBudgetExhausted`. A ``None`` deadline (no limit) falls
614 back to a bounded attempt count so the retry can't spin forever.
615 """
617 def _attempt(client: LlamaServerClient) -> str:
618 remaining = max(0.0, deadline - time.monotonic()) if deadline is not None else None
619 if remaining == 0.0:
620 raise _PageBudgetExhausted
621 return _vision_call(client, messages, remaining)
623 return retry_on_busy(
624 lambda: _dispatch_vision(pool, _attempt),
625 retries=_VISION_BUSY_RETRIES,
626 deadline=deadline,
627 )
630def _pdf_drain_budget(total_pages: int, per_page_timeout_s: float | None) -> float | None:
631 """Total OCR wall-clock budget = pages*per_page + load grace, or None for no cap.
633 Mirrors the in-process drain budget: one document-wide deadline rather than a
634 per-page cap, so a slow page borrows from fast ones and the vision model's cold
635 first-inference is absorbed by the grace instead of tripping a fixed page limit.
636 """
637 from lilbee.core.config import cfg
639 if not per_page_timeout_s or per_page_timeout_s <= 0:
640 return None
641 return total_pages * per_page_timeout_s + cfg.vision_load_budget_s
644def _ocr_deadline(per_page_timeout_s: float | None) -> float | None:
645 """Absolute monotonic deadline for one image OCR, or None when uncapped.
647 An image is a one-page document, so it gets the same budget as a PDF page:
648 the per-page timeout plus the cold-load grace, spanning queue wait and
649 generation together.
650 """
651 budget = _pdf_drain_budget(1, per_page_timeout_s)
652 return None if budget is None else time.monotonic() + budget
655_ROLE_TO_MODEL_FIELD = {role: field for field, role in MODEL_FIELD_TO_ROLE.items()}
658def _configured_model_for(role: WorkerRole) -> str:
659 """The cfg model ref for *role*, empty when the role is unset."""
660 field = _ROLE_TO_MODEL_FIELD.get(role)
661 return getattr(cfg, field) or "" if field else ""
664def _unusable_engine_reason() -> str | None:
665 """Why no server can start on this host, or None once an engine resolves.
667 Planning drops an engine-less host to serving nothing and says so only at
668 debug, so by the time a surface has an empty pool the engine is the one cause
669 it cannot see. Re-resolving here is also what keeps an engine installed
670 mid-session from being reported as still missing.
671 """
672 try:
673 resolve_llama_server()
674 except ProviderError as exc:
675 return str(exc)
676 return None
679def _chat_needs_local_engine() -> bool:
680 """Whether the configured chat model is one this host has to serve itself.
682 A chat ref routed to an SDK backend runs without any local engine, so a
683 missing one is not its failure and must not be stamped on its warm.
684 """
685 ref = _configured_model_for(WorkerRole.CHAT)
686 return bool(ref) and not parse_model_ref(ref).is_remote
689def _no_server_message(role: WorkerRole) -> str:
690 """User-facing reason *role* has no server, engine state first.
692 A missing engine and a model that never placed both arrive as an empty pool,
693 and reading the second onto the first sends the reader to a model
694 configuration that is already correct.
695 """
696 reason = _unusable_engine_reason()
697 if reason is not None:
698 return f"No {role.value} model server is running: {reason}"
699 return (
700 f"No {role.value} model server is running. Make sure the {role.value} "
701 "model is installed and configured, then try again."
702 )
705class _EngineDemand(NamedTuple):
706 """What this process needs an engine to serve: pairs plus its chat window."""
708 pairs: set[tuple[WorkerRole, str]]
709 # Per-slot chat tokens this process needs; 0 demands nothing.
710 chat_ctx: int
711 # Configured roles the plan skipped because their model is not installed.
712 # Carried out of the demand plan so the warm tracker can name the missing
713 # model even when the ladder never reaches _plan_and_spawn (zero installed
714 # models fail _can_build_engine first) or binds an existing engine.
715 skipped_not_installed: dict[WorkerRole, str]
716 # Launches the demand plan refused for an unusable window (role -> reason);
717 # recorded even when the ladder binds an engine or never builds one.
718 skipped_unusable_ctx: dict[WorkerRole, str]
721def _placeable_demand() -> _EngineDemand:
722 """Configured (role, model) pairs a fresh plan would serve, and the chat window.
724 A configured role is wanted only when the planner would place it. The plan
725 is the co-placement authority: a role that fits alone but cannot co-tenant (a
726 unified-memory box past its budget) gets no launch, so it is dropped here too,
727 and bind matches a running engine instead of judging it a partial cover and
728 restarting the shared engine on every process start. The per-role check stays
729 as the cheap gate for the reasons in its own docstring. Empty when no engine
730 binary resolves: nothing is placeable, so the ladder serves nothing.
731 """
732 from lilbee.providers.fleet.planning import (
733 placeable_total_vram,
734 plan_all_launches,
735 role_model_placeable,
736 )
738 try:
739 plan = plan_all_launches()
740 except ProviderError as exc:
741 if exc.kind is ProviderErrorKind.NOT_FOUND:
742 return _EngineDemand(set(), 0, {}, {})
743 raise
744 placed = {launch.role for launch in plan.launches}
745 total_vram = placeable_total_vram()
746 pairs = {
747 (role, model)
748 for role in WorkerRole
749 if role in placed
750 and (model := _configured_model_for(role))
751 and role_model_placeable(role, model, total_vram)
752 }
753 return _EngineDemand(
754 pairs,
755 _demanded_chat_ctx(plan.launches, pairs),
756 dict(plan.skipped_not_installed),
757 dict(plan.skipped_unusable_ctx),
758 )
761def _demanded_chat_ctx(
762 launches: Iterable[InstanceLaunch], pairs: set[tuple[WorkerRole, str]]
763) -> int:
764 """Per-slot chat window this process needs an engine to serve; 0 for none.
766 The cfg target (a ``num_ctx`` pin, else ``chat_n_ctx_target``) capped by
767 this process's own planned chat window: a window the plan itself cannot
768 reach (model ceiling, hardware) is not a demand a rebuild could satisfy,
769 so capping keeps the fit check from rebuilding the engine in a loop.
771 The cap applies only to a single-device chat plan, whose window is sized
772 against device totals and holds regardless of what is resident. A
773 tensor-split plan is sized against live free VRAM, which a resident
774 incumbent deflates; capping by it would shrink the demand to whatever the
775 incumbent left free and let the fit check pass vacuously.
776 """
777 if not any(role is WorkerRole.CHAT for role, _model in pairs):
778 return 0
779 chat_launches = [launch for launch in launches if launch.role is WorkerRole.CHAT]
780 planned = max((launch.ctx for launch in chat_launches), default=0)
781 if planned <= 0:
782 return 0
783 # Always positive: num_ctx validates ge=1 and chat_n_ctx_target ge=512.
784 target = cfg.num_ctx if cfg.num_ctx is not None else cfg.chat_n_ctx_target
785 split = any(len(launch.est_vram_by_device) > 1 for launch in chat_launches)
786 return target if split else min(target, planned)
789def _can_build_engine(wanted: set[tuple[WorkerRole, str]]) -> bool:
790 """Preconditions for a viable build, checked before stopping a warm engine.
792 A process that can serve nothing (no placeable model, an unresolvable engine
793 binary) must not stop an engine another setup left warm and then spawn nothing.
794 Probing the engine here resolves the binary AND enumerates devices, so a wedged
795 GPU probe or an unusable CUDA runtime raises loud at this point -- before the
796 caller stops a replaceable incumbent. Were the probe left to run only inside
797 ``_plan_and_spawn`` (after the stop), that raise would kill an engine other
798 members still hold and then skip the overflow build, leaving zero engines. This
799 takes no memory snapshot (device enumeration reads no residency); the clean-box
800 sizing snapshot is captured by ``_plan_and_spawn`` after the stop.
801 """
802 from lilbee.providers.fleet import planning
804 if not wanted:
805 return False
806 try:
807 planning.assert_engine_probeable()
808 except ProviderError as exc:
809 # A genuinely-missing engine binary keeps the quiet serve-nothing path;
810 # every other probe failure must propagate (fail loud) rather than be read
811 # as "cannot build" and silently stand down.
812 if exc.kind is not ProviderErrorKind.NOT_FOUND:
813 raise
814 return False
815 except OSError:
816 return False
817 return True
820def _warm_ttl_seconds(*, hold_warm_for_session: bool = False) -> int:
821 """llama-swap idle-unload timer in seconds for the spawned fleet.
823 A ttl of 0 keeps weights resident until the engine is stopped; otherwise an
824 idle engine releases its weights after ``engine_idle_ttl_minutes`` and reloads
825 transparently on the next prompt. The timer is held off (ttl 0) whenever
826 someone is actively depending on an instant response: *hold_warm_for_session*
827 is set for a provider serving an interactive session, which owns the process
828 for its whole lifetime (close lilbee to release the engine); a bulk ingest
829 holds the fleet resident for its run so an unevenly loaded replica cannot
830 idle-unload and reload cold mid-run; and ``keep_engine_warm`` pins the weights
831 for a process meant to stay ready.
832 """
833 if hold_warm_for_session or ingest_keep_warm() or cfg.keep_engine_warm:
834 return 0
835 return cfg.engine_idle_ttl_minutes * 60
838class FleetProvider:
839 """Routes every role to the managed llama-server fleet (a fleet-of-one on one box)."""
841 def __init__(self, *, hold_warm: bool = False) -> None:
842 # An interactive session (the TUI) owns this process for its whole
843 # lifetime, so its fleet stays resident instead of idle-unloading under a
844 # user who is still in the app; closing lilbee releases it. Set by the
845 # container that built this provider, never mutated afterwards.
846 self._hold_warm_for_session = hold_warm
847 # One llama-swap per placed group, so restarting one group's servers (a
848 # placement or per-role model change) never unloads another group's. A
849 # co-tenant group holds chat and vision, which evict each other on load.
850 self._swaps: dict[SwapGroup, SwapManager] = {}
851 # The group each placed role runs in, so a role's clients and its swap
852 # process can be reached from the role alone.
853 self._role_group: dict[WorkerRole, SwapGroup] = {}
854 # The launches each running group was started with, kept so a reload can
855 # diff the fresh plan against what is running and restart only the groups
856 # whose launches actually changed. Launch argv is port-free (ports are
857 # injected at config render), so the comparison is stable across starts.
858 self._launches: dict[SwapGroup, tuple[InstanceLaunch, ...]] = {}
859 # Engine dirs this provider holds membership in (machine slot and/or
860 # the private overflow), and the dir each running group lives in.
861 self._engine_holds: dict[Path, UserLockHold] = {}
862 self._group_dirs: dict[SwapGroup, Path] = {}
863 # Latched once shutdown runs. A discarded provider (reset_services swaps
864 # in a new one) can still have an in-flight warm-up or reload daemon
865 # thread; without this latch that thread could start a llama-swap after
866 # shutdown already ran, leaving a process no live provider owns.
867 # _ensure_fleet checks it under the build lock so a post-shutdown build is
868 # refused (the swap_manager reaper is the backstop if one slips through).
869 self._shut_down = False
870 # A pool of OpenAI clients per placed role (one per data-parallel replica),
871 # all pointed at the llama-swap endpoint and routed by replica model id;
872 # rebuilt whenever the swap process (re)starts. Requests round-robin the pool.
873 self._clients: dict[WorkerRole, list[LlamaServerClient]] = {}
874 # Clients retired by a reload, awaiting close. A reload's old clients may
875 # still be held by an in-flight reader, so they are closed at the *next*
876 # reload (by when those readers have finished) or at shutdown, never while
877 # potentially in use. See _retire_clients.
878 self._retiring_clients: list[LlamaServerClient] = []
879 # Chat batching slots and per-slot context from the chat launch, surfaced to
880 # the concurrency gate and clients; defaults until the chat group is up.
881 self._chat_slots = 1
882 self._chat_ctx: int | None = None
883 # Single-flight guard: the HTTP/MCP servers route concurrently, so two
884 # first-requests must not each start a swap (double GPU allocation) or
885 # tear one down mid-route. Reentrant: invalidate_load_cache nests calls.
886 self._lock = threading.RLock()
887 # Serializes the slow startup (GPU probe + GGUF parse + llama-swap spawn)
888 # across concurrent callers, so the off-thread warm-up and an on-demand call
889 # can't start two swaps. Held only during startup, NOT while routing.
890 self._build_lock = threading.Lock()
891 # Spawn-lifecycle listeners (set by the TUI via add_spawn_listener). Stored
892 # so warm-up can report per-role progress as it pre-loads each upstream.
893 self._on_spawning: Callable[[WorkerRole], None] | None = None
894 self._on_spawned: Callable[[WorkerRole], None] | None = None
895 # Granular cold-load progress for the chat role, streamed to a launcher so
896 # the user sees real read/engine-load progress instead of a frozen spinner.
897 self._warm_tracker = WarmProgressTracker()
898 # The engine's own reason a role's model failed to warm, so the launcher and
899 # the TUI report the real cause instead of a generic "did not load".
900 self._warm_errors: dict[WorkerRole, str] = {}
901 # Configured roles the last plan left unplaced because their model isn't
902 # installed (role -> ref). The warm finalizer reads it to fail a not-installed
903 # chat with a named reason instead of clearing to a silent "not ready" retry.
904 self._skipped_not_installed: dict[WorkerRole, str] = {}
905 # Launches the last plan refused for an unusable window (role -> reason);
906 # read by the warm finalizer and _require_clients.
907 self._skipped_unusable_ctx: dict[WorkerRole, str] = {}
908 # Single-flight guard for the off-thread warm-up: True from the moment a
909 # warm thread is dispatched until it finishes, so a second warm_up_pool
910 # never starts a second swap and double-allocates GPU memory.
911 self._warming = False
912 # Single-flight guard for the off-thread reload: a second reload_role
913 # while one is in flight sets the pending flag instead of dispatching,
914 # and the in-flight thread re-runs the plan loop once per pending flag.
915 self._reloading = False
916 # Set when a reload arrives mid-reload: the in-flight pass may have
917 # already snapshotted its plan, so the change must be re-applied.
918 self._reload_pending = False
919 # Notified when ``_reloading`` clears, so a ``reload_role(wait=True)`` caller
920 # can block until the reload it requested (or the in-flight one that will
921 # run its pending pass) has finished.
922 self._reload_done = threading.Condition(self._lock)
924 def _ensure_fleet(self) -> bool:
925 """Start one llama-swap per placed role exactly once across concurrent callers.
927 Returns whether any role group is running afterwards; ``False`` when no
928 role is configured and installed (nothing to serve), leaving no process
929 spawned. The startup runs under ``_build_lock`` (not the routing lock),
930 so the off-thread warm-up and an on-demand call can't start two fleets --
931 which would double-allocate GPU and parse the same GGUF twice. A second
932 caller blocks on the build lock and reuses the groups the first one
933 started. A group failing to start tears down the groups already started
934 in this build, so a partial fleet never leaks past the failure.
935 """
936 with self._lock:
937 if self._swaps:
938 return True
939 with self._build_lock:
940 with self._lock:
941 if self._swaps:
942 return True
943 if self._shut_down:
944 # Provider was shut down (and likely discarded by reset_services)
945 # while this warm-up/reload thread was in flight; do not spawn a
946 # llama-swap no live provider would ever reap.
947 return False
949 return self._acquire_engine(cfg.data_root)
951 def _acquire_engine(self, config_root: Path) -> bool:
952 """The acquisition ladder: bind to a compatible engine, else build one.
954 Machine slot first. An incumbent is replaced in place when no live
955 user holds it, or when it is this contract's own engine (pin-equal,
956 serving only wanted models) left partially dead or partially covering:
957 its members are waiting for exactly that rebuild, and overflowing
958 around it would load duplicate weights. Only a live incompatible
959 engine in active use sends the build to the config root's private
960 overflow dir. Per-dir the step is all-or-nothing: every configured
961 (role, model) pair bound, or built fresh. Runs under the cross-process
962 build lock, so two starts never both build and stop-if-last never
963 races an arrival.
964 """
965 pin = engine_pin()
966 demand = _placeable_demand()
967 # Record the demand plan's skips before walking the ladder: a bind or an
968 # early serve-nothing exit never reaches _plan_and_spawn, and the warm
969 # tracker must still be able to say "chat model X is not installed"
970 # rather than a retryable not-ready.
971 self._skipped_not_installed = dict(demand.skipped_not_installed)
972 self._skipped_unusable_ctx = dict(demand.skipped_unusable_ctx)
973 machine_dir = machine_engine_dir()
974 if kernel_arbitrates_locks(machine_dir):
975 machine = self._acquire_in_dir(machine_dir, pin, demand, is_overflow=False)
976 if machine is not None:
977 return machine
978 else:
979 # Without kernel-arbitrated locks the membership refcount cannot be
980 # trusted, and sharing is exactly what needs it: a probe would
981 # destroy a live member's lock, so the slot would look free while
982 # another setup is serving from it. Keep to our own dir instead.
983 log.warning(
984 "Engine dir %s is on a filesystem without working file locks; "
985 "using a private engine instead of the shared one. Set %s to a "
986 "path on a local filesystem to share one engine across lilbees.",
987 machine_dir,
988 ENGINE_DIR_ENV,
989 )
990 # The machine slot holds a live incompatible engine in active use: overflow
991 # to this config root's private dir rather than evict another model setup.
992 private = private_engine_dir(config_root)
993 return self._acquire_in_dir(private, pin, demand, is_overflow=True) or False
995 def _acquire_in_dir(
996 self, engine_dir: Path, pin: str, demand: _EngineDemand, *, is_overflow: bool
997 ) -> bool | None:
998 """Bind or build one engine dir; ``None`` on the slot means overflow next.
1000 Binds a compatible running engine. Whether an incumbent may be replaced
1001 is decided by kernel-refcounted membership, not the proxy HTTP probe: an
1002 engine with a live user is never reaped or stopped, so a transient probe
1003 failure (fd exhaustion, host thrash) cannot kill a busy engine. Replace in
1004 place only when no live user holds it or it is this contract's own engine
1005 (pin-equal, serving only wanted models). A live incompatible engine in
1006 active use is never evicted or stacked on: on the machine slot it returns
1007 ``None`` (overflow), and in the overflow dir it serves nothing rather than
1008 duplicate weights beside it. Before building, any recorded engine is cleared
1009 -- keyed on
1010 the state file, not the probe -- so an unprobeable incumbent is stopped
1011 rather than double-built beside. The stop is gated on ``_can_build_engine``
1012 so a process that can serve nothing never destroys a warm engine it can't
1013 replace. Held under the cross-process build lock.
1014 """
1015 wanted = demand.pairs
1016 with build_lock(engine_dir):
1017 states = _healthy_states(engine_dir)
1018 if wanted and self._bind_all_in_dir(engine_dir, states, pin, demand):
1019 self._hold_membership(engine_dir)
1020 return True
1021 replaceable = not live_users_exist(engine_dir) or _healthy_groups_ours(
1022 states, pin, wanted
1023 )
1024 if not replaceable:
1025 # A live engine another setup is actively using is never evicted or
1026 # stacked on. On the machine slot that means overflow (None); in the
1027 # overflow dir there is nowhere further to go, so serve nothing rather
1028 # than kill the incumbent or load a second fleet's weights beside it
1029 # (an OOM on a small-VRAM box).
1030 return None if not is_overflow else False
1031 if not _can_build_engine(wanted):
1032 return False
1033 # No live user holds this dir now (or it is ours to rebuild): reap dead
1034 # leftovers and stop any recorded engine so planning sees true free VRAM
1035 # and the build never lands beside an unprobeable incumbent.
1036 reap_stale(engine_dir)
1037 if engine_record_exists(engine_dir):
1038 stop_engine(engine_dir)
1039 if self._plan_and_spawn(engine_dir):
1040 self._hold_membership(engine_dir)
1041 return True
1042 return False
1044 def _bind_all_in_dir(
1045 self,
1046 engine_dir: Path,
1047 states: dict[SwapGroup, SwapState],
1048 pin: str,
1049 demand: _EngineDemand,
1050 ) -> bool:
1051 """Bind every group needed to cover the demanded pairs; False leaves nothing bound.
1053 Binding never touches groups serving models outside the demand; the dir
1054 matches only when healthy, pin-equal groups cover every wanted pair and
1055 the served chat window covers the demanded per-slot ctx. (Whether an
1056 unmatched dir's engine is then replaced or overflowed around is the
1057 ladder's call, based on live users.)
1058 """
1059 wanted = demand.pairs
1060 candidates: list[tuple[SwapGroup, SwapState, list[InstanceLaunch]]] = []
1061 covered: set[tuple[WorkerRole, str]] = set()
1062 for group, state in states.items():
1063 found = _bindable_group(state, pin, wanted)
1064 if found is None:
1065 continue
1066 bindable, launches, pairs = found
1067 if not chat_ctx_covers(launches, demand.chat_ctx):
1068 # The live chat window is smaller than this process needs.
1069 return False
1070 candidates.append((group, bindable, launches))
1071 covered |= pairs
1072 if covered != wanted:
1073 return False
1074 bound: dict[SwapGroup, tuple[SwapManager, list[InstanceLaunch]]] = {}
1075 for group, state, launches in candidates:
1076 swap = SwapManager(engine_dir, group)
1077 if not swap.bind(state):
1078 for prior, _launches in bound.values():
1079 prior.shutdown()
1080 return False
1081 bound[group] = (swap, launches)
1082 with self._lock:
1083 for group, (swap, launches) in bound.items():
1084 self._adopt_group(group, swap, launches)
1085 self._group_dirs[group] = engine_dir
1086 log.info("Bound to the running engine at %s", engine_dir)
1087 return True
1089 def _reload_dir(self) -> Path:
1090 """The engine dir a reload rebuilds into: where our groups already live.
1092 All this provider's groups share one dir by construction (the ladder is
1093 all-or-nothing per dir); an empty provider rebuilds into the machine slot.
1094 """
1095 with self._lock:
1096 dirs = set(self._group_dirs.values())
1097 return next(iter(dirs)) if dirs else machine_engine_dir()
1099 def _hold_membership(self, engine_dir: Path) -> None:
1100 """Record this process as a user of *engine_dir*'s engine.
1102 The single point every acquisition passes through, bind and build alike,
1103 so it is also where this user's persistence opt-in is recorded against
1104 the engine. Marking on bind (not only on build) is what makes the
1105 setting mean what it says on a shared slot: a user who asked for a warm
1106 engine keeps it warm even when a default-config sibling is last out.
1107 """
1108 from lilbee.core.config import cfg
1110 if engine_dir not in self._engine_holds:
1111 self._engine_holds[engine_dir] = hold_user_lock(engine_dir)
1112 if cfg.keep_engine_warm:
1113 request_keep_warm(engine_dir, cfg.data_root)
1115 def _plan_and_spawn(self, data_dir: Path) -> bool:
1116 """Plan placement against the clean box and start one swap per group.
1118 Caller holds the build lock. False when the engine binary is missing or
1119 nothing is installed/configured, so the provider serves nothing.
1120 """
1121 try:
1122 # Snapshot the clean box; this plan and every later reload size
1123 # ctx, slots, and budgets against it (a live probe under a loaded
1124 # fleet would report our own residency as unavailable). Inside the
1125 # try: capturing resolves the engine binary, and a binary-less
1126 # host must serve nothing, not raise. Every other planning failure
1127 # (a wedged GPU probe, an unusable CUDA runtime) propagates so the
1128 # warm tracker and the caller report the real reason instead of a
1129 # silent never-ready fleet.
1130 planning.capture_plan_probe()
1131 plan = planning.plan_all_launches()
1132 except ProviderError as exc:
1133 # Only a genuinely-missing engine binary keeps the quiet no-fleet path;
1134 # any other planning failure (a wedged GPU probe, an unusable CUDA
1135 # runtime) must surface to the warm tracker and on-demand callers
1136 # rather than silently serving nothing (#540).
1137 if exc.kind is not ProviderErrorKind.NOT_FOUND:
1138 raise
1139 log.debug("Engine binary unavailable; no swap started")
1140 plan = None
1141 # plan None (no engine binary) keeps the demand-time record from
1142 # _acquire_engine instead of wiping it.
1143 if plan is not None:
1144 self._skipped_not_installed = dict(plan.skipped_not_installed)
1145 self._skipped_unusable_ctx = dict(plan.skipped_unusable_ctx)
1146 if plan is None or not plan.launches:
1147 # No engine binary, or no installed/configured model: serve nothing.
1148 return False
1149 by_group = _launches_by_group(plan)
1150 started: dict[SwapGroup, SwapManager] = {}
1151 try:
1152 for group, group_launches in by_group.items():
1153 swap = SwapManager(data_dir, group)
1154 swap.start(
1155 list(group_launches),
1156 ttl_seconds=_warm_ttl_seconds(
1157 hold_warm_for_session=self._hold_warm_for_session
1158 ),
1159 bind_lifetime=not cfg.keep_engine_warm,
1160 )
1161 started[group] = swap
1162 except BaseException:
1163 for swap in started.values():
1164 swap.shutdown()
1165 raise
1166 with self._lock:
1167 for group, swap in started.items():
1168 self._adopt_group(group, swap, list(by_group[group]))
1169 self._group_dirs[group] = data_dir
1170 return True
1172 def _adopt_group(
1173 self, group: SwapGroup, swap: SwapManager, launches: list[InstanceLaunch]
1174 ) -> None:
1175 """Record *group*'s freshly started swap and build a client pool per role.
1177 Caller holds ``self._lock``. Each launch (one per replica) becomes a client
1178 keyed by its replica model id against this group's own proxy endpoint;
1179 the chat launch carries the slots/ctx so the capacity and served context
1180 come from the launch, not a probe.
1181 """
1182 self._swaps[group] = swap
1183 self._launches[group] = tuple(launches)
1184 endpoint = swap.endpoint()
1185 for role, role_launches in _by_role(launches).items():
1186 # Retire the role's previous clients (a reload re-adopts over an existing
1187 # pool): closing them now would error a reader still mid-call on an old
1188 # client snapshot, and never closing leaks an httpx pool per replica.
1189 old_clients = list(self._clients.get(role, []))
1190 self._role_group[role] = group
1191 # token_cap truncates oversize embed/rerank inputs to the per-slot context
1192 # (the in-process backstop); the longer timeout covers a cold upstream load.
1193 self._clients[role] = [
1194 LlamaServerClient(
1195 endpoint,
1196 launch.model_id,
1197 token_cap=launch.token_cap,
1198 timeout=_request_timeout_s(launch.weights_bytes),
1199 rerank_mode=launch.rerank_mode,
1200 inline_reasoning=role is WorkerRole.CHAT,
1201 # A cold embed replica 429s bulk ingest until its slots load; wait
1202 # out the same cold-load budget llama-swap keeps it alive for so a
1203 # burst never drops files while the server is legitimately warming.
1204 embed_busy_deadline_s=(
1205 cold_load_timeout_s(launch.weights_bytes)
1206 if role is WorkerRole.EMBED
1207 else None
1208 ),
1209 )
1210 for launch in role_launches
1211 ]
1212 if role is WorkerRole.CHAT:
1213 self._chat_slots = role_launches[0].slots
1214 self._chat_ctx = role_launches[0].ctx
1215 self._retire_clients(old_clients)
1217 def _swap_for(self, role: WorkerRole) -> SwapManager | None:
1218 """The swap process serving *role*, or None when the role has no server.
1220 Caller holds ``self._lock``.
1221 """
1222 group = self._role_group.get(role)
1223 return None if group is None else self._swaps.get(group)
1225 def _role_launches(self, role: WorkerRole) -> tuple[InstanceLaunch, ...]:
1226 """*role*'s launch snapshot, empty when it has no server.
1228 A co-tenant group holds more than one role's launches, so the group's
1229 snapshot is filtered down to this role's replicas.
1230 """
1231 group = self._role_group.get(role)
1232 if group is None:
1233 return ()
1234 return tuple(launch for launch in self._launches.get(group, ()) if launch.role is role)
1236 def _drop_group(self, group: SwapGroup) -> SwapManager | None:
1237 """Forget *group*'s swap/launches and every member role's clients.
1239 Caller holds ``self._lock``. Member clients are retired (closed at a later
1240 reload or shutdown, never while a reader could still hold one) and the chat
1241 capacity falls back to its defaults when chat's group is dropped.
1242 """
1243 swap = self._swaps.pop(group, None)
1244 self._launches.pop(group, None)
1245 # Prune the dir map with the group: a stale entry outliving its group makes
1246 # _reload_dir see two dirs and pick one arbitrarily, splitting the provider.
1247 self._group_dirs.pop(group, None)
1248 for role in [r for r, g in self._role_group.items() if g is group]:
1249 del self._role_group[role]
1250 self._retire_clients(self._clients.pop(role, []))
1251 if role is WorkerRole.CHAT:
1252 self._chat_slots = 1
1253 self._chat_ctx = None
1254 return swap
1256 def _retire_clients(self, old_clients: list[LlamaServerClient]) -> None:
1257 """Close the previously-retired clients, then retire *old_clients*.
1259 Caller holds ``self._lock``. Retired clients are never handed to new
1260 readers (they are out of ``self._clients``), so by this reload any reader
1261 that held one from a prior reload has finished; an ``in_flight == 0``
1262 check confirms it before close, and any still-busy client stays retired
1263 for the next reload. This closes idle reloaded-away pools without ever
1264 closing one a reader could still use. Shutdown closes whatever remains.
1265 """
1266 still_busy: list[LlamaServerClient] = []
1267 for client in self._retiring_clients:
1268 if client.in_flight == 0:
1269 client.close()
1270 else:
1271 still_busy.append(client)
1272 self._retiring_clients = still_busy + old_clients
1274 def _require_clients(self, role: WorkerRole) -> list[LlamaServerClient]:
1275 """The client pool for *role*, or a user-facing error when it has no server.
1277 A configured, placeable role gets one or more replica clients; their absence
1278 means the role is unconfigured or did not fit memory. llama-swap loads each
1279 upstream on its first request, so a returned client may still be cold. No
1280 in-process fallback, so a missing pool is a hard error.
1282 When the pool is empty but a swap was previously built and its process has
1283 since exited (detected via ``is_live()``), a one-shot rebuild is attempted
1284 before raising so a transient llama-swap restart recovers transparently.
1285 """
1286 self._ensure_fleet()
1287 with self._lock:
1288 clients = self._clients.get(role)
1289 swap = self._swap_for(role)
1290 if not clients and swap is not None and not swap.is_live():
1291 self._rebuild_role(role)
1292 with self._lock:
1293 clients = self._clients.get(role)
1294 if not clients:
1295 # A refused launch's recorded reason wins over the generic line.
1296 reason = self._skipped_unusable_ctx.get(role)
1297 if reason is not None:
1298 raise ProviderError(reason, provider=_PROVIDER_NAME)
1299 raise ProviderError(_no_server_message(role), provider=_PROVIDER_NAME)
1300 return list(clients)
1302 def _with_rediscover(self, call: Callable[[], _T], *, role: WorkerRole | None = None) -> _T:
1303 """Run *call*; on a connection-kind or load-capacity failure, retry once.
1305 A vanished engine (its last user left on a config change, or it died)
1306 surfaces as ProviderErrorKind.CONNECTION, or as a raw httpx transport
1307 error when the proxy itself is gone (nothing listening to answer with
1308 a status). Membership is still held, so dropping the swap refs and
1309 retrying sends the call through _ensure_fleet, which rediscovers the
1310 new proxy ports or rebuilds. One retry only; a second failure surfaces
1311 to the caller.
1313 A ProviderErrorKind.CAPACITY failure is the engine dying on load because
1314 the estimate was too optimistic. Retrying it unchanged respawns the same
1315 launch into the same death, so *role*'s auto context steps down first and
1316 the role is rebuilt against the smaller plan. When there is no step left
1317 to take (a user-pinned context, or already at the floor) the failure
1318 surfaces instead: a retry that asks for the same thing is a crash loop.
1319 """
1320 try:
1321 return call()
1322 except (ProviderError, httpx.TransportError) as err:
1323 if is_rebuildable_failure(err) and role is not None:
1324 return self._retry_rebuilt(call, role, err)
1325 if not is_connection_failure(err):
1326 raise
1327 log.info("Engine unreachable; rediscovering before one retry")
1328 self._drop_swap_refs()
1329 self._release_holds()
1330 return call()
1332 def _retry_rebuilt(self, call: Callable[[], _T], role: WorkerRole, err: BaseException) -> _T:
1333 """Rebuild *role* so the retry is a different launch, and run *call* again.
1335 A held port just needs the rebuild, which picks a new one. A memory
1336 shortfall needs the plan to come back smaller too, so the context steps
1337 down first; when there is no step left to take, *err* is re-raised
1338 untouched rather than rebuilding into the same death.
1339 """
1340 from lilbee.providers.fleet.planning import record_ctx_downshift
1342 if is_load_capacity_failure(err):
1343 if not record_ctx_downshift(role):
1344 log.warning(
1345 "%s ran out of device memory on load and its context cannot be "
1346 "reduced further; lower num_ctx or use a smaller model",
1347 role.value,
1348 )
1349 raise err
1350 log.warning(
1351 "%s ran out of device memory on load; re-planning it with a smaller "
1352 "context where its window has room to give",
1353 role.value,
1354 )
1355 else:
1356 log.warning("%s could not claim its port; rebuilding it on a new one", role.value)
1357 self._rebuild_role(role)
1358 return call()
1360 def _rebuild_role(self, role: WorkerRole) -> None:
1361 """Restart just *role*'s dead group (new port) from a fresh plan.
1363 Other roles' groups keep serving; only the dead group is torn down and
1364 respawned. Runs the same diff-driven pass as a reload, forcing *role*
1365 into the restart set so an unchanged plan still replaces its dead swap.
1366 """
1367 self._reload_pass(force=frozenset((role,)))
1369 def role_ready(self, role: WorkerRole) -> bool:
1370 """Whether *role*'s upstream is loaded and ready, without starting the swap.
1372 A read-only probe for surfaces (HTTP status, SSE warming event) that want
1373 to report cold-start state without triggering a load. False before the swap
1374 is up or while the role's upstream is still loading.
1375 """
1376 with self._lock:
1377 swap = self._swap_for(role)
1378 return swap is not None and swap.role_ready(role)
1380 def max_concurrent_chats(self) -> int:
1381 """The chat server's batching-slot capacity, so the gate admits that many.
1383 Falls back to ``1`` before the chat group is up, so chat is serialized
1384 until the slot count is known (the launcher warms the engine before a
1385 client connects, so the real capacity is in effect by the first chat).
1386 """
1387 with self._lock:
1388 if WorkerRole.CHAT not in self._role_group:
1389 return 1
1390 return self._chat_slots
1392 def served_chat_ctx(self) -> int | None:
1393 """Per-slot context the chat server runs with, or None if not up."""
1394 with self._lock:
1395 return self._chat_ctx if WorkerRole.CHAT in self._role_group else None
1397 def warm_pending(self) -> bool:
1398 """Whether a requested warm is still running.
1400 The tracker only stamps a phase once the chat role starts loading, which is
1401 seconds after the swap is spawned, so ``warm_progress`` alone cannot tell a
1402 not-yet-started warm from no warm at all.
1403 """
1404 with self._lock:
1405 return self._warming
1407 def warm_progress(self) -> WarmProgress | None:
1408 """Live cold-load progress for the chat role, or None before warm begins."""
1409 return self._warm_tracker.snapshot()
1411 def _shutdown_swap(self, *, latch: bool = True) -> None:
1412 """Release this process's engine use; ``latch=False`` keeps the provider reusable.
1414 Terminal ``shutdown()`` latches ``_shut_down`` so a discarded provider's
1415 in-flight warm/reload thread can't spawn an orphan swap, then releases
1416 membership: the engine stops only when this was the last user and
1417 persistence was not opted into. The cache-drop paths
1418 (``invalidate_load_cache``, ``drop_loaded_models_async``) pass
1419 ``latch=False``: a config change restarts the shared engine for every
1420 user (they rediscover), and this provider rebuilds on next use.
1421 """
1422 # Latched before the lock, not inside it: every _shut_down check runs after
1423 # acquiring the build lock, so a warm or reload thread queued behind us can
1424 # only bail early if the flag is already set when its turn comes.
1425 if latch:
1426 with self._lock:
1427 self._shut_down = True
1428 # The build lock serializes shutdown against a concurrent reload/build:
1429 # both mutate self._swaps and the llama-swap processes, so an unserialized
1430 # loser would overwrite the winner's state and leak a live llama-swap.
1431 # Bounded, because a wedged engine start holds this lock and an unbounded
1432 # wait would hang process exit outright. On timeout the teardown proceeds
1433 # anyway: whatever the builder leaves behind is recorded in the engine
1434 # dir's state files, so the next start's reap finds it by record, while a
1435 # shutdown that never returns cannot be recovered from at all.
1436 acquired = self._build_lock.acquire(timeout=_SHUTDOWN_BUILD_LOCK_WAIT_S)
1437 if not acquired:
1438 log.warning(
1439 "Engine build still in progress after %.0fs; shutting down without "
1440 "waiting for it. Leftovers are reaped from their records on the next start.",
1441 _SHUTDOWN_BUILD_LOCK_WAIT_S,
1442 )
1443 try:
1444 # Terminal shutdown closes every client; a config-change teardown
1445 # retires them so an in-flight reader is never severed.
1446 self._drop_swap_refs(close_all=latch)
1447 self._release_engines(config_changed=not latch)
1448 finally:
1449 if acquired:
1450 self._build_lock.release()
1452 def _release_engines(self, *, config_changed: bool = False) -> None:
1453 """Drop membership in every used engine dir; stop each engine we leave last.
1455 Runs under each dir's cross-process build lock so a departing last user
1456 can never race an arriving binder: the arrival either sees the engine
1457 (and its bind holds it live) or sees the slot empty and builds.
1459 Whether the engine outlives us is the union of every user's opt-in, not
1460 just the exiting process's config: the machine slot is shared by
1461 installations that configure it differently, and which one leaves last
1462 is arbitrary.
1464 *config_changed* is the cache-drop path, where this provider's settings
1465 or model changed. That makes the running engine stale for us, so no
1466 persistence opt-in preserves it -- but it says nothing about the peers
1467 still serving requests against it, so a shared engine is left running
1468 and the next use re-runs the ladder, binding it if it happens to match
1469 and overflowing to a private dir if it does not.
1471 The hold map is cleared either way. Leaving a stale hold behind is not
1472 benign: after a lazy rebuild overflows to a private dir (a foreign
1473 process having claimed the machine slot in the gap), the next release
1474 would iterate the stale machine hold and stop that foreign engine
1475 mid-use, and the stale flock would keep live_users_exist true so the
1476 foreign engine's real last user could never reap it.
1477 """
1478 from lilbee.core.config import cfg
1480 for engine_dir, hold in list(self._engine_holds.items()):
1481 with build_lock(engine_dir, best_effort=True):
1482 last = hold.release_and_check_last()
1483 # A flip after binding never re-acquires, so reconcile our own mark
1484 # here. Skipped on a config change, which stops and clears regardless.
1485 if not config_changed:
1486 if cfg.keep_engine_warm:
1487 request_keep_warm(engine_dir, cfg.data_root)
1488 else:
1489 withdraw_keep_warm(engine_dir, cfg.data_root)
1490 # Any remaining opt-in keeps the engine, including a peer's.
1491 keep = not config_changed and keep_warm_requested(engine_dir)
1492 if last and not keep:
1493 stop_engine(engine_dir)
1494 log.info("Engine stopped at %s (last user out)", engine_dir)
1495 elif last:
1496 log.info("Engine left warm at %s (last user out)", engine_dir)
1497 elif config_changed:
1498 log.info("Engine left running at %s (still in use by peers)", engine_dir)
1499 self._engine_holds = {}
1501 def _release_holds(self) -> None:
1502 """Drop this process's engine memberships without stopping any engine.
1504 The rediscover retry re-runs the acquisition ladder. A retained membership
1505 would make this process count itself as a live user of the machine slot, so
1506 the ladder would judge the slot in use and overflow to a private engine
1507 instead of rebinding a recovered engine or rebuilding a dead one -- N
1508 private engines and N times the VRAM after the shared engine first dies.
1509 Nothing is stopped here: a live engine is rebound by the retry, and a dead
1510 one is cleared by the rebuild that retry triggers.
1511 """
1512 for hold in list(self._engine_holds.values()):
1513 hold.release_and_check_last()
1514 self._engine_holds = {}
1516 def _drop_swap_refs(self, *, close_all: bool = False) -> None:
1517 """Clear every group's swap/clients and the chat capacity so the next call rebuilds.
1519 Live pools are RETIRED through the ``in_flight`` check rather than closed
1520 outright: ``_with_rediscover`` reaches this on any connection blip, so a
1521 chat proxy hiccup must not sever the client another thread is mid-embed or
1522 mid-stream on (a streamed response is handed out after the retry returns,
1523 and failures past the first frame are not retried). Idle clients close now,
1524 busy ones stay retired for a later pass.
1526 *close_all* is the terminal-shutdown path, where nothing will read again and
1527 whatever remains must actually be closed.
1528 """
1529 doomed: list[LlamaServerClient] = []
1530 with self._lock:
1531 live = [client for pool in self._clients.values() for client in pool]
1532 self._swaps = {}
1533 self._launches = {}
1534 self._role_group = {}
1535 self._group_dirs = {}
1536 self._clients = {}
1537 if close_all:
1538 doomed = live + self._retiring_clients
1539 self._retiring_clients = []
1540 else:
1541 self._retire_clients(live)
1542 self._chat_slots = 1
1543 self._chat_ctx = None
1544 # A torn-down fleet's load failures describe servers that no longer
1545 # exist; the next warm records its own.
1546 self._warm_errors = {}
1547 # Full teardown: the next build starts from a clean box, so it must
1548 # re-snapshot memory rather than plan against this boot's probe.
1549 planning.clear_plan_probe()
1550 for client in doomed:
1551 client.close()
1553 def _drop_dead_swaps(self) -> None:
1554 """Drop the refs of groups whose process is gone so the next call rebuilds them.
1556 A no-op for groups still running (e.g. the failure was in planning), so
1557 a live engine is never abandoned unstopped.
1558 """
1559 with self._build_lock, self._lock:
1560 for group in [g for g, swap in self._swaps.items() if not swap.running]:
1561 self._drop_group(group)
1563 def _require_configured_model(
1564 self, model: str | None, configured: str, role: WorkerRole
1565 ) -> None:
1566 """Reject a per-call model that differs from the server's configured one.
1568 The fleet serves the configured model for each role; switching models is
1569 a config change that respawns the server (via ``invalidate_load_cache``),
1570 not a per-call override. An empty/None ``model`` means "use the configured
1571 one" and is always accepted.
1572 """
1573 if model and model != configured:
1574 raise ProviderError(
1575 configured_model_message(role, configured, model),
1576 provider=_PROVIDER_NAME,
1577 kind=ProviderErrorKind.BAD_REQUEST,
1578 )
1580 @overload
1581 def chat(
1582 self,
1583 messages: list[ChatMessage],
1584 *,
1585 stream: Literal[False] = False,
1586 options: dict[str, Any] | None = None,
1587 model: str | None = None,
1588 tools: list[dict[str, Any]] | None = None,
1589 tool_choice: str | dict[str, Any] | None = None,
1590 ) -> ChatResult: ...
1592 @overload
1593 def chat(
1594 self,
1595 messages: list[ChatMessage],
1596 *,
1597 stream: Literal[True],
1598 options: dict[str, Any] | None = None,
1599 model: str | None = None,
1600 tools: list[dict[str, Any]] | None = None,
1601 tool_choice: str | dict[str, Any] | None = None,
1602 ) -> ClosableIterator[ChatStreamItem]: ...
1604 def chat(
1605 self,
1606 messages: list[ChatMessage],
1607 *,
1608 stream: bool = False,
1609 options: dict[str, Any] | None = None,
1610 model: str | None = None,
1611 tools: list[dict[str, Any]] | None = None,
1612 tool_choice: str | dict[str, Any] | None = None,
1613 ) -> ChatResult | ClosableIterator[ChatStreamItem]:
1614 """Route a chat turn to the least-busy chat server.
1616 Non-streaming returns a :class:`ChatResult` (text, tool calls, finish
1617 reason); streaming yields :data:`ChatStreamItem` frames. ``--jinja`` on
1618 the server parses native tool calls, so tool support needs no per-family
1619 parser here.
1620 """
1621 from lilbee.providers.engine_params import chat_options_to_kwargs
1623 self._require_configured_model(model, str(cfg.chat_model), WorkerRole.CHAT)
1624 self._require_clients(WorkerRole.CHAT)
1625 messages = self._fit_chat_context(messages, tools, options, model or str(cfg.chat_model))
1626 # Translate options exactly as the in-process path did (validate via
1627 # LLMOptions, num_predict -> max_tokens, drop num_ctx) so the server
1628 # honors the same generation settings; a raw passthrough would drop
1629 # num_predict and leak the load-only num_ctx.
1630 server_options = chat_options_to_kwargs(options) or None
1631 if stream:
1632 # The first frame is pulled eagerly so a dead proxy fails inside
1633 # _with_rediscover; a failure past the first frame surfaces to the
1634 # caller as a retry error, and rediscovery covers the next call.
1635 return self._with_rediscover(
1636 lambda: _primed_stream(
1637 _least_in_flight(self._require_clients(WorkerRole.CHAT)).chat_stream_items(
1638 messages, tools=tools, tool_choice=tool_choice, options=server_options
1639 )
1640 ),
1641 role=WorkerRole.CHAT,
1642 )
1643 return self._with_rediscover(
1644 lambda: _least_in_flight(self._require_clients(WorkerRole.CHAT)).chat_result(
1645 messages, tools=tools, tool_choice=tool_choice, options=server_options
1646 ),
1647 role=WorkerRole.CHAT,
1648 )
1650 def chat_with_tools(
1651 self,
1652 messages: list[ChatMessage],
1653 *,
1654 tools: list[dict[str, Any]],
1655 tool_choice: str | dict[str, Any] | None = None,
1656 options: dict[str, Any] | None = None,
1657 model: str | None = None,
1658 ) -> ChatToolResult:
1659 """Route a tool-enabled chat turn to the least-busy chat server."""
1660 from lilbee.providers.engine_params import chat_options_to_kwargs
1662 self._require_configured_model(model, str(cfg.chat_model), WorkerRole.CHAT)
1663 self._require_clients(WorkerRole.CHAT)
1664 messages = self._fit_chat_context(messages, tools, options, model or str(cfg.chat_model))
1665 server_options = chat_options_to_kwargs(options) or None
1666 return self._with_rediscover(
1667 lambda: _least_in_flight(self._require_clients(WorkerRole.CHAT)).chat_tools(
1668 messages, tools=tools, tool_choice=tool_choice, options=server_options
1669 ),
1670 role=WorkerRole.CHAT,
1671 )
1673 def _fit_chat_context(
1674 self,
1675 messages: list[ChatMessage],
1676 tools: list[dict[str, Any]] | None,
1677 options: dict[str, Any] | None,
1678 model: str,
1679 ) -> list[ChatMessage]:
1680 """Drop oldest turns so the prompt fits the served context.
1682 A ``num_predict`` reservation larger than the default generation room is
1683 capped to it, so an agent client that over-reserves keeps its history
1684 instead of having it evicted; llama-server stops at the context edge
1685 anyway. A smaller reservation is honored as-is and widens the prompt.
1686 Raises ``ProviderError(CONTEXT_OVERFLOW)`` only when system messages,
1687 tools, and the final turn exceed the window even with the capped
1688 reserve (mapped to a 400 by the chat-completions route).
1689 """
1690 # 0/None means the served context is unknown (no chat launch adopted yet);
1691 # a real per-slot context is always positive, so skip windowing.
1692 if not self._chat_ctx:
1693 return messages
1694 # An output reservation only ever buys the prompt MORE room, never less:
1695 # a num_predict past the default is a ceiling on what the model may
1696 # generate, not a claim on prompt space, and llama-server stops at the
1697 # context edge regardless. Capping it here rather than retrying after a
1698 # failed fit is the difference between a policy and a rescue -- an agent
1699 # reserving most of the window leaves a budget of a few dozen tokens, in
1700 # which the final turn still "fits" while the whole conversation is
1701 # silently evicted.
1702 requested = (options or {}).get("num_predict")
1703 reserve = min(requested, GENERATION_RESERVE_TOKENS) if requested else None
1704 result = window_messages(messages, tools, prompt_token_budget(self._chat_ctx, reserve))
1705 if not result.fits:
1706 raise ProviderError(
1707 f"Prompt of about {result.prompt_tokens} tokens exceeds the "
1708 f"{self._chat_ctx}-token context window for {model!r}. Shorten the "
1709 "conversation or the system prompt.",
1710 provider=_PROVIDER_NAME,
1711 kind=ProviderErrorKind.CONTEXT_OVERFLOW,
1712 )
1713 return result.messages
1715 def embed(self, texts: list[str]) -> list[Vector]:
1716 return self._with_rediscover(lambda: self._embed_once(texts), role=WorkerRole.EMBED)
1718 def _embed_once(self, texts: list[str]) -> list[Vector]:
1719 clients = self._require_clients(WorkerRole.EMBED)
1720 return _call_with_failover(clients, lambda client: client.embed(texts))
1722 def count_tokens(self, text: str) -> int:
1723 """Exact token count of *text* under the embedding model's tokenizer.
1725 Routes to the embed server's ``/tokenize`` so chunk sizing counts the same
1726 tokens the embedder will consume. Raises ``ProviderError`` when no embed
1727 server is configured; callers on the chunk-sizing path degrade to an
1728 estimate rather than propagate it.
1729 """
1730 clients = self._require_clients(WorkerRole.EMBED)
1731 return _call_with_failover(clients, lambda client: client.count_tokens(text))
1733 def vision_ocr(
1734 self, png_bytes: bytes, model: str, prompt: str = "", *, timeout: float | None = None
1735 ) -> str:
1736 from lilbee.vision import build_vision_messages, resolve_ocr_prompt
1738 self._require_configured_model(model, str(cfg.vision_model), WorkerRole.VISION)
1739 pool = self._vision_pool()
1740 effective = model or str(cfg.vision_model)
1741 messages = build_vision_messages(prompt or resolve_ocr_prompt(effective), png_bytes)
1742 try:
1743 return _ocr_dispatch(pool, messages, _ocr_deadline(timeout))
1744 except _PageBudgetExhausted:
1745 raise ProviderError(
1746 "Vision OCR timed out waiting for a free vision slot.",
1747 provider=_PROVIDER_NAME,
1748 ) from None
1750 def vision_slot_capacity(self) -> int | None:
1751 """Total fitted ``--parallel`` slots across the running vision replicas.
1753 ``None`` before the fleet is up (no launch snapshot yet), so the ingest
1754 fan-out keeps its own estimate until real capacity is known. A modest
1755 card that fit fewer slots than requested reports the smaller real number,
1756 so the fan-out never queues more pages than the servers can serve.
1757 """
1758 launches = self._role_launches(WorkerRole.VISION)
1759 if not launches:
1760 return None
1761 return max(1, sum(launch.slots for launch in launches))
1763 def _vision_pool(self) -> list[_VisionReplica]:
1764 """Each vision replica paired with its fitted ``--parallel`` slot count.
1766 The fitted count can be lower than ``vision_ocr_concurrency`` when memory
1767 forced a smaller fit; dispatching at the configured ceiling instead
1768 over-subscribes that server. The configured ceiling applies per replica
1769 only when no matching launch snapshot exists (a reload can momentarily
1770 drop it between two reads).
1771 """
1772 clients = self._require_clients(WorkerRole.VISION)
1773 launches = self._role_launches(WorkerRole.VISION)
1774 if launches and len(launches) == len(clients):
1775 return [
1776 _VisionReplica(client, max(1, launch.slots))
1777 for client, launch in zip(clients, launches, strict=True)
1778 ]
1779 fallback_slots = max(1, cfg.vision_ocr_concurrency)
1780 return [_VisionReplica(client, fallback_slots) for client in clients]
1782 # PDF/image OCR now runs inside xberg via the registered lilbee-vision
1783 # backend (see data.extract.backends.vision_ocr); this provider only exposes
1784 # single-image vision_ocr, which that backend calls.
1786 def rerank(self, query: str, candidates: list[str]) -> list[float]:
1787 return self._with_rediscover(
1788 lambda: self._rerank_once(query, candidates), role=WorkerRole.RERANK
1789 )
1791 def _rerank_once(self, query: str, candidates: list[str]) -> list[float]:
1792 clients = self._require_clients(WorkerRole.RERANK)
1793 return _call_with_failover(clients, lambda client: client.rerank(query, candidates))
1795 # --- model management: registry / GGUF reads, no running server needed ---
1797 def supports_rerank(self) -> bool:
1798 """Serve a cross-encoder (rank pooling) or an LLM reranker (yes/no logprob)."""
1799 return True
1801 def list_models(self) -> list[str]:
1802 """List installed models from the registry."""
1803 from lilbee.app.services import get_services
1805 registry = get_services().registry
1806 return sorted(m.ref for m in registry.list_installed())
1808 def list_chat_models(self, provider: str) -> list[str]:
1809 """The local engine has no frontier-provider catalog; always ``[]``."""
1810 del provider
1811 return []
1813 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:
1814 """Not supported directly: ``lilbee.catalog`` handles GGUF downloads."""
1815 del on_progress
1816 raise NotImplementedError(
1817 f"The local engine cannot pull model {model!r}. "
1818 "Download GGUF files through the catalog or 'lilbee model pull'."
1819 )
1821 def show_model(self, model: str) -> dict[str, Any] | None:
1822 """Return model metadata from GGUF headers, or ``None`` if unresolved."""
1823 from lilbee.providers.engine_params import resolve_model_path
1824 from lilbee.providers.gguf_meta import read_gguf_metadata
1826 try:
1827 path = resolve_model_path(model)
1828 except ProviderError:
1829 return None
1830 return read_gguf_metadata(path)
1832 def get_capabilities(self, model: str) -> list[str]:
1833 """Detect capabilities from the local GGUF files.
1835 Cross-encoder rerank GGUFs report ``["rerank"]`` (they cannot generate);
1836 other models report ``"completion"`` plus ``"vision"`` when an mmproj
1837 sidecar is present.
1838 """
1839 from lilbee.catalog import is_rerank_ref
1840 from lilbee.providers.engine_params import resolve_model_path
1841 from lilbee.providers.gguf_meta import find_mmproj_for_model
1843 if model and is_rerank_ref(model):
1844 return ["rerank"]
1845 caps = ["completion"]
1846 try:
1847 path = resolve_model_path(model)
1848 except ProviderError:
1849 return caps
1850 try:
1851 find_mmproj_for_model(path)
1852 caps.append("vision")
1853 except ProviderError:
1854 pass
1855 return caps
1857 def supports_tools(self, model_ref: str) -> bool:
1858 """True iff *model_ref*'s GGUF chat template references tool tokens.
1860 The server parses native tool calls via ``--jinja``; a template that
1861 declares tools is the signal that the model was trained to emit them.
1862 Cached on ``(path, mtime)`` so a tool-bearing chat doesn't re-read the
1863 GGUF header each request; a re-quantised file at the same path
1864 invalidates because its mtime changes.
1865 """
1866 from lilbee.providers.engine_params import resolve_model_path
1868 try:
1869 path = resolve_model_path(model_ref)
1870 except (ProviderError, OSError):
1871 log.debug("supports_tools: resolve_model_path failed for %s", model_ref, exc_info=True)
1872 return False
1873 try:
1874 mtime_ns = path.stat().st_mtime_ns
1875 except OSError:
1876 mtime_ns = 0
1877 return _supports_tools_cached(str(path), mtime_ns)
1879 def warm_up_pool(self) -> None:
1880 """Pre-load every configured role off the caller's thread (idempotent).
1882 Starting the swap and loading each role's model (seconds on a cold large
1883 model) runs on a background thread and this returns at once: the eager-start
1884 at TUI mount must not freeze the UI. The spawn listeners fire per role as it
1885 loads, so the UI shows progress. A second call while warm-up is in flight
1886 (or once the fleet is up) is a no-op.
1887 """
1888 with self._lock:
1889 if self._warming:
1890 return
1891 fleet_up = bool(self._swaps)
1892 # A live swap whose model llama-swap idle-unloaded (its ttl stops only the
1893 # llama-server child, leaving the swap handle in _swaps) reports its role
1894 # cold. Re-warm so a prompt sent into that gap drives llama-swap's
1895 # on-demand reload; bailing on "swaps exist" alone stranded every later
1896 # prompt on a stale not-ready. A fully-loaded fleet still short-circuits.
1897 # The probe runs off the lock (role_ready may hit the proxy).
1898 if fleet_up and self._roles_ready():
1899 return
1900 with self._lock:
1901 if self._warming:
1902 return
1903 self._warming = True
1904 threading.Thread(
1905 target=self._warm_up_blocking,
1906 name="fleet-warm-up",
1907 daemon=True,
1908 ).start()
1910 def _roles_ready(self) -> bool:
1911 """Whether every configured role's upstream is loaded (fleet fully warm)."""
1912 with self._lock:
1913 roles = list(self._role_group)
1914 return bool(roles) and all(self.role_ready(role) for role in roles)
1916 def _warm_up_blocking(self) -> None:
1917 """Start the fleet and pre-load every role on a background thread.
1919 Runs on a daemon thread with no caller to catch failures, so a startup
1920 error is logged and swallowed: a role that can't load surfaces a
1921 user-facing ProviderError on the next call, not a thread traceback.
1923 The tracker is stamped STARTING before the fleet spawn so surfaces
1924 show the engine coming up from the first moment (spawn plus health
1925 check takes seconds and previously reported nothing), and stamped
1926 ERROR with the real reason when the warm fails before the chat warm
1927 proper begins.
1928 """
1929 try:
1930 self._warm_tracker.begin(str(cfg.chat_model))
1931 self._ensure_fleet()
1932 self._preload_roles()
1933 self._finalize_warm_if_chat_never_ran()
1934 except Exception as exc:
1935 if isinstance(exc, RuntimeError) and sys.is_finalizing():
1936 # A fast CLI exit can tear down the interpreter while this daemon
1937 # thread is still warming; pool submission then raises "cannot
1938 # schedule new futures after interpreter shutdown". The process is
1939 # leaving anyway, so drop it quietly instead of stack-tracing.
1940 log.debug("Engine warm-up abandoned during interpreter shutdown: %s", exc)
1941 else:
1942 # A warm-up failure is handled (roles lazy-load on first use), so
1943 # keep the full traceback at debug: a WARNING carrying exc_info
1944 # reads like a crash for a condition the next real call recovers
1945 # from.
1946 log.warning("Engine warm-up failed: %s", exc)
1947 log.debug("Engine warm-up failure detail.", exc_info=True)
1948 self._fail_warm_unless_ready(str(exc))
1949 finally:
1950 with self._lock:
1951 self._warming = False
1953 def _fail_warm_unless_ready(self, message: str) -> None:
1954 """Stamp the warm tracker ERROR unless the chat warm already finished.
1956 A failure in a later role's preload must not clobber a chat warm that
1957 reached READY; every earlier failure leaves the tracker mid-phase,
1958 where surfaces would spin forever and the prompt path could not name
1959 the reason.
1960 """
1961 snapshot = self._warm_tracker.snapshot()
1962 if snapshot is None or snapshot.phase is not WarmPhase.READY:
1963 self._warm_tracker.fail(message)
1965 def _finalize_warm_if_chat_never_ran(self) -> None:
1966 """Terminate the early STARTING stamp when no chat instance was placed.
1968 ``_warm_chat_role`` always ends in READY or ERROR when it runs, so a
1969 snapshot still on STARTING after a successful preload means the plan had
1970 no chat instance. A chat model that isn't installed, one whose launch the
1971 plan refused for an unusable window, and one with no engine to run it all
1972 fail the warm with a user-facing reason so the prompt path renders
1973 "failed to load" instead of spinning a "not ready" retry that can never
1974 succeed; any other reason (a remote-routed chat has no local server to
1975 warm) clears the stamp.
1976 """
1977 snapshot = self._warm_tracker.snapshot()
1978 if snapshot is None or snapshot.phase is not WarmPhase.STARTING:
1979 return
1980 missing = self._skipped_not_installed.get(WorkerRole.CHAT)
1981 if missing is not None:
1982 self._warm_tracker.fail(f"chat model {clean_display_name(missing)} is not installed")
1983 return
1984 unusable = self._skipped_unusable_ctx.get(WorkerRole.CHAT)
1985 if unusable is not None:
1986 self._warm_tracker.fail(unusable)
1987 return
1988 if _chat_needs_local_engine() and (reason := _unusable_engine_reason()) is not None:
1989 self._warm_tracker.fail(reason)
1990 return
1991 self._warm_tracker.clear()
1993 def _preload_roles(self, roles: frozenset[WorkerRole] | None = None) -> None:
1994 """Issue a cheap request per replica so llama-swap loads each upstream now.
1996 llama-swap starts an upstream on its first request, so warming sends a
1997 minimal call to every replica of every role (firing the spawn listeners
1998 around each role). A per-replica failure is logged and skipped; that replica
1999 still loads on its first real use. The chat role routes through
2000 :meth:`_warm_chat_role` so a launcher gets granular progress. *roles*
2001 narrows the warm to just those roles (a reload warms only what restarted).
2003 Roles on separate devices warm concurrently: chat is the long pole (a
2004 large model's load dominates), so the light roles load alongside it
2005 instead of before it. Roles whose launches pin overlapping devices warm
2006 one at a time instead, chat last: two engines loading into the same
2007 card at once race each other for VRAM, and the loser's first load can
2008 OOM even though both fit once settled.
2009 """
2010 with self._lock:
2011 pools = {
2012 role: list(clients)
2013 for role, clients in self._clients.items()
2014 if roles is None or role in roles
2015 }
2016 on_spawning, on_spawned = self._on_spawning, self._on_spawned
2017 device_sets = _role_device_sets(
2018 launch for launches in self._launches.values() for launch in launches
2019 )
2021 if not pools:
2022 return
2023 listeners = (on_spawning, on_spawned)
2024 chains = _warm_chains(list(pools), device_sets)
2025 with ThreadPoolExecutor(
2026 max_workers=len(chains), thread_name_prefix="fleet-preload"
2027 ) as pool:
2028 futures = [pool.submit(self._warm_chain, chain, pools, listeners) for chain in chains]
2029 for future in futures:
2030 future.result()
2032 def _warm_chain(
2033 self,
2034 chain: list[WorkerRole],
2035 pools: dict[WorkerRole, list[LlamaServerClient]],
2036 listeners: tuple[Callable[[WorkerRole], None] | None, Callable[[WorkerRole], None] | None],
2037 ) -> None:
2038 """Warm *chain*'s roles one at a time; every role gets its attempt.
2040 An unexpected error warming one role (a listener blowing up) must not rob
2041 the roles behind it of their warm, so the first error is re-raised only
2042 after the chain finishes.
2043 """
2044 on_spawning, on_spawned = listeners
2045 first_exc: Exception | None = None
2046 for role in chain:
2047 try:
2048 if on_spawning is not None:
2049 on_spawning(role)
2050 if role is WorkerRole.CHAT:
2051 self._warm_chat_role(pools[role])
2052 else:
2053 self._warm_role_clients(role, pools[role])
2054 if on_spawned is not None:
2055 on_spawned(role)
2056 except Exception as exc:
2057 first_exc = first_exc or exc
2058 if first_exc is not None:
2059 raise first_exc
2061 def _warm_role_clients(self, role: WorkerRole, clients: list[LlamaServerClient]) -> bool:
2062 """Warm every replica of *role*; return whether at least one loaded.
2064 A replica that fails to load is reported at warning level with the engine's
2065 own message (an unsupported architecture, a corrupt file). Warm-up stays
2066 best-effort, but the failure must not be silent: the role then serves
2067 nothing, and a caller that never reaches it would otherwise see only an
2068 unexplained empty answer.
2069 """
2070 warmed = False
2071 self._warm_errors.pop(role, None)
2072 for client in clients:
2073 try:
2074 _warm_role(role, client)
2075 client.mark_healthy()
2076 warmed = True
2077 except Exception as exc:
2078 # A replica that cannot load is not routable. Marking it takes it
2079 # out of the pool so calls go to a sibling on a device that works,
2080 # instead of every request picking the dead one again. It is a
2081 # device fault as often as a model one: an adapter that enumerates
2082 # but cannot allocate fails here and nowhere else. The health flag
2083 # carries its own cool-down, so a device that recovers rejoins
2084 # without anything having to remember it was bad.
2085 client.mark_unhealthy()
2086 self._warm_errors[role] = str(exc)
2087 log.warning(
2088 "The %s model failed to load: %s",
2089 role.value,
2090 exc,
2091 exc_info=log.isEnabledFor(logging.DEBUG),
2092 )
2093 if warmed:
2094 self._warm_errors.pop(role, None)
2095 return warmed
2097 def _warm_chat_role(self, clients: list[LlamaServerClient]) -> None:
2098 """Warm the chat role, driving the tracker through read -> load -> ready/fail.
2100 Readiness is decided by whether a warm request actually returned, not by
2101 re-probing llama-swap (which can transiently report empty right after a
2102 successful load). The terminal phase is stamped in ``finally`` so an
2103 unexpected error mid-warm still ends the launcher's progress stream.
2104 """
2105 self._warm_tracker.begin(str(cfg.chat_model))
2106 warmed = False
2107 try:
2108 self._prewarm_chat_weights()
2109 self._warm_tracker.loading_engine()
2110 warmed = self._warm_role_clients(WorkerRole.CHAT, clients)
2111 finally:
2112 if warmed:
2113 self._warm_tracker.ready()
2114 else:
2115 self._warm_tracker.fail(self._chat_load_failure())
2117 def _chat_load_failure(self) -> str:
2118 """The engine's own reason the chat model did not load, when it gave one."""
2119 reason = self._warm_errors.get(WorkerRole.CHAT)
2120 if not reason:
2121 return "The chat model did not finish loading."
2122 return f"The chat model did not load: {reason}"
2124 def _prewarm_chat_weights(self) -> None:
2125 """Page the chat model's GGUF shards into the OS cache, reporting byte progress.
2127 Reading the shards before llama-swap loads them does two things: it gives a
2128 true read-phase percentage for the warm tracker, and it warms the page cache
2129 so the engine's mmap faults hit memory (a large win on a network filesystem,
2130 where random mmap faults stalled cold loads). Best-effort: any failure to
2131 resolve or size the shards (unregistered ref, cache miss, I/O error) is
2132 skipped, and the model still loads on the warm request.
2133 """
2134 try:
2135 shards = ModelRegistry(cfg.models_dir).shard_paths(str(cfg.chat_model))
2136 total = sum(shard.stat().st_size for shard in shards)
2137 except Exception:
2138 log.debug("Prewarm skipped; could not resolve chat shards.", exc_info=True)
2139 return
2140 if total <= 0:
2141 return
2142 keys = [_prewarm_key(shard) for shard in shards]
2143 if all(key in _PREWARMED_SHARDS for key in keys):
2144 # Already paged in this boot (e.g. a placement rebuild); the cache is hot.
2145 self._warm_tracker.reading(total, total)
2146 return
2147 done = 0
2148 self._warm_tracker.reading(0, total)
2149 chunk = bytearray(_PREWARM_CHUNK_BYTES)
2150 for index, (shard, key) in enumerate(zip(shards, keys, strict=True)):
2151 detail = f"shard {index + 1}/{len(shards)}" if len(shards) > 1 else None
2152 try:
2153 with shard.open("rb", buffering=0) as handle:
2154 while True:
2155 read = handle.readinto(chunk)
2156 if not read:
2157 break
2158 done += read
2159 self._warm_tracker.reading(done, total, detail=detail)
2160 _PREWARMED_SHARDS.add(key)
2161 except OSError:
2162 # A partial/locked shard just shortens the read bar; the engine load
2163 # surfaces any real fault as a user-facing error on the warm request.
2164 log.debug("Prewarm read of %s stopped early.", shard, exc_info=True)
2166 def cancel_inference(self) -> None:
2167 """Sever every in-flight chat stream so its blocked reader unwinds.
2169 A cooperative worker cancel cannot reach a thread blocked in a socket
2170 read, and the reader's own close runs only when its worker unwinds, so
2171 the disconnect must happen here. Retired clients are swept too: a
2172 model-swap reload retires a busy client before the cancel lands.
2173 """
2174 with self._lock:
2175 clients = [*self._clients.get(WorkerRole.CHAT, ()), *self._retiring_clients]
2176 for client in clients:
2177 client.abort_streams()
2179 def reload_role(self, role: WorkerRole, *, wait: bool = False) -> None:
2180 """Apply a model/settings change for *role* with current cfg.
2182 The whole fleet is re-planned, but only the roles whose launches changed
2183 restart, so the other roles' loaded models stay resident (*role* names
2184 the change for the thread label; the diff decides what restarts).
2185 """
2186 self._dispatch_reload(f"fleet-reload-{role.value}", wait=wait)
2188 def reload_placement(self, *, wait: bool = False) -> None:
2189 """Apply a placement change with current cfg, restarting only moved roles.
2191 The fresh plan is diffed per role against the running fleet: a role whose
2192 devices (and so its launch argv) did not change keeps serving through the
2193 change -- moving the embedder never unloads a 100GB chat model. When no
2194 fleet is up, the next use plans fresh, so this returns at once.
2195 """
2196 self._dispatch_reload("fleet-reload-placement", wait=wait)
2198 def _dispatch_reload(self, thread_name: str, *, wait: bool) -> None:
2199 """Run the diff-driven reload once, off-thread unless *wait*.
2201 Dispatched to a background thread because the slow restart (rewrite config +
2202 respawn + wait-ready) must not block the settings/model-picker callback.
2203 If no group is up yet, the next use starts the fleet with current cfg.
2204 Single-flight: a reload while one is in flight sets the pending flag (the
2205 in-flight pass may have already snapshotted its plan), and the in-flight
2206 thread runs one more pass per pending flag so the change is applied, not
2207 dropped.
2209 ``wait=True`` runs the reload in the caller's thread and returns only once
2210 the restart (and any reload already in flight that will run the pending
2211 pass) has finished and the proxies are healthy again, so a caller already
2212 off the event loop gets a real completion signal. A restarted role's model
2213 still loads lazily (the reload kicks an off-thread warm). It propagates a
2214 reload failure as an exception.
2215 """
2216 with self._lock:
2217 if not self._swaps:
2218 return
2219 if self._reloading:
2220 self._reload_pending = True
2221 if wait:
2222 while self._reloading:
2223 self._reload_done.wait()
2224 return
2225 self._reloading = True
2226 self._reload_pending = False
2227 if wait:
2228 self._reload_blocking()
2229 return
2230 threading.Thread(
2231 target=self._reload_blocking,
2232 name=thread_name,
2233 daemon=True,
2234 ).start()
2236 def _reload_blocking(self) -> None:
2237 """Run reload passes until no further reload arrived mid-pass.
2239 A failed pass with the pending flag set still runs the pending pass (the
2240 fresh plan may succeed under the new cfg); only the final pass's failure
2241 propagates, after dropping the refs to a dead swap so the next call can
2242 rebuild. The pending check and the guard release happen under one lock
2243 acquisition, so a reload_role landing between them cannot be acknowledged
2244 and dropped.
2245 """
2246 while True:
2247 try:
2248 self._reload_pass()
2249 except BaseException:
2250 with self._lock:
2251 rerun = self._reload_pending
2252 self._reload_pending = False
2253 if not rerun:
2254 self._reloading = False
2255 self._reload_done.notify_all()
2256 if rerun:
2257 log.warning(
2258 "Engine reload failed; retrying with the pending change.", exc_info=True
2259 )
2260 continue
2261 self._drop_dead_swaps()
2262 raise
2263 with self._lock:
2264 if not self._reload_pending:
2265 self._reloading = False
2266 self._reload_done.notify_all()
2267 return
2268 self._reload_pending = False
2270 def _rebind_or_overflow(self) -> list[WorkerRole]:
2271 """Re-acquire a bound engine after a config change; caller holds the build lock.
2273 A provider that rode another process's engine owns none of its groups and
2274 cannot restart them. Restarting "in place" would spawn a duplicate fleet
2275 into the shared slot (a bound manager's ``shutdown`` only detaches, leaving
2276 the incumbent resident) and size it blind against VRAM the incumbent still
2277 holds. Instead drop every binding and this process's membership, then re-run
2278 the acquisition ladder: it rebinds to the reconfigured shared engine, builds
2279 fresh in the machine slot if we were its last user, or overflows to a private
2280 engine sized against a fresh probe. Returns the roles now served (to preload).
2281 """
2282 from lilbee.core.config import cfg
2284 with self._lock:
2285 groups = list(self._swaps)
2286 for group in groups:
2287 with self._lock:
2288 swap = self._drop_group(group) # also prunes _group_dirs
2289 if swap is not None:
2290 swap.shutdown() # bound: detaches; the shared engine keeps running
2291 self._release_engines() # a shared engine's builder keeps it live; no stop here
2292 if not self._acquire_engine(cfg.data_root):
2293 return []
2294 with self._lock:
2295 return list(self._role_group)
2297 def _reload_pass(self, force: frozenset[WorkerRole] = frozenset()) -> None:
2298 """One re-plan from current cfg, restarting only the groups that changed.
2300 The fresh plan is diffed per swap group against the launches each running
2301 group was started with; a group restarts only when its launches differ
2302 (covers added and removed groups too), so an untouched group's loaded model
2303 stays resident through a placement or per-role model change. *force* adds a
2304 role's group to the restart set even when its plan is unchanged (dead-swap
2305 recovery). Changed groups stop before the new ones start, so the planned
2306 VRAM is actually free when the new servers spawn. Runs under the build lock
2307 so a racing shutdown/build can't interleave with the restart and leak a live
2308 llama-swap holding GPU memory.
2309 """
2311 restarted: list[WorkerRole] = []
2312 with self._build_lock:
2313 # The device list is structural and was captured once at boot, so a
2314 # card that has since left keeps being planned onto. The memory
2315 # figures beside it are deliberately not re-taken: this fleet is
2316 # resident, and charging it against itself is what the snapshot exists
2317 # to prevent.
2318 planning.refresh_plan_devices()
2319 with self._lock:
2320 if self._shut_down:
2321 # Terminal shutdown landed while this reload was queued; a
2322 # rebuild here would spawn a fleet no live provider owns.
2323 return
2324 running = set(self._swaps)
2325 old = dict(self._launches)
2326 # All groups share one dir and one ownership by construction, so any
2327 # bound manager means this provider rides another process's engine.
2328 bound = any(swap.bound for swap in self._swaps.values())
2329 if bound:
2330 # Cannot restart a shared engine's groups in place (that duplicates
2331 # the fleet into the slot); drop the bindings and re-acquire.
2332 restarted = self._rebind_or_overflow()
2333 self._preload_restarted(restarted)
2334 return
2335 # Reap dead engines in our dirs before re-planning, as a build would.
2336 reload_dir = self._reload_dir()
2337 # Serialize the reload against peer acquisitions with the same
2338 # cross-process lock a build takes. Without it, this reap_stale can kill
2339 # a peer's swap that is spawned but not yet answering its proxy, and the
2340 # stop-then-spawn gap lets a peer's ladder see a half-stopped slot and
2341 # build a second fleet into the same dir, double-allocating VRAM.
2342 with build_lock(reload_dir):
2343 reap_stale(reload_dir)
2344 try:
2345 if not running:
2346 # Nothing loaded (a resurrect after a failed pass): the box is
2347 # clean, so refresh the plan snapshot like a first build would.
2348 planning.capture_plan_probe()
2349 plan = planning.plan_all_launches()
2350 except ProviderError as exc:
2351 # Same policy as the initial build: a genuinely-missing engine
2352 # binary aborts the reload quietly (nothing to serve), while any
2353 # other planning failure (a wedged GPU probe, an unusable CUDA
2354 # runtime) propagates to fail loud. The raise lands before the
2355 # stop phase, so a running fleet is left intact rather than half
2356 # torn down.
2357 if exc.kind is not ProviderErrorKind.NOT_FOUND:
2358 raise
2359 log.debug("Engine binary unavailable; reload left the fleet as-is")
2360 return
2361 # Keep the skip reasons in step with the fresh plan.
2362 self._skipped_not_installed = dict(plan.skipped_not_installed)
2363 self._skipped_unusable_ctx = dict(plan.skipped_unusable_ctx)
2364 new = _launches_by_group(plan)
2365 # A group restarts when its launches changed OR its running/planned
2366 # presence disagrees (covers a group the new plan drops or adds).
2367 changed = {
2368 group
2369 for group in running | set(new)
2370 if (group in running) != (group in new)
2371 or old.get(group, ()) != new.get(group, ())
2372 }
2373 changed |= {group_for(role, plan.co_tenants) for role in force}
2374 # Stop phase: free the changed groups' VRAM before their replacements
2375 # (or another group's grown plan) spawn against it.
2376 for group in sorted(changed, key=lambda g: g.value):
2377 with self._lock:
2378 swap = self._drop_group(group)
2379 if swap is not None:
2380 swap.shutdown()
2381 # Start phase: spawn the changed groups present in the new plan.
2382 for group in sorted(changed & set(new), key=lambda g: g.value):
2383 group_launches = list(new[group])
2384 swap = SwapManager(reload_dir, group)
2385 swap.start(
2386 group_launches,
2387 ttl_seconds=_warm_ttl_seconds(
2388 hold_warm_for_session=self._hold_warm_for_session
2389 ),
2390 bind_lifetime=not cfg.keep_engine_warm,
2391 )
2392 with self._lock:
2393 self._adopt_group(group, swap, group_launches)
2394 self._group_dirs[group] = reload_dir
2395 restarted.extend(_by_role(group_launches))
2396 self._preload_restarted(restarted)
2398 def _preload_restarted(self, restarted: list[WorkerRole]) -> None:
2399 """Load the restarted roles' models off-thread (a no-op for none).
2401 llama-swap spawns an upstream on its first request, so the reload returns
2402 once the proxies answer and the UI's spawn listeners track the model loads.
2403 """
2404 if not restarted:
2405 return
2406 threading.Thread(
2407 target=self._preload_roles,
2408 kwargs={"roles": frozenset(restarted)},
2409 name="fleet-reload-warm",
2410 daemon=True,
2411 ).start()
2413 def add_spawn_listener(
2414 self,
2415 *,
2416 on_spawning: Callable[[WorkerRole], None] | None = None,
2417 on_spawned: Callable[[WorkerRole], None] | None = None,
2418 ) -> None:
2419 """Store spawn-lifecycle callbacks; warm-up fires them as each role loads."""
2420 with self._lock:
2421 self._on_spawning = on_spawning
2422 self._on_spawned = on_spawned
2424 def invalidate_load_cache(self, model_path: Path | None = None) -> None:
2425 """A model or settings change restarts the engine: drop the swap."""
2426 del model_path # the whole engine restarts on next use; no per-model scope.
2427 self._shutdown_swap(latch=False)
2429 def drop_loaded_models_async(self) -> None:
2430 """Drop the swap off the caller's thread; next use restarts with current cfg.
2432 ``_shutdown_swap`` stops llama-swap and waits on its process group, so a
2433 role-agnostic load-key change (num_ctx, kv_cache_type) routes here rather
2434 than blocking the settings callback. A no-op when no swap is up.
2435 """
2436 with self._lock:
2437 if not self._swaps:
2438 return
2439 threading.Thread(
2440 target=lambda: self._shutdown_swap(latch=False),
2441 name="fleet-drop",
2442 daemon=True,
2443 ).start()
2445 def shutdown(self) -> None:
2446 self._shutdown_swap()