Coverage for src/lilbee/providers/fleet/provider.py: 100%
1255 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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 dataclasses
16import functools
17import json
18import logging
19import re
20import sys
21import threading
22import time
23from contextlib import contextmanager
24from pathlib import Path
25from typing import TYPE_CHECKING, Any, Literal, NamedTuple, NoReturn, TypeVar, overload
27import httpx
29from lilbee.catalog import clean_display_name
30from lilbee.core.config import cfg
31from lilbee.core.health_warnings import HealthWarning
32from lilbee.core.settings import read_persisted
33from lilbee.core.vectors import Vector
34from lilbee.modelhub.registry import ModelRegistry
35from lilbee.providers import engine_params
36from lilbee.providers.base import (
37 GENERATION_RESERVE_TOKENS,
38 ProviderError,
39 ProviderErrorKind,
40 prompt_token_budget,
41)
42from lilbee.providers.fleet import planning
43from lilbee.providers.fleet.binary import engine_pin, resolve_llama_server
44from lilbee.providers.fleet.client import (
45 ChatDeadlineError,
46 LlamaServerClient,
47 is_connection_failure,
48 is_load_capacity_failure,
49 is_rebuildable_failure,
50 retry_on_gateway_error,
51)
52from lilbee.providers.fleet.contract import (
53 NO_SETTINGS,
54 chat_ctx_covers,
55 contract_matches,
56 decoded_launches,
57 vision_slots_cover,
58)
59from lilbee.providers.fleet.groups import SwapGroup, group_for
60from lilbee.providers.fleet.ingest_warmth import ingest_keep_warm
61from lilbee.providers.fleet.launch import InstanceLaunch
62from lilbee.providers.fleet.readback import engine_log_path, writes_engine_log
63from lilbee.providers.fleet.swap_config import cold_load_timeout_s
64from lilbee.providers.fleet.swap_manager import (
65 EngineLifetime,
66 PlanOrigin,
67 SwapManager,
68 SwapState,
69 engine_record_exists,
70 ensure_engine_reaper,
71 find_live_state,
72 reap_stale,
73 state_is_healthy,
74 stop_engine,
75)
76from lilbee.providers.fleet.windowing import window_messages
77from lilbee.providers.model_ref import parse_model_ref
78from lilbee.providers.reload_keys import FLEET_RELOAD_KEYS
79from lilbee.providers.roles import (
80 MODEL_FIELD_TO_ROLE,
81 WorkerRole,
82 configured_model_message,
83)
84from lilbee.providers.warm_progress import (
85 WarmPhase,
86 WarmProgress,
87 WarmProgressTracker,
88 is_active_warm,
89)
90from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError
91from lilbee.runtime.daemon_call import DaemonCall
92from lilbee.runtime.engine_lock import (
93 ENGINE_DIR_ENV,
94 UserLockHold,
95 build_lock,
96 hold_user_lock,
97 keep_warm_requested,
98 kernel_arbitrates_locks,
99 live_users_exist,
100 machine_engine_dir,
101 private_engine_dir,
102 request_keep_warm,
103 withdraw_keep_warm,
104)
106log = logging.getLogger(__name__)
108# How long a shutdown waits for an in-flight build to finish before tearing down
109# regardless. Generous against a legitimate llama-swap spawn (a 30 s boot budget)
110# and far short of any supervisor's patience for a process that will not exit.
111_SHUTDOWN_BUILD_LOCK_WAIT_S = 45.0
113# The thread that releases the engine holds no group uses once a retired client is idle.
114_HOLD_SWEEP_THREAD = "fleet-hold-sweep"
116# How long a request waits for a reload that is restarting its role's server.
117_RELOAD_GAP_WAIT_S = 60.0
119if TYPE_CHECKING:
120 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
122 from lilbee.providers.base import (
123 ChatMessage,
124 ChatResult,
125 ChatStreamItem,
126 ChatToolResult,
127 ClosableIterator,
128 )
130# User-facing name for this engine in error messages.
131_PROVIDER_NAME = "llama-server"
132# No-healthy-replica message for a server that stopped answering at the transport.
133_NOT_RESPONDING_MESSAGE = (
134 "The model server is not responding and no healthy replica is available. "
135 "It may be restarting; try again in a moment."
136)
137# Why a process has no fleet when the shared and the private engine both have a live user.
138_ENGINES_IN_USE_MESSAGE = (
139 "The shared and private engines are in use by other lilbee processes, so this "
140 "process has no model server. Close the other lilbee processes, or set "
141 f"{ENGINE_DIR_ENV} to a separate engine directory, then try again."
142)
143# Tokens held back from the served context for the model's own generation when the
144# request does not cap it, plus a margin for chat-template overhead and estimate drift.
145# Minimal input used to pre-load a role's upstream during warm-up (llama-swap
146# starts an upstream on its first request, so warming issues one cheap call).
147_WARM_PROMPT = "warm"
148_WARM_MAX_TOKENS = 1
149# Read size for paging chat shards into the page cache during warm; large enough
150# to keep sequential reads efficient without holding much resident at once.
151_PREWARM_CHUNK_BYTES = 8 * 1024 * 1024
152# Shards fully paged in this boot, keyed on (path, size, mtime_ns); a fleet
153# rebuild (e.g. a placement change) skips re-reading a hot cache. Module-level so
154# it survives reset_services() replacing the provider instance.
155_PREWARMED_SHARDS: set[tuple[str, int, int]] = set()
156# Per-role client request budget: the first request covers the lazy cold load plus
157# generation, so the weights-scaled cold-load budget plus the margin raises this floor.
158# The embedder generates nothing, so its budget is the cold load plus the margin alone.
159_REQUEST_TIMEOUT_FLOOR_S = 900.0
160_REQUEST_TIMEOUT_GENERATION_MARGIN_S = 120.0
161# How long a peer waits for a live builder of its config root to rebuild on a newer
162# config.toml before it overflows, and how often it looks at the slot meanwhile.
163_OWNER_CATCH_UP_WAIT_S = 30.0
164_OWNER_CATCH_UP_POLL_S = 0.5
165# Jinja chat templates flag tool support by referencing one of these names as an
166# identifier inside a ``{% ... %}`` / ``{{ ... }}`` block (not free-text prose).
167# The server parses tool calls natively via ``--jinja``; this probe only decides
168# whether to offer tools to a given model at all.
169_TOOL_TEMPLATE_PATTERN = re.compile(r"\{[%{][^}]*\b(?:tools|tool_calls|functions|function_calls)\b")
170# Attempt cap for the gateway retry only when a page has no deadline (ocr_timeout=0,
171# "no limit"): it backstops the retry so a replica that never comes back can't spin
172# forever. A page with a deadline retries until that deadline instead (see
173# _ocr_dispatch), so the count doesn't bound the common case.
174_VISION_GATEWAY_RETRIES = 18
175# llama.cpp's repeat penalty on every vision OCR request: it damps tokens seen in
176# the last window, which stops a page looping one line to the token cap.
177_VISION_REPEAT_PENALTY = 1.1
178_VISION_REPEAT_LAST_N = 64
179# How often a waiter blocked on full replicas re-polls their health: an
180# unhealthy replica re-admits itself by cool-down expiry, which notifies nobody.
181_DISPATCH_HEALTH_RECHECK_S = 0.5
182_T = TypeVar("_T")
185def _prewarm_key(shard: Path) -> tuple[str, int, int]:
186 """The prewarm identity of *shard*: same path, size, and mtime -> same pages."""
187 stat = shard.stat()
188 return (str(shard), stat.st_size, stat.st_mtime_ns)
191def _request_timeout_s(launch: InstanceLaunch) -> float:
192 """Per-client request budget: the role's cold load plus margin, floored when it generates."""
193 budget = (
194 cold_load_timeout_s(launch.weights_bytes, launch.role)
195 + _REQUEST_TIMEOUT_GENERATION_MARGIN_S
196 )
197 if launch.role is WorkerRole.EMBED:
198 return budget
199 return max(_REQUEST_TIMEOUT_FLOOR_S, budget)
202def _launches_by_group(
203 plan: planning.FleetPlan,
204) -> dict[SwapGroup, tuple[InstanceLaunch, ...]]:
205 """Group a plan's launches by swap group, replica order preserved within each group.
207 Co-tenant roles land in one group, so llama-swap evicts between them rather
208 than holding both resident.
209 """
210 grouped: dict[SwapGroup, list[InstanceLaunch]] = {}
211 for launch in plan.launches:
212 grouped.setdefault(group_for(launch.role, plan.co_tenants), []).append(launch)
213 return {group: tuple(group_launches) for group, group_launches in grouped.items()}
216def _by_role(launches: list[InstanceLaunch]) -> dict[WorkerRole, list[InstanceLaunch]]:
217 """Split one group's launches per role, replica order preserved."""
218 grouped: dict[WorkerRole, list[InstanceLaunch]] = {}
219 for launch in launches:
220 grouped.setdefault(launch.role, []).append(launch)
221 return grouped
224def _least_in_flight(clients: list[LlamaServerClient]) -> LlamaServerClient:
225 """Pick the healthy client with the fewest in-flight requests.
227 Falls back to the full pool when every client is marked unhealthy, so a
228 fully-dead pool still gets a call (which surfaces the error and lets a
229 recovered replica mark itself healthy again).
230 """
231 healthy = [client for client in clients if client.healthy]
232 return min(healthy or clients, key=lambda c: c.in_flight)
235# Serializes pick-and-reserve so concurrent routers see each other's assignment.
236# Held only for the O(replicas) selection, never across the request itself.
237_ROUTE_LOCK = threading.Lock()
240def _reserve_least_in_flight(clients: list[LlamaServerClient]) -> LlamaServerClient:
241 """Atomically pick the least-loaded healthy client and reserve a slot on it.
243 Selection and reservation are one critical section: without it, concurrent
244 callers all read the same idlest replica before any of them increments its
245 counter and route there together (a thundering herd that starves the rest of
246 the fleet). The caller must :meth:`~LlamaServerClient.release` the slot.
247 """
248 with _ROUTE_LOCK:
249 client = _least_in_flight(clients)
250 client.reserve()
251 return client
254def _launch_changes(old: tuple[InstanceLaunch, ...], new: tuple[InstanceLaunch, ...]) -> str:
255 """The launch fields that differ, as ``field before -> after``."""
256 if len(old) != len(new):
257 return f"replicas {len(old)} -> {len(new)}"
258 moved = [
259 f"{field.name} {getattr(before, field.name)!r} -> {getattr(after, field.name)!r}"
260 for before, after in zip(old, new, strict=True)
261 for field in dataclasses.fields(InstanceLaunch)
262 if getattr(before, field.name) != getattr(after, field.name)
263 ]
264 return ", ".join(moved) or "no launch field moved"
267def _healthy_states(engine_dir: Path) -> dict[SwapGroup, SwapState]:
268 """One probe pass over *engine_dir*: the recorded, answering group states.
270 The ladder's single view of a dir. Bind eligibility and the replaceability
271 check both read this snapshot, so they cannot disagree about an engine that
272 died between them, and one wedged proxy port is paid for once per ladder
273 pass rather than once per decision -- all of it under the build lock, which
274 every other lilbee start is waiting on.
275 """
276 found: dict[SwapGroup, SwapState] = {}
277 for group in SwapGroup:
278 state = find_live_state(engine_dir, group)
279 if state is not None and state_is_healthy(state):
280 found[group] = state
281 return found
284class _SlotOwed(NamedTuple):
285 """A ladder pass found the slot's builder owes config.toml a rebuild; *states* is the slot."""
287 states: dict[SwapGroup, SwapState]
290def _plan_origin() -> PlanOrigin:
291 """This config root and what its config.toml sets for a fleet plan, as a state records them.
293 Only the keys whose change makes the fleet owner rebuild count, so a write of any
294 other key never makes a fleet look older than the file. A torn or unreadable file
295 gives no view.
296 """
297 _stamp, values = read_persisted(cfg.data_root)
298 view = None
299 if values is not None:
300 picked = {key: value for key, value in values.items() if key in FLEET_RELOAD_KEYS}
301 view = json.dumps(picked, sort_keys=True, default=str)
302 return PlanOrigin(str(cfg.data_root), view)
305def _owner_owes_the_file(states: dict[SwapGroup, SwapState]) -> bool:
306 """True when a fleet of this config root was planned before config.toml's last plan write.
308 Its builder has still to take the write and rebuild, so a peer on the new value waits
309 for that rather than load a second fleet beside it.
310 """
311 now = _plan_origin()
312 return now.fleet_view is not None and any(
313 state.origin is not None
314 and state.origin.config_root == now.config_root
315 and state.origin.fleet_view not in (None, now.fleet_view)
316 for state in states.values()
317 )
320def _await_slot_change(
321 engine_dir: Path, before: dict[SwapGroup, SwapState], deadline: float
322) -> None:
323 """Wait, holding no lock, until *engine_dir*'s fleet changes or loses its users.
325 The wait ends at *deadline* (a ``time.monotonic`` value) at the latest.
326 """
327 while time.monotonic() < deadline:
328 time.sleep(_OWNER_CATCH_UP_POLL_S)
329 if not live_users_exist(engine_dir, exclude_self=True):
330 return
331 if _healthy_states(engine_dir) != before:
332 return
335def _bindable_group(
336 state: SwapState, pin: str, demand: _EngineDemand
337) -> tuple[SwapState, list[InstanceLaunch], set[tuple[WorkerRole, str]]] | None:
338 """*state*'s launches and the wanted pairs it covers, or ``None``.
340 ``None`` for every reason an already-healthy group is not bindable by us:
341 a foreign pin, other settings, an undecodable contract, or serving nothing
342 we want.
343 """
344 if not contract_matches(state, (), pin, demand.settings):
345 return None
346 launches = decoded_launches(state)
347 if launches is None:
348 return None
349 pairs = {(launch.role, launch.model) for launch in launches} & demand.pairs
350 return (state, launches, pairs) if pairs else None
353def _log_adopted_launches(
354 candidates: list[tuple[SwapGroup, SwapState, list[InstanceLaunch]]],
355) -> None:
356 """Name the engine every bound instance now runs on, and the pid that owns it."""
357 for _group, state, launches in candidates:
358 for launch in launches:
359 planning.log_engine_launch(launch, owner_pid=state.pid)
362class _PrimedStream:
363 """A stream re-fronted with its eagerly-pulled first frame.
365 close() always reaches the source stream, even before any iteration, so a
366 caller that truncates immediately still releases the fleet's in-flight
367 request slot (an unstarted chaining generator would silently drop it).
368 """
370 def __init__(self, first: ChatStreamItem, source: ClosableIterator[ChatStreamItem]) -> None:
371 self._first: list[ChatStreamItem] = [first]
372 self._source = source
374 def __iter__(self) -> _PrimedStream:
375 return self
377 def __next__(self) -> ChatStreamItem:
378 if self._first:
379 return self._first.pop()
380 return next(self._source)
382 def close(self) -> None:
383 self._source.close()
386def _primed_stream(items: ClosableIterator[ChatStreamItem]) -> ClosableIterator[ChatStreamItem]:
387 """Pull the first frame of *items* now, so a dead engine raises to the caller.
389 The stream connects lazily on first iteration; without priming, a proxy
390 that died raises only inside the consumer's loop, past any rediscovery.
391 """
392 try:
393 first = next(items)
394 except StopIteration:
395 return items # already exhausted; still closable
396 return _PrimedStream(first, items)
399def _call_with_failover(
400 clients: list[LlamaServerClient],
401 call: Callable[[LlamaServerClient], _T],
402) -> _T:
403 """Run *call* on the least-busy healthy client, retrying once on another replica.
405 The client is reserved at selection so concurrent ingest threads spread
406 across replicas. A connection-level failure marks the client unhealthy and
407 retries once on a different replica; with no other replica the failure
408 surfaces. The reservation is released once the call resolves.
409 """
410 client = _reserve_least_in_flight(clients)
411 try:
412 result = call(client)
413 except Exception as exc:
414 if not is_connection_failure(exc):
415 raise
416 client.mark_unhealthy()
417 return _retry_on_other_replica(clients, client, call, exc)
418 else:
419 client.mark_healthy()
420 return result
421 finally:
422 client.release()
425def _retry_on_other_replica(
426 clients: list[LlamaServerClient],
427 failed: LlamaServerClient,
428 call: Callable[[LlamaServerClient], _T],
429 cause: Exception,
430) -> _T:
431 """Retry *call* once on a replica other than *failed*, marking its health."""
432 others = [c for c in clients if c is not failed]
433 if not others:
434 _raise_no_healthy_replica(cause)
435 retry = _reserve_least_in_flight(others)
436 try:
437 retry_result = call(retry)
438 except Exception as retry_exc:
439 if is_connection_failure(retry_exc):
440 retry.mark_unhealthy()
441 raise
442 else:
443 retry.mark_healthy()
444 return retry_result
445 finally:
446 retry.release()
449def _engine_log_of(launch: InstanceLaunch, log_dir: Path) -> Path | None:
450 """The file *launch*'s engine writes its own output to, or None when it writes none."""
451 return engine_log_path(log_dir, launch.model_id) if writes_engine_log(launch.argv) else None
454def _raise_no_healthy_replica(cause: Exception) -> NoReturn:
455 """Raise the user-facing error for a call with no healthy replica left to retry on.
457 A ``ProviderError`` cause is the client's own plain line for an engine that
458 died (model name, cause, engine log), so it surfaces unchanged. A transport
459 error has no such line and reads as an unresponsive server.
460 """
461 if isinstance(cause, ProviderError):
462 raise cause
463 raise ProviderError(
464 _NOT_RESPONDING_MESSAGE, provider=_PROVIDER_NAME, kind=ProviderErrorKind.CONNECTION
465 ) from cause
468# Env vars a launch pins its devices with, one per backend (Metal has none).
469_VISIBLE_DEVICE_ENV_VARS = (
470 "CUDA_VISIBLE_DEVICES",
471 "ROCR_VISIBLE_DEVICES",
472 "HIP_VISIBLE_DEVICES",
473 "GGML_VK_VISIBLE_DEVICES",
474 "ONEAPI_DEVICE_SELECTOR",
475)
478def _role_device_sets(
479 launches: Iterable[InstanceLaunch],
480) -> dict[WorkerRole, frozenset[str]]:
481 """Backend-qualified device tokens each role's launches pin, by role.
483 A role's set is the union across its replicas. Roles whose launches carry
484 no visibility env (Metal, or an unpinned backend) are absent: without
485 pinning there is no proof of sharing, so they keep the concurrent warm.
486 """
487 sets: dict[WorkerRole, set[str]] = {}
488 for launch in launches:
489 for var in _VISIBLE_DEVICE_ENV_VARS:
490 value = launch.env_overrides.get(var)
491 if value:
492 sets.setdefault(launch.role, set()).update(
493 f"{var}={part.strip()}" for part in value.split(",")
494 )
495 return {role: frozenset(tokens) for role, tokens in sets.items()}
498def _warm_chains(
499 warm_roles: list[WorkerRole], device_sets: dict[WorkerRole, frozenset[str]]
500) -> list[list[WorkerRole]]:
501 """Group *warm_roles* into chains warmed sequentially; chains run in parallel.
503 Roles with overlapping device sets land in one chain, merged transitively.
504 Within a chain chat goes last: it sizes its KV against the headroom the
505 settled residents leave, so it must not race their loads. A role with no
506 device set shares nothing provable and gets its own chain.
507 """
508 chains: list[tuple[set[str], list[WorkerRole]]] = []
509 ordered = sorted(warm_roles, key=lambda r: (r is WorkerRole.CHAT, list(WorkerRole).index(r)))
510 for role in ordered:
511 tokens = device_sets.get(role)
512 if not tokens:
513 chains.append((set(), [role]))
514 continue
515 merged_tokens, merged_roles = set(tokens), [role]
516 kept: list[tuple[set[str], list[WorkerRole]]] = []
517 for chain_tokens, chain_roles in chains:
518 if chain_tokens & merged_tokens:
519 merged_tokens |= chain_tokens
520 merged_roles = chain_roles + merged_roles
521 else:
522 kept.append((chain_tokens, chain_roles))
523 kept.append((merged_tokens, merged_roles))
524 chains = kept
525 return [roles for _tokens, roles in chains]
528def _warm_role(role: WorkerRole, client: LlamaServerClient) -> None:
529 """Send the cheapest request that loads *role*'s upstream behind llama-swap.
531 Vision is skipped (its load is heavy and it warms on the first OCR); chat,
532 embed, and rerank each issue a minimal call to trigger the upstream start.
533 """
534 if role is WorkerRole.CHAT:
535 client.chat(
536 [{"role": "user", "content": _WARM_PROMPT}],
537 options={"max_tokens": _WARM_MAX_TOKENS},
538 stream=False,
539 )
540 elif role is WorkerRole.EMBED:
541 client.embed([_WARM_PROMPT])
542 elif role is WorkerRole.RERANK:
543 client.rerank(_WARM_PROMPT, [_WARM_PROMPT])
546@functools.lru_cache(maxsize=32)
547def _supports_tools_cached(path_str: str, _mtime_ns: int) -> bool:
548 """Memoised tool-template probe keyed on the GGUF's path + mtime.
550 The mtime arg participates in the cache key only; a re-quantised file at the
551 same path invalidates automatically because its mtime changes.
552 """
553 from lilbee.providers.gguf_meta import read_gguf_metadata
555 meta = read_gguf_metadata(Path(path_str))
556 if not isinstance(meta, dict):
557 return False
558 template = meta.get("chat_template")
559 if not isinstance(template, str):
560 return False
561 return _TOOL_TEMPLATE_PATTERN.search(template) is not None
564class _VisionReplica(NamedTuple):
565 """One vision server paired with its fitted ``--parallel`` slot count."""
567 client: LlamaServerClient
568 slots: int
571class _PageBudgetExhausted(Exception): # noqa: N818 - internal control flow, not an error API
572 """A page's document-wide OCR budget ran out before its slot came up."""
575class _VisionDispatcher:
576 """Process-wide per-replica slot assignment for vision requests.
578 The ingest file fan-out runs many OCR requests at once; each request is
579 assigned one specific replica and only while that replica has a free
580 continuous-batching slot, so lilbee's own traffic never queues past a
581 vision server's slots (an aggregate cap plus racy least-busy routing
582 can). Requests past capacity wait in-process until any usable replica
583 frees a slot; unhealthy replicas take no new work until their half-open
584 cool-down re-admits them.
585 """
587 def __init__(self) -> None:
588 self._cond = threading.Condition()
589 self._assigned: dict[LlamaServerClient, int] = {}
591 @contextmanager
592 def slot(
593 self, pool: Sequence[_VisionReplica], cancel: CancelSignal | None = None
594 ) -> Iterator[LlamaServerClient]:
595 """Hold one batching slot on the pool's best replica; yields that client.
597 A set *cancel* raises ``TaskCancelledError`` while the request still waits.
598 """
599 client = self._acquire(pool, cancel)
600 try:
601 yield client
602 finally:
603 self._release(client)
605 def _acquire(
606 self, pool: Sequence[_VisionReplica], cancel: CancelSignal | None
607 ) -> LlamaServerClient:
608 with self._cond:
609 while True:
610 if cancel is not None and cancel.is_set():
611 raise TaskCancelledError
612 client = self._pick(pool)
613 if client is not None:
614 self._assigned[client] = self._assigned.get(client, 0) + 1
615 return client
616 # The timed wait re-polls health: a replica can become routable
617 # again by cool-down expiry alone, which notifies no waiter.
618 self._cond.wait(timeout=_DISPATCH_HEALTH_RECHECK_S)
620 def _pick(self, pool: Sequence[_VisionReplica]) -> LlamaServerClient | None:
621 """The usable replica with the most free slots, or None while all are full.
623 Falls back to the full pool when every replica is unhealthy (mirrors
624 ``_least_in_flight``), so a dead pool surfaces the error instead of
625 queueing forever.
626 """
627 usable = [replica for replica in pool if replica.client.healthy] or list(pool)
628 best = max(usable, key=self._free_slots)
629 return best.client if self._free_slots(best) > 0 else None
631 def _free_slots(self, replica: _VisionReplica) -> int:
632 return replica.slots - self._assigned.get(replica.client, 0)
634 def _release(self, client: LlamaServerClient) -> None:
635 with self._cond:
636 remaining = self._assigned.get(client, 0) - 1
637 if remaining <= 0:
638 self._assigned.pop(client, None)
639 else:
640 self._assigned[client] = remaining
641 self._cond.notify_all()
644_VISION_DISPATCHER = _VisionDispatcher()
647def _dispatch_vision(
648 pool: Sequence[_VisionReplica],
649 call: Callable[[LlamaServerClient], _T],
650 cancel: CancelSignal | None = None,
651) -> _T:
652 """Run *call* on a replica with a free batching slot, failing over once.
654 Blocks until a slot frees rather than racing requests at a full server. A
655 connection-level failure marks the replica unhealthy and retries once on
656 another replica's slot; with no other replica the failure surfaces.
657 """
658 with _VISION_DISPATCHER.slot(pool, cancel) as client:
659 try:
660 result = call(client)
661 except Exception as exc:
662 if not is_connection_failure(exc):
663 raise
664 client.mark_unhealthy()
665 failed, cause = client, exc
666 else:
667 client.mark_healthy()
668 return result
669 others = [replica for replica in pool if replica.client is not failed]
670 if not others:
671 _raise_no_healthy_replica(cause)
672 with _VISION_DISPATCHER.slot(others, cancel) as retry_client:
673 try:
674 retry_result = call(retry_client)
675 except Exception as retry_exc:
676 if is_connection_failure(retry_exc):
677 retry_client.mark_unhealthy()
678 raise
679 retry_client.mark_healthy()
680 return retry_result
683def _vision_call(
684 client: LlamaServerClient,
685 messages: Sequence[Mapping[str, Any]],
686 timeout: float | None,
687 cancel: CancelSignal | None = None,
688) -> str:
689 """Run a vision chat on *client*, bounded by a total *timeout* (``None`` or 0: no limit).
691 The repeat penalty stops a page looping one line; ``cfg.vision_ocr_max_tokens``
692 caps generation if a loop still escapes it. A timeout surfaces as a
693 ``ProviderError`` so the page-level OCR caller can fail just that page.
694 Callers hold a dispatcher slot, so queue time isn't billed against the timeout.
695 A set *cancel* aborts the request at once, before or after its first token.
696 """
697 options = {
698 "max_tokens": cfg.vision_ocr_max_tokens,
699 "repeat_penalty": _VISION_REPEAT_PENALTY,
700 "repeat_last_n": _VISION_REPEAT_LAST_N,
701 }
702 deadline_s = timeout if timeout and timeout > 0 else None
703 try:
704 return client.chat_abortable(
705 messages, options=options, deadline_s=deadline_s, cancel=cancel
706 )
707 except ChatDeadlineError:
708 raise ProviderError(
709 f"Vision OCR timed out after {timeout:g}s.",
710 provider=_PROVIDER_NAME,
711 ) from None
714def _ocr_dispatch(
715 pool: Sequence[_VisionReplica],
716 messages: Sequence[Mapping[str, Any]],
717 deadline: float | None,
718 cancel: CancelSignal | None = None,
719) -> str:
720 """OCR *messages* on a free replica slot, retrying gateway errors until *deadline*.
722 Backpressure (the dispatcher blocking until a slot frees) keeps lilbee's own
723 traffic within the slots; a gateway error is a replica restarting mid-run,
724 and a 429 is foreign overload that is reported, not re-sent. The retry is
725 deadline-bound rather than attempt-bound so a page keeps riding out a
726 restart until its own budget passes instead of dropping after a fixed count. Each attempt
727 is bounded by the budget remaining before *deadline*; an exhausted budget
728 raises :class:`_PageBudgetExhausted`. A ``None`` deadline (no limit) falls
729 back to a bounded attempt count so the retry can't spin forever.
730 """
732 def _attempt(client: LlamaServerClient) -> str:
733 remaining = max(0.0, deadline - time.monotonic()) if deadline is not None else None
734 if remaining == 0.0:
735 raise _PageBudgetExhausted
736 return _vision_call(client, messages, remaining, cancel)
738 return retry_on_gateway_error(
739 lambda: _dispatch_vision(pool, _attempt, cancel),
740 retries=_VISION_GATEWAY_RETRIES,
741 deadline=deadline,
742 )
745def _pdf_drain_budget(total_pages: int, per_page_timeout_s: float | None) -> float | None:
746 """Total OCR wall-clock budget = pages*per_page + load grace, or None for no cap.
748 Mirrors the in-process drain budget: one document-wide deadline rather than a
749 per-page cap, so a slow page borrows from fast ones and the vision model's cold
750 first-inference is absorbed by the grace instead of tripping a fixed page limit.
751 """
752 from lilbee.core.config import cfg
754 if not per_page_timeout_s or per_page_timeout_s <= 0:
755 return None
756 return total_pages * per_page_timeout_s + cfg.vision_load_budget_s
759def _ocr_deadline(per_page_timeout_s: float | None) -> float | None:
760 """Absolute monotonic deadline for one image OCR, or None when uncapped.
762 An image is a one-page document, so it gets the same budget as a PDF page:
763 the per-page timeout plus the cold-load grace, spanning queue wait and
764 generation together.
765 """
766 budget = _pdf_drain_budget(1, per_page_timeout_s)
767 return None if budget is None else time.monotonic() + budget
770_ROLE_TO_MODEL_FIELD = {role: field for field, role in MODEL_FIELD_TO_ROLE.items()}
773def _configured_model_for(role: WorkerRole) -> str:
774 """The cfg model ref for *role*, empty when the role is unset."""
775 field = _ROLE_TO_MODEL_FIELD.get(role)
776 return getattr(cfg, field) or "" if field else ""
779def _unusable_engine_reason() -> str | None:
780 """Why no server can start on this host, or None once an engine resolves.
782 Planning drops an engine-less host to serving nothing and says so only at
783 debug, so by the time a surface has an empty pool the engine is the one cause
784 it cannot see. Re-resolving here is also what keeps an engine installed
785 mid-session from being reported as still missing.
786 """
787 try:
788 resolve_llama_server()
789 except ProviderError as exc:
790 return str(exc)
791 return None
794def _chat_needs_local_engine() -> bool:
795 """Whether the configured chat model is one this host has to serve itself.
797 A chat ref routed to an SDK backend runs without any local engine, so a
798 missing one is not its failure and must not be stamped on its warm.
799 """
800 ref = _configured_model_for(WorkerRole.CHAT)
801 return bool(ref) and not parse_model_ref(ref).is_remote
804def _no_server_message(role: WorkerRole) -> str:
805 """User-facing reason *role* has no server, engine state first.
807 A missing engine and a model that never placed both arrive as an empty pool,
808 and reading the second onto the first sends the reader to a model
809 configuration that is already correct.
810 """
811 reason = _unusable_engine_reason()
812 if reason is not None:
813 return f"No {role.value} model server is running: {reason}"
814 ref = _configured_model_for(role)
815 if role is WorkerRole.VISION and ref and planning.vision_mmproj(ref) is None:
816 return (
817 f"No vision model server is running: {ref} has no mmproj (CLIP projector) "
818 "file beside it. Put the projector next to the model file, then try again."
819 )
820 return (
821 f"No {role.value} model server is running. Make sure the {role.value} "
822 "model is installed and configured, then try again."
823 )
826def _reload_timeout_message(role: WorkerRole) -> str:
827 """User-facing reason a request gave up waiting for *role*'s restarting server."""
828 return (
829 f"The {role.value} model server is still restarting after a settings change. "
830 "Wait for the restart to finish, then try again."
831 )
834class _EngineDemand(NamedTuple):
835 """What this process needs an engine to serve: pairs plus its chat window."""
837 pairs: set[tuple[WorkerRole, str]]
838 # Per-slot chat tokens this process needs; 0 demands nothing.
839 chat_ctx: int
840 # Configured roles the plan skipped because their model is not installed.
841 # Carried out of the demand plan so the warm tracker can name the missing
842 # model even when the ladder never reaches _plan_and_spawn (zero installed
843 # models fail _can_build_engine first) or binds an existing engine.
844 skipped_not_installed: dict[WorkerRole, str]
845 # Launches the demand plan refused for an unusable window (role -> reason);
846 # recorded even when the ladder binds an engine or never builds one.
847 skipped_unusable_ctx: dict[WorkerRole, str]
848 # Vision slots this process needs at once; 0 demands nothing.
849 vision_slots: int = 0
850 # The settings the demand plan gives each role's launch.
851 settings: Mapping[WorkerRole, str] = NO_SETTINGS
854def _placeable_demand() -> _EngineDemand:
855 """Configured (role, model) pairs a fresh plan would serve, and the chat window.
857 A configured role is wanted only when the planner would place it. The plan
858 is the co-placement authority: a role that fits alone but cannot co-tenant (a
859 unified-memory box past its budget) gets no launch, so it is dropped here too,
860 and bind matches a running engine instead of judging it a partial cover and
861 restarting the shared engine on every process start. The per-role check stays
862 as the cheap gate for the reasons in its own docstring. Empty when no engine
863 binary resolves: nothing is placeable, so the ladder serves nothing.
864 """
865 from lilbee.providers.fleet.planning import (
866 placeable_total_vram,
867 plan_all_launches,
868 role_model_placeable,
869 )
871 try:
872 plan = plan_all_launches()
873 except ProviderError as exc:
874 if exc.kind is ProviderErrorKind.NOT_FOUND:
875 return _EngineDemand(set(), 0, {}, {})
876 raise
877 placed = {launch.role for launch in plan.launches}
878 total_vram = placeable_total_vram()
879 pairs = {
880 (role, model)
881 for role in WorkerRole
882 if role in placed
883 and (model := _configured_model_for(role))
884 and role_model_placeable(role, model, total_vram)
885 }
886 return _EngineDemand(
887 pairs,
888 _demanded_chat_ctx(plan.launches, pairs),
889 dict(plan.skipped_not_installed),
890 dict(plan.skipped_unusable_ctx),
891 _demanded_vision_slots(pairs),
892 {launch.role: launch.settings for launch in plan.launches},
893 )
896def _demanded_vision_slots(pairs: set[tuple[WorkerRole, str]]) -> int:
897 """Vision slots this process needs an engine to serve at once; 0 for none."""
898 if not any(role is WorkerRole.VISION for role, _ in pairs):
899 return 0
900 return max(1, cfg.vision_ocr_concurrency)
903def _demanded_chat_ctx(
904 launches: Iterable[InstanceLaunch], pairs: set[tuple[WorkerRole, str]]
905) -> int:
906 """Per-slot chat window this process needs an engine to serve; 0 for none.
908 The cfg target (a ``num_ctx`` pin, else ``chat_n_ctx_target``) capped by
909 this process's own planned chat window: a window the plan itself cannot
910 reach (model ceiling, hardware) is not a demand a rebuild could satisfy,
911 so capping keeps the fit check from rebuilding the engine in a loop.
913 The cap applies only to a single-device chat plan, whose window is sized
914 against device totals and holds regardless of what is resident. A
915 tensor-split plan is sized against live free VRAM, which a resident
916 incumbent deflates; capping by it would shrink the demand to whatever the
917 incumbent left free and let the fit check pass vacuously.
918 """
919 if not any(role is WorkerRole.CHAT for role, _model in pairs):
920 return 0
921 chat_launches = [launch for launch in launches if launch.role is WorkerRole.CHAT]
922 planned = max((launch.ctx for launch in chat_launches), default=0)
923 if planned <= 0:
924 return 0
925 # Always positive: num_ctx validates ge=1 and chat_n_ctx_target ge=512.
926 target = cfg.num_ctx if cfg.num_ctx is not None else cfg.chat_n_ctx_target
927 split = any(len(launch.est_vram_by_device) > 1 for launch in chat_launches)
928 return target if split else min(target, planned)
931def _can_build_engine(wanted: set[tuple[WorkerRole, str]]) -> bool:
932 """Preconditions for a viable build, checked before stopping a warm engine.
934 A process that can serve nothing (no placeable model, an unresolvable engine
935 binary) must not stop an engine another setup left warm and then spawn nothing.
936 Probing the engine here resolves the binary AND enumerates devices, so a wedged
937 GPU probe or an unusable CUDA runtime raises loud at this point -- before the
938 caller stops a replaceable incumbent. Were the probe left to run only inside
939 ``_plan_and_spawn`` (after the stop), that raise would kill an engine other
940 members still hold and then skip the overflow build, leaving zero engines. This
941 takes no memory snapshot (device enumeration reads no residency); the clean-box
942 sizing snapshot is captured by ``_plan_and_spawn`` after the stop.
943 """
944 from lilbee.providers.fleet import planning
946 if not wanted:
947 return False
948 try:
949 planning.assert_engine_probeable()
950 except ProviderError as exc:
951 # A genuinely-missing engine binary keeps the quiet serve-nothing path;
952 # every other probe failure must propagate (fail loud) rather than be read
953 # as "cannot build" and silently stand down.
954 if exc.kind is not ProviderErrorKind.NOT_FOUND:
955 raise
956 return False
957 except OSError:
958 return False
959 return True
962def _engine_lifetime(engine_dir: Path) -> EngineLifetime:
963 """What ends an engine this process builds in *engine_dir*.
965 A keep-warm engine outlives lilbee; its reaper only unloads its weights.
966 Otherwise shared ownership needs kernel-arbitrated user locks; without them
967 the engine dies with its builder.
968 """
969 if cfg.keep_engine_warm:
970 return EngineLifetime.DETACHED
971 if kernel_arbitrates_locks(engine_dir):
972 return EngineLifetime.USERS
973 return EngineLifetime.BUILDER
976def _clear_engine_dir(engine_dir: Path) -> None:
977 """Reap dead leftovers in *engine_dir* and stop its recorded engine; caller holds its lock."""
978 reap_stale(engine_dir)
979 if engine_record_exists(engine_dir):
980 stop_engine(engine_dir)
983def _watch_bound_engine(engine_dir: Path) -> None:
984 """Make sure a reaper watches the engine this process just bound in *engine_dir*.
986 A reaper can die (killed, or a failed start) while its engine serves. Every
987 bind repairs that, so the users of a live engine leave a reaper behind them:
988 it stops a shared engine, or unloads a kept-warm one's weights after the idle time.
989 """
990 if not kernel_arbitrates_locks(engine_dir):
991 return
992 if not ensure_engine_reaper(engine_dir):
993 log.warning(
994 "No reaper watches the engine at %s; it stops at a clean last exit only.", engine_dir
995 )
998def _warm_ttl_seconds(
999 lifetime: EngineLifetime, engine_dir: Path, *, hold_warm_for_session: bool = False
1000) -> int:
1001 """llama-swap idle-unload timer in seconds for the fleet spawned in *engine_dir*.
1003 A ttl of 0 keeps weights resident until the engine is stopped; otherwise an
1004 idle engine releases its weights after ``engine_idle_ttl_minutes`` and reloads
1005 transparently on the next prompt. The timer is held off (ttl 0) for an engine
1006 bound to its builder (*lifetime* ``BUILDER``) while the builder depends on an
1007 instant response: an interactive session (*hold_warm_for_session*), or a bulk
1008 ingest whose unevenly loaded replicas must not idle-unload and reload cold.
1010 An engine that outlives its builder takes no such hold: llama-swap's ttl is
1011 fixed at launch, so a hold baked into it would keep the weights loaded after
1012 the session or ingest ended. A kept-warm engine (this process's setting or
1013 any user's marker) also gets ttl 0 when its users can be counted and a reaper
1014 is confirmed running, because that reaper unloads the weights after the last
1015 user leaves. Without a confirmed reaper the idle ttl stays.
1016 """
1017 if lifetime is EngineLifetime.BUILDER and (hold_warm_for_session or ingest_keep_warm()):
1018 return 0
1019 if _reaper_unloads_weights(engine_dir) and ensure_engine_reaper(engine_dir):
1020 return 0
1021 return cfg.engine_idle_ttl_minutes * 60
1024def _reaper_unloads_weights(engine_dir: Path) -> bool:
1025 """Whether a reaper, not llama-swap's timer, unloads the weights of *engine_dir*'s engine."""
1026 kept_warm = cfg.keep_engine_warm or keep_warm_requested(engine_dir)
1027 return kept_warm and kernel_arbitrates_locks(engine_dir)
1030class FleetProvider:
1031 """Routes every role to the managed llama-server fleet (a fleet-of-one on one box)."""
1033 def __init__(
1034 self, *, hold_warm: bool = False, config_refresh: Callable[[], None] | None = None
1035 ) -> None:
1036 # Takes config.toml changes now. The process that follows the file injects it,
1037 # so a retry plans on the settings an owner restarted its fleet with.
1038 self._config_refresh = config_refresh
1039 # An interactive session (the TUI) owns this process for its whole
1040 # lifetime, so a fleet bound to it stays resident instead of idle-unloading
1041 # under a user who is still in the app; closing lilbee releases it. A
1042 # fleet that outlives the session (shared or keep-warm) takes no hold, so
1043 # _warm_ttl_seconds ignores it there. Set by the container that built this provider, never
1044 # mutated afterwards.
1045 self._hold_warm_for_session = hold_warm
1046 # One llama-swap per placed group, so restarting one group's servers (a
1047 # placement or per-role model change) never unloads another group's. A
1048 # co-tenant group holds chat and vision, which evict each other on load.
1049 self._swaps: dict[SwapGroup, SwapManager] = {}
1050 # The group each placed role runs in, so a role's clients and its swap
1051 # process can be reached from the role alone.
1052 self._role_group: dict[WorkerRole, SwapGroup] = {}
1053 # The launches each running group was started with, kept so a reload can
1054 # diff the fresh plan against what is running and restart only the groups
1055 # whose launches actually changed. Launch argv is port-free (ports are
1056 # injected at config render), so the comparison is stable across starts.
1057 self._launches: dict[SwapGroup, tuple[InstanceLaunch, ...]] = {}
1058 # Engine dirs this provider holds membership in (machine slot and/or
1059 # the private overflow), and the dir each running group lives in.
1060 self._engine_holds: dict[Path, UserLockHold] = {}
1061 self._group_dirs: dict[SwapGroup, Path] = {}
1062 # Latched once shutdown runs. A discarded provider (reset_services swaps
1063 # in a new one) can still have an in-flight warm-up or reload daemon
1064 # thread; without this latch that thread could start a llama-swap after
1065 # shutdown already ran, leaving a process no live provider owns.
1066 # _ensure_fleet checks it under the build lock so a post-shutdown build is
1067 # refused (the swap_manager reaper is the backstop if one slips through).
1068 self._shut_down = False
1069 # A pool of OpenAI clients per placed role (one per data-parallel replica),
1070 # all pointed at the llama-swap endpoint and routed by replica model id;
1071 # rebuilt whenever the swap process (re)starts. Requests round-robin the pool.
1072 self._clients: dict[WorkerRole, list[LlamaServerClient]] = {}
1073 # Clients retired by a reload, awaiting close. A reload's old clients may
1074 # still be held by an in-flight reader, so they are closed at the *next*
1075 # reload (by when those readers have finished) or at shutdown, never while
1076 # potentially in use. See _retire_clients.
1077 self._retiring_clients: list[LlamaServerClient] = []
1078 # Rediscover retries between their drop and their end; the hold sweep waits for zero.
1079 self._open_rediscovers = 0
1080 # Chat batching slots and per-slot context from the chat launch, surfaced to
1081 # the concurrency gate and clients; defaults until the chat group is up.
1082 self._chat_slots = 1
1083 self._chat_ctx: int | None = None
1084 # Latest chat prefill progress reported by a streaming client, cleared
1085 # by the same stream when generation starts or the stream ends.
1086 self._chat_prefill: tuple[int, int] | None = None
1087 # Single-flight guard: the HTTP/MCP servers route concurrently, so two
1088 # first-requests must not each start a swap (double GPU allocation) or
1089 # tear one down mid-route.
1090 self._lock = threading.RLock()
1091 # Serializes the slow startup (GPU probe + GGUF parse + llama-swap spawn)
1092 # across concurrent callers, so the off-thread warm-up and an on-demand call
1093 # can't start two swaps. Held only during startup, NOT while routing.
1094 self._build_lock = threading.Lock()
1095 # Single-flight for starting vision on request: concurrent OCR pages wait
1096 # for the first page's re-plan instead of each re-planning.
1097 self._vision_request_lock = threading.Lock()
1098 # Spawn-lifecycle listeners (set by the TUI via add_spawn_listener). Stored
1099 # so warm-up can report per-role progress as it pre-loads each upstream.
1100 self._on_spawning: Callable[[WorkerRole], None] | None = None
1101 self._on_spawned: Callable[[WorkerRole], None] | None = None
1102 # Granular cold-load progress for the chat role, streamed to a launcher so
1103 # the user sees real read/engine-load progress instead of a frozen spinner.
1104 self._warm_tracker = WarmProgressTracker()
1105 # The engine's own reason a role's model failed to warm, so the launcher and
1106 # the TUI report the real cause instead of a generic "did not load".
1107 self._warm_errors: dict[WorkerRole, str] = {}
1108 # Configured roles the last plan left unplaced because their model isn't
1109 # installed (role -> ref). The warm finalizer reads it to fail a not-installed
1110 # chat with a named reason instead of clearing to a silent "not ready" retry.
1111 self._skipped_not_installed: dict[WorkerRole, str] = {}
1112 # Launches the last plan refused for an unusable window (role -> reason);
1113 # read by the warm finalizer and _require_clients.
1114 self._skipped_unusable_ctx: dict[WorkerRole, str] = {}
1115 # Why the last acquisition left this process without a fleet because the
1116 # private dir has a live user; shown before any generic no-server line.
1117 self._engine_refusal: str | None = None
1118 # Single-flight guard for the off-thread warm-up: True from the moment a
1119 # warm thread is dispatched until it finishes, so a second warm_up_pool
1120 # never starts a second swap and double-allocates GPU memory.
1121 self._warming = False
1122 # Request-triggered lazy-warm single-flight: one request begins the warm
1123 # cycle, concurrent siblings share it instead of each resetting.
1124 self._lazy_warming = False
1125 # Single-flight guard for the off-thread reload: a second reload_role
1126 # while one is in flight sets the pending flag instead of dispatching,
1127 # and the in-flight thread re-runs the plan loop once per pending flag.
1128 self._reloading = False
1129 # Set when a reload arrives mid-reload: the in-flight pass may have
1130 # already snapshotted its plan, so the change must be re-applied.
1131 self._reload_pending = False
1132 # Roles the running reload may still bring a server back for: every role
1133 # from the start of a reload until its plan is known, then the planned
1134 # restarts. A request for any other role does not wait for the reload.
1135 self._reload_owed: frozenset[WorkerRole] = frozenset()
1136 # Notified when ``_reloading`` clears, when a reload publishes its plan and
1137 # when it adopts a group, so a ``reload_role(wait=True)`` caller and a
1138 # request waiting for its role's server wake at the change they wait for.
1139 self._reload_done = threading.Condition(self._lock)
1141 def _ensure_fleet(self) -> bool:
1142 """Start one llama-swap per placed role exactly once across concurrent callers.
1144 Returns whether any role group is running afterwards; ``False`` when no
1145 role is configured and installed (nothing to serve), leaving no process
1146 spawned. The startup runs under ``_build_lock`` (not the routing lock),
1147 so the off-thread warm-up and an on-demand call can't start two fleets --
1148 which would double-allocate GPU and parse the same GGUF twice. A second
1149 caller blocks on the build lock and reuses the groups the first one
1150 started. A group failing to start tears down the groups already started
1151 in this build, so a partial fleet never leaks past the failure.
1152 """
1153 with self._lock:
1154 if self._swaps:
1155 return True
1156 with self._build_lock:
1157 with self._lock:
1158 if self._swaps:
1159 return True
1160 if self._shut_down:
1161 # Provider was shut down (and likely discarded by reset_services)
1162 # while this warm-up/reload thread was in flight; do not spawn a
1163 # llama-swap no live provider would ever reap.
1164 return False
1166 return self._acquire_engine(cfg.data_root)
1168 def _acquire_engine(self, config_root: Path) -> bool:
1169 """The acquisition ladder: bind to a compatible engine, else build one.
1171 Machine slot first. An incumbent is replaced in place only when no live
1172 user holds it, and a matching fleet in the private dir is bound before
1173 any build in a free slot. An engine with a live user that this process cannot bind
1174 sends the build to the config root's private overflow dir. Per-dir the
1175 step is all-or-nothing: every configured (role, model) pair bound, or
1176 built fresh. Runs under the cross-process build lock, so two starts
1177 never both build and stop-if-last never races an arrival.
1178 """
1179 pin = engine_pin()
1180 demand = _placeable_demand()
1181 # Record the demand plan's skips before walking the ladder: a bind or an
1182 # early serve-nothing exit never reaches _plan_and_spawn, and the warm
1183 # tracker must still be able to say "chat model X is not installed"
1184 # rather than a retryable not-ready.
1185 self._skipped_not_installed = dict(demand.skipped_not_installed)
1186 self._skipped_unusable_ctx = dict(demand.skipped_unusable_ctx)
1187 self._engine_refusal = None
1188 machine_dir = machine_engine_dir()
1189 private = private_engine_dir(config_root)
1190 # One bound for the attempt: a wait on the slot leaves the private dir only the rest.
1191 wait_until = time.monotonic() + _OWNER_CATCH_UP_WAIT_S
1192 if kernel_arbitrates_locks(machine_dir):
1193 machine = self._acquire_in_dir(
1194 machine_dir, pin, demand, overflow_dir=private, wait_until=wait_until
1195 )
1196 if machine is not None:
1197 return machine
1198 log.warning(
1199 "The engine in %s serves another setup; starting a private engine for this one.",
1200 machine_dir,
1201 )
1202 else:
1203 # Without kernel-arbitrated locks the membership refcount cannot be
1204 # trusted, and sharing is exactly what needs it: a probe would
1205 # destroy a live member's lock, so the slot would look free while
1206 # another setup is serving from it. Keep to our own dir instead.
1207 log.warning(
1208 "Engine dir %s is on a filesystem without working file locks; "
1209 "using a private engine instead of the shared one. Set %s to a "
1210 "path on a local filesystem to share one engine across lilbees.",
1211 machine_dir,
1212 ENGINE_DIR_ENV,
1213 )
1214 # The machine slot is out; use this config root's private engine dir.
1215 return (
1216 self._acquire_in_dir(private, pin, demand, overflow_dir=None, wait_until=wait_until)
1217 or False
1218 )
1220 def _acquire_in_dir(
1221 self,
1222 engine_dir: Path,
1223 pin: str,
1224 demand: _EngineDemand,
1225 *,
1226 overflow_dir: Path | None,
1227 wait_until: float,
1228 ) -> bool | None:
1229 """Bind or build one engine dir; ``None`` on the slot means overflow next.
1231 When the live fleet there was planned for this config root before config.toml's
1232 last plan write, its builder rebuilds on that write soon. The ladder then lets the
1233 build lock go, waits until the fleet changes or *wait_until*, and runs once more,
1234 so this process does not load a second fleet beside the rebuilt one. The caller
1235 holds this provider's in-process build lock across the wait, so this process's
1236 other starts wait the same bound.
1237 """
1238 outcome = self._acquire_in_dir_locked(
1239 engine_dir,
1240 pin,
1241 demand,
1242 overflow_dir=overflow_dir,
1243 may_wait=time.monotonic() < wait_until,
1244 )
1245 if isinstance(outcome, _SlotOwed):
1246 _await_slot_change(engine_dir, outcome.states, wait_until)
1247 outcome = self._acquire_in_dir_locked(
1248 engine_dir, pin, demand, overflow_dir=overflow_dir, may_wait=False
1249 )
1250 return None if isinstance(outcome, _SlotOwed) else outcome
1252 def _acquire_in_dir_locked(
1253 self,
1254 engine_dir: Path,
1255 pin: str,
1256 demand: _EngineDemand,
1257 *,
1258 overflow_dir: Path | None,
1259 may_wait: bool,
1260 ) -> bool | None | _SlotOwed:
1261 """One ladder pass over one engine dir under its build lock; see ``_acquire_in_dir``.
1263 *overflow_dir* is the private dir behind the machine slot, and ``None`` when
1264 *engine_dir* is that private dir. Before a build in the slot, a matching
1265 fleet in *overflow_dir* is bound instead, so the card never holds the
1266 same models twice.
1268 Binds a compatible running engine. Whether an incumbent may be replaced
1269 is decided by kernel-refcounted membership, not the proxy HTTP probe: an
1270 engine with a live user is never reaped or stopped, so a transient probe
1271 failure (fd exhaustion, host thrash) cannot kill a busy engine. Replace in
1272 place only when no live user holds it. An engine with a live user that this
1273 process cannot bind is never evicted or stacked on: on the machine slot it
1274 returns ``None`` (overflow), and in the overflow dir it serves nothing rather
1275 than duplicate weights beside it. Before building, any recorded engine is
1276 cleared -- keyed on the state file, not the probe -- so an unprobeable
1277 incumbent is stopped rather than double-built beside. The stop is gated on
1278 ``_can_build_engine`` so a process that can serve nothing never destroys a
1279 warm engine it can't replace. Held under the cross-process build lock.
1280 """
1281 wanted = demand.pairs
1282 with build_lock(engine_dir):
1283 # One probe pass serves the bind and the owed-rebuild check.
1284 states = _healthy_states(engine_dir) if wanted else {}
1285 if self._bind_held_dir(engine_dir, pin, demand, states):
1286 return True
1287 # This process's own hold is no reason to overflow: the ladder runs only
1288 # when this provider serves no group, so the hold is all that is left.
1289 if live_users_exist(engine_dir, exclude_self=True):
1290 if may_wait and _owner_owes_the_file(states):
1291 return _SlotOwed(states)
1292 return self._refuse_live_engine(engine_dir, overflow_dir)
1293 if overflow_dir is not None:
1294 # Lock order: the slot's build lock, then the private dir's. No
1295 # path takes the slot's build lock while it holds a private one.
1296 with build_lock(overflow_dir):
1297 bound_private = self._bind_held_dir(overflow_dir, pin, demand)
1298 if bound_private:
1299 # No live user holds the slot, so its engine is stopped, as
1300 # before a build, and its weights do not stay beside the private fleet.
1301 _clear_engine_dir(engine_dir)
1302 return True
1303 if not _can_build_engine(wanted):
1304 return False
1305 # No live user holds this dir now: reap dead leftovers and stop any
1306 # recorded engine so planning sees true free VRAM and the build never
1307 # lands beside an unprobeable incumbent.
1308 _clear_engine_dir(engine_dir)
1309 built = self._plan_and_spawn(engine_dir)
1310 if built:
1311 self._hold_membership(engine_dir)
1312 _watch_bound_engine(engine_dir)
1313 return built
1315 def _refuse_live_engine(self, engine_dir: Path, overflow_dir: Path | None) -> bool | None:
1316 """Leave a live engine this process cannot bind: overflow (None) or serve nothing.
1318 A live engine is never evicted or stacked on. On the machine slot that means
1319 overflow; in the overflow dir there is nowhere further to go, so serve nothing
1320 rather than kill the incumbent or load a second fleet's weights beside it (an
1321 OOM on a small-VRAM box).
1322 """
1323 log.info("The engine at %s has a live user and cannot serve this setup", engine_dir)
1324 if overflow_dir is None:
1325 self._engine_refusal = _ENGINES_IN_USE_MESSAGE
1326 return False
1327 return None
1329 def _bind_held_dir(
1330 self,
1331 engine_dir: Path,
1332 pin: str,
1333 demand: _EngineDemand,
1334 states: dict[SwapGroup, SwapState] | None = None,
1335 ) -> bool:
1336 """Bind *engine_dir*'s running engine and join its users; caller holds its build lock.
1338 The user lock taken here is what a peer's last-user stop counts, so a
1339 bound engine is never stopped under this process. *states* is the caller's
1340 probe of the dir, when it has one.
1341 """
1342 if not demand.pairs:
1343 return False
1344 if states is None:
1345 states = _healthy_states(engine_dir)
1346 if not self._bind_all_in_dir(engine_dir, states, pin, demand):
1347 return False
1348 self._hold_membership(engine_dir)
1349 _watch_bound_engine(engine_dir)
1350 return True
1352 def _bind_all_in_dir(
1353 self,
1354 engine_dir: Path,
1355 states: dict[SwapGroup, SwapState],
1356 pin: str,
1357 demand: _EngineDemand,
1358 ) -> bool:
1359 """Bind every group needed to cover the demanded pairs; False leaves nothing bound.
1361 Binding never touches groups serving models outside the demand; the dir
1362 matches only when healthy, contract-equal groups cover every wanted pair and
1363 the served chat window covers the demanded per-slot ctx. (Whether an
1364 unmatched dir's engine is then replaced or overflowed around is the
1365 ladder's call, based on live users.)
1366 """
1367 wanted = demand.pairs
1368 candidates: list[tuple[SwapGroup, SwapState, list[InstanceLaunch]]] = []
1369 covered: set[tuple[WorkerRole, str]] = set()
1370 for group, state in states.items():
1371 found = _bindable_group(state, pin, demand)
1372 if found is None:
1373 continue
1374 bindable, launches, pairs = found
1375 if not chat_ctx_covers(launches, demand.chat_ctx) or not vision_slots_cover(
1376 launches, demand.vision_slots
1377 ):
1378 # The live chat window or vision slots are below this process's need.
1379 return False
1380 candidates.append((group, bindable, launches))
1381 covered |= pairs
1382 if covered != wanted:
1383 return False
1384 bound: dict[SwapGroup, tuple[SwapManager, list[InstanceLaunch]]] = {}
1385 for group, state, launches in candidates:
1386 swap = SwapManager(engine_dir, group)
1387 if not swap.bind(state):
1388 for prior, _launches in bound.values():
1389 prior.shutdown()
1390 return False
1391 bound[group] = (swap, launches)
1392 with self._lock:
1393 for group, (swap, launches) in bound.items():
1394 self._adopt_group(group, swap, launches)
1395 self._group_dirs[group] = engine_dir
1396 log.info("Bound to the running engine at %s", engine_dir)
1397 _log_adopted_launches(candidates)
1398 return True
1400 def _reload_dir(self) -> Path:
1401 """The engine dir a reload rebuilds into: where our groups already live.
1403 All this provider's groups share one dir by construction (the ladder is
1404 all-or-nothing per dir); an empty provider rebuilds into the machine slot.
1405 """
1406 with self._lock:
1407 dirs = set(self._group_dirs.values())
1408 return next(iter(dirs)) if dirs else machine_engine_dir()
1410 def _hold_membership(self, engine_dir: Path) -> None:
1411 """Record this process as a user of *engine_dir*'s engine.
1413 The single point every acquisition passes through, bind and build alike,
1414 so it is also where this user's persistence opt-in is recorded against
1415 the engine. Marking on bind (not only on build) is what makes the
1416 setting mean what it says on a shared slot: a user who asked for a warm
1417 engine keeps it warm even when a default-config sibling is last out.
1418 """
1419 from lilbee.core.config import cfg
1421 if engine_dir not in self._engine_holds:
1422 self._engine_holds[engine_dir] = hold_user_lock(engine_dir)
1423 if cfg.keep_engine_warm:
1424 request_keep_warm(engine_dir, cfg.data_root, cfg.engine_idle_ttl_minutes)
1426 def _plan_and_spawn(self, data_dir: Path) -> bool:
1427 """Plan placement against the clean box and start one swap per group.
1429 Caller holds the build lock. False when the engine binary is missing or
1430 nothing is installed/configured, so the provider serves nothing.
1431 """
1432 try:
1433 # Snapshot the clean box; this plan and every later reload size
1434 # ctx, slots, and budgets against it (a live probe under a loaded
1435 # fleet would report our own residency as unavailable). Inside the
1436 # try: capturing resolves the engine binary, and a binary-less
1437 # host must serve nothing, not raise. Every other planning failure
1438 # (a wedged GPU probe, an unusable CUDA runtime) propagates so the
1439 # warm tracker and the caller report the real reason instead of a
1440 # silent never-ready fleet.
1441 origin = _plan_origin()
1442 planning.capture_plan_probe()
1443 plan = planning.plan_all_launches()
1444 except ProviderError as exc:
1445 # Only a genuinely-missing engine binary keeps the quiet no-fleet path;
1446 # any other planning failure (a wedged GPU probe, an unusable CUDA
1447 # runtime) must surface to the warm tracker and on-demand callers
1448 # rather than silently serving nothing (#540).
1449 if exc.kind is not ProviderErrorKind.NOT_FOUND:
1450 raise
1451 log.debug("Engine binary unavailable; no swap started")
1452 plan = None
1453 # plan None (no engine binary) keeps the demand-time record from
1454 # _acquire_engine instead of wiping it.
1455 if plan is not None:
1456 self._skipped_not_installed = dict(plan.skipped_not_installed)
1457 self._skipped_unusable_ctx = dict(plan.skipped_unusable_ctx)
1458 if plan is None or not plan.launches:
1459 # No engine binary, or no installed/configured model: serve nothing.
1460 return False
1461 by_group = _launches_by_group(plan)
1462 started: dict[SwapGroup, SwapManager] = {}
1463 try:
1464 lifetime = _engine_lifetime(data_dir)
1465 ttl_seconds = _warm_ttl_seconds(
1466 lifetime, data_dir, hold_warm_for_session=self._hold_warm_for_session
1467 )
1468 for group, group_launches in by_group.items():
1469 swap = SwapManager(data_dir, group)
1470 swap.start(
1471 list(group_launches), ttl_seconds=ttl_seconds, lifetime=lifetime, origin=origin
1472 )
1473 started[group] = swap
1474 except BaseException:
1475 for swap in started.values():
1476 swap.shutdown()
1477 raise
1478 with self._lock:
1479 for group, swap in started.items():
1480 self._adopt_group(group, swap, list(by_group[group]))
1481 self._group_dirs[group] = data_dir
1482 return True
1484 def _adopt_group(
1485 self, group: SwapGroup, swap: SwapManager, launches: list[InstanceLaunch]
1486 ) -> None:
1487 """Record *group*'s freshly started swap and build a client pool per role.
1489 Caller holds ``self._lock``. Each launch (one per replica) becomes a client
1490 keyed by its replica model id against this group's own proxy endpoint;
1491 the chat launch carries the slots/ctx so the capacity and served context
1492 come from the launch, not a probe.
1493 """
1494 self._swaps[group] = swap
1495 self._launches[group] = tuple(launches)
1496 endpoint = swap.endpoint()
1497 for role, role_launches in _by_role(launches).items():
1498 # Retire the role's previous clients (a reload re-adopts over an existing
1499 # pool): closing them now would error a reader still mid-call on an old
1500 # client snapshot, and never closing leaks an httpx pool per replica.
1501 old_clients = list(self._clients.get(role, []))
1502 self._role_group[role] = group
1503 # token_cap truncates oversize embed/rerank inputs to the per-slot context
1504 # (the in-process backstop); the longer timeout covers a cold upstream load.
1505 self._clients[role] = [
1506 LlamaServerClient(
1507 endpoint,
1508 launch.model_id,
1509 token_cap=launch.token_cap,
1510 timeout=_request_timeout_s(launch),
1511 rerank_mode=launch.rerank_mode,
1512 inline_reasoning=role is WorkerRole.CHAT,
1513 on_prefill=self._record_chat_prefill if role is WorkerRole.CHAT else None,
1514 engine_log=_engine_log_of(launch, swap.log_path.parent),
1515 model_ref=launch.model,
1516 )
1517 for launch in role_launches
1518 ]
1519 if role is WorkerRole.CHAT:
1520 self._chat_slots = role_launches[0].slots
1521 self._chat_ctx = role_launches[0].ctx
1522 # Every serving process passes through adoption (fresh launch,
1523 # reload, guest bind), so the warning fires in each of them.
1524 planning.warn_when_chat_downsized(role_launches[0])
1525 if role is WorkerRole.EMBED:
1526 planning.warn_when_embed_window_below_chunk(role_launches[0])
1527 self._retire_clients(old_clients)
1528 self._reload_done.notify_all()
1530 def _swap_for(self, role: WorkerRole) -> SwapManager | None:
1531 """The swap process serving *role*, or None when the role has no server.
1533 Caller holds ``self._lock``.
1534 """
1535 group = self._role_group.get(role)
1536 return None if group is None else self._swaps.get(group)
1538 def _role_launches(self, role: WorkerRole) -> tuple[InstanceLaunch, ...]:
1539 """*role*'s launch snapshot, empty when it has no server.
1541 A co-tenant group holds more than one role's launches, so the group's
1542 snapshot is filtered down to this role's replicas.
1543 """
1544 group = self._role_group.get(role)
1545 if group is None:
1546 return ()
1547 return tuple(launch for launch in self._launches.get(group, ()) if launch.role is role)
1549 def _drop_group(self, group: SwapGroup) -> SwapManager | None:
1550 """Forget *group*'s swap/launches and every member role's clients.
1552 Caller holds ``self._lock``. Member clients are retired (closed at a later
1553 reload or shutdown, never while a reader could still hold one) and the chat
1554 capacity falls back to its defaults when chat's group is dropped.
1555 """
1556 swap = self._swaps.pop(group, None)
1557 self._launches.pop(group, None)
1558 # Prune the dir map with the group: a stale entry outliving its group makes
1559 # _reload_dir see two dirs and pick one arbitrarily, splitting the provider.
1560 self._group_dirs.pop(group, None)
1561 for role in [r for r, g in self._role_group.items() if g is group]:
1562 del self._role_group[role]
1563 self._retire_clients(self._clients.pop(role, []))
1564 if role is WorkerRole.CHAT:
1565 self._chat_slots = 1
1566 self._chat_ctx = None
1567 return swap
1569 def _retire_clients(self, old_clients: list[LlamaServerClient]) -> None:
1570 """Close the previously-retired clients, then retire *old_clients*.
1572 Caller holds ``self._lock``. Retired clients are never handed to new
1573 readers (they are out of ``self._clients``), so by this reload any reader
1574 that held one from a prior reload has finished; an ``in_flight == 0``
1575 check confirms it before close, and any still-busy client stays retired
1576 for the next reload. This closes idle reloaded-away pools without ever
1577 closing one a reader could still use. Shutdown closes whatever remains.
1578 """
1579 still_busy: list[LlamaServerClient] = []
1580 for client in self._retiring_clients:
1581 if client.in_flight == 0:
1582 client.close()
1583 else:
1584 still_busy.append(client)
1585 for client in old_clients:
1586 # A hold the sweep kept for this client's stream is released when the
1587 # stream ends; see _start_hold_sweep.
1588 client.notify_when_idle(self._start_hold_sweep)
1589 self._retiring_clients = still_busy + old_clients
1591 def _require_clients(self, role: WorkerRole) -> list[LlamaServerClient]:
1592 """The client pool for *role*, or a user-facing error when it has no server.
1594 A configured, placeable role gets one or more replica clients; their absence
1595 means the role is unconfigured or did not fit memory. llama-swap loads each
1596 upstream on its first request, so a returned client may still be cold. No
1597 in-process fallback, so a missing pool is a hard error.
1599 When the pool is empty but a swap was previously built and its process has
1600 since exited (detected via ``is_live()``), a one-shot rebuild is attempted
1601 before raising so a transient llama-swap restart recovers transparently.
1602 A request that lands while a reload restarts the role's group waits, up to
1603 ``_RELOAD_GAP_WAIT_S``, until the role has clients or the reload will not
1604 bring the role back. A teardown on another thread can empty the fleet after
1605 the check below passes, so a fleet that was running and is now empty goes
1606 through the check once more. A fleet that never started does not.
1607 """
1608 had_fleet = self._ensure_fleet()
1609 with self._lock:
1610 # The reload dropped this role's group and has not adopted the new one.
1611 settled = self._reload_done.wait_for(
1612 lambda: (
1613 bool(self._clients.get(role))
1614 or not self._reloading
1615 or role not in self._reload_owed
1616 ),
1617 _RELOAD_GAP_WAIT_S,
1618 )
1619 clients = self._clients.get(role)
1620 if not clients and not settled:
1621 raise ProviderError(_reload_timeout_message(role), provider=_PROVIDER_NAME)
1622 swap = self._swap_for(role)
1623 fleet_is_empty = not self._swaps
1624 if not clients and fleet_is_empty and had_fleet:
1625 self._ensure_fleet()
1626 with self._lock:
1627 clients = self._clients.get(role)
1628 if not clients and swap is not None and not swap.is_live():
1629 self._rebuild_role(role)
1630 with self._lock:
1631 clients = self._clients.get(role)
1632 if not clients:
1633 # A refused launch's recorded reason wins over the generic line.
1634 # BAD_REQUEST so the HTTP surfaces return this deterministic,
1635 # actionable refusal in a 4xx body instead of a generic 500.
1636 reason = self._engine_refusal or self._skipped_unusable_ctx.get(role)
1637 if reason is not None:
1638 raise ProviderError(
1639 reason, provider=_PROVIDER_NAME, kind=ProviderErrorKind.BAD_REQUEST
1640 )
1641 raise ProviderError(_no_server_message(role), provider=_PROVIDER_NAME)
1642 return list(clients)
1644 def _with_rediscover(self, call: Callable[[], _T], *, role: WorkerRole | None = None) -> _T:
1645 """Run *call*; on a connection-kind or load-capacity failure, retry once.
1647 A vanished engine (it died, or its owner restarted its group) surfaces as
1648 ProviderErrorKind.CONNECTION, or as a raw httpx transport error when the
1649 proxy itself is gone (nothing listening to answer with a status). Membership
1650 is still held, so dropping the swap refs and retrying sends the call through
1651 _ensure_fleet, which rediscovers the new proxy ports or rebuilds. One retry
1652 only; a second failure surfaces to the caller.
1654 A ProviderErrorKind.CAPACITY failure is the engine dying on load because
1655 the estimate was too optimistic. Retrying it unchanged respawns the same
1656 launch into the same death, so *role*'s auto context steps down first and
1657 the role is rebuilt against the smaller plan. When there is no step left
1658 to take (a user-pinned context, or already at the floor) the failure
1659 surfaces instead: a retry that asks for the same thing is a crash loop.
1660 """
1661 try:
1662 return call()
1663 except (ProviderError, httpx.TransportError) as err:
1664 if is_rebuildable_failure(err) and role is not None:
1665 return self._retry_rebuilt(call, role, err)
1666 if not is_connection_failure(err):
1667 raise
1668 log.info("Engine unreachable; rediscovering before one retry")
1669 # The holds stay through the retry: a stream on a retired client and a
1670 # second thread's rebuild both need the dir held, and the ladder does not
1671 # count this process's own hold as a live user. The open count keeps any
1672 # sweep from releasing a hold before this retry ends. The count goes up
1673 # inside the try, so an exception at any point after it brings it down.
1674 opened = False
1675 try:
1676 with self._lock:
1677 self._open_rediscovers += 1
1678 opened = True
1679 self._drop_swap_refs()
1680 # An owner that restarted its fleet for new settings wrote them first.
1681 # With no group held, the replan the take starts returns at once, and
1682 # the retry plans on the new settings, so it binds the restarted fleet.
1683 # When the retry reaches the slot before the owner rebuilds, the ladder
1684 # waits for that rebuild (see _acquire_in_dir) instead of overflowing.
1685 if self._config_refresh is not None:
1686 self._config_refresh()
1687 return call()
1688 finally:
1689 if opened:
1690 with self._lock:
1691 self._open_rediscovers -= 1
1692 try:
1693 self._release_unused_holds()
1694 except Exception:
1695 # A failed sweep must not replace the retry's own result or error.
1696 log.warning("Hold sweep after a rediscover failed", exc_info=True)
1698 def _retry_rebuilt(self, call: Callable[[], _T], role: WorkerRole, err: BaseException) -> _T:
1699 """Rebuild *role* so the retry is a different launch, and run *call* again.
1701 A held port just needs the rebuild, which picks a new one. A memory
1702 shortfall needs the plan to come back smaller too, so the context steps
1703 down first; when there is no step left to take, *err* is re-raised
1704 untouched rather than rebuilding into the same death.
1705 """
1706 from lilbee.providers.fleet.planning import record_ctx_downshift
1708 if is_load_capacity_failure(err):
1709 if not record_ctx_downshift(role):
1710 log.warning(
1711 "%s ran out of device memory on load and its context cannot be "
1712 "reduced further; lower num_ctx or use a smaller model",
1713 role.value,
1714 )
1715 raise err
1716 log.warning(
1717 "%s ran out of device memory on load; re-planning it with a smaller "
1718 "context where its window has room to give",
1719 role.value,
1720 )
1721 else:
1722 log.warning("%s could not claim its port; rebuilding it on a new one", role.value)
1723 self._rebuild_role(role)
1724 return call()
1726 def _rebuild_role(self, role: WorkerRole) -> None:
1727 """Restart just *role*'s dead group (new port) from a fresh plan.
1729 Other roles' groups keep serving; only the dead group is torn down and
1730 respawned. Runs the same diff-driven pass as a reload, forcing *role*
1731 into the restart set so an unchanged plan still replaces its dead swap.
1732 """
1733 self._reload_pass(force=frozenset((role,)))
1735 def role_ready(self, role: WorkerRole) -> bool:
1736 """Whether *role*'s upstream is loaded and ready, without starting the swap.
1738 A read-only probe for surfaces (HTTP status, SSE warming event) that want
1739 to report cold-start state without triggering a load. False before the swap
1740 is up or while the role's upstream is still loading.
1741 """
1742 with self._lock:
1743 swap = self._swap_for(role)
1744 return swap is not None and swap.role_ready(role)
1746 def max_concurrent_chats(self) -> int:
1747 """The chat server's batching-slot capacity, so the gate admits that many.
1749 Falls back to ``1`` before the chat group is up, so chat is serialized
1750 until the slot count is known (the launcher warms the engine before a
1751 client connects, so the real capacity is in effect by the first chat).
1752 """
1753 with self._lock:
1754 if WorkerRole.CHAT not in self._role_group:
1755 return 1
1756 return self._chat_slots
1758 def served_chat_ctx(self) -> int | None:
1759 """Per-slot context the chat server runs with, or None if not up."""
1760 with self._lock:
1761 return self._chat_ctx if WorkerRole.CHAT in self._role_group else None
1763 def served_chat_slots(self) -> int | None:
1764 """Batching slots the chat server runs with, or None if not up."""
1765 with self._lock:
1766 return self._chat_slots if WorkerRole.CHAT in self._role_group else None
1768 def embed_token_cap(self) -> int | None:
1769 """The served embed launch's token cap, or the planned one before it is up."""
1770 with self._lock:
1771 launches = self._role_launches(WorkerRole.EMBED)
1772 if launches:
1773 return launches[0].token_cap
1774 return planning.planned_embed_token_cap(cfg.embedding_model)
1776 def health_warnings(self) -> list[HealthWarning]:
1777 """Serving degradations: the embed window and any placement divergences."""
1778 warnings: list[HealthWarning] = []
1779 cap = self.embed_token_cap()
1780 if cap is not None:
1781 warning = engine_params.embed_window_warning(cap)
1782 if warning is not None:
1783 warnings.append(warning)
1784 with self._lock:
1785 swaps = list(self._swaps.values())
1786 for swap in swaps:
1787 warnings.extend(swap.health_warnings())
1788 return warnings
1790 def chat_prefill_progress(self) -> tuple[int, int] | None:
1791 """``(processed, total)`` of a chat prefill in flight, or None when idle."""
1792 with self._lock:
1793 return self._chat_prefill
1795 def _record_chat_prefill(self, progress: tuple[int, int] | None) -> None:
1796 """Store a stream's prefill reading. Latest writer wins across concurrent
1797 streams; a live prefill rewrites itself on its next engine batch."""
1798 with self._lock:
1799 self._chat_prefill = progress
1801 def warm_pending(self) -> bool:
1802 """Whether a requested warm is still running.
1804 The tracker only stamps a phase once the chat role starts loading, which is
1805 seconds after the swap is spawned, so ``warm_progress`` alone cannot tell a
1806 not-yet-started warm from no warm at all.
1807 """
1808 with self._lock:
1809 return self._warming
1811 def warm_progress(self) -> WarmProgress | None:
1812 """Live cold-load progress for the chat role, or None before warm begins."""
1813 return self._warm_tracker.snapshot()
1815 def _shutdown_swap(self) -> None:
1816 """Release this process's engine use for good.
1818 Latches ``_shut_down`` so a discarded provider's in-flight warm/reload
1819 thread can't spawn an orphan swap, then releases membership: the engine
1820 stops only when this was the last user and persistence was not opted into.
1821 """
1822 # Latched before the lock, not inside it: every _shut_down check runs after
1823 # acquiring the build lock, so a warm or reload thread queued behind us can
1824 # only bail early if the flag is already set when its turn comes.
1825 with self._lock:
1826 self._shut_down = True
1827 # The build lock serializes shutdown against a concurrent reload/build:
1828 # both mutate self._swaps and the llama-swap processes, so an unserialized
1829 # loser would overwrite the winner's state and leak a live llama-swap.
1830 # Bounded, because a wedged engine start holds this lock and an unbounded
1831 # wait would hang process exit outright. On timeout the teardown proceeds
1832 # anyway: whatever the builder leaves behind is recorded in the engine
1833 # dir's state files, so the next start's reap finds it by record, while a
1834 # shutdown that never returns cannot be recovered from at all.
1835 acquired = self._build_lock.acquire(timeout=_SHUTDOWN_BUILD_LOCK_WAIT_S)
1836 if not acquired:
1837 log.warning(
1838 "Engine build still in progress after %.0fs; shutting down without "
1839 "waiting for it. Leftovers are reaped from their records on the next start.",
1840 _SHUTDOWN_BUILD_LOCK_WAIT_S,
1841 )
1842 try:
1843 self._drop_swap_refs(close_all=True)
1844 self._release_engines()
1845 finally:
1846 if acquired:
1847 self._build_lock.release()
1849 def _release_engines(self) -> None:
1850 """Drop membership in every used engine dir; stop each engine we leave last.
1852 Runs under each dir's cross-process build lock so a departing last user
1853 can never race an arriving binder: the arrival either sees the engine
1854 (and its bind holds it live) or sees the slot empty and builds.
1856 Whether the engine outlives us is the union of every user's opt-in, not
1857 just the exiting process's config: the machine slot is shared by
1858 installations that configure it differently, and which one leaves last
1859 is arbitrary.
1861 The hold map is cleared either way. Leaving a stale hold behind is not
1862 benign: after a lazy rebuild overflows to a private dir (a foreign
1863 process having claimed the machine slot in the gap), the next release
1864 would iterate the stale machine hold and stop that foreign engine
1865 mid-use, and the stale flock would keep live_users_exist true so the
1866 foreign engine's real last user could never reap it.
1867 """
1868 from lilbee.core.config import cfg
1870 for engine_dir, hold in list(self._engine_holds.items()):
1871 with build_lock(engine_dir, best_effort=True):
1872 last = hold.release_and_check_last()
1873 # A flip after binding never re-acquires, so reconcile our own mark
1874 # here. A withdrawn opt-in turns a detached engine into one its
1875 # remaining users share, so a reaper must watch it as after a bind.
1876 if cfg.keep_engine_warm:
1877 request_keep_warm(engine_dir, cfg.data_root, cfg.engine_idle_ttl_minutes)
1878 else:
1879 withdraw_keep_warm(engine_dir, cfg.data_root)
1880 if not last:
1881 _watch_bound_engine(engine_dir)
1882 # Any remaining opt-in keeps the engine, including a peer's.
1883 if last and not keep_warm_requested(engine_dir):
1884 stop_engine(engine_dir)
1885 log.info("Engine stopped at %s (last user out)", engine_dir)
1886 elif last:
1887 _watch_bound_engine(engine_dir)
1888 log.info("Engine left warm at %s (last user out)", engine_dir)
1889 self._engine_holds = {}
1891 def _release_unused_holds(self) -> None:
1892 """Drop each membership whose dir no group uses, once no retired client is busy.
1894 After a rediscover retry overflows, the old dir's hold is no longer needed,
1895 and a stale hold keeps that engine's real last user from stopping it. A
1896 retired client still streams from the old engine, so while any retired
1897 client is in flight the holds stay; that client's last request starts the
1898 sweep again (see _start_hold_sweep). While a rediscover is open its dropped
1899 groups have no dir yet, so the holds stay; its own retry ends with a sweep.
1900 The build lock keeps the sweep from releasing a dir that a build on another
1901 thread is about to hold. A first check without the build lock returns at once
1902 when there is nothing to release, so the common sweep never waits on a build.
1903 Nothing is stopped here.
1904 """
1905 with self._lock:
1906 if not self._unused_hold_dirs():
1907 return
1908 with self._build_lock:
1909 with self._lock:
1910 holds = [self._engine_holds.pop(path) for path in self._unused_hold_dirs()]
1911 for hold in holds:
1912 hold.release_and_check_last()
1914 def _unused_hold_dirs(self) -> list[Path]:
1915 """The held dirs that no group uses, or none while a sweep must wait. Caller holds _lock."""
1916 if self._open_rediscovers or any(client.in_flight for client in self._retiring_clients):
1917 return []
1918 used = set(self._group_dirs.values())
1919 return [path for path in self._engine_holds if path not in used]
1921 def _start_hold_sweep(self) -> None:
1922 """Run the hold sweep on a new thread: a retired client ended its last request.
1924 The thread that ended the request belongs to the caller and can hold any of
1925 its locks. An abandoned stream also closes from garbage collection, inside
1926 whatever frame runs then, including one that holds the build lock. A sweep
1927 on that thread could wait on a lock the thread itself holds. A new thread
1928 holds no lock and takes them in the sweep's own order. At interpreter exit no
1929 new thread can start; the sweep is skipped then, as the process ends anyway.
1930 """
1931 try:
1932 threading.Thread(
1933 target=self._release_unused_holds,
1934 name=_HOLD_SWEEP_THREAD,
1935 daemon=True,
1936 ).start()
1937 except RuntimeError:
1938 log.debug("Hold sweep not started: no new thread can start", exc_info=True)
1940 def _drop_swap_refs(self, *, close_all: bool = False) -> None:
1941 """Clear every group's swap/clients and the chat capacity so the next call rebuilds.
1943 Live pools are RETIRED through the ``in_flight`` check rather than closed
1944 outright: ``_with_rediscover`` reaches this on any connection blip, so a
1945 chat proxy hiccup must not sever the client another thread is mid-embed or
1946 mid-stream on (a streamed response is handed out after the retry returns,
1947 and failures past the first frame are not retried). Idle clients close now,
1948 busy ones stay retired for a later pass.
1950 *close_all* is the terminal-shutdown path, where nothing will read again and
1951 whatever remains must actually be closed.
1952 """
1953 doomed: list[LlamaServerClient] = []
1954 with self._lock:
1955 live = [client for pool in self._clients.values() for client in pool]
1956 self._swaps = {}
1957 self._launches = {}
1958 self._role_group = {}
1959 self._group_dirs = {}
1960 self._clients = {}
1961 if close_all:
1962 doomed = live + self._retiring_clients
1963 self._retiring_clients = []
1964 else:
1965 self._retire_clients(live)
1966 self._chat_slots = 1
1967 self._chat_ctx = None
1968 # A torn-down fleet's load failures describe servers that no longer
1969 # exist; the next warm records its own.
1970 self._warm_errors = {}
1971 # Full teardown: the next build starts from a clean box, so it must
1972 # re-snapshot memory rather than plan against this boot's probe.
1973 planning.clear_plan_probe()
1974 for client in doomed:
1975 client.close()
1977 def _drop_dead_swaps(self) -> None:
1978 """Drop the refs of groups whose process is gone so the next call rebuilds them.
1980 A no-op for groups still running (e.g. the failure was in planning), so
1981 a live engine is never abandoned unstopped.
1982 """
1983 with self._build_lock, self._lock:
1984 for group in [g for g, swap in self._swaps.items() if not swap.running]:
1985 self._drop_group(group)
1987 def _require_configured_model(
1988 self, model: str | None, configured: str, role: WorkerRole
1989 ) -> None:
1990 """Reject a per-call model that differs from the server's configured one.
1992 The fleet serves the configured model for each role; switching models is
1993 a config change that respawns the server, not a per-call override. An
1994 empty/None ``model`` means "use the configured one" and is always accepted.
1995 """
1996 if model and model != configured:
1997 raise ProviderError(
1998 configured_model_message(role, configured, model),
1999 provider=_PROVIDER_NAME,
2000 kind=ProviderErrorKind.BAD_REQUEST,
2001 )
2003 @overload
2004 def chat(
2005 self,
2006 messages: list[ChatMessage],
2007 *,
2008 stream: Literal[False] = False,
2009 options: dict[str, Any] | None = None,
2010 model: str | None = None,
2011 tools: list[dict[str, Any]] | None = None,
2012 tool_choice: str | dict[str, Any] | None = None,
2013 ) -> ChatResult: ...
2015 @overload
2016 def chat(
2017 self,
2018 messages: list[ChatMessage],
2019 *,
2020 stream: Literal[True],
2021 options: dict[str, Any] | None = None,
2022 model: str | None = None,
2023 tools: list[dict[str, Any]] | None = None,
2024 tool_choice: str | dict[str, Any] | None = None,
2025 ) -> ClosableIterator[ChatStreamItem]: ...
2027 def chat(
2028 self,
2029 messages: list[ChatMessage],
2030 *,
2031 stream: bool = False,
2032 options: dict[str, Any] | None = None,
2033 model: str | None = None,
2034 tools: list[dict[str, Any]] | None = None,
2035 tool_choice: str | dict[str, Any] | None = None,
2036 ) -> ChatResult | ClosableIterator[ChatStreamItem]:
2037 """Route a chat turn to the least-busy chat server.
2039 Non-streaming returns a :class:`ChatResult` (text, tool calls, finish
2040 reason); streaming yields :data:`ChatStreamItem` frames. ``--jinja`` on
2041 the server parses native tool calls, so tool support needs no per-family
2042 parser here.
2043 """
2044 from lilbee.providers.engine_params import chat_options_to_kwargs
2046 self._require_configured_model(model, str(cfg.chat_model), WorkerRole.CHAT)
2047 with self._lazy_warm_scope():
2048 self._require_clients(WorkerRole.CHAT)
2049 messages = self._fit_chat_context(
2050 messages, tools, options, model or str(cfg.chat_model)
2051 )
2052 # Translate options exactly as the in-process path did (validate via
2053 # LLMOptions, num_predict -> max_tokens, drop num_ctx) so the server
2054 # honors the same generation settings; a raw passthrough would drop
2055 # num_predict and leak the load-only num_ctx.
2056 server_options = chat_options_to_kwargs(options) or None
2057 if stream:
2058 # The first frame is pulled eagerly so a dead proxy fails inside
2059 # _with_rediscover; a failure past the first frame surfaces to the
2060 # caller as a retry error, and rediscovery covers the next call.
2061 return self._with_rediscover(
2062 lambda: _primed_stream(
2063 _least_in_flight(self._require_clients(WorkerRole.CHAT)).chat_stream_items(
2064 messages, tools=tools, tool_choice=tool_choice, options=server_options
2065 )
2066 ),
2067 role=WorkerRole.CHAT,
2068 )
2069 return self._with_rediscover(
2070 lambda: _least_in_flight(self._require_clients(WorkerRole.CHAT)).chat_result(
2071 messages, tools=tools, tool_choice=tool_choice, options=server_options
2072 ),
2073 role=WorkerRole.CHAT,
2074 )
2076 def chat_with_tools(
2077 self,
2078 messages: list[ChatMessage],
2079 *,
2080 tools: list[dict[str, Any]],
2081 tool_choice: str | dict[str, Any] | None = None,
2082 options: dict[str, Any] | None = None,
2083 model: str | None = None,
2084 ) -> ChatToolResult:
2085 """Route a tool-enabled chat turn to the least-busy chat server."""
2086 from lilbee.providers.engine_params import chat_options_to_kwargs
2088 self._require_configured_model(model, str(cfg.chat_model), WorkerRole.CHAT)
2089 with self._lazy_warm_scope():
2090 self._require_clients(WorkerRole.CHAT)
2091 messages = self._fit_chat_context(
2092 messages, tools, options, model or str(cfg.chat_model)
2093 )
2094 server_options = chat_options_to_kwargs(options) or None
2095 return self._with_rediscover(
2096 lambda: _least_in_flight(self._require_clients(WorkerRole.CHAT)).chat_tools(
2097 messages, tools=tools, tool_choice=tool_choice, options=server_options
2098 ),
2099 role=WorkerRole.CHAT,
2100 )
2102 def _fit_chat_context(
2103 self,
2104 messages: list[ChatMessage],
2105 tools: list[dict[str, Any]] | None,
2106 options: dict[str, Any] | None,
2107 model: str,
2108 ) -> list[ChatMessage]:
2109 """Drop oldest turns so the prompt fits the served context.
2111 A ``num_predict`` reservation larger than the default generation room is
2112 capped to it, so an agent client that over-reserves keeps its history
2113 instead of having it evicted; llama-server stops at the context edge
2114 anyway. A smaller reservation is honored as-is and widens the prompt.
2115 Raises ``ProviderError(CONTEXT_OVERFLOW)`` only when system messages,
2116 tools, and the final turn exceed the window even with the capped
2117 reserve (mapped to a 400 by the chat-completions route).
2118 """
2119 # 0/None means the served context is unknown (no chat launch adopted yet);
2120 # a real per-slot context is always positive, so skip windowing.
2121 if not self._chat_ctx:
2122 return messages
2123 # An output reservation only ever buys the prompt MORE room, never less:
2124 # a num_predict past the default is a ceiling on what the model may
2125 # generate, not a claim on prompt space, and llama-server stops at the
2126 # context edge regardless. Capping it here rather than retrying after a
2127 # failed fit is the difference between a policy and a rescue -- an agent
2128 # reserving most of the window leaves a budget of a few dozen tokens, in
2129 # which the final turn still "fits" while the whole conversation is
2130 # silently evicted.
2131 requested = (options or {}).get("num_predict")
2132 reserve = min(requested, GENERATION_RESERVE_TOKENS) if requested else None
2133 budget = prompt_token_budget(self._chat_ctx, reserve)
2134 result = window_messages(messages, tools, budget)
2135 if not result.fits:
2136 raise ProviderError(
2137 f"Prompt of about {result.prompt_tokens} tokens exceeds the "
2138 f"{budget}-token prompt budget for {model!r} "
2139 f"({self._chat_ctx}-token window minus reserve and margin). "
2140 "Shorten the conversation or the system prompt.",
2141 provider=_PROVIDER_NAME,
2142 kind=ProviderErrorKind.CONTEXT_OVERFLOW,
2143 )
2144 return result.messages
2146 def embed(self, texts: list[str]) -> list[Vector]:
2147 return self._with_rediscover(lambda: self._embed_once(texts), role=WorkerRole.EMBED)
2149 def _embed_once(self, texts: list[str]) -> list[Vector]:
2150 clients = self._require_clients(WorkerRole.EMBED)
2151 return _call_with_failover(clients, lambda client: client.embed(texts))
2153 def count_tokens(self, text: str) -> int:
2154 """Exact token count of *text* under the embedding model's tokenizer.
2156 Routes to the embed server's ``/tokenize`` so chunk sizing counts the same
2157 tokens the embedder consumes, with the same rediscovery retry as embedding.
2158 """
2159 return self._with_rediscover(lambda: self._count_once(text), role=WorkerRole.EMBED)
2161 def _count_once(self, text: str) -> int:
2162 clients = self._require_clients(WorkerRole.EMBED)
2163 return _call_with_failover(clients, lambda client: client.count_tokens(text))
2165 def count_chat_prompt_tokens(
2166 self,
2167 messages: list[ChatMessage],
2168 *,
2169 options: dict[str, Any] | None = None,
2170 model: str | None = None,
2171 tools: list[dict[str, Any]] | None = None,
2172 tool_choice: str | dict[str, Any] | None = None,
2173 ) -> int:
2174 """Tokens the chat server prefills for this prompt, template applied.
2176 Takes the arguments :meth:`chat` takes and sends the same body. The
2177 conversation is not fitted to the served window first, so the count
2178 covers the whole prompt the caller asked about.
2179 """
2180 from lilbee.providers.engine_params import chat_options_to_kwargs
2182 self._require_configured_model(model, str(cfg.chat_model), WorkerRole.CHAT)
2183 server_options = chat_options_to_kwargs(options) or None
2184 with self._lazy_warm_scope():
2185 return self._with_rediscover(
2186 lambda: _least_in_flight(
2187 self._require_clients(WorkerRole.CHAT)
2188 ).count_chat_prompt_tokens(
2189 messages, tools=tools, tool_choice=tool_choice, options=server_options
2190 ),
2191 role=WorkerRole.CHAT,
2192 )
2194 def vision_ocr(
2195 self,
2196 png_bytes: bytes,
2197 model: str,
2198 prompt: str = "",
2199 *,
2200 timeout: float | None = None,
2201 cancel: CancelSignal | None = None,
2202 ) -> str:
2203 from lilbee.vision import build_vision_messages, resolve_ocr_prompt
2205 self._require_configured_model(model, str(cfg.vision_model), WorkerRole.VISION)
2206 pool = self._vision_pool()
2207 effective = model or str(cfg.vision_model)
2208 messages = build_vision_messages(prompt or resolve_ocr_prompt(effective), png_bytes)
2209 try:
2210 return _ocr_dispatch(pool, messages, _ocr_deadline(timeout), cancel)
2211 except _PageBudgetExhausted:
2212 raise ProviderError(
2213 "Vision OCR timed out waiting for a free vision slot.",
2214 provider=_PROVIDER_NAME,
2215 ) from None
2217 def vision_slot_capacity(self) -> int | None:
2218 """Total fitted ``--parallel`` slots across the running vision replicas.
2220 ``None`` before the fleet is up (no launch snapshot yet), so the ingest
2221 fan-out keeps its own estimate until real capacity is known. A modest
2222 card that fit fewer slots than requested reports the smaller real number,
2223 so the fan-out never queues more pages than the servers can serve.
2224 """
2225 launches = self._role_launches(WorkerRole.VISION)
2226 if not launches:
2227 return None
2228 return max(1, sum(launch.slots for launch in launches))
2230 def _vision_pool(self) -> list[_VisionReplica]:
2231 """Each vision replica paired with its fitted ``--parallel`` slot count.
2233 The fitted count can be lower than ``vision_ocr_concurrency`` when memory
2234 forced a smaller fit; dispatching at the configured ceiling instead
2235 over-subscribes that server. The configured ceiling applies per replica
2236 only when no matching launch snapshot exists (a reload can momentarily
2237 drop it between two reads).
2238 """
2239 self._serve_vision_on_request()
2240 clients = self._require_clients(WorkerRole.VISION)
2241 launches = self._role_launches(WorkerRole.VISION)
2242 if launches and len(launches) == len(clients):
2243 return [
2244 _VisionReplica(client, max(1, launch.slots))
2245 for client, launch in zip(clients, launches, strict=True)
2246 ]
2247 fallback_slots = max(1, cfg.vision_ocr_concurrency)
2248 return [_VisionReplica(client, fallback_slots) for client in clients]
2250 def _serve_vision_on_request(self) -> None:
2251 """Re-plan the fleet with vision for an OCR call that ``ocr`` off left out.
2253 Only a request that turned OCR on reaches vision OCR while the setting is
2254 off, so the call itself is the demand. The re-plan is the diff-driven pass
2255 a vision model change runs: it restarts the groups whose launches change,
2256 which includes chat where chat and vision share one group. A failed
2257 re-plan drops the grant, so the next page tries again.
2258 """
2259 ref = str(cfg.vision_model)
2260 if not ref:
2261 return
2262 with self._vision_request_lock:
2263 if planning.vision_role_wanted(ref):
2264 return
2265 planning.grant_vision_on_request()
2266 try:
2267 self._dispatch_reload("fleet-vision-on-request", wait=True)
2268 except BaseException:
2269 planning.revoke_vision_on_request()
2270 raise
2272 # PDF/image OCR now runs inside xberg via the registered lilbee-vision
2273 # backend (see data.extract.backends.vision_ocr); this provider only exposes
2274 # single-image vision_ocr, which that backend calls.
2276 def rerank(self, query: str, candidates: list[str]) -> list[float]:
2277 return self._with_rediscover(
2278 lambda: self._rerank_once(query, candidates), role=WorkerRole.RERANK
2279 )
2281 def _rerank_once(self, query: str, candidates: list[str]) -> list[float]:
2282 clients = self._require_clients(WorkerRole.RERANK)
2283 return _call_with_failover(clients, lambda client: client.rerank(query, candidates))
2285 # --- model management: registry / GGUF reads, no running server needed ---
2287 def supports_rerank(self) -> bool:
2288 """Serve a cross-encoder (rank pooling) or an LLM reranker (yes/no logprob)."""
2289 return True
2291 def list_models(self) -> list[str]:
2292 """List installed models from the registry."""
2293 from lilbee.app.services import get_services
2295 registry = get_services().registry
2296 return sorted(m.ref for m in registry.list_installed())
2298 def list_chat_models(self, provider: str) -> list[str]:
2299 """The local engine has no frontier-provider catalog; always ``[]``."""
2300 del provider
2301 return []
2303 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:
2304 """Not supported directly: ``lilbee.catalog`` handles GGUF downloads."""
2305 del on_progress
2306 raise NotImplementedError(
2307 f"The local engine cannot pull model {model!r}. "
2308 "Download GGUF files through the catalog or 'lilbee model pull'."
2309 )
2311 def show_model(self, model: str) -> dict[str, Any] | None:
2312 """Return model metadata from GGUF headers, or ``None`` if unresolved."""
2313 from lilbee.providers.engine_params import resolve_model_path
2314 from lilbee.providers.gguf_meta import read_gguf_metadata
2316 try:
2317 path = resolve_model_path(model)
2318 except ProviderError:
2319 return None
2320 return read_gguf_metadata(path)
2322 def get_capabilities(self, model: str) -> list[str]:
2323 """Detect capabilities from the local GGUF files.
2325 Cross-encoder rerank GGUFs report ``["rerank"]`` (they cannot generate);
2326 other models report ``"completion"`` plus ``"vision"`` when an mmproj
2327 sidecar is present.
2328 """
2329 from lilbee.catalog import is_rerank_ref
2330 from lilbee.providers.engine_params import resolve_model_path
2331 from lilbee.providers.gguf_meta import find_mmproj_for_model
2333 if model and is_rerank_ref(model):
2334 return ["rerank"]
2335 caps = ["completion"]
2336 try:
2337 path = resolve_model_path(model)
2338 except ProviderError:
2339 return caps
2340 try:
2341 find_mmproj_for_model(path)
2342 caps.append("vision")
2343 except ProviderError:
2344 pass
2345 return caps
2347 def supports_tools(self, model_ref: str) -> bool:
2348 """True iff *model_ref*'s GGUF chat template references tool tokens.
2350 The server parses native tool calls via ``--jinja``; a template that
2351 declares tools is the signal that the model was trained to emit them.
2352 Cached on ``(path, mtime)`` so a tool-bearing chat doesn't re-read the
2353 GGUF header each request; a re-quantised file at the same path
2354 invalidates because its mtime changes.
2355 """
2356 from lilbee.providers.engine_params import resolve_model_path
2358 try:
2359 path = resolve_model_path(model_ref)
2360 except (ProviderError, OSError):
2361 log.debug("supports_tools: resolve_model_path failed for %s", model_ref, exc_info=True)
2362 return False
2363 try:
2364 mtime_ns = path.stat().st_mtime_ns
2365 except OSError:
2366 mtime_ns = 0
2367 return _supports_tools_cached(str(path), mtime_ns)
2369 def warm_up_pool(self) -> None:
2370 """Pre-load every configured role off the caller's thread (idempotent).
2372 Starting the swap and loading each role's model (seconds on a cold large
2373 model) runs on a background thread and this returns at once: the eager-start
2374 at TUI mount must not freeze the UI. The spawn listeners fire per role as it
2375 loads, so the UI shows progress. A second call while warm-up is in flight
2376 (or once the fleet is up) is a no-op.
2377 """
2378 with self._lock:
2379 if self._warming:
2380 return
2381 fleet_up = bool(self._swaps)
2382 # A live swap whose model llama-swap idle-unloaded (its ttl stops only the
2383 # llama-server child, leaving the swap handle in _swaps) reports its role
2384 # cold. Re-warm so a prompt sent into that gap drives llama-swap's
2385 # on-demand reload; bailing on "swaps exist" alone stranded every later
2386 # prompt on a stale not-ready. A fully-loaded fleet still short-circuits.
2387 # The probe runs off the lock (role_ready may hit the proxy).
2388 if fleet_up and self._roles_ready():
2389 return
2390 with self._lock:
2391 if self._warming:
2392 return
2393 self._warming = True
2394 threading.Thread(
2395 target=self._warm_up_blocking,
2396 name="fleet-warm-up",
2397 daemon=True,
2398 ).start()
2400 def _roles_ready(self) -> bool:
2401 """Whether every configured role's upstream is loaded (fleet fully warm)."""
2402 with self._lock:
2403 roles = list(self._role_group)
2404 return bool(roles) and all(self.role_ready(role) for role in roles)
2406 def _warm_up_blocking(self) -> None:
2407 """Start the fleet and pre-load every role on a background thread.
2409 Runs on a daemon thread with no caller to catch failures, so a startup
2410 error is logged and swallowed: a role that can't load surfaces a
2411 user-facing ProviderError on the next call, not a thread traceback.
2413 The tracker is stamped STARTING before the fleet spawn so surfaces
2414 show the engine coming up from the first moment (spawn plus health
2415 check takes seconds and previously reported nothing), and stamped
2416 ERROR with the real reason when the warm fails before the chat warm
2417 proper begins.
2418 """
2419 try:
2420 self._warm_tracker.begin(str(cfg.chat_model))
2421 self._ensure_fleet()
2422 self._preload_roles()
2423 self._finalize_warm_if_chat_never_ran()
2424 except Exception as exc:
2425 if isinstance(exc, RuntimeError) and sys.is_finalizing():
2426 # A fast CLI exit can tear down the interpreter while this daemon
2427 # thread is still warming; the process is leaving, so drop it quietly.
2428 log.debug("Engine warm-up abandoned during interpreter shutdown: %s", exc)
2429 else:
2430 # A warm-up failure is handled (roles lazy-load on first use), so
2431 # keep the full traceback at debug: a WARNING carrying exc_info
2432 # reads like a crash for a condition the next real call recovers
2433 # from.
2434 log.warning("Engine warm-up failed: %s", exc)
2435 log.debug("Engine warm-up failure detail.", exc_info=True)
2436 self._fail_warm_unless_ready(str(exc))
2437 finally:
2438 with self._lock:
2439 self._warming = False
2441 def _fail_warm_unless_ready(self, message: str) -> None:
2442 """Stamp the warm tracker ERROR unless the chat warm already finished.
2444 A failure in a later role's preload must not clobber a chat warm that
2445 reached READY; every earlier failure leaves the tracker mid-phase,
2446 where surfaces would spin forever and the prompt path could not name
2447 the reason.
2448 """
2449 snapshot = self._warm_tracker.snapshot()
2450 if snapshot is None or snapshot.phase is not WarmPhase.READY:
2451 self._warm_tracker.fail(message)
2453 def _finalize_warm_if_chat_never_ran(self) -> None:
2454 """Terminate the early STARTING stamp when no chat instance was placed.
2456 ``_warm_chat_role`` always ends in READY or ERROR when it runs, so a
2457 snapshot still on STARTING after a successful preload means the plan had
2458 no chat instance. A chat model that isn't installed, one whose launch the
2459 plan refused for an unusable window, and one with no engine to run it all
2460 fail the warm with a user-facing reason so the prompt path renders
2461 "failed to load" instead of spinning a "not ready" retry that can never
2462 succeed; any other reason (a remote-routed chat has no local server to
2463 warm) clears the stamp.
2464 """
2465 snapshot = self._warm_tracker.snapshot()
2466 if snapshot is None or snapshot.phase is not WarmPhase.STARTING:
2467 return
2468 missing = self._skipped_not_installed.get(WorkerRole.CHAT)
2469 if missing is not None:
2470 self._warm_tracker.fail(f"chat model {clean_display_name(missing)} is not installed")
2471 return
2472 needs_engine = _chat_needs_local_engine()
2473 unusable = (self._engine_refusal if needs_engine else None) or (
2474 self._skipped_unusable_ctx.get(WorkerRole.CHAT)
2475 )
2476 if unusable is not None:
2477 self._warm_tracker.fail(unusable)
2478 return
2479 if needs_engine and (reason := _unusable_engine_reason()) is not None:
2480 self._warm_tracker.fail(reason)
2481 return
2482 self._warm_tracker.clear()
2484 def _preload_roles(self, roles: frozenset[WorkerRole] | None = None) -> None:
2485 """Issue a cheap request per replica so llama-swap loads each upstream now.
2487 llama-swap starts an upstream on its first request, so warming sends a
2488 minimal call to every replica of every role (firing the spawn listeners
2489 around each role). A per-replica failure is logged and skipped; that replica
2490 still loads on its first real use. The chat role routes through
2491 :meth:`_warm_chat_role` so a launcher gets granular progress. *roles*
2492 narrows the warm to just those roles (a reload warms only what restarted).
2494 Roles on separate devices warm concurrently: chat is the long pole (a
2495 large model's load dominates), so the light roles load alongside it
2496 instead of before it. Roles whose launches pin overlapping devices warm
2497 one at a time instead, chat last: two engines loading into the same
2498 card at once race each other for VRAM, and the loser's first load can
2499 OOM even though both fit once settled.
2500 """
2501 with self._lock:
2502 pools = {
2503 role: list(clients)
2504 for role, clients in self._clients.items()
2505 if roles is None or role in roles
2506 }
2507 on_spawning, on_spawned = self._on_spawning, self._on_spawned
2508 device_sets = _role_device_sets(
2509 launch for launches in self._launches.values() for launch in launches
2510 )
2512 listeners = (on_spawning, on_spawned)
2513 calls = [
2514 DaemonCall(
2515 functools.partial(self._warm_chain, chain, pools, listeners),
2516 name=f"fleet-preload-{'+'.join(role.value for role in chain)}",
2517 )
2518 for chain in _warm_chains(list(pools), device_sets)
2519 ]
2520 # Every chain finishes before the first chain's error, in chain order, is raised.
2521 for call in calls:
2522 call.wait()
2523 for call in calls:
2524 call.result()
2526 def _warm_chain(
2527 self,
2528 chain: list[WorkerRole],
2529 pools: dict[WorkerRole, list[LlamaServerClient]],
2530 listeners: tuple[Callable[[WorkerRole], None] | None, Callable[[WorkerRole], None] | None],
2531 ) -> None:
2532 """Warm *chain*'s roles one at a time; every role gets its attempt.
2534 An unexpected error warming one role (a listener blowing up) must not rob
2535 the roles behind it of their warm, so the first error is re-raised only
2536 after the chain finishes.
2537 """
2538 on_spawning, on_spawned = listeners
2539 first_exc: Exception | None = None
2540 for role in chain:
2541 try:
2542 if on_spawning is not None:
2543 on_spawning(role)
2544 if role is WorkerRole.CHAT:
2545 self._warm_chat_role(pools[role])
2546 else:
2547 self._warm_role_clients(role, pools[role])
2548 if on_spawned is not None:
2549 on_spawned(role)
2550 except Exception as exc:
2551 first_exc = first_exc or exc
2552 if first_exc is not None:
2553 raise first_exc
2555 def _warm_role_clients(self, role: WorkerRole, clients: list[LlamaServerClient]) -> bool:
2556 """Warm every replica of *role*; return whether at least one loaded.
2558 A replica that fails to load is reported at warning level with the engine's
2559 own message (an unsupported architecture, a corrupt file). Warm-up stays
2560 best-effort, but the failure must not be silent: the role then serves
2561 nothing, and a caller that never reaches it would otherwise see only an
2562 unexplained empty answer.
2563 """
2564 warmed = False
2565 self._warm_errors.pop(role, None)
2566 for client in clients:
2567 try:
2568 _warm_role(role, client)
2569 client.mark_healthy()
2570 warmed = True
2571 except Exception as exc:
2572 # Shutdown closes the client under an in-flight warm request. The
2573 # model did not fail, so there is nothing to report.
2574 if self._shut_down:
2575 return warmed
2576 # A replica that cannot load is not routable. Marking it takes it
2577 # out of the pool so calls go to a sibling on a device that works,
2578 # instead of every request picking the dead one again. It is a
2579 # device fault as often as a model one: an adapter that enumerates
2580 # but cannot allocate fails here and nowhere else. The health flag
2581 # carries its own cool-down, so a device that recovers rejoins
2582 # without anything having to remember it was bad.
2583 client.mark_unhealthy()
2584 self._warm_errors[role] = str(exc)
2585 log.warning(
2586 "The %s model failed to load: %s",
2587 role.value,
2588 exc,
2589 exc_info=log.isEnabledFor(logging.DEBUG),
2590 )
2591 if warmed:
2592 self._warm_errors.pop(role, None)
2593 return warmed
2595 def _warm_chat_role(self, clients: list[LlamaServerClient]) -> None:
2596 """Warm the chat role, driving the tracker through read -> load -> ready/fail.
2598 Readiness is decided by whether a warm request actually returned, not by
2599 re-probing llama-swap (which can transiently report empty right after a
2600 successful load). The terminal phase is stamped in ``finally`` so an
2601 unexpected error mid-warm still ends the launcher's progress stream.
2602 """
2603 self._warm_tracker.begin(str(cfg.chat_model))
2604 warmed = False
2605 try:
2606 self._prewarm_chat_weights()
2607 self._warm_tracker.loading_engine()
2608 warmed = self._warm_role_clients(WorkerRole.CHAT, clients)
2609 finally:
2610 if warmed:
2611 self._warm_tracker.ready()
2612 else:
2613 self._warm_tracker.fail(self._chat_load_failure())
2615 @contextmanager
2616 def _lazy_warm_scope(self) -> Iterator[None]:
2617 """Report a request-triggered chat load on the warm tracker.
2619 The request that finds the chat role cold with no warm in flight owns
2620 the warm: it stamps STARTING on entry and READY or ERROR on exit, on
2621 every exit path. Concurrent requests share the load and touch nothing.
2622 """
2623 owner = self._begin_lazy_warm()
2624 try:
2625 yield
2626 except Exception as exc:
2627 if owner:
2628 self._end_lazy_warm(str(exc))
2629 raise
2630 if owner:
2631 self._end_lazy_warm(None)
2633 def _begin_lazy_warm(self) -> bool:
2634 """Begin a warm when the chat role is cold; True when this call owns it."""
2635 with self._lock:
2636 if self._lazy_warming or self._warming:
2637 return False
2638 if is_active_warm(self._warm_tracker.snapshot()) or self.role_ready(WorkerRole.CHAT):
2639 return False
2640 with self._lock:
2641 if self._lazy_warming:
2642 return False
2643 self._lazy_warming = True
2644 self._warm_tracker.begin(str(cfg.chat_model))
2645 return True
2647 def _end_lazy_warm(self, error: str | None) -> None:
2648 """Stamp the owned warm READY when the role serves, else ERROR with *error*."""
2649 with self._lock:
2650 self._lazy_warming = False
2651 if error is None or self.role_ready(WorkerRole.CHAT):
2652 self._warm_tracker.ready()
2653 else:
2654 self._warm_tracker.fail(error)
2656 def _chat_load_failure(self) -> str:
2657 """The engine's own reason the chat model did not load, when it gave one."""
2658 reason = self._warm_errors.get(WorkerRole.CHAT)
2659 if not reason:
2660 return "The chat model did not finish loading."
2661 return f"The chat model did not load: {reason}"
2663 def _prewarm_chat_weights(self) -> None:
2664 """Page the chat model's GGUF shards into the OS cache, reporting byte progress.
2666 Reading the shards before llama-swap loads them does two things: it gives a
2667 true read-phase percentage for the warm tracker, and it warms the page cache
2668 so the engine's mmap faults hit memory (a large win on a network filesystem,
2669 where random mmap faults stalled cold loads). Best-effort: any failure to
2670 resolve or size the shards (unregistered ref, cache miss, I/O error) is
2671 skipped, and the model still loads on the warm request.
2672 """
2673 try:
2674 shards = ModelRegistry(cfg.models_dir).shard_paths(str(cfg.chat_model))
2675 total = sum(shard.stat().st_size for shard in shards)
2676 except Exception:
2677 log.debug("Prewarm skipped; could not resolve chat shards.", exc_info=True)
2678 return
2679 if total <= 0:
2680 return
2681 keys = [_prewarm_key(shard) for shard in shards]
2682 if all(key in _PREWARMED_SHARDS for key in keys):
2683 # Already paged in this boot (e.g. a placement rebuild); the cache is hot.
2684 self._warm_tracker.reading(total, total)
2685 return
2686 done = 0
2687 self._warm_tracker.reading(0, total)
2688 chunk = bytearray(_PREWARM_CHUNK_BYTES)
2689 for index, (shard, key) in enumerate(zip(shards, keys, strict=True)):
2690 detail = f"shard {index + 1}/{len(shards)}" if len(shards) > 1 else None
2691 try:
2692 with shard.open("rb", buffering=0) as handle:
2693 while True:
2694 read = handle.readinto(chunk)
2695 if not read:
2696 break
2697 done += read
2698 self._warm_tracker.reading(done, total, detail=detail)
2699 _PREWARMED_SHARDS.add(key)
2700 except OSError:
2701 # A partial/locked shard just shortens the read bar; the engine load
2702 # surfaces any real fault as a user-facing error on the warm request.
2703 log.debug("Prewarm read of %s stopped early.", shard, exc_info=True)
2705 def cancel_inference(self) -> None:
2706 """Sever every in-flight chat stream so its blocked reader unwinds.
2708 A cooperative worker cancel cannot reach a thread blocked in a socket
2709 read, and the reader's own close runs only when its worker unwinds, so
2710 the disconnect must happen here. Retired clients are swept too: a
2711 model-swap reload retires a busy client before the cancel lands.
2712 """
2713 with self._lock:
2714 clients = [*self._clients.get(WorkerRole.CHAT, ()), *self._retiring_clients]
2715 for client in clients:
2716 client.abort_streams()
2718 def reload_role(self, role: WorkerRole, *, wait: bool = False) -> None:
2719 """Apply a model/settings change for *role* with current cfg.
2721 The whole fleet is re-planned, but only the roles whose launches changed
2722 restart, so the other roles' loaded models stay resident (*role* names
2723 the change for the thread label; the diff decides what restarts). A vision
2724 setting change drops any on-request vision, so the plan follows the setting.
2725 """
2726 if role is WorkerRole.VISION:
2727 planning.revoke_vision_on_request()
2728 self._dispatch_reload(f"fleet-reload-{role.value}", wait=wait)
2730 def replan_fleet(self, *, wait: bool = False) -> None:
2731 """Apply a placement or load-setting change with current cfg, restarting only moved roles.
2733 The fresh plan is diffed per role against the running fleet: a role whose
2734 launch did not change keeps serving through the change -- moving the
2735 embedder never unloads a 100GB chat model. When no fleet is up, the next
2736 use plans fresh, so this returns at once.
2737 """
2738 self._dispatch_reload("fleet-replan", wait=wait)
2740 def _dispatch_reload(self, thread_name: str, *, wait: bool) -> None:
2741 """Run the diff-driven reload once, off-thread unless *wait*.
2743 Dispatched to a background thread because the slow restart (rewrite config +
2744 respawn + wait-ready) must not block the settings/model-picker callback.
2745 If no group is up yet, the next use starts the fleet with current cfg.
2746 Single-flight: a reload while one is in flight sets the pending flag (the
2747 in-flight pass may have already snapshotted its plan), and the in-flight
2748 thread runs one more pass per pending flag so the change is applied, not
2749 dropped.
2751 ``wait=True`` runs the reload in the caller's thread and returns only once
2752 the restart (and any reload already in flight that will run the pending
2753 pass) has finished and the proxies are healthy again, so a caller already
2754 off the event loop gets a real completion signal. A restarted role's model
2755 still loads lazily (the reload kicks an off-thread warm). It propagates a
2756 reload failure as an exception.
2757 """
2758 with self._lock:
2759 if not self._swaps:
2760 return
2761 if self._reloading:
2762 self._reload_pending = True
2763 if wait:
2764 while self._reloading:
2765 self._reload_done.wait()
2766 return
2767 self._reloading = True
2768 self._reload_owed = frozenset(WorkerRole)
2769 self._reload_pending = False
2770 if wait:
2771 self._reload_blocking()
2772 return
2773 try:
2774 threading.Thread(
2775 target=self._reload_blocking,
2776 name=thread_name,
2777 daemon=True,
2778 ).start()
2779 except RuntimeError:
2780 with self._lock:
2781 self._reloading = False
2782 self._reload_owed = frozenset()
2783 self._reload_done.notify_all()
2784 raise
2786 def _reload_blocking(self) -> None:
2787 """Run reload passes until no further reload arrived mid-pass.
2789 A failed pass with the pending flag set still runs the pending pass (the
2790 fresh plan may succeed under the new cfg); only the final pass's failure
2791 propagates, after dropping the refs to a dead swap so the next call can
2792 rebuild. The pending check and the guard release happen under one lock
2793 acquisition, so a reload_role landing between them cannot be acknowledged
2794 and dropped.
2795 """
2796 while True:
2797 try:
2798 self._reload_pass()
2799 except BaseException:
2800 with self._lock:
2801 rerun = self._reload_pending
2802 self._reload_pending = False
2803 if rerun:
2804 self._reload_owed = frozenset(WorkerRole)
2805 else:
2806 self._reloading = False
2807 self._reload_owed = frozenset()
2808 self._reload_done.notify_all()
2809 if rerun:
2810 log.warning(
2811 "Engine reload failed; retrying with the pending change.", exc_info=True
2812 )
2813 continue
2814 self._drop_dead_swaps()
2815 raise
2816 with self._lock:
2817 if not self._reload_pending:
2818 self._reloading = False
2819 self._reload_owed = frozenset()
2820 self._reload_done.notify_all()
2821 return
2822 self._reload_pending = False
2823 self._reload_owed = frozenset(WorkerRole)
2825 def _rebind_or_overflow(self) -> list[WorkerRole]:
2826 """Re-acquire a bound engine after a config change; caller holds the build lock.
2828 A provider that rode another process's engine owns none of its groups and
2829 cannot restart them. Restarting "in place" would spawn a duplicate fleet
2830 into the shared slot (a bound manager's ``shutdown`` only detaches, leaving
2831 the incumbent resident) and size it blind against VRAM the incumbent still
2832 holds. Instead drop every binding and this process's membership, then re-run
2833 the acquisition ladder: it rebinds to the reconfigured shared engine, builds
2834 fresh in the machine slot if we were its last user, or overflows to a private
2835 engine sized against a fresh probe. Returns the roles now served (to preload).
2836 """
2837 from lilbee.core.config import cfg
2839 with self._lock:
2840 groups = list(self._swaps)
2841 for group in groups:
2842 with self._lock:
2843 swap = self._drop_group(group) # also prunes _group_dirs
2844 if swap is not None:
2845 swap.shutdown() # bound: detaches; the shared engine keeps running
2846 self._release_engines() # a shared engine's builder keeps it live; no stop here
2847 if not self._acquire_engine(cfg.data_root):
2848 return []
2849 with self._lock:
2850 return list(self._role_group)
2852 def _reload_pass(self, force: frozenset[WorkerRole] = frozenset()) -> None:
2853 """One re-plan from current cfg, restarting only the groups that changed.
2855 The fresh plan is diffed per swap group against the launches each running
2856 group was started with; a group restarts only when its launches differ
2857 (covers added and removed groups too), so an untouched group's loaded model
2858 stays resident through a placement or per-role model change. *force* adds a
2859 role's group to the restart set even when its plan is unchanged (dead-swap
2860 recovery). Changed groups stop before the new ones start, so the planned
2861 VRAM is actually free when the new servers spawn. Runs under the build lock
2862 so a racing shutdown/build can't interleave with the restart and leak a live
2863 llama-swap holding GPU memory.
2864 """
2866 restarted: list[WorkerRole] = []
2867 with self._build_lock:
2868 # The device list is structural and was captured once at boot, so a
2869 # card that has since left keeps being planned onto. The memory
2870 # figures beside it are deliberately not re-taken: this fleet is
2871 # resident, and charging it against itself is what the snapshot exists
2872 # to prevent.
2873 planning.refresh_plan_devices()
2874 with self._lock:
2875 if self._shut_down:
2876 # Terminal shutdown landed while this reload was queued; a
2877 # rebuild here would spawn a fleet no live provider owns.
2878 return
2879 running = set(self._swaps)
2880 old = dict(self._launches)
2881 # All groups share one dir and one ownership by construction, so any
2882 # bound manager means this provider rides another process's engine.
2883 bound = any(swap.bound for swap in self._swaps.values())
2884 if bound:
2885 # Cannot restart a shared engine's groups in place (that duplicates
2886 # the fleet into the slot); drop the bindings and re-acquire.
2887 restarted = self._rebind_or_overflow()
2888 self._preload_restarted(restarted)
2889 return
2890 # Reap dead engines in our dirs before re-planning, as a build would.
2891 reload_dir = self._reload_dir()
2892 # Serialize the reload against peer acquisitions with the same
2893 # cross-process lock a build takes. Without it, this reap_stale can kill
2894 # a peer's swap that is spawned but not yet answering its proxy, and the
2895 # stop-then-spawn gap lets a peer's ladder see a half-stopped slot and
2896 # build a second fleet into the same dir, double-allocating VRAM.
2897 with build_lock(reload_dir):
2898 reap_stale(reload_dir)
2899 origin = _plan_origin()
2900 try:
2901 if not running:
2902 # Nothing loaded (a resurrect after a failed pass): the box is
2903 # clean, so refresh the plan snapshot like a first build would.
2904 planning.capture_plan_probe()
2905 plan = planning.plan_all_launches()
2906 except ProviderError as exc:
2907 # Same policy as the initial build: a genuinely-missing engine
2908 # binary aborts the reload quietly (nothing to serve), while any
2909 # other planning failure (a wedged GPU probe, an unusable CUDA
2910 # runtime) propagates to fail loud. The raise lands before the
2911 # stop phase, so a running fleet is left intact rather than half
2912 # torn down.
2913 if exc.kind is not ProviderErrorKind.NOT_FOUND:
2914 raise
2915 log.debug("Engine binary unavailable; reload left the fleet as-is")
2916 return
2917 # Keep the skip reasons in step with the fresh plan.
2918 self._skipped_not_installed = dict(plan.skipped_not_installed)
2919 self._skipped_unusable_ctx = dict(plan.skipped_unusable_ctx)
2920 new = _launches_by_group(plan)
2921 # A group restarts when its launches changed OR its running/planned
2922 # presence disagrees (covers a group the new plan drops or adds).
2923 changed = {
2924 group
2925 for group in running | set(new)
2926 if (group in running) != (group in new)
2927 or old.get(group, ()) != new.get(group, ())
2928 }
2929 changed |= {group_for(role, plan.co_tenants) for role in force}
2930 with self._lock:
2931 self._reload_owed = frozenset(
2932 role for group in changed & set(new) for role in _by_role(list(new[group]))
2933 )
2934 self._reload_done.notify_all()
2935 self._log_dropped_groups(changed, old, new)
2936 # Stop phase: free the changed groups' VRAM before their replacements
2937 # (or another group's grown plan) spawn against it.
2938 for group in sorted(changed, key=lambda g: g.value):
2939 with self._lock:
2940 swap = self._drop_group(group)
2941 if swap is not None:
2942 swap.shutdown()
2943 # Start phase: spawn the changed groups present in the new plan.
2944 for group in sorted(changed & set(new), key=lambda g: g.value):
2945 group_launches = list(new[group])
2946 swap = SwapManager(reload_dir, group)
2947 lifetime = _engine_lifetime(reload_dir)
2948 swap.start(
2949 group_launches,
2950 ttl_seconds=_warm_ttl_seconds(
2951 lifetime,
2952 reload_dir,
2953 hold_warm_for_session=self._hold_warm_for_session,
2954 ),
2955 lifetime=lifetime,
2956 origin=origin,
2957 )
2958 with self._lock:
2959 self._adopt_group(group, swap, group_launches)
2960 self._group_dirs[group] = reload_dir
2961 restarted.extend(_by_role(group_launches))
2962 self._restamp_kept(running - changed, origin)
2963 self._preload_restarted(restarted)
2965 def _restamp_kept(self, kept: set[SwapGroup], origin: PlanOrigin) -> None:
2966 """Record *origin* on the *kept* groups: their plan matches the file, so no peer waits."""
2967 with self._lock:
2968 swaps = [self._swaps[group] for group in kept if group in self._swaps]
2969 for swap in swaps:
2970 try:
2971 swap.restamp(origin)
2972 except OSError as exc:
2973 # Best effort: the restamp only lets waiting peers stop sooner. On Windows a
2974 # peer that holds the state file open makes the write fail.
2975 log.warning(
2976 "Could not record the new config on a running engine group (%s); "
2977 "waiting peers use their time limit instead.",
2978 exc,
2979 )
2981 @staticmethod
2982 def _log_dropped_groups(
2983 changed: set[SwapGroup],
2984 old: dict[SwapGroup, tuple[InstanceLaunch, ...]],
2985 new: dict[SwapGroup, tuple[InstanceLaunch, ...]],
2986 ) -> None:
2987 """Log what moved in each restarted group's launches, before and after."""
2988 for group in sorted(changed, key=lambda g: g.value):
2989 log.debug(
2990 "Reload restarts group %s: %s",
2991 group.value,
2992 _launch_changes(old.get(group, ()), new.get(group, ())),
2993 )
2995 def _preload_restarted(self, restarted: list[WorkerRole]) -> None:
2996 """Load the restarted roles' models off-thread (a no-op for none).
2998 llama-swap spawns an upstream on its first request, so the reload returns
2999 once the proxies answer and the UI's spawn listeners track the model loads.
3000 """
3001 if not restarted:
3002 return
3003 threading.Thread(
3004 target=self._preload_roles,
3005 kwargs={"roles": frozenset(restarted)},
3006 name="fleet-reload-warm",
3007 daemon=True,
3008 ).start()
3010 def add_spawn_listener(
3011 self,
3012 *,
3013 on_spawning: Callable[[WorkerRole], None] | None = None,
3014 on_spawned: Callable[[WorkerRole], None] | None = None,
3015 ) -> None:
3016 """Store spawn-lifecycle callbacks; warm-up fires them as each role loads."""
3017 with self._lock:
3018 self._on_spawning = on_spawning
3019 self._on_spawned = on_spawned
3021 def shutdown(self) -> None:
3022 # The on-request vision grant is process-wide; the next provider starts from the setting.
3023 planning.revoke_vision_on_request()
3024 self._shutdown_swap()