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