Coverage for src/lilbee/providers/fleet/planning.py: 100%
1047 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Launch planning for the fleet: device probe, VRAM estimate, placement, argv."""
3from __future__ import annotations
5import contextvars
6import logging
7import re
8import threading
9import time
10from contextlib import contextmanager
11from dataclasses import dataclass, field, replace
12from pathlib import Path
13from typing import TYPE_CHECKING, TypeAlias
15from lilbee.core.config.enums import KV_CACHE_TYPE_BYTES, KvCacheType
16from lilbee.core.system import is_network_path
17from lilbee.providers import engine_params, model_cache
18from lilbee.providers.base import ProviderError
19from lilbee.providers.fleet import ctx as fleet_ctx
20from lilbee.providers.fleet.adapters import (
21 FLAG_FIT_TARGET,
22 FLAG_KV_UNIFIED_PER_SLOT,
23 LLM_RERANK_CONCURRENCY,
24 ROLE_SPECS,
25 RoleServerSpec,
26 build_server_argv,
27 embed_spec,
28 rerank_spec,
29 resolve_rerank_mode,
30)
31from lilbee.providers.fleet.binary import (
32 engine_binary_identity,
33 engine_build_id,
34 llama_server_runtime_env,
35 resolve_llama_server,
36)
37from lilbee.providers.fleet.devices import (
38 MIB,
39 VULKAN_BACKEND,
40 FleetDevice,
41 host_lacks_nvlink,
42 probe_devices,
43 visible_env,
44)
45from lilbee.providers.fleet.launch import InstanceLaunch
46from lilbee.providers.fleet.placement import (
47 InstancePlan,
48 ModelPlacementInput,
49 PeakEstimator,
50 Placement,
51 SplitCtxFitter,
52 placement_from_spec,
53 plan_placement,
54)
55from lilbee.providers.fleet.placement_spec import PlacementError, PlacementSpec
56from lilbee.providers.fleet.readback import engine_advertises, supports_memory_readback
57from lilbee.providers.fleet.replicas import resolve_replica_count
58from lilbee.providers.fleet.vram import estimate_instance_footprint, usable_vram_fraction
59from lilbee.providers.model_cache import free_system_memory, total_system_memory
60from lilbee.providers.model_ref import parse_model_ref
61from lilbee.providers.roles import ROLE_REGISTRY, EngineBackend, RerankMode, WorkerRole
62from lilbee.runtime.progress.types import OcrBackendUsed
64log = logging.getLogger(__name__)
66if TYPE_CHECKING:
67 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
69# Fleet-only concurrency: continuous-batching slots (--parallel) per server.
70_CHAT_SLOTS = 4
71# The free memory llama-server's own fit leaves on a card when ``--fit-target`` is not given.
72ENGINE_FIT_MARGIN_MIB = 1024
73# Roles the warm-up loads before chat on a card they share (provider._warm_chains).
74# Their memory is gone from the card's free memory when chat's fit runs.
75_LOADED_BEFORE_CHAT = frozenset({WorkerRole.EMBED, WorkerRole.RERANK})
76# Stand-in in the launch log for a host whose devices could not be read.
77_NO_DEVICES = "none"
80def log_engine_launch(launch: InstanceLaunch, *, owner_pid: int | None = None) -> None:
81 """Log the binary, build, backend and devices serving *launch*.
83 *owner_pid* is the engine's owner when this process adopted it rather than
84 spawned it.
85 """
86 if owner_pid is None:
87 log.info("Launched %s serving %s on %s", launch.model_id, launch.model, _engine_id(launch))
88 return
89 log.info(
90 "Adopted %s serving %s from engine pid %d on %s",
91 launch.model_id,
92 launch.model,
93 owner_pid,
94 _engine_id(launch),
95 )
98def _engine_id(launch: InstanceLaunch) -> str:
99 """*launch*'s binary with the engine build, backend, and probed devices."""
100 devices = probed_devices()
101 names = ", ".join(f"{d.backend}{d.index}: {d.name}" for d in devices) or _NO_DEVICES
102 return (
103 f"{launch.binary} (build {engine_build_id()}, "
104 f"backend {engine_backend().value}, devices: {names})"
105 )
108def warn_when_embed_window_below_chunk(launch: InstanceLaunch) -> None:
109 """Log when the embed engine's input cap is below the configured chunk budget.
111 Runs at engine adoption so the serving process's own log names the window
112 it serves; the health route carries the same warning for clients.
113 """
114 if launch.token_cap is None:
115 return
116 warning = engine_params.embed_window_warning(launch.token_cap)
117 if warning is not None:
118 log.warning("%s", " ".join(filter(None, (warning.message, warning.remedy))))
121def warn_when_chat_downsized(launch: InstanceLaunch) -> None:
122 """Log when a chat engine's granted shape ends below the requested one.
124 Runs at engine adoption so the warning lands on the serving process's
125 own log, where the operator of that engine reads it.
126 """
127 below_ctx = launch.built_ctx_target > 0 and launch.ctx < launch.built_ctx_target
128 below_slots = launch.slots < _CHAT_SLOTS
129 if not (below_ctx or below_slots):
130 return
131 # A pre-field launch record carries built_ctx_target=0; show the granted
132 # window as the requested one rather than "x 0 context". The remedy names
133 # both target knobs because the recorded target does not say which of
134 # num_ctx (pin) or chat_n_ctx_target produced it.
135 requested_ctx = launch.built_ctx_target or launch.ctx
136 log.warning(
137 "Chat engine downsized: serving %d slot(s) x %d context "
138 "(requested: up to %d slots x %d context). Requests beyond %d at once "
139 "wait for a free slot, so parallel agents can look stalled. "
140 "A smaller model or a lower context target (num_ctx or "
141 "chat_n_ctx_target) frees room for more slots.",
142 launch.slots,
143 launch.ctx,
144 _CHAT_SLOTS,
145 requested_ctx,
146 launch.slots,
147 )
150# Slots the PLACEMENT estimate reserves KV for on a tensor-split chat: one full
151# window, the minimum any split must hold. The launch may serve more than this
152# (see _resolve_split_chat_slots) when the cards measurably have room for several
153# full windows, so this is a planning floor, not the served slot count.
154_SPLIT_CHAT_SLOTS = 1
155# Floor context the PLACEMENT estimate reserves KV for, so a large model is never
156# single-carded into a KV corner too small for real use (a 17GB model on a 24GB
157# card leaves ~no KV room -> n_ctx collapses to a few hundred tokens). Sizing the
158# placement reserve against this floor forces a tensor-split when one card cannot
159# hold weights + a usable context; the served ctx is then grown by resolve_chat_ctx
160# (single) / fit_split_ctx (split) toward the cards' real headroom, with each
161# sequence capped at the working-context target. A split may then serve several
162# such sequences, so the served total can exceed the single-window reserve; the
163# per-device headroom test in fit_split_ctx is what bounds it. A fit still at
164# its floor after the forced split is refused by plan_launches
165# (_unusable_chat_ctx_reason).
166_MIN_USABLE_CHAT_CTX = 8192
167# Embed and cross-encoder rerank serve one request at a time. Raising it was
168# tried and measured worse: on 8xA40 with an 8B Q8 embedder and one ~100-token
169# passage per request, --parallel 1 gave 133 docs/sec at 81% SM while
170# --parallel 8 gave 100 at 63%. At one slot the card is already busy, so there is
171# no stall for slot-batching to reclaim and the extra slots only add
172# continuous-batching and KV-fragmentation overhead. Batch on the request side
173# (embed_batch_sequences) instead.
174_AUX_SLOTS = 1
175# A tensor-split needs at least this many GPUs; below it the chat context objective
176# (a gguf read) is pointless because the model can only single-card or stay unplaced.
177_MIN_SPLIT_GPUS = 2
178# Pooled single-slot search roles (embed/cross-encoder rerank) whose whole input
179# batches in one pass; derived from the role registry.
180_EMBED_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.pooled)
181# Roles whose loaders offload every layer regardless of cfg.n_gpu_layers; only
182# chat honors cfg.n_gpu_layers.
183_ALL_LAYER_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.offload_all_layers)
184_FLASH_ON = "on"
185_FLASH_OFF = "off"
186_FLASH_AUTO = "auto"
187# llama-server's documented way to say "offload nothing": --device none.
188_NO_DEVICE = "none"
189# Backends pinned by the name the engine printed rather than through an env var,
190# because their variables index a different space than --list-devices reports.
191_NAME_PINNED_BACKENDS = frozenset({VULKAN_BACKEND, "SYCL"})
192# Backends whose flash-attention coverage in llama.cpp is complete enough to ask
193# for it outright. Vulkan and SYCL are behind CUDA's and have been incomplete on
194# Intel's mesa driver, so those are left to the engine's own auto, which enables
195# flash attention only where the backend really supports it.
196_TRUSTED_FLASH_BACKENDS = frozenset({"CUDA", "ROCm", "HIP", "MTL", "Metal"})
197# Roles to which flash attention applies; embed/rerank run without it.
198_FLASH_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.flash_attn)
201# Cap vision's own KV footprint at this fraction of usable VRAM when sizing its
202# batching slots, leaving room for the weights and any co-located role.
203_VISION_VRAM_FRACTION = 0.5
205# Cap an LLM reranker's footprint at this fraction of usable VRAM when sizing its
206# slots; its per-slot ctx is tiny, so a normal GPU fits the full fan-out and a
207# small one steps down toward 1.
208_LLM_RERANK_VRAM_FRACTION = 0.5
210# RAM kept free for the OS when placing against system memory (no discrete GPU):
211# a quarter of total RAM, capped at 4 GiB. A fixed 4 GiB floor leaves a small
212# host (7-8 GB) with no budget at all, refusing to serve even tiny models.
213_SYSTEM_MEMORY_FLOOR_DIVISOR = 4
214# A GPU driver still initializing at boot answers with no devices. Ask again
215# before letting that decide the daemon's whole run; two extra probes cost a
216# couple of seconds only on a host that has a card the engine could not see.
217_PROBE_RETRIES = 2
218_PROBE_RETRY_DELAY_S = 1.0
220# A network filesystem makes mmap dangerous (page faults served over the wire can
221# wedge the loader in uninterruptible I/O), so the chat server loads its weights
222# into a malloc'd host copy (the no-mmap flag the binary lists, --load-mode none
223# or --no-mmap) whenever that copy fits in this fraction of total system RAM.
224# Local disk keeps mmap: its lazy paging gives a faster first token on a cold
225# cache -- the common desktop first launch -- and the no-mmap buffered full read
226# only wins on an already-hot cache (#474: 33s vs 43s for a 112GB model on 3
227# GPUs) while pessimizing cold start. Keyed on TOTAL memory
228# (stable), not free (fluctuates), so replans do not flap the launch argv. The
229# exact ceiling is tuned on a network-volume host.
230_NO_MMAP_NETWORK_RAM_FRACTION = 0.85
232# llama.cpp split-GGUF shard naming ("%s-%05d-of-%05d.gguf"); the cold-load
233# timeout must scale with the SUM of the shards, not the first file alone.
234_SPLIT_GGUF_NAME = re.compile(r"^(?P<prefix>.+)-(?P<index>\d{5})-of-(?P<total>\d{5})\.gguf$")
237def _weights_bytes(model_path: Path) -> int:
238 """Total weights size on disk; a split GGUF sums every sibling shard."""
239 match = _SPLIT_GGUF_NAME.fullmatch(model_path.name)
240 if match is None:
241 return model_path.stat().st_size
242 return sum(
243 sibling.stat().st_size
244 for sibling in model_path.parent.iterdir()
245 if _is_sibling_shard(sibling.name, match)
246 )
249def _is_sibling_shard(name: str, match: re.Match[str]) -> bool:
250 """Whether *name* is a shard of the same split GGUF as *match*."""
251 shard = _SPLIT_GGUF_NAME.fullmatch(name)
252 return (
253 shard is not None
254 and shard["prefix"] == match["prefix"]
255 and shard["total"] == match["total"]
256 )
259def _slots_for(
260 role: WorkerRole,
261 model_path: Path,
262 ctx: int,
263 *,
264 mmproj_path: Path | None = None,
265 unified_budget: int | None = None,
266 chat_reservation: int = 0,
267 rerank_mode: RerankMode | None = None,
268 device: FleetDevice | None = None,
269) -> int:
270 """Continuous-batching slots (--parallel) for a role's server.
272 Chat batches concurrent turns; vision batches concurrent OCR pages since a
273 one-page decode underutilizes the GPU; an LLM reranker batches its per-candidate
274 chat requests; embed and cross-encoder rerank are single-slot (their batching is
275 request-side). The memory-aware roles drop toward 1 on a small or shared host
276 instead of overcommitting. ``unified_budget`` caps sizing against free system RAM
277 with no discrete GPU; ``chat_reservation`` is the search-role footprint held back
278 from chat; ``device`` is the card the role was placed on, whose memory the
279 budget comes from once placement has chosen one.
280 """
281 if role is WorkerRole.CHAT:
282 return _resolve_chat_slots(
283 model_path,
284 ctx,
285 mmproj_path=mmproj_path,
286 unified_budget=unified_budget,
287 chat_reservation=chat_reservation,
288 device=device,
289 )
290 if role is WorkerRole.VISION:
291 return _resolve_vision_slots(
292 model_path, ctx, mmproj_path=mmproj_path, unified_budget=unified_budget, device=device
293 )
294 if role is WorkerRole.RERANK and rerank_mode is RerankMode.LLM:
295 return _resolve_llm_rerank_slots(
296 model_path, ctx, unified_budget=unified_budget, device=device
297 )
298 return _AUX_SLOTS
301def _resolve_split_chat_slots(fit_fn: Callable[[int], int]) -> tuple[int, int]:
302 """Largest split-chat slot count whose sequences each keep the full window.
304 ``fit_fn(n)`` is the per-slot context that fits when serving ``n`` sequences
305 (``fit_split_ctx``, capped at the working target and verified against real
306 per-card headroom). More slots divide the KV, so a split whose cards hold
307 several full windows can serve that many agents concurrently instead of one.
308 Returns ``(slots, per_slot_ctx)``, falling to one slot when only one full
309 window fits (or the fit degenerated to the floor), which preserves the
310 max-context single-sequence behaviour on a tight card.
312 Found by bisection rather than a scan because every ``fit_fn`` call is a
313 complete binary search whose probes each shell out to gguf-parser, and the
314 whole thing runs while this process holds the cross-process build lock that
315 every other lilbee start waits on without a deadline. A descending scan paid
316 for all of ``_CHAT_SLOTS - 1`` searches in exactly the tight-card case where
317 none of them fit. Bisection is sound here because the fit is non-increasing
318 in the slot count: more sequences divide the same headroom, so once a count
319 fails no larger one can succeed.
320 """
321 full = fit_fn(1)
322 if full <= model_cache._DYNAMIC_CTX_FLOOR:
323 return 1, full
324 low, high = 1, _CHAT_SLOTS
325 while low < high:
326 mid = (low + high + 1) // 2
327 if fit_fn(mid) >= full:
328 low = mid
329 else:
330 high = mid - 1
331 return low, full
334def _resolve_chat_slots(
335 model_path: Path,
336 ctx: int,
337 *,
338 mmproj_path: Path | None = None,
339 unified_budget: int | None = None,
340 chat_reservation: int = 0,
341 device: FleetDevice | None = None,
342) -> int:
343 """Largest chat slot count (<= ``_CHAT_SLOTS``) whose footprint fits the budget
344 after reserving the search roles; steps to 1 when none fit.
346 The budget is the whole serve budget less the search roles' measured
347 footprint. ``cfg.gpu_memory_fraction`` is already the margin held back from
348 the card, and the room for co-located embed/rerank is ``chat_reservation``,
349 which is what those servers were sized at rather than a flat share. A second
350 fraction on top charged that room twice and left a fifth of the card unused
351 on a box with no search roles at all.
352 """
353 budget = _slot_budget(unified_budget, device) - chat_reservation
354 return _fit_slots(
355 _CHAT_SLOTS,
356 WorkerRole.CHAT,
357 model_path,
358 ctx,
359 mmproj_path=mmproj_path,
360 unified=unified_budget is not None,
361 budget=budget,
362 )
365def _resolve_vision_slots(
366 model_path: Path,
367 ctx: int,
368 *,
369 mmproj_path: Path | None = None,
370 unified_budget: int | None = None,
371 device: FleetDevice | None = None,
372) -> int:
373 """Largest OCR batching slot count (<= ``cfg.vision_ocr_concurrency``) that fits
374 the memory budget; 1 when the ceiling is 1 or nothing larger fits."""
375 from lilbee.core.config import cfg
377 ceiling = max(1, cfg.vision_ocr_concurrency)
378 if ceiling == 1:
379 return 1
380 return _fit_slots(
381 ceiling,
382 WorkerRole.VISION,
383 model_path,
384 ctx,
385 mmproj_path=mmproj_path,
386 unified=unified_budget is not None,
387 budget=_slot_budget(unified_budget, device, vram_fraction=_VISION_VRAM_FRACTION),
388 )
391def _resolve_llm_rerank_slots(
392 model_path: Path,
393 ctx: int,
394 *,
395 unified_budget: int | None = None,
396 device: FleetDevice | None = None,
397) -> int:
398 """Largest LLM-reranker slot count (<= ``LLM_RERANK_CONCURRENCY``) that fits the
399 memory budget; 1 when nothing larger fits. Matches the client's request fan-out."""
400 return _fit_slots(
401 LLM_RERANK_CONCURRENCY,
402 WorkerRole.RERANK,
403 model_path,
404 ctx,
405 mmproj_path=None,
406 unified=unified_budget is not None,
407 budget=_slot_budget(unified_budget, device, vram_fraction=_LLM_RERANK_VRAM_FRACTION),
408 rerank_mode=RerankMode.LLM,
409 )
412def _slot_budget(
413 unified_budget: int | None,
414 device: FleetDevice | None = None,
415 *,
416 vram_fraction: float = 1.0,
417) -> int:
418 """Memory budget for slot sizing: the usable memory on *device* (the fleet's
419 smallest when placement has not chosen one yet), capped by ``unified_budget``
420 (free system RAM) when there is no discrete GPU so the count steps down to fit
421 free memory instead of overcommitting.
423 *vram_fraction* takes less than the whole for a role that shares its card by
424 design. Chat takes all of it and subtracts what the search roles were sized
425 at, which is the same room stated as a measurement rather than a share."""
426 budget = int(plan_sizing_budget(device) * vram_fraction)
427 if unified_budget is not None:
428 budget = min(budget, unified_budget)
429 return budget
432def _fit_slots(
433 ceiling: int,
434 role: WorkerRole,
435 model_path: Path,
436 ctx: int,
437 *,
438 mmproj_path: Path | None,
439 unified: bool,
440 budget: int,
441 rerank_mode: RerankMode | None = None,
442) -> int:
443 """Largest slot count in ``1..ceiling`` whose instance footprint fits *budget*;
444 1 when none larger fit."""
445 from lilbee.providers.base import ProviderError
447 for slots in range(ceiling, 1, -1):
448 try:
449 est = estimate_instance_footprint(
450 model_path,
451 ctx=ctx,
452 slots=slots,
453 gpu_layers=_role_gpu_layers(role),
454 flash_attn=_role_flash(role, rerank_mode),
455 kv_cache_type=_role_kv_cache_type(role),
456 mmproj_path=mmproj_path,
457 expert_offload=_role_expert_offload(model_path),
458 )
459 except (ProviderError, OSError):
460 # An unsizable model runs a single slot; the load decides the rest.
461 return 1
462 if est.footprint(unified=unified) <= budget:
463 return slots
464 return 1
467def fit_chat_ctx(
468 model_path: Path,
469 meta: dict[str, str] | None,
470 *,
471 available_bytes: int,
472 ctx_ceiling: int,
473) -> int:
474 """The chat window gguf-parser says *available_bytes* backs at the configured offload.
476 Sizes one slot, as :func:`_slots_for` then grows the slot count against
477 what the window leaves. A window below a grounded prompt is raised to one
478 when the engine's own fit will move layers and experts off the card to
479 hold it (:func:`_engine_fits_chat_offload`). Raises when the estimator
480 cannot answer, which sends :func:`engine_params.resolve_chat_ctx` to its
481 header-math fallback.
482 """
483 ctx = fleet_ctx.fit_single_ctx(
484 model_path,
485 meta=meta,
486 slots=1,
487 available_bytes=available_bytes,
488 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
489 flash_attn=_role_flash(WorkerRole.CHAT),
490 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
491 unified=plan_sizing_is_unified(),
492 ctx_ceiling=ctx_ceiling,
493 expert_offload=_role_expert_offload(model_path),
494 )
495 needed = engine_params.min_usable_chat_ctx()
496 if (
497 ctx >= needed
498 or ctx_ceiling < needed
499 or not _engine_fits_chat_offload(model_path, _engine_binary())
500 ):
501 return ctx
502 quantum = model_cache._DYNAMIC_CTX_QUANTUM
503 return min(ctx_ceiling, -(-needed // quantum) * quantum)
506def _engine_fits_chat_offload(model_path: Path, binary: Path | None) -> bool:
507 """Whether llama-server's own fit places the chat layers at load.
509 It does for an auto layer count with no expert overrides, on a card with
510 memory of its own, when *binary* lists the fit flags in its ``--help``; a
511 numeric ``n_gpu_layers`` or tensor override stops it.
512 """
513 return (
514 _role_gpu_layers(WorkerRole.CHAT) == engine_params.N_GPU_LAYERS_AUTO
515 and not _role_expert_offload(model_path)
516 and not plan_sizing_is_unified()
517 and binary is not None
518 and engine_advertises(binary, FLAG_FIT_TARGET)
519 )
522def _engine_binary() -> Path | None:
523 """The llama-server lilbee launches, or None on a host without one."""
524 try:
525 return resolve_llama_server()
526 except ProviderError:
527 return None
530def _chat_host_estimate(model_path: Path, ctx: int, slots: int) -> int:
531 """Host bytes gguf-parser expects chat to hold at its configured offload; 0 when unsized."""
532 try:
533 estimate = estimate_instance_footprint(
534 model_path,
535 ctx=ctx,
536 slots=slots,
537 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
538 flash_attn=_role_flash(WorkerRole.CHAT),
539 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
540 expert_offload=_role_expert_offload(model_path),
541 )
542 except (ProviderError, OSError):
543 return 0
544 return estimate.ram_bytes
547def _role_ctx(
548 role: WorkerRole,
549 model_path: Path,
550 meta: dict[str, str] | None,
551 device: FleetDevice | None = None,
552) -> int:
553 """Per-slot context for a role, derived as the in-process loader does.
555 Embed/rerank use the embedding model's training context; vision uses the
556 vision loader's training-context picker; chat honors ``cfg.num_ctx`` then
557 falls back to the single-GPU dynamic chat-ctx picker, sized against *device*
558 once placement has chosen one. A tensor-split chat is sized against its
559 per-device headroom instead (see :func:`fit_split_ctx`).
560 """
561 from lilbee.core.config import cfg
563 if role is WorkerRole.EMBED:
564 return engine_params.resolve_embed_ctx(meta, model_path)
565 if role is WorkerRole.RERANK:
566 if _rerank_mode_for(meta) is RerankMode.LLM:
567 return engine_params.resolve_llm_rerank_ctx(meta, model_path)
568 return engine_params.resolve_embed_ctx(meta, model_path)
569 if role is WorkerRole.VISION:
570 return engine_params.resolve_vision_ctx(model_path)
571 if cfg.num_ctx is not None:
572 return _pinned_chat_ctx(model_path, meta)
573 return engine_params.resolve_chat_ctx(
574 model_path, meta, available_bytes=plan_sizing_budget(device)
575 )
578def _pinned_chat_ctx(model_path: Path, meta: dict[str, str] | None) -> int:
579 """``cfg.num_ctx``, clamped to what the model was trained for.
581 Every unpinned resolver already clamps, and both docstrings here claimed the
582 pin did too. It did not, so a pin past the trained window was passed straight
583 to the engine, which clamps it silently and serves a different number than
584 every budget was sized for.
586 Only against a window that is actually known. A GGUF whose header cannot be
587 read falls back to a default that is a guess, and contradicting an explicit
588 pin with a guess would break the hosts where the header is the thing that is
589 broken.
590 """
591 from lilbee.core.config import cfg
593 pinned = cfg.num_ctx
594 assert pinned is not None # noqa: S101 - callers check; this documents the contract
595 ceiling = _known_chat_ceiling(model_path, meta)
596 if ceiling is None or pinned <= ceiling:
597 return pinned
598 log.warning(
599 "num_ctx is set to %d but %s was trained for %d, so %d is what will be served. "
600 "Lower num_ctx to stop planning against a window this model does not have.",
601 pinned,
602 model_path.name,
603 ceiling,
604 ceiling,
605 )
606 return ceiling
609def _known_chat_ceiling(model_path: Path, meta: dict[str, str] | None) -> int | None:
610 """The largest chat window this model is known to support, or ``None``.
612 ``None`` when the GGUF header gave no usable context length and the user set
613 no ``cfg.num_ctx_max``: there is then no measured ceiling, only a default.
614 """
615 from lilbee.core.config import cfg
616 from lilbee.providers.gguf_meta import train_ctx_from_meta
618 sentinel = -1
619 trained = train_ctx_from_meta(meta, fallback=sentinel, model_path=model_path)
620 known = [value for value in (trained, cfg.num_ctx_max) if value is not None and value > 0]
621 return min(known) if known else None
624def _rerank_mode_for(meta: dict[str, str] | None) -> RerankMode:
625 """Resolve the RERANK serving mode from cfg + the reranker GGUF arch."""
626 from lilbee.core.config import cfg
628 arch = meta.get("architecture") if meta else None
629 return resolve_rerank_mode(cfg.reranker_type, arch)
632def _role_rerank_mode(role: WorkerRole, meta: dict[str, str] | None) -> RerankMode | None:
633 """The RERANK serving mode for *role*, or ``None`` for every other role."""
634 return _rerank_mode_for(meta) if role is WorkerRole.RERANK else None
637def _server_spec(
638 role: WorkerRole, rerank_mode: RerankMode | None, meta: dict[str, str] | None
639) -> RoleServerSpec:
640 """The llama-server spec for a launch: rerank mode, decoder-aware embed pooling,
641 or the role default. EMBED forces ``--pooling last`` for decoder-only archs."""
642 if rerank_mode is not None:
643 return rerank_spec(rerank_mode)
644 if role is WorkerRole.EMBED:
645 return embed_spec(meta)
646 return ROLE_SPECS[role]
649def _pooled_batch_size(role: WorkerRole, rerank_mode: RerankMode | None, ctx: int) -> int | None:
650 """The ``--batch-size``/``--ubatch-size`` the launch raises for pooled
651 embed/cross-encoder rerank (the full context), or ``None`` for other roles."""
652 if role in _EMBED_ROLES and rerank_mode is not RerankMode.LLM:
653 return ctx
654 return None
657def _role_gpu_layers(role: WorkerRole) -> int:
658 """GPU-layer offload: chat honors ``cfg.n_gpu_layers``, others offload all layers."""
660 return engine_params.resolve_n_gpu_layers(embedding=role in _ALL_LAYER_ROLES)
663def _flash_enabled() -> bool:
664 """Flash attention is on unless ``cfg.flash_attention`` is explicitly ``False``."""
665 from lilbee.core.config import cfg
667 return cfg.flash_attention is not False
670def probed_devices() -> tuple[FleetDevice, ...]:
671 """Devices the engine enumerated, empty when they could not be read.
673 Prefers the plan snapshot so a whole planning pass answers consistently, and
674 falls back to the short-TTL read cache rather than a fresh probe.
675 """
676 probe = _current_plan_probe()
677 if probe is not None:
678 return probe.devices
679 try:
680 return tuple(_read_device_cache.get(resolve_llama_server()).devices)
681 except (ProviderError, OSError):
682 return ()
685def engine_backend() -> EngineBackend:
686 """The backend the engine selected, UNKNOWN when it could not be asked.
688 Prefers the plan snapshot so a whole planning pass answers consistently, and
689 reads it exactly as :func:`probed_devices` does, through the restating reader:
690 a snapshot taken against another engine binary re-probes here too, so the
691 answer does not depend on which reader a caller happens to reach first.
693 The placement route reports the backend of the reading its device list came
694 from instead (:func:`resolve_placement_plan`), so one payload never pairs a
695 live device list with a snapshot's backend.
696 """
697 probe = _current_plan_probe()
698 if probe is not None:
699 return probe.backend
700 try:
701 return _read_device_cache.get(resolve_llama_server()).backend
702 except (ProviderError, OSError):
703 return EngineBackend.UNKNOWN
706def _fleet_backend() -> str | None:
707 """The engine backend this host plans onto, or ``None`` when unknown."""
708 return next((device.backend for device in probed_devices()), None)
711def _flash_attention_is_trusted() -> bool:
712 """Whether to ask for flash attention outright rather than let the engine decide.
714 Unknown backends answer yes, which keeps every host that works today on the
715 argv it has now; only the backends known to lag get the engine's own auto.
716 """
717 backend = _fleet_backend()
718 return backend is None or backend in _TRUSTED_FLASH_BACKENDS
721def flash_attn_flag() -> str:
722 """``--flash-attn`` argv value for chat and vision."""
723 if not _flash_enabled():
724 return _FLASH_OFF
725 return _FLASH_ON if _flash_attention_is_trusted() else _FLASH_AUTO
728def _role_launches_with_flash(role: WorkerRole, rerank_mode: RerankMode | None = None) -> bool:
729 """Whether the launch asks the engine for flash attention on *role*.
731 The one place that answers this. The registry marks RERANK as a non-flash
732 role because a cross-encoder pools in one batch, but an LLM reranker is
733 generative and launches exactly like chat, so the mode decides there.
734 """
735 if role is WorkerRole.RERANK:
736 return rerank_mode is RerankMode.LLM
737 return role in _FLASH_ROLES
740def _role_flash(role: WorkerRole, rerank_mode: RerankMode | None = None) -> bool:
741 """Whether the estimate may assume flash attention for *role*.
743 The launch's own answer, narrowed to a definite ``on``. Under ``auto`` the
744 engine decides at load time, and assuming it would size the KV cache below
745 what the launch may need.
746 """
747 return _role_launches_with_flash(role, rerank_mode) and flash_attn_flag() == _FLASH_ON
750def _role_kv_cache_type(role: WorkerRole) -> KvCacheType:
751 """The one KV cache type *role* launches with, for both K and V.
753 Chat carries ``cfg.kv_cache_type`` only under flash attention ``on``; every
754 other role, and chat under ``auto`` or ``off``, is f16. llama-server refuses
755 K != V on MLA models, forces flash attention on for a quantized V under
756 ``auto``, and refuses a quantized V under ``off``.
757 """
758 from lilbee.core.config import cfg
760 if role is WorkerRole.CHAT and flash_attn_flag() == _FLASH_ON:
761 return cfg.kv_cache_type
762 return KvCacheType.F16
765def launch_settings(
766 role: WorkerRole,
767 rerank_mode: RerankMode | None,
768 meta: dict[str, str] | None,
769 model_path: Path,
770) -> str:
771 """The settings a launch of *role* bakes in, as the text two processes compare.
773 Only what this role's launch asks for, resolved from config and the model's
774 metadata, never from a device probe. The text records the request: chat carries
775 ``kv_cache_type`` whenever flash attention is not switched off, because whether
776 the argv applies it depends on the probed backend. A cache type change under
777 flash ``auto`` can restart a worker whose argv stays the same.
778 Chat and vision ctx sizing is not here
779 (see contract.chat_ctx_covers). The model is not here: the bind matches
780 (role, model) pairs.
781 """
782 from lilbee.core.config import cfg
784 spec = PlacementSpec.from_json(cfg.placement) if cfg.placement else None
785 flash = _flash_enabled() if _role_launches_with_flash(role, rerank_mode) else None
786 values: dict[str, object] = {
787 "gpu_devices": cfg.gpu_devices,
788 "placement": spec.roles.get(role) if spec else None,
789 "n_gpu_layers": _role_gpu_layers(role),
790 "cpu_moe": expert_offload_all(meta),
791 "n_cpu_moe": expert_offload_layers(meta),
792 "flash_attention": flash,
793 "kv_cache_type": cfg.kv_cache_type if role is WorkerRole.CHAT and flash else None,
794 "rerank_mode": rerank_mode,
795 "embed_window": _role_ctx(role, model_path, meta) if role in _EMBED_ROLES else None,
796 }
797 return ";".join(f"{key}={value}" for key, value in values.items())
800def _replica_count(role: WorkerRole, device_count: int) -> int:
801 """Requested data-parallel instances for *role* via the shared resolver."""
802 return resolve_replica_count(role, device_count)
805def chat_cache_type_flag() -> str | None:
806 """``--cache-type-k`` and ``--cache-type-v`` for chat; ``None`` leaves the f16 default."""
807 kind = _role_kv_cache_type(WorkerRole.CHAT)
808 return None if kind is KvCacheType.F16 else kind.value
811def vision_mmproj(model_ref: str) -> Path | None:
812 """Resolve a vision model's mmproj sidecar, or ``None`` if absent."""
813 from lilbee.providers.base import ProviderError
814 from lilbee.providers.gguf_meta import find_mmproj_for_model
816 try:
817 return find_mmproj_for_model(engine_params.resolve_model_path(model_ref))
818 except (ProviderError, OSError, ValueError, KeyError):
819 return None
822def _estimate_role(
823 role: WorkerRole,
824 model_ref: str,
825 *,
826 slots: int | None = None,
827 unified_budget: int | None = None,
828 chat_reservation: int = 0,
829 device_count: int = 0,
830) -> ModelPlacementInput:
831 """Estimate one role-model's footprint via gguf-parser (+ mmproj for vision).
833 ``slots`` defaults to the role's resolved batching slots (chat and vision are
834 memory-aware); ``chat_reservation`` shrinks chat to leave room for the search
835 roles; ``device_count`` resolves an auto (0) replica knob to one per GPU.
836 Charges the unified footprint with no discrete GPU, else the VRAM one.
837 """
838 from lilbee.providers.gguf_meta import read_gguf_metadata
840 path = engine_params.resolve_model_path(model_ref)
841 mmproj = vision_mmproj(model_ref) if role is WorkerRole.VISION else None
842 meta = read_gguf_metadata(path)
843 # Size the single-instance footprint against the placement reserve (a usable
844 # KV floor for chat), so a model that fits weights-only on one card but not
845 # weights + a usable context falls through to a tensor-split instead of being
846 # single-carded into a tiny n_ctx. Non-chat roles keep their launch ctx.
847 ctx = _placement_estimate_ctx(role, path, meta)
848 rerank_mode = _role_rerank_mode(role, meta)
849 if slots is None:
850 slots = _slots_for(
851 role,
852 path,
853 ctx,
854 mmproj_path=mmproj,
855 unified_budget=unified_budget,
856 chat_reservation=chat_reservation,
857 rerank_mode=rerank_mode,
858 )
859 est = estimate_instance_footprint(
860 path,
861 ctx=ctx,
862 slots=slots,
863 gpu_layers=_role_gpu_layers(role),
864 flash_attn=_role_flash(role, rerank_mode),
865 kv_cache_type=_role_kv_cache_type(role),
866 mmproj_path=mmproj,
867 batch_size=_pooled_batch_size(role, rerank_mode, ctx),
868 expert_offload=_role_expert_offload(path),
869 )
870 fp = est.footprint(unified=unified_budget is not None)
871 if role is WorkerRole.CHAT and unified_budget is None:
872 fp = _chat_serve_budget_footprint(fp)
873 return ModelPlacementInput(
874 role=role,
875 est_vram_bytes=fp,
876 replicas=_replica_count(role, device_count),
877 est_ram_bytes=est.ram_bytes,
878 )
881def _chat_serve_budget_footprint(footprint: int) -> int:
882 """Charge a chat instance against the serve budget, not the placement headroom.
884 The planner fits instances within ``cfg.usable_vram_fraction`` of a card, but a
885 single-card chat then sizes its KV cache against the smaller
886 ``cfg.gpu_memory_fraction`` budget (``resolve_chat_ctx``). A model that fills a
887 card at 0.9 leaves no room for KV at 0.75 and collapses to a few hundred tokens,
888 so scale its placement footprint by the budget ratio: it then needs a
889 tensor-split (pooling VRAM across cards) whenever single-carding it would starve
890 its context. Small models are unaffected -- they fit the serve budget with KV
891 room to spare.
892 """
893 from lilbee.core.config import cfg
895 # Never below 1.0. The ratio only compensates while the serve budget is the
896 # smaller of the two; a gpu_memory_fraction raised past the usable fraction
897 # inverts it, and the same line that exists to charge chat more starts
898 # charging it less than the model takes.
899 return int(footprint * max(1.0, usable_vram_fraction() / cfg.gpu_memory_fraction))
902def _placement_estimate_ctx(role: WorkerRole, model_path: Path, meta: dict[str, str] | None) -> int:
903 """Per-slot context the placement estimate sizes a role against.
905 For chat this reserves KV for a usable floor (``_MIN_USABLE_CHAT_CTX``, or the
906 user's ``cfg.num_ctx`` pin), capped by the model's trained ceiling -- not the
907 single-GPU dynamic ctx (which shrinks to fit one card and then confirms a
908 single-card placement) nor the full trained ceiling (which over-reserves). A
909 model that cannot hold weights + this floor on one card is tensor-split.
910 """
911 from lilbee.core.config import cfg
913 if role is WorkerRole.CHAT:
914 if cfg.num_ctx is not None:
915 return _pinned_chat_ctx(model_path, meta)
916 return apply_ctx_downshift(
917 role,
918 min(
919 engine_params.chat_ctx_ceiling(meta, model_path),
920 max(cfg.chat_n_ctx_target, _MIN_USABLE_CHAT_CTX),
921 ),
922 )
923 return apply_ctx_downshift(role, _role_ctx(role, model_path, meta))
926def _placement_estimate_slots(role: WorkerRole, meta: dict[str, str] | None) -> int:
927 """The slot count the placement estimate reserves KV for.
929 A tensor-split chat reserves one full-context sequence here: a conservative
930 floor for the card-count decision. The launch then fills the placed cards'
931 real headroom with as many full-context slots as fit (``_resolve_split_chat_slots``),
932 never exceeding what those cards hold, so a larger launch count can't OOM.
933 """
934 from lilbee.core.config import cfg
936 if role is WorkerRole.CHAT:
937 return _SPLIT_CHAT_SLOTS
938 if role is WorkerRole.VISION:
939 return max(1, cfg.vision_ocr_concurrency)
940 if role is WorkerRole.RERANK and _rerank_mode_for(meta) is RerankMode.LLM:
941 return LLM_RERANK_CONCURRENCY
942 return _AUX_SLOTS
945def _peak_estimator(model_refs: dict[WorkerRole, str]) -> PeakEstimator:
946 """Per-device VRAM-vector estimator for the planner, bound to the configured models.
948 Estimates each role at its launch ceiling (ctx x slots) with the candidate
949 tensor-split ratio, so the planner reserves enough cards for the busiest one.
950 """
951 from lilbee.providers.gguf_meta import read_gguf_metadata
953 def estimate_peak(role: WorkerRole, ratio: tuple[int, ...]) -> tuple[int, ...]:
954 path = engine_params.resolve_model_path(model_refs[role])
955 meta = read_gguf_metadata(path)
956 mmproj = vision_mmproj(model_refs[role]) if role is WorkerRole.VISION else None
957 slots = _placement_estimate_slots(role, meta)
958 ctx = _placement_estimate_ctx(role, path, meta)
959 rerank_mode = _role_rerank_mode(role, meta)
960 est = estimate_instance_footprint(
961 path,
962 ctx=ctx,
963 slots=slots,
964 gpu_layers=_role_gpu_layers(role),
965 flash_attn=_role_flash(role, rerank_mode),
966 kv_cache_type=_role_kv_cache_type(role),
967 mmproj_path=mmproj,
968 tensor_split=ratio,
969 batch_size=_pooled_batch_size(role, rerank_mode, ctx),
970 expert_offload=_role_expert_offload(path),
971 )
972 return est.per_device_vram
974 return estimate_peak
977def _chat_split_ctx_objective(
978 model_refs: dict[WorkerRole, str],
979) -> tuple[SplitCtxFitter | None, int]:
980 """The chat split's context fitter and target, or ``(None, 0)`` with no chat model.
982 The fitter sizes a candidate shard's served context exactly as the launch does
983 (:func:`fit_split_ctx`), so the planner widens chat onto idle cards only when a
984 tighter shard would starve KV below the target. See docs/architecture.md.
985 """
986 if WorkerRole.CHAT not in model_refs:
987 return None, 0
988 from lilbee.providers.gguf_meta import read_gguf_metadata
990 path = engine_params.resolve_model_path(model_refs[WorkerRole.CHAT])
991 meta = read_gguf_metadata(path)
992 target = _placement_estimate_ctx(WorkerRole.CHAT, path, meta)
994 def fit(ratio: tuple[int, ...], per_device_free_bytes: Sequence[int]) -> int:
995 return fleet_ctx.fit_split_ctx(
996 path,
997 meta=meta,
998 slots=_SPLIT_CHAT_SLOTS,
999 ratio=ratio,
1000 per_device_free_bytes=per_device_free_bytes,
1001 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
1002 flash_attn=_role_flash(WorkerRole.CHAT),
1003 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
1004 ctx_ceiling=target,
1005 expert_offload=_role_expert_offload(path),
1006 )
1008 return fit, target
1011def _search_reservation(inputs: dict[WorkerRole, ModelPlacementInput]) -> int:
1012 """Total footprint of the placed search roles (all replicas), held back ahead
1013 of chat."""
1014 return sum(
1015 inputs[role].est_vram_bytes * inputs[role].replicas
1016 for role in _EMBED_ROLES
1017 if role in inputs
1018 )
1021def _role_weights_bytes(role: WorkerRole, ref: str) -> int:
1022 """The model's weight bytes on disk (plus the mmproj for vision): a
1023 ground-truth lower bound on residency. 0 when the file cannot be resolved."""
1024 from lilbee.providers.base import ProviderError
1026 try:
1027 size = _weights_bytes(engine_params.resolve_model_path(ref))
1028 if role is WorkerRole.VISION:
1029 mmproj = vision_mmproj(ref)
1030 if mmproj is not None:
1031 size += int(mmproj.stat().st_size)
1032 except (ProviderError, OSError):
1033 return 0
1034 return size
1037def _is_moe(meta: dict[str, str] | None) -> bool:
1038 """Whether the GGUF declares routed experts, so its experts can be offloaded."""
1039 count = (meta or {}).get("expert_count")
1040 try:
1041 return int(count) > 0 if count is not None else False
1042 except ValueError:
1043 return False
1046def expert_offload_all(meta: dict[str, str] | None) -> bool:
1047 """Whether to keep every layer's experts in system memory; MoE models only."""
1048 from lilbee.core.config import cfg
1050 return bool(cfg.cpu_moe) and _is_moe(meta)
1053def expert_offload_layers(meta: dict[str, str] | None) -> int | None:
1054 """How many layers' experts to keep in system memory, or None for no split.
1056 A non-positive ``n_cpu_moe`` offloads nothing (it would emit a no-op
1057 ``--n-cpu-moe 0``), so it reads as unset.
1058 """
1059 from lilbee.core.config import cfg
1061 if cfg.n_cpu_moe is None or cfg.n_cpu_moe < 1 or not _is_moe(meta):
1062 return None
1063 return cfg.n_cpu_moe
1066def _role_expert_offload(model_path: Path) -> tuple[str, ...]:
1067 """Expert patterns the launch will offload, for sizing the same way it runs.
1069 Reads the GGUF (cached) rather than taking metadata as an argument so every
1070 estimate site charges the same tensors the launch moves off the GPU.
1071 """
1072 from lilbee.providers.fleet.adapters import expert_offload_patterns
1073 from lilbee.providers.gguf_meta import read_gguf_metadata
1075 meta = read_gguf_metadata(model_path)
1076 return expert_offload_patterns(
1077 cpu_moe=expert_offload_all(meta), n_cpu_moe=expert_offload_layers(meta)
1078 )
1081def _expert_offload_configured() -> bool:
1082 """Whether the user asked for expert offload that would actually take effect.
1084 A non-positive ``n_cpu_moe`` offloads nothing, so it does not count.
1085 """
1086 from lilbee.core.config import cfg
1088 return bool(cfg.cpu_moe) or (cfg.n_cpu_moe is not None and cfg.n_cpu_moe >= 1)
1091def _weights_exceed_everything(size: int, *, total_vram: int, total_ram: int) -> bool:
1092 """True when a model's weights fit neither the GPUs nor system memory.
1094 File size is ground truth, not an estimate, so this bound cannot repeat the
1095 false-refusal class: no estimator error makes a 40 GiB file fit a 1 GiB box.
1096 Past both pools there is nowhere for a layer to go and no launch can win, so
1097 saying so beats a load that thrashes and then dies.
1098 """
1099 ceiling = total_vram + total_ram
1100 return ceiling > 0 and size > ceiling
1103def _weights_exceed_hardware(size: int, total_vram: int, *, is_moe: bool) -> bool:
1104 """True when this model cannot be served on this machine at all.
1106 Exceeding VRAM alone is not that. The engine chooses how many layers fit and
1107 keeps the rest in system memory, so a model larger than every card is a
1108 partial offload and lilbee's job is to launch it and say what will happen.
1109 Refusing there meant the fit never ran and the role was skipped, which left
1110 the user hand-tuning n_gpu_layers to get back what the engine does by itself.
1112 What still refuses is a model past VRAM and system memory together, where no
1113 arrangement of layers exists. A user-set n_gpu_layers or expert offload keeps
1114 standing the bound down entirely, since the user has said where the weights
1115 should go.
1116 """
1117 from lilbee.core.config import cfg
1119 if cfg.n_gpu_layers is not None:
1120 return False
1121 if is_moe and _expert_offload_configured():
1122 return False
1123 return _weights_exceed_everything(
1124 size, total_vram=total_vram, total_ram=model_cache.total_system_memory()
1125 )
1128def _vision_without_mmproj(role: WorkerRole, ref: str) -> bool:
1129 """True (with a warning) for a configured vision model whose mmproj is missing.
1131 The skip would silently disable OCR; the warning names the cause and the fix.
1132 """
1133 if role is not WorkerRole.VISION or vision_mmproj(ref) is not None:
1134 return False
1135 log.warning(
1136 "Vision model %s has no mmproj (CLIP projector); OCR is disabled. "
1137 "Re-run 'lilbee model pull %s' to fetch the projector.",
1138 ref,
1139 ref,
1140 )
1141 return True
1144def _estimate_or_fallback(
1145 role: WorkerRole,
1146 ref: str,
1147 *,
1148 unified_budget: int | None,
1149 chat_reservation: int,
1150 device_count: int,
1151 total_vram: int,
1152 skipped_not_installed: dict[WorkerRole, str],
1153 host_committed: int = 0,
1154) -> ModelPlacementInput | None:
1155 """Size *role* for placement, degrading rather than refusing.
1157 A missing model is skipped and recorded; a sizing failure on an installed
1158 model falls back to the analytic floor; weights alone exceeding the physical
1159 VRAM refuse with a plain message (ground truth, not an estimate).
1160 """
1161 from lilbee.providers.base import ProviderError, ProviderErrorKind
1163 try:
1164 estimate = _estimate_role(
1165 role,
1166 ref,
1167 unified_budget=unified_budget,
1168 chat_reservation=chat_reservation,
1169 device_count=device_count,
1170 )
1171 except (ProviderError, OSError) as exc:
1172 if isinstance(exc, ProviderError) and exc.kind is ProviderErrorKind.NOT_FOUND:
1173 log.warning("Skipping %s server: model %r is not installed.", role.value, ref)
1174 skipped_not_installed[role] = ref
1175 return None
1176 return _sizing_failure_fallback(
1177 role,
1178 ref,
1179 exc,
1180 device_count=device_count,
1181 total_vram=total_vram,
1182 host_committed=host_committed,
1183 )
1184 return _admit_estimate(
1185 _floor_implausible_estimate(estimate, role, ref),
1186 role,
1187 ref,
1188 total_vram=total_vram,
1189 ram_bytes=estimate.est_ram_bytes,
1190 host_committed=host_committed,
1191 )
1194def _admit_estimate(
1195 estimate: ModelPlacementInput,
1196 role: WorkerRole,
1197 ref: str,
1198 *,
1199 total_vram: int,
1200 ram_bytes: int,
1201 host_committed: int = 0,
1202) -> ModelPlacementInput | None:
1203 """*estimate*, or ``None`` when this model cannot load on this machine.
1205 Two hardware bounds, one per kind of memory: the weights must fit the GPUs
1206 unless something offloads, and whatever offloading puts in system memory must
1207 fit the system.
1208 """
1209 weights = _role_weights_bytes(role, ref)
1210 if _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref)):
1211 _warn_weights_exceed(role, ref, weights, total_vram)
1212 return None
1213 if total_vram > 0 and weights > total_vram:
1214 _warn_weights_spill(role, ref, weights, total_vram)
1215 if _host_memory_refuses(role, ref, ram_bytes, host_committed):
1216 return None
1217 return estimate
1220def _analytic_footprint_floor(
1221 weights: int, *, role: WorkerRole, meta: dict[str, str] | None, ctx: int, slots: int
1222) -> int:
1223 """The least this instance can occupy: weights, its KV cache, and overhead.
1225 Used when the estimator cannot answer. Charging weight bytes alone was a
1226 knowing under-charge: the engine allocates a KV cache sized by context and
1227 slot count, plus compute buffers, and omitting all of it lets placement fit a
1228 model that cannot fit. The comment said the load would decide, and it did, by
1229 running out of memory.
1231 A floor rather than an estimate. It is derived from the header the same way
1232 the in-process sizing path derives it, and it is deliberately the smallest
1233 defensible number, because refusing a model that would have fit is its own
1234 failure. Without a readable header the per-token fallback still applies:
1235 zero is the one answer that is certainly wrong.
1236 """
1238 kv_bytes = (
1239 model_cache.kv_bytes_per_token(
1240 meta,
1241 KV_CACHE_TYPE_BYTES[_role_kv_cache_type(role)],
1242 KV_CACHE_TYPE_BYTES[_role_kv_cache_type(role)],
1243 )
1244 * ctx
1245 * max(slots, 1)
1246 )
1247 overhead = int(weights * model_cache._BUFFER_OVERHEAD_FRACTION)
1248 return weights + kv_bytes + overhead
1251def _estimate_is_implausible(*, estimated: int, floor: int) -> bool:
1252 """Whether *estimated* describes a load that cannot exist.
1254 Below the bytes the card must hold there is no arrangement of memory that
1255 serves the model. A floor of zero means nothing could be computed to compare
1256 against, and a guess is not grounds to discard the only measurement there is.
1257 """
1258 return floor > 0 and 0 < estimated < floor
1261def _floor_implausible_estimate(
1262 estimate: ModelPlacementInput, role: WorkerRole, ref: str
1263) -> ModelPlacementInput:
1264 """*estimate*, or the model's weight bytes when it reports less than those.
1266 The bound is the weights alone, not the analytic footprint. That figure
1267 sizes a KV cache as though every layer ran dense attention over the whole
1268 window, which linear-attention, sliding-window and MLA models do not, so it
1269 sits far above what they hold. Weights are architecture-independent, which
1270 makes an estimate under them the one answer the planner can call impossible
1271 without modelling attention.
1273 Offload lifts the bound: the layers in system memory are weight bytes the
1274 card never holds.
1275 """
1276 if _cpu_offload_in_play():
1277 return estimate
1278 weights = _role_weights_bytes(role, ref)
1279 if not _estimate_is_implausible(estimated=estimate.est_vram_bytes, floor=weights):
1280 return estimate
1281 log.warning(
1282 "The estimator sized the %s model %s at %.1f GiB, below the %.1f GiB of "
1283 "weights the card has to hold. Charging the weights instead.",
1284 role.value,
1285 ref,
1286 estimate.est_vram_bytes / 1024**3,
1287 weights / 1024**3,
1288 )
1289 return replace(estimate, est_vram_bytes=weights)
1292def _sizing_failure_fallback(
1293 role: WorkerRole,
1294 ref: str,
1295 exc: Exception,
1296 *,
1297 device_count: int,
1298 total_vram: int,
1299 host_committed: int = 0,
1300) -> ModelPlacementInput | None:
1301 """Analytic-floor placement input for an installed model the estimator cannot
1302 size; ``None`` skips the role (the file is unresolvable, its weights alone
1303 exceed the hardware, or offloading it would exceed system memory).
1305 The host bound applies here too. Charging the whole floor to VRAM and
1306 skipping it let an unsizable model past a check every sized model faces."""
1307 weights = _role_weights_bytes(role, ref)
1308 if weights == 0:
1309 log.warning("Skipping %s server: could not size model %r (%s).", role.value, ref, exc)
1310 return None
1311 if _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref)):
1312 _warn_weights_exceed(role, ref, weights, total_vram)
1313 return None
1314 floor = _fallback_floor_for(role, ref, weights)
1315 log.warning(
1316 "Could not size the %s model %s (%s). Charging %.1f GiB, its weights plus the "
1317 "cache and buffers it will allocate, which is a floor rather than an estimate: "
1318 "the load may still need more.",
1319 role.value,
1320 ref,
1321 exc,
1322 floor / 1024**3,
1323 )
1324 if _host_memory_refuses(role, ref, floor, host_committed):
1325 return None
1326 return ModelPlacementInput(
1327 role=role, est_vram_bytes=floor, replicas=_replica_count(role, device_count)
1328 )
1331def _fallback_floor_for(role: WorkerRole, ref: str, weights: int) -> int:
1332 """:func:`_analytic_footprint_floor` for *role*, reading what metadata it can."""
1333 from lilbee.providers.base import ProviderError
1334 from lilbee.providers.gguf_meta import read_gguf_metadata
1336 try:
1337 path = engine_params.resolve_model_path(ref)
1338 meta = read_gguf_metadata(path)
1339 except (ProviderError, OSError, ValueError):
1340 meta = None
1341 path = None
1342 ctx = _placement_estimate_ctx(role, path, meta) if path is not None else _MIN_USABLE_CHAT_CTX
1343 return _analytic_footprint_floor(
1344 weights, role=role, meta=meta, ctx=ctx, slots=_placement_estimate_slots(role, meta)
1345 )
1348def planned_embed_token_cap(ref: str) -> int | None:
1349 """The token cap an embed launch of *ref* gets, or None when *ref* cannot be read.
1351 Computed from the model's metadata the way the launch is, so the chunker can
1352 bound itself to the cap before the engine is up.
1353 """
1354 from lilbee.providers.base import ProviderError
1355 from lilbee.providers.gguf_meta import read_gguf_metadata
1357 try:
1358 path = engine_params.resolve_model_path(ref)
1359 meta = read_gguf_metadata(path)
1360 except (ProviderError, OSError, ValueError):
1361 return None
1362 return engine_params.embed_token_cap(
1363 apply_ctx_downshift(WorkerRole.EMBED, _role_ctx(WorkerRole.EMBED, path, meta))
1364 )
1367def _ref_is_moe(ref: str) -> bool:
1368 """Whether *ref*'s GGUF declares routed experts; False when it cannot be read."""
1369 from lilbee.providers.base import ProviderError
1370 from lilbee.providers.gguf_meta import read_gguf_metadata
1372 try:
1373 return _is_moe(read_gguf_metadata(engine_params.resolve_model_path(ref)))
1374 except (ProviderError, OSError):
1375 return False
1378def _cpu_offload_in_play() -> bool:
1379 """Whether this configuration puts any of a model's weights in system memory.
1381 Expert offload moves the experts, a partial ``n_gpu_layers`` moves whole
1382 layers, and zero moves the model. Without one of these the engine keeps
1383 everything on the card and the estimator's host figure describes memory
1384 nobody will allocate.
1385 """
1386 from lilbee.core.config import cfg
1388 return _expert_offload_configured() or cfg.n_gpu_layers is not None
1391def _host_bytes_must_be_resident(role: WorkerRole, ref: str) -> bool:
1392 """Whether *role*'s host bytes have to fit RAM rather than page in and out.
1394 The estimator's host figure counts mmap pages, and llama.cpp maps CPU-side
1395 weights over that mapping instead of allocating them, so with mmap they are
1396 evictable page cache: a model far larger than RAM streams from disk and
1397 serves, which is a practiced setup for a large mixture-of-experts. Only
1398 the no-mmap load turns them into a buffered read that must be resident, and the
1399 single path that asks for it is a chat model on a network filesystem.
1401 Anything this cannot determine counts as mappable, because a false refusal
1402 here has no override and costs the user a model that would have run.
1403 """
1404 if role is not WorkerRole.CHAT:
1405 return False
1406 try:
1407 path = engine_params.resolve_model_path(ref)
1408 except (ProviderError, OSError, ValueError):
1409 return False
1410 if not is_network_path(path):
1411 return False
1412 return _chat_no_mmap(_role_weights_bytes(role, ref), on_network_fs=True)
1415def _host_committed(admitted: Mapping[WorkerRole, ModelPlacementInput]) -> int:
1416 """System-memory bytes the roles already admitted to this plan will hold."""
1417 return sum(inp.est_ram_bytes for inp in admitted.values())
1420def _host_memory_refuses(role: WorkerRole, ref: str, ram_bytes: int, committed: int) -> bool:
1421 """Whether *role*'s system-memory half is too big for this machine to load.
1423 Charged only when something actually offloads, and only when the bytes must
1424 be resident: refusing a mapped model that would have streamed from disk is a
1425 false refusal with no override, which is worse than a slow load.
1427 Measured against the whole plan, not this role alone. Every role was
1428 previously compared to the entire machine on its own, so two roles that each
1429 fit and together do not were both admitted.
1430 """
1431 if not _cpu_offload_in_play() or ram_bytes <= 0:
1432 return False
1433 wanted = committed + ram_bytes
1434 total = total_system_memory()
1435 if total and wanted > total and _host_bytes_must_be_resident(role, ref):
1436 log.warning(
1437 "The %s model %s cannot load: this plan puts %.1f GiB in system memory, which "
1438 "cannot be paged out here, and the machine has %.1f GiB in total. Use a smaller "
1439 "model, or offload less.",
1440 role.value,
1441 ref,
1442 wanted / 1024**3,
1443 total / 1024**3,
1444 )
1445 return True
1446 free = free_system_memory()
1447 if free and wanted > free:
1448 log.warning(
1449 "Offloading the %s model %s brings this plan to %.1f GiB in system memory and "
1450 "only %.1f GiB is free. It will still load; close other programs if it swaps "
1451 "or runs slowly.",
1452 role.value,
1453 ref,
1454 wanted / 1024**3,
1455 free / 1024**3,
1456 )
1457 return False
1460def _warn_weights_exceed(role: WorkerRole, ref: str, weights: int, total_vram: int) -> None:
1461 log.warning(
1462 "The %s model %s cannot load: its weights are %.1f GiB and this machine has "
1463 "%.1f GiB of GPU memory and %.1f GiB of system memory, so there is nowhere "
1464 "for its layers to go. Use a smaller model or a smaller quantization.",
1465 role.value,
1466 ref,
1467 weights / 1024**3,
1468 total_vram / 1024**3,
1469 model_cache.total_system_memory() / 1024**3,
1470 )
1473def _warn_weights_spill(role: WorkerRole, ref: str, weights: int, total_vram: int) -> None:
1474 """Say that a model larger than the GPUs will run partly in system memory."""
1475 log.warning(
1476 "The %s model %s is %.1f GiB and this machine has %.1f GiB of GPU memory, so "
1477 "the engine will keep the layers that fit on the GPU and the rest in system "
1478 "memory. It will run, and it will be slower than a model that fits.",
1479 role.value,
1480 ref,
1481 weights / 1024**3,
1482 total_vram / 1024**3,
1483 )
1486def placeable_total_vram() -> int:
1487 """Physical VRAM across all cards, for the weights-exceed placeability bound.
1489 Physical total is box-state-independent (a running incumbent doesn't skew
1490 it), so it is safe to read without a clean box. Reuses the plan probe when
1491 one is captured; otherwise probes best-effort and returns ``0`` on failure,
1492 which disables only the weights-exceed filter (its own ``total > 0`` guard).
1493 """
1494 probe = _current_plan_probe()
1495 if probe is not None:
1496 return sum(d.total_bytes for d in probe.devices)
1497 from lilbee.providers.base import ProviderError
1499 try:
1500 return sum(d.total_bytes for d in resolve_devices(resolve_llama_server()))
1501 except (ProviderError, OSError):
1502 return 0
1505def role_model_placeable(role: WorkerRole, ref: str, total_vram: int) -> bool:
1506 """Whether a fresh plan would actually serve *role* on *ref*.
1508 Mirrors the planner's own drop conditions (SDK-routed role, vision without a
1509 projector, model not installed, weights exceeding physical VRAM) using the
1510 same primitives, so the acquisition ladder binds and replaces against what
1511 an engine can serve rather than the raw config. Without this a
1512 configured-but-unplaceable role keeps bind from ever matching a running
1513 engine and restarts the shared engine on every process start.
1514 """
1515 if parse_model_ref(ref).is_remote or _vision_without_mmproj(role, ref):
1516 return False
1517 weights = _role_weights_bytes(role, ref) # 0 when not installed / unresolvable
1518 if weights == 0:
1519 return False
1520 return not _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref))
1523def _server_model_inputs(
1524 roles: tuple[WorkerRole, ...] | None = None,
1525 *,
1526 unified_budget: int | None = None,
1527 device_count: int = 0,
1528 total_vram: int = 0,
1529) -> tuple[list[ModelPlacementInput], dict[WorkerRole, str], int, dict[WorkerRole, str]]:
1530 """Build placement inputs for the configured server roles.
1532 The search and vision roles are estimated first; chat is then sized against the
1533 budget minus the search footprint (the ``reservation``) so a large chat cannot
1534 starve embed/rerank on a shared-memory host. ``device_count`` resolves an auto
1535 replica knob to one per GPU. When *roles* is given, only those are considered.
1536 Skips an unconfigured optional role, a vision model with no resolvable mmproj
1537 projector, a role whose model is not installed on disk (returned as
1538 ``skipped_not_installed`` so a surface can say so), and a model whose weight
1539 bytes alone exceed ``total_vram`` (physically unloadable under all-GPU layers).
1540 A model the estimator cannot size is enrolled at its file size instead of
1541 skipped, so the load, not the estimator, decides.
1542 """
1543 from lilbee.core.config import cfg
1545 inputs: dict[WorkerRole, ModelPlacementInput] = {}
1546 model_refs: dict[WorkerRole, str] = {}
1547 skipped_not_installed: dict[WorkerRole, str] = {}
1549 def consider(role: WorkerRole, *, chat_reservation: int = 0) -> None:
1550 if roles is not None and role not in roles:
1551 return
1552 # Any role may be "" (unconfigured) -> skipped, so that role has no
1553 # server and no not-installed complaint.
1554 ref = str(getattr(cfg, ROLE_REGISTRY[role].config_field))
1555 if not ref:
1556 return # unconfigured optional role -> no server
1557 if parse_model_ref(ref).is_remote:
1558 return # SDK-routed role: no local server to plan, not a missing install
1559 if _vision_without_mmproj(role, ref):
1560 return # no projector -> vision can't run on a server
1561 estimate = _estimate_or_fallback(
1562 role,
1563 ref,
1564 unified_budget=unified_budget,
1565 chat_reservation=chat_reservation,
1566 device_count=device_count,
1567 total_vram=total_vram,
1568 skipped_not_installed=skipped_not_installed,
1569 host_committed=_host_committed(inputs),
1570 )
1571 if estimate is None:
1572 return
1573 inputs[role] = estimate
1574 model_refs[role] = ref
1576 # Estimate every non-chat role first so the search footprint is known, then size
1577 # chat against the remainder. The reservation only applies on a shared-memory
1578 # host; discrete GPUs pin each role to its own VRAM and pack independently.
1579 for role in ROLE_REGISTRY:
1580 if role is not WorkerRole.CHAT:
1581 consider(role)
1582 reservation = _search_reservation(inputs) if unified_budget is not None else 0
1583 consider(WorkerRole.CHAT, chat_reservation=reservation)
1585 ordered = [inputs[role] for role in ROLE_REGISTRY if role in inputs]
1586 return ordered, model_refs, reservation, skipped_not_installed
1589def _non_chat_reservation(
1590 instances: Sequence[InstancePlan],
1591 inputs: Sequence[ModelPlacementInput],
1592 co_tenants: frozenset[WorkerRole] = frozenset(),
1593) -> dict[int, int]:
1594 """Per-device VRAM the non-chat role servers occupy, keyed by device index.
1596 A tensor-split chat shard must size its KV against the headroom left after the
1597 embed/rerank/vision servers on the same card, not the card's raw free VRAM, or
1598 it over-commits and OOMs at launch. Chat is excluded because it sizes its own
1599 weights. Chat's own swap-group siblings are excluded too: they are evicted while
1600 chat is resident, so their VRAM is chat's to use. That only holds when chat is
1601 itself a co-tenant; a co-tenant group that does not include chat runs behind its
1602 own swap process and can be resident beside chat, so it is charged normally.
1603 Non-chat roles are single-device, so each charges its full footprint (once per
1604 replica) to its card.
1605 """
1606 chat_siblings = co_tenants if WorkerRole.CHAT in co_tenants else frozenset()
1607 charge_by_role = {inp.role: inp.est_vram_bytes for inp in inputs}
1608 reserved: dict[int, int] = {}
1609 for inst in instances:
1610 if inst.role is WorkerRole.CHAT or inst.role in chat_siblings:
1611 continue
1612 charge = charge_by_role[inst.role]
1613 for device in inst.devices:
1614 reserved[device] = reserved.get(device, 0) + charge
1615 return reserved
1618def _charge_by_device(
1619 chosen: tuple[FleetDevice, ...], ratio: tuple[int, ...], total: int
1620) -> dict[str, int]:
1621 """What each of *chosen* was charged, keyed by the name the engine prints.
1623 A single-card instance carries the whole charge. A split carries it in the
1624 proportions it launches with, which is what the planner decided and therefore
1625 what the engine's own report should be compared against.
1626 """
1627 from lilbee.providers.fleet.readback import device_label
1629 if total <= 0 or not chosen:
1630 return {}
1631 if len(chosen) == 1:
1632 return {device_label(chosen[0]): total}
1633 weights = ratio if len(ratio) == len(chosen) else (1,) * len(chosen)
1634 denominator = sum(weights) or len(chosen)
1635 return {
1636 device_label(device): total * weight // denominator
1637 for device, weight in zip(chosen, weights, strict=True)
1638 }
1641def _chat_fit_target_mib(
1642 binary: Path, chosen: tuple[FleetDevice, ...], reserved: Mapping[int, int]
1643) -> tuple[int, ...]:
1644 """Memory, in MiB, the engine's fit leaves free on each of chat's cards.
1646 The engine's own margin plus what *reserved* holds on the card for roles
1647 that load after chat. Empty for an engine without ``--fit-target`` and on
1648 shared memory, where a layer moved off the device frees nothing.
1649 """
1650 if any(device.unified for device in chosen) or not engine_advertises(binary, FLAG_FIT_TARGET):
1651 return ()
1652 return tuple(
1653 ENGINE_FIT_MARGIN_MIB + -(-reserved.get(device.index, 0) // MIB) for device in chosen
1654 )
1657def _engine_shares_kv(binary: Path) -> bool:
1658 """Whether *binary* can give its slots one KV pool with a per-slot cap."""
1659 return engine_advertises(binary, FLAG_KV_UNIFIED_PER_SLOT)
1662def _launch_for(
1663 plan: InstancePlan,
1664 model_ref: str,
1665 binary: Path,
1666 by_index: dict[int, FleetDevice],
1667 *,
1668 unified_budget: int | None = None,
1669 chat_reservation: int = 0,
1670 reserved_by_device: dict[int, int] | None = None,
1671 reserved_after_chat: dict[int, int] | None = None,
1672 est_vram_bytes: int = 0,
1673 model_path: Path | None = None,
1674) -> InstanceLaunch:
1675 """Build the launch spec (argv + device-pinning env) for one planned instance."""
1676 from lilbee.providers.gguf_meta import read_gguf_metadata
1678 # The self-check holds a downloaded file rather than a configured reference,
1679 # so it hands the path over instead of asking for one to be resolved.
1680 model_path = model_path or engine_params.resolve_model_path(model_ref)
1681 weights_bytes = _weights_bytes(model_path)
1682 meta = read_gguf_metadata(model_path)
1683 from lilbee.core.config import cfg
1685 chosen = tuple(by_index[i] for i in plan.devices)
1686 # ctx and slots are sized against the card this role landed on, not the fleet's
1687 # smallest, which is all the pre-placement estimate had to go on. A role spread
1688 # over several cards has no one budget; the split chat, the only such role today,
1689 # sizes against per-device headroom below.
1690 placed_device = chosen[0] if len(chosen) == 1 else None
1691 is_chat = plan.role is WorkerRole.CHAT
1692 is_vision = plan.role is WorkerRole.VISION
1693 mmproj = vision_mmproj(model_ref) if is_vision else None
1694 chat_on_network_fs = is_chat and is_network_path(model_path)
1695 if chat_on_network_fs and not _chat_no_mmap(weights_bytes, on_network_fs=True):
1696 log.warning(
1697 "Chat model %s is served from a network filesystem and is too large to load "
1698 "into host RAM; mmap over the network can stall the load in uninterruptible "
1699 "I/O. Stage it on local disk for a reliable load.",
1700 model_ref,
1701 )
1702 # A tensor-split chat serves one full-context sequence sized against the busiest
1703 # card's headroom. A cfg.num_ctx pin overrides the fit (handled by _role_ctx).
1704 multi_card_chat = is_chat and len(chosen) > 1
1705 split_chat = multi_card_chat and cfg.num_ctx is None
1706 if multi_card_chat and host_lacks_nvlink():
1707 log.warning(
1708 "Chat model %s is tensor-split across GPUs %s on a host without NVLink; "
1709 "generation is PCIe all-reduce bound and can be very slow. A model that fits "
1710 "on fewer cards will generate faster.",
1711 model_ref,
1712 list(plan.devices),
1713 )
1714 split_slots = _SPLIT_CHAT_SLOTS
1715 reserved = reserved_by_device or {}
1716 if split_chat:
1717 # Headroom left after the embed/rerank servers on each shared card, not the
1718 # card's raw free VRAM, so the chat KV doesn't over-commit.
1719 per_device_free = [max(0, d.free_bytes - reserved.get(d.index, 0)) for d in chosen]
1721 def _split_fit(slots: int) -> int:
1722 return fleet_ctx.fit_split_ctx(
1723 model_path,
1724 meta=meta,
1725 slots=slots,
1726 ratio=plan.tensor_split,
1727 per_device_free_bytes=per_device_free,
1728 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
1729 flash_attn=_role_flash(WorkerRole.CHAT),
1730 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
1731 ctx_ceiling=_placement_estimate_ctx(WorkerRole.CHAT, model_path, meta),
1732 expert_offload=_role_expert_offload(model_path),
1733 )
1735 split_slots, ctx = _resolve_split_chat_slots(_split_fit)
1736 else:
1737 # Downshifted here and not only in the estimate: the role resolvers are
1738 # pure functions of model and config, so without this the retry after a
1739 # load OOM re-emits a byte-identical argv and dies the same way. The
1740 # split branch above already inherits it through its ctx_ceiling.
1741 ctx = apply_ctx_downshift(plan.role, _role_ctx(plan.role, model_path, meta, placed_device))
1742 rerank_mode = _role_rerank_mode(plan.role, meta)
1743 is_llm_rerank = rerank_mode is RerankMode.LLM
1744 # A multi-card chat runs as many full-context slots as its cards' KV headroom
1745 # holds (split_slots, one when a num_ctx pin skips the fit); other roles size
1746 # --parallel against the budget the same way the estimator did.
1747 slots = (
1748 split_slots
1749 if multi_card_chat
1750 else _slots_for(
1751 plan.role,
1752 model_path,
1753 ctx,
1754 mmproj_path=mmproj,
1755 unified_budget=unified_budget,
1756 chat_reservation=chat_reservation,
1757 rerank_mode=rerank_mode,
1758 device=placed_device,
1759 )
1760 )
1761 spec = _server_spec(plan.role, rerank_mode, meta)
1762 # Where the engine's fit can move chat layers off the card, the readback
1763 # compares the host memory it reports with this to say so.
1764 fit_moves_chat = (
1765 is_chat and placed_device is not None and _engine_fits_chat_offload(model_path, binary)
1766 )
1767 # Cross-encoder embed/rerank pools the whole input in one batch; an LLM reranker
1768 # is generative and uses the default batching plus flash attention.
1769 cross_encoder_pooled = plan.role in _EMBED_ROLES and not is_llm_rerank
1770 argv = build_server_argv(
1771 binary=binary,
1772 spec=spec,
1773 model_path=model_path,
1774 devices=plan.devices,
1775 n_gpu_layers=_role_gpu_layers(plan.role),
1776 slots=slots,
1777 ctx_per_slot=ctx,
1778 tensor_split=plan.tensor_split,
1779 mmproj=mmproj,
1780 flash_attn=flash_attn_flag() if _role_launches_with_flash(plan.role, rerank_mode) else None,
1781 cache_type=chat_cache_type_flag() if is_chat else None,
1782 batch_size=_pooled_batch_size(plan.role, rerank_mode, ctx),
1783 no_mmap=is_chat and _chat_no_mmap(weights_bytes, on_network_fs=chat_on_network_fs),
1784 cpu_moe=expert_offload_all(meta),
1785 n_cpu_moe=expert_offload_layers(meta),
1786 device_names=_device_names(chosen) or _cpu_pin_when_every_device_was_refused(),
1787 memory_endpoint=supports_memory_readback(binary),
1788 fit_target_mib=(
1789 _chat_fit_target_mib(binary, chosen, reserved_after_chat or {}) if is_chat else ()
1790 ),
1791 kv_unified_per_slot=ctx if is_chat and _engine_shares_kv(binary) else None,
1792 )
1793 return InstanceLaunch(
1794 role=plan.role,
1795 argv=argv,
1796 env_overrides={**visible_env(chosen), **llama_server_runtime_env()},
1797 model=model_ref,
1798 # token_cap drives cross-encoder/embed input truncation; the LLM rerank path
1799 # doesn't truncate (it relies on the per-slot ctx headroom), so leave it None.
1800 token_cap=engine_params.embed_token_cap(ctx) if cross_encoder_pooled else None,
1801 # Weights size scales the cold-load ready timeout (larger model = longer).
1802 weights_bytes=weights_bytes,
1803 # Slots is the chat concurrency the gate admits; ctx is what a client fits to.
1804 slots=slots,
1805 ctx=ctx,
1806 built_ctx_target=(
1807 (cfg.num_ctx if cfg.num_ctx is not None else cfg.chat_n_ctx_target) if is_chat else 0
1808 ),
1809 built_slots_target=(
1810 max(1, cfg.vision_ocr_concurrency) if plan.role is WorkerRole.VISION else 0
1811 ),
1812 replica=plan.replica,
1813 rerank_mode=rerank_mode,
1814 # What placement charged this instance, for the post-launch check against
1815 # the engine's own report of what it really allocated.
1816 est_vram_bytes=est_vram_bytes,
1817 est_vram_by_device=_charge_by_device(chosen, plan.tensor_split, est_vram_bytes),
1818 est_unreported_bytes=_unreported_bytes(plan.role, mmproj),
1819 est_host_bytes=_chat_host_estimate(model_path, ctx, slots) if fit_moves_chat else 0,
1820 settings=launch_settings(plan.role, rerank_mode, meta, model_path),
1821 )
1824def build_single_role_launch(role: WorkerRole, model_path: Path) -> InstanceLaunch:
1825 """The launch the fleet would build for *role* serving *model_path*, alone.
1827 One construction path. The self-check used to assemble its own beside this
1828 one and the two disagreed on slot count, on the context that follows from it,
1829 on device pinning and on the tensor split, so a green check proved nothing
1830 about the launch serving actually performs, and a red one could be a
1831 configuration serving would never have chosen.
1833 Placement is the planner's, on the devices the plan snapshot holds, so the
1834 check runs on the card the role would really land on.
1835 """
1836 binary = resolve_llama_server()
1837 with _one_engine_per_pass():
1838 devices = _plan_devices(binary)
1839 by_index = {d.index: d for d in devices}
1840 # The whole machine, since nothing else is resident during a self-check.
1841 placed = (min(by_index),) if by_index else ()
1842 plan = InstancePlan(role=role, devices=placed)
1843 return _launch_for(
1844 plan,
1845 str(model_path),
1846 binary,
1847 by_index,
1848 unified_budget=_unified_memory_budget(devices),
1849 model_path=model_path,
1850 )
1853def resolve_devices(binary: Path) -> list[FleetDevice]:
1854 """Enumerate devices in the binary's index space, or the Vulkan VRAM probe."""
1855 return _read_devices(binary).devices
1858# The visibility variable each vendor's runtime reads, named in the warning so
1859# the reader checks the one that applies to the card they actually have.
1860_VENDOR_VISIBILITY_HINT = {
1861 "NVIDIA": "CUDA_VISIBLE_DEVICES",
1862 "AMD": "ROCR_VISIBLE_DEVICES / HIP_VISIBLE_DEVICES",
1863 "Intel": "ONEAPI_DEVICE_SELECTOR",
1864}
1867def _warn_gpu_present_but_unenumerated(binary: Path) -> None:
1868 """Say so when the host has a GPU the engine did not list.
1870 Previously asked only whether an NVIDIA card was present, so an AMD or Intel
1871 host whose engine enumerated nothing produced the identical symptom, a fleet
1872 quietly planned for CPU, and said nothing. The vendor lookup is the same one
1873 the Vulkan ICD rules use, and it works on Windows as well as Linux.
1874 """
1875 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
1876 from lilbee.providers.fleet.gpu_select import PCIVendorID
1878 present = installed_gpu_vendor_ids()
1879 names = sorted(
1880 v.name.title() if v.name == "INTEL" else v.name for v in PCIVendorID if v in present
1881 )
1882 if not names:
1883 return
1884 hints = sorted(
1885 {_VENDOR_VISIBILITY_HINT[name] for name in names if name in _VENDOR_VISIBILITY_HINT}
1886 )
1887 log.warning(
1888 "This host has a %s GPU but the engine's device probe (%s --list-devices) "
1889 "reported none; placement is falling back to shared-memory mode with unpinned "
1890 "GPUs. Check the GPU driver, %s, and that this llama-server build supports "
1891 "that GPU.",
1892 " and ".join(names),
1893 binary,
1894 " / ".join(hints) if hints else "the vendor's visibility variable",
1895 )
1898@dataclass(frozen=True)
1899class DeviceReading:
1900 """What one ``--list-devices`` run answered, as the whole fleet reads it.
1902 The device list and the backend are independent answers. When the engine never
1903 spoke the protocol but the host's Vulkan loader supplies devices, the reading
1904 carries those devices with an ``unknown`` backend, so a client can see a
1905 non-empty GPU list beside ``engine_backend: "unknown"``.
1906 """
1908 devices: list[FleetDevice]
1909 # The backend the engine selected, UNKNOWN when it never answered. Not
1910 # derivable from ``devices``: an empty list is a CPU host and a failed probe
1911 # alike, and only the engine can say which.
1912 backend: EngineBackend
1913 # The engine listed GPUs and lilbee rejected all of them, so the plan is
1914 # CPU-shaped while the engine would still choose one of those devices.
1915 refused_all: bool = False
1918def _read_devices(binary: Path) -> DeviceReading:
1919 """Enumerate devices, the selected backend, and whether every GPU was refused.
1921 One function because all three answers come from one ``--list-devices`` run,
1922 and that run costs a subprocess against a driver that may be wedged. Asking
1923 twice would pay it twice.
1925 The binary's ``--list-devices`` is authoritative, including when it lists
1926 nothing: it prints every non-CPU device it can use, so an empty list means
1927 the engine has no usable GPU rather than that we failed to look. The Vulkan
1928 VRAM probe is consulted only when the binary produced no output at all, and
1929 it reports the same index space. A
1930 probe that times out raises instead (a wedged GPU driver); falling through
1931 to the in-process Vulkan probe there could hang this thread unkillably.
1932 """
1933 from lilbee.providers.fleet.cuda_runtime import assert_cuda_devices_usable
1934 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
1935 from lilbee.providers.fleet.gpu_select import enumerate_gpu_vram
1936 from lilbee.providers.fleet.rocm_runtime import assert_rocm_devices_usable
1938 probe = probe_devices(binary)
1939 # An engine that answered but listed no device while a GPU is physically
1940 # present may be hitting a transient init error (the card momentarily held by
1941 # another process, e.g. an embedder served alongside -- the "ggml_cuda_init:
1942 # initialization error" symptom). Re-probe before treating the empty list as
1943 # fatal; a persistently empty list still hits the fail-loud asserts below.
1944 # The engine identity a plan snapshot carries does not cover this: the binary
1945 # is the same one across a transient init error, so only asking again tells it
1946 # apart from a host whose engine has no usable GPU.
1947 for _ in range(_DEVICE_PROBE_EMPTY_RETRIES):
1948 if probe.devices or not probe.spoke_protocol or not installed_gpu_vendor_ids():
1949 break
1950 time.sleep(_DEVICE_PROBE_EMPTY_RETRY_DELAY_S)
1951 probe = probe_devices(binary)
1952 devices = probe.devices
1953 # A GPU build that links a runtime it cannot serve the host's GPU with must
1954 # fail loud, not silently fall back to CPU (the Vulkan VRAM probe below
1955 # would mask it). Only when the engine actually answered, though: a binary
1956 # that does not support --list-devices enumerated nothing because it was
1957 # never asked, and accusing its driver of failing would be wrong and fatal.
1958 if probe.spoke_protocol:
1959 assert_cuda_devices_usable(binary, devices, probe.output)
1960 assert_rocm_devices_usable(binary, devices, probe.output)
1961 if not devices and probe.spoke_protocol:
1962 _warn_gpu_present_but_unenumerated(binary)
1963 if not devices and not probe.spoke_protocol:
1964 # Only when the binary never answered the question. An engine that ran
1965 # and listed nothing is reporting a fact, not a gap: believing the host
1966 # loader instead invents devices the engine cannot see. A CPU-only build
1967 # on a desktop with mesa is the clearest case, and the cost is not merely
1968 # a wrong device list. The fleet is planned onto GPUs, the pins are
1969 # no-ops, the shared-RAM guard is off because devices looked non-empty,
1970 # and every role then loads its full weights into system RAM while
1971 # running on the CPU anyway.
1972 #
1973 # Keyed on the exit code and the header rather than on there being no
1974 # output at all: the probe merges stderr into stdout, so a build that
1975 # predates --list-devices prints usage text and would otherwise be read
1976 # as an authoritative "no GPUs here".
1977 from lilbee.providers.fleet.gpu_select import integrated_vulkan_indices
1979 integrated = integrated_vulkan_indices()
1980 devices = [
1981 FleetDevice(
1982 VULKAN_BACKEND, idx, "", vram, free, unified=idx in integrated, from_loader=True
1983 )
1984 for idx, vram, free in (enumerate_gpu_vram() or [])
1985 ]
1986 if devices:
1987 log.warning(
1988 "The engine's device probe returned nothing, so placement is using "
1989 "the host's Vulkan loader instead and found %d device(s). If the "
1990 "engine has no Vulkan backend it will run on CPU regardless; set %s "
1991 "to override the engine location if that is wrong.",
1992 len(devices),
1993 "LILBEE_ENGINE_DIR",
1994 )
1995 return DeviceReading(devices, probe.backend, probe.refused_all)
1998_DEVICE_PROBE_TTL_S = 2.0
1999# A failed probe is cached much longer than a good one: each retry against a
2000# wedged GPU driver costs a full probe timeout, so a per-poll retry would stall
2001# every placement read for a minute at a time.
2002_DEVICE_PROBE_FAILURE_TTL_S = 60.0
2003# How long a plan read holds off re-probing an engine whose restate probe raised.
2004# The same order as the read cache's failure TTL and for the same reason (the
2005# retry ladder costs a full probe timeout), but a separate knob: the read cache
2006# bounds how long a wedged driver goes un-re-probed by view reads, this one how
2007# long a stale plan snapshot stays stale.
2008_PLAN_RESTATE_FAILURE_WAIT_S = 60.0
2009# An engine that lists no device on a GPU host may be hitting a transient GPU-init
2010# error (the card momentarily held by another process); re-probe before treating
2011# the empty list as fatal.
2012_DEVICE_PROBE_EMPTY_RETRIES = 3
2013_DEVICE_PROBE_EMPTY_RETRY_DELAY_S = 0.5
2014# Stands in for the engine binary's identity on a host that has none, so "no
2015# engine yet" and "an engine at last" are different identities.
2016_NO_ENGINE_BINARY = "no-engine"
2019class _ReadDeviceCache:
2020 """Short-TTL device-probe cache for the read/view path.
2022 Not a ``cachetools.TTLCache``: it caches the *failure* too, under its own
2023 longer TTL, and re-raises it. A memoizing cache stores return values only,
2024 so a failing probe would re-spawn the subprocess on every placement read.
2026 Inspecting placement (GET placement/gpus, preview, ``placement show``)
2027 resolves devices on every call, which spawns a ``llama-server --list-devices``
2028 subprocess; a brief TTL collapses a burst of reads onto one probe. A probe
2029 failure is cached too (with its own TTL) and re-raised to every read in the
2030 window. The launch path is never served from here -- it sizes against the
2031 clean-box plan snapshot below (captured after stale-server reaping).
2033 An entry also carries the identity of the binary that answered, so it cannot
2034 describe an engine that has since been installed or upgraded in place.
2035 """
2037 def __init__(self, ttl_s: float, failure_ttl_s: float) -> None:
2038 self._ttl_s = ttl_s
2039 self._failure_ttl_s = failure_ttl_s
2040 self._lock = threading.Lock()
2041 self._at: float | None = None
2042 self._engine: str | None = None
2043 self._reading: DeviceReading | None = None
2044 self._failure: ProviderError | None = None
2046 def _is_fresh(self, engine: str) -> bool:
2047 """Whether the entry in hand answers for *engine* and is still inside its TTL."""
2048 ttl = self._ttl_s if self._failure is None else self._failure_ttl_s
2049 return self._at is not None and self._engine == engine and time.monotonic() - self._at < ttl
2051 def get(self, binary: Path) -> DeviceReading:
2052 engine = engine_binary_identity(binary)
2053 with self._lock:
2054 fresh = self._is_fresh(engine)
2055 if fresh and self._failure is not None:
2056 raise self._failure
2057 if self._reading is None or not fresh:
2058 self._at = time.monotonic()
2059 self._engine = engine
2060 try:
2061 self._reading = _read_devices(binary)
2062 except ProviderError as exc:
2063 self._reading = None
2064 self._failure = exc
2065 raise
2066 self._failure = None
2067 return self._reading
2069 def clear(self) -> None:
2070 with self._lock:
2071 self._at = None
2072 self._engine = None
2073 self._reading = None
2074 self._failure = None
2077_read_device_cache = _ReadDeviceCache(_DEVICE_PROBE_TTL_S, _DEVICE_PROBE_FAILURE_TTL_S)
2080def clear_read_device_cache() -> None:
2081 """Drop the read-path device probe cache (e.g. after the fleet is reconfigured).
2083 Also drops what the host's Vulkan loader told us about device types, which is
2084 otherwise held for the process lifetime and would survive a driver reload or
2085 an eGPU being plugged in.
2086 """
2087 from lilbee.providers.fleet.gpu_select import (
2088 integrated_vulkan_indices,
2089 vulkan_device_types_by_name,
2090 )
2092 _read_device_cache.clear()
2093 vulkan_device_types_by_name.cache_clear()
2094 integrated_vulkan_indices.cache_clear()
2097@dataclass(frozen=True)
2098class _PlanProbe:
2099 """Clean-box memory snapshot every plan is sized against.
2101 Captured once, right after stale-server reaping and before the first build,
2102 when nothing lilbee owns is loaded. Reloads re-plan against this same
2103 snapshot instead of re-probing: a live probe under a loaded fleet reports
2104 our own residency as unavailable, which would shrink chat context and slot
2105 counts, widen splits, and (on a unified-memory host) evict roles outright.
2106 Launches stay a pure function of config + hardware + this snapshot, so the
2107 reload diff restarts only real changes. Cleared on full fleet teardown so
2108 the next boot probes the clean box afresh.
2109 """
2111 devices: tuple[FleetDevice, ...]
2112 # What one role may size its ctx and slots against, already scaled by
2113 # cfg.gpu_memory_fraction. System memory only on a host with no GPU.
2114 sizing_budget: int
2115 free_system: int
2116 # The engine binary that answered the probe. A snapshot is only an answer
2117 # about the binary it was taken against.
2118 engine: str
2119 # The engine listed GPUs and lilbee rejected all of them, so the plan is
2120 # CPU-shaped while the engine would still choose one of those devices.
2121 engine_devices_all_refused: bool = False
2122 # The backend the engine selected when this snapshot was taken.
2123 backend: EngineBackend = EngineBackend.UNKNOWN
2126@dataclass(frozen=True)
2127class _FailedProbe:
2128 """An engine identity whose restate probe raised, and when it raised."""
2130 engine: str
2131 at: float
2134# Re-reads a plan snapshot against the hardware, or answers ``None`` when the
2135# probe could not run.
2136_Restate: TypeAlias = "Callable[[_PlanProbe], _PlanProbe | None]"
2139class _PlanProbeStore:
2140 """Holds the captured plan snapshot; a single instance below (no bare global)."""
2142 def __init__(self) -> None:
2143 self._lock = threading.Lock()
2144 self._probe: _PlanProbe | None = None
2145 self._failed: _FailedProbe | None = None
2146 # Held across the probe a restate runs, so a burst of reads after an
2147 # engine swap pays one probe. Private, and never taken outside this
2148 # class, so no caller can hold it across a probe of its own. Separate
2149 # from _lock, which must never be held across a subprocess.
2150 self._restate_lock = threading.Lock()
2152 def set(self, probe: _PlanProbe) -> None:
2153 with self._lock:
2154 self._probe = probe
2155 self._failed = None
2157 def get(self) -> _PlanProbe | None:
2158 with self._lock:
2159 return self._probe
2161 def note_failed_probe(self, engine: str) -> None:
2162 """Record that the restate against *engine* could not probe."""
2163 with self._lock:
2164 self._failed = _FailedProbe(engine, time.monotonic())
2166 def probe_failed_recently(self, engine: str) -> bool:
2167 """Whether a restate against *engine* raised inside the failure wait."""
2168 with self._lock:
2169 failed = self._failed
2170 return (
2171 failed is not None
2172 and failed.engine == engine
2173 and time.monotonic() - failed.at < _PLAN_RESTATE_FAILURE_WAIT_S
2174 )
2176 def restate(self, restate: _Restate, engine: str) -> None:
2177 """Re-read the snapshot with *restate*, whatever binary it answers for now."""
2178 with self._restate_lock:
2179 probe = self.get()
2180 if probe is not None:
2181 self._store_restated(probe, restate, engine)
2183 def restate_if_stale(self, restate: _Restate, engine: str) -> _PlanProbe | None:
2184 """The snapshot, re-read with *restate* first when it answers for another binary.
2186 The staleness and the failure wait are re-checked under the lock, so a
2187 burst of stale reads pays one probe and the rest take its answer.
2188 """
2189 with self._restate_lock:
2190 probe = self.get()
2191 if probe is None or probe.engine == engine or self.probe_failed_recently(engine):
2192 return probe
2193 return self._store_restated(probe, restate, engine)
2195 def _store_restated(self, probe: _PlanProbe, restate: _Restate, engine: str) -> _PlanProbe:
2196 """Store *probe* restated, or keep it and hold off re-probing *engine* for a while.
2198 A probe that raised is not an answer about *engine*, so the snapshot keeps
2199 the identity that did answer for it and the failure is recorded instead.
2200 """
2201 fresh = restate(probe)
2202 if fresh is None:
2203 self.note_failed_probe(engine)
2204 return probe
2205 self.set(fresh)
2206 return fresh
2208 def clear(self) -> None:
2209 with self._lock:
2210 self._probe = None
2211 self._failed = None
2214_plan_probe_store = _PlanProbeStore()
2217# Where the ladder stops. Sized for chat, below which the answers are too short
2218# to be useful, so a role that still will not load here has a real problem the
2219# planner cannot size its way out of and the failure should surface. Roles whose
2220# window already sits under it (a small embedding context) are left alone rather
2221# than raised to meet it, so for them the ladder is a no-op and the failure
2222# surfaces after the one retry.
2223MIN_DOWNSHIFT_CTX = 4096
2226class _CtxDownshiftStore:
2227 """How many halvings each role's auto context has taken after a load OOM.
2229 An estimate that was too optimistic is only recoverable if the retry asks
2230 for something different. Halving the auto context does that, and keeping the
2231 count here rather than in the launch means the whole plan is re-predicted
2232 against the smaller number, including the placement it implies.
2233 """
2235 def __init__(self) -> None:
2236 self._lock = threading.Lock()
2237 self._steps: dict[WorkerRole, int] = {}
2238 # The last unshifted context each role was sized from, recorded as it is
2239 # applied. Deciding whether another halving would change anything needs
2240 # the number being halved, and this is the only place that sees it.
2241 self._base: dict[WorkerRole, int] = {}
2243 def steps(self, role: WorkerRole) -> int:
2244 with self._lock:
2245 return self._steps.get(role, 0)
2247 def note_base(self, role: WorkerRole, ctx: int) -> None:
2248 with self._lock:
2249 self._base[role] = ctx
2251 def base(self, role: WorkerRole) -> int | None:
2252 with self._lock:
2253 return self._base.get(role)
2255 def step(self, role: WorkerRole) -> int:
2256 with self._lock:
2257 taken = self._steps.get(role, 0) + 1
2258 self._steps[role] = taken
2259 return taken
2261 def clear(self, role: WorkerRole | None = None) -> None:
2262 with self._lock:
2263 if role is None:
2264 self._steps.clear()
2265 self._base.clear()
2266 return
2267 self._steps.pop(role, None)
2268 self._base.pop(role, None)
2271_ctx_downshift_store = _CtxDownshiftStore()
2274def apply_ctx_downshift(role: WorkerRole, ctx: int) -> int:
2275 """*ctx* halved once per downshift step recorded for *role*, floored.
2277 Never more than *ctx*. The floor is a stopping point, not a target: applied
2278 to a context already below it (a small embedding window, a model trained for
2279 2048 tokens) a bare floor would hand back a larger number, and the retry
2280 after a load OOM would ask for more memory than the launch that just ran out
2281 of it. Such a role simply has nothing to give back, and its failure surfaces
2282 after the one retry instead.
2284 A user's ``cfg.num_ctx`` pin is returned untouched: serving a window smaller
2285 than the one that was asked for, without being asked, is worse than failing
2286 to load and saying so.
2287 """
2288 from lilbee.core.config import cfg
2290 if role is WorkerRole.CHAT and cfg.num_ctx is not None:
2291 return ctx
2292 _ctx_downshift_store.note_base(role, ctx)
2293 return _shifted(ctx, _ctx_downshift_store.steps(role))
2296def _shifted(ctx: int, steps: int) -> int:
2297 """*ctx* halved *steps* times, never below the floor and never above *ctx*."""
2298 return min(ctx, max(MIN_DOWNSHIFT_CTX, ctx >> steps)) if steps else ctx
2301def record_ctx_downshift(role: WorkerRole) -> bool:
2302 """Take one downshift step for *role*; False when there is none left to take.
2304 False means the retry would ask for the same thing again, so the caller must
2305 surface the load failure instead of respawning an identical launch.
2306 """
2307 from lilbee.core.config import cfg
2309 if role is WorkerRole.CHAT and cfg.num_ctx is not None:
2310 return False
2311 base = _ctx_downshift_store.base(role)
2312 if base is None:
2313 # Nothing has been sized for this role yet, so there is no number to
2314 # decide against. Allow one step rather than trusting that a plan always
2315 # runs first: an unbounded grant here would let a caller that never
2316 # sizes anything loop forever.
2317 if _ctx_downshift_store.steps(role):
2318 return False
2319 _ctx_downshift_store.step(role)
2320 return True
2321 steps = _ctx_downshift_store.steps(role)
2322 if _shifted(base, steps + 1) == _shifted(base, steps):
2323 return False
2324 _ctx_downshift_store.step(role)
2325 return True
2328def clear_ctx_downshift(role: WorkerRole | None = None) -> None:
2329 """Forget *role*'s recorded downshift, or every role's, back to full size.
2331 Called when a role's engine reports ready, which is proof the reduced plan
2332 loaded: keeping the reduction after that would carry a shrunken window into
2333 a machine that has since freed memory, or into a smaller model the user
2334 switched to, and would then refuse on its first failure with a budget it
2335 had already spent.
2336 """
2337 _ctx_downshift_store.clear(role)
2340def _probe_engine_devices() -> DeviceReading:
2341 """Resolve the binary and enumerate devices.
2343 This is the wedge point: a missing binary raises NOT_FOUND, and a CUDA build
2344 that cannot init a GPU (a broken-runtime host) raises loud from resolve_devices
2345 rather than silently degrading. Device enumeration reads no residency, so it is
2346 safe to run while an incumbent engine is still up.
2347 """
2348 binary = resolve_llama_server()
2349 reading = _read_devices(binary)
2350 if reading.devices:
2351 return reading
2352 return _reprobe_while_a_gpu_is_installed(binary, reading)
2355def _reprobe_while_a_gpu_is_installed(binary: Path, reading: DeviceReading) -> DeviceReading:
2356 """Ask again when the host has a GPU the engine did not list.
2358 The plan snapshot is taken once, on a clean box, and is not retaken until a
2359 full teardown, so an empty first answer decides the whole run. A GPU driver
2360 that is still initializing when the daemon starts, which is ordinary under
2361 systemd or right after a container gains a device, would leave a GPU host
2362 serving on CPU until someone noticed and restarted it.
2364 Only where a card is actually installed. A host with no GPU answers empty
2365 every time and must not pay a retry for it on every start.
2366 """
2367 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
2369 if not installed_gpu_vendor_ids():
2370 return reading
2371 for attempt in range(1, _PROBE_RETRIES + 1):
2372 log.info(
2373 "The engine listed no GPU on a host that has one; asking again in %.1fs "
2374 "(attempt %d of %d) in case the driver is still initializing.",
2375 _PROBE_RETRY_DELAY_S,
2376 attempt,
2377 _PROBE_RETRIES,
2378 )
2379 time.sleep(_PROBE_RETRY_DELAY_S)
2380 clear_read_device_cache()
2381 reading = _read_devices(binary)
2382 if reading.devices:
2383 return reading
2384 return reading
2387def assert_engine_probeable() -> None:
2388 """Raise if the engine cannot be probed; capture no snapshot.
2390 A build precondition that must run BEFORE stopping a replaceable incumbent:
2391 it surfaces a wedged GPU probe or an unusable CUDA runtime without taking the
2392 residency-dependent memory snapshot (that belongs on the clean box, after the
2393 stop, in capture_plan_probe). resolve_devices caches within its TTL, so the
2394 follow-up capture reuses this enumeration rather than re-probing the hardware.
2395 """
2396 _probe_engine_devices()
2399def _engine_identity() -> str:
2400 """Identity of the engine binary a plan snapshot has to agree with.
2402 A host with no engine binary has an identity of its own, so a snapshot taken
2403 before the engine arrived never reads as one taken after.
2404 """
2405 try:
2406 return engine_binary_identity(resolve_llama_server())
2407 except (ProviderError, OSError):
2408 return _NO_ENGINE_BINARY
2411# The binary one planning pass answers about. A pass reads the snapshot half a
2412# dozen times; resolving the identity per read let an engine that landed between
2413# two of them size the plan against one snapshot and place it against another.
2414_pass_engine: contextvars.ContextVar[str | None] = contextvars.ContextVar(
2415 "lilbee_plan_pass_engine", default=None
2416)
2419@contextmanager
2420def _one_engine_per_pass() -> Iterator[None]:
2421 """Pin the engine identity every snapshot read in this planning pass answers about.
2423 Re-entrant: a nested pass keeps the outer one's binary, so a pass that calls
2424 another still asks about one engine. A snapshot that is stale when the pass
2425 opens is restated once, at the first read.
2426 """
2427 if _pass_engine.get() is not None:
2428 yield
2429 return
2430 token = _pass_engine.set(_engine_identity())
2431 try:
2432 yield
2433 finally:
2434 _pass_engine.reset(token)
2437def _pass_engine_identity() -> str:
2438 """The engine this read answers about: the pass's binary, or the live one outside a pass."""
2439 pinned = _pass_engine.get()
2440 return pinned if pinned is not None else _engine_identity()
2443def _probe_engine_devices_and_identity() -> tuple[str, DeviceReading]:
2444 """The engine's reading, tagged with the identity read before the probe ran.
2446 The identity is read first so a binary replaced while the probe runs tags the
2447 answer with the binary that gave it, and the next read restates. The other
2448 order tags the outgoing binary's answer with the incoming binary.
2449 """
2450 engine = _engine_identity()
2451 return engine, _probe_engine_devices()
2454def capture_plan_probe() -> None:
2455 """Snapshot devices and memory for planning; call only on a clean box."""
2456 engine, reading = _probe_engine_devices_and_identity()
2457 _plan_probe_store.set(
2458 _PlanProbe(
2459 devices=tuple(reading.devices),
2460 sizing_budget=_device_sizing_budget(reading.devices),
2461 free_system=model_cache.free_system_memory(),
2462 engine=engine,
2463 engine_devices_all_refused=reading.refused_all,
2464 backend=reading.backend,
2465 )
2466 )
2469def _structural(devices: Iterable[FleetDevice]) -> tuple[FleetDevice, ...]:
2470 """*devices* with the volatile free reading zeroed, so equality is structural.
2472 This probe runs while the fleet is resident, so its free figures are deflated
2473 by the very models a reload is about to stop; comparing or keying on them
2474 would read every reload as a hardware change.
2475 """
2476 return tuple(replace(d, free_bytes=0) for d in devices)
2479def refresh_plan_devices() -> None:
2480 """Re-read which devices exist, keeping the clean-box memory figures.
2482 The snapshot is captured once and only a full teardown clears it, so an eGPU
2483 unplugged, a driver reset, or a VM hot-remove left the fleet pinning a device
2484 that is no longer there and every rebuild replanning onto it.
2486 Only the structural half is restated. The memory figures are what make a
2487 reload plan the way the boot did, and re-taking them while the fleet is
2488 resident would charge it against itself, which is the whole reason the
2489 snapshot exists. So a card that survives the refresh keeps the snapshot's
2490 free figure; only a genuinely new card contributes a fresh one.
2492 A probe that cannot run keeps the previous device list: the last known one is
2493 a better answer than none, and the loud paths for an unreachable engine live
2494 in the build, not here. A reload asks to look again, so it probes even inside
2495 the failure wait that holds the read path off.
2496 """
2497 _plan_probe_store.restate(_restated, _engine_identity())
2500def _restated(probe: _PlanProbe) -> _PlanProbe | None:
2501 """*probe* re-read against the hardware, or ``None`` when the probe could not run."""
2502 clear_read_device_cache()
2503 try:
2504 engine, reading = _probe_engine_devices_and_identity()
2505 except (ProviderError, OSError) as exc:
2506 log.debug("Device rediscovery could not run, keeping the previous list: %s", exc)
2507 return None
2508 devices = reading.devices
2509 if _structural(devices) == _structural(probe.devices):
2510 # The refusal answer is carried too: an engine that lists only devices
2511 # lilbee refuses probes as an empty list, which is structurally the same
2512 # as no GPU at all. Keeping the old binary's refusal here would leave the
2513 # CPU pin off and let ggml fall back onto the adapter just refused.
2514 return replace(
2515 probe,
2516 engine=engine,
2517 engine_devices_all_refused=reading.refused_all,
2518 backend=reading.backend,
2519 )
2520 log.info(
2521 "The set of GPUs changed since this fleet was planned (%d device(s) now, %d before); "
2522 "replanning against the ones that are here.",
2523 len(devices),
2524 len(probe.devices),
2525 )
2526 snapshot_free = {replace(d, free_bytes=0): d.free_bytes for d in probe.devices}
2527 merged = tuple(
2528 replace(d, free_bytes=snapshot_free.get(replace(d, free_bytes=0), d.free_bytes))
2529 for d in devices
2530 )
2531 return _PlanProbe(
2532 devices=merged,
2533 sizing_budget=_device_sizing_budget(merged),
2534 free_system=probe.free_system,
2535 engine=engine,
2536 engine_devices_all_refused=reading.refused_all,
2537 backend=reading.backend,
2538 )
2541def _current_plan_probe() -> _PlanProbe | None:
2542 """The plan snapshot, restated first when another engine binary is now in place.
2544 The snapshot carries the identity of the binary that answered, so an engine
2545 installed or upgraded under a running serve re-probes on the next plan read.
2546 """
2547 probe = _plan_probe_store.get()
2548 if probe is None:
2549 return None
2550 engine = _pass_engine_identity()
2551 return probe if probe.engine == engine else _restate_plan_probe(engine)
2554def _restate_plan_probe(engine: str) -> _PlanProbe | None:
2555 """Restate the snapshot for *engine*, once per burst of stale reads.
2557 An engine whose probe keeps raising is re-probed at most once per failure
2558 wait on this read path, so a broken binary costs the retry ladder
2559 occasionally rather than on every read. A reload restates through
2560 refresh_plan_devices instead, which probes every time.
2562 The wait is keyed on the binary's identity, which is a digest of its bytes,
2563 so a binary repaired in place is retried without a reload as soon as its
2564 content changes, whatever its stat fields do.
2565 """
2566 return _plan_probe_store.restate_if_stale(_restated, engine)
2569def clear_plan_probe() -> None:
2570 """Drop the plan snapshot (full fleet teardown); the next build re-captures."""
2571 _plan_probe_store.clear()
2574def _cpu_pin_when_every_device_was_refused() -> tuple[str, ...]:
2575 """``("none",)`` when the engine offered GPUs that lilbee refused, else empty.
2577 Dropping a device from lilbee's view does not stop the engine using it. With
2578 no pin at all, ggml applies its own selection, and its fallback takes the
2579 first non-CPU adapter, which is exactly the paravirtual device just refused;
2580 with every layer offloaded by default the model then runs on it while
2581 placement budgeted against system RAM. Naming no device keeps the engine on
2582 the CPU the plan was shaped for.
2583 """
2584 probe = _current_plan_probe()
2585 if probe is None or not probe.engine_devices_all_refused:
2586 return ()
2587 log.warning(
2588 "The engine listed GPU devices that lilbee will not plan onto, so it is being "
2589 "run on the CPU. Serving from one of them would be slower than the CPU or fail "
2590 "outright, and placement has been sized for system RAM."
2591 )
2592 return (_NO_DEVICE,)
2595def _plan_devices(binary: Path) -> list[FleetDevice]:
2596 """Devices the plan paths size against: the snapshot, else a live probe."""
2597 probe = _current_plan_probe()
2598 return list(probe.devices) if probe is not None else resolve_devices(binary)
2601def plan_sizing_budget(device: FleetDevice | None = None) -> int:
2602 """Usable memory for ctx/slot sizing: *device*'s own, else the snapshot, else live."""
2603 from lilbee.core.config import cfg
2605 if device is not None:
2606 return int(device.total_bytes * cfg.gpu_memory_fraction)
2607 probe = _current_plan_probe()
2608 if probe is not None:
2609 return probe.sizing_budget
2610 return _device_sizing_budget(_live_sizing_devices())
2613def plan_sizing_is_unified() -> bool:
2614 """Whether ctx sizing charges the shared-memory footprint rather than VRAM.
2616 True when no device has memory of its own -- an integrated GPU, an Apple
2617 Silicon Mac, or no GPU at all -- because there a model's would-be VRAM and
2618 its host bytes are the same memory, and charging only the VRAM half
2619 under-reports the load by everything it maps. The same test decides the pool
2620 placement charges against (:func:`_unified_memory_budget`).
2621 """
2622 probe = _current_plan_probe()
2623 devices = list(probe.devices) if probe is not None else _live_sizing_devices()
2624 return all(dev.unified for dev in devices)
2627def _device_sizing_budget(devices: Sequence[FleetDevice]) -> int:
2628 """Memory one role may size its ctx and slots against, in bytes.
2630 Read from the engine's own device report, which ran under the environment the
2631 servers will run under and states each device's memory whatever the backend.
2632 A host-memory read answers with system RAM on every host without an NVIDIA
2633 card, which gave a 24 GiB AMD card a budget the size of the machine, and on
2634 Apple Silicon it ignores that Metal will not allocate past
2635 ``recommendedMaxWorkingSetSize``, which is the figure the probe carries.
2637 The smallest device, since this is asked before placement has picked one;
2638 :func:`_launch_for` re-sizes against the card the role actually landed on.
2639 System memory only when the engine reports no device at all, where the fleet
2640 runs on the CPU and system memory is the budget.
2641 """
2642 from lilbee.core.config import cfg
2644 if devices:
2645 return int(min(d.total_bytes for d in devices) * cfg.gpu_memory_fraction)
2646 return int(model_cache.total_system_memory() * cfg.gpu_memory_fraction)
2649def _live_sizing_devices() -> list[FleetDevice]:
2650 """Devices to size against with no plan snapshot; empty when none can be read."""
2651 try:
2652 return _read_device_cache.get(resolve_llama_server()).devices
2653 except (ProviderError, OSError):
2654 return []
2657def _plan_free_system_memory() -> int:
2658 """Free system RAM for the unified-memory budget: the snapshot, else live."""
2659 probe = _current_plan_probe()
2660 return probe.free_system if probe is not None else model_cache.free_system_memory()
2663def _unreported_bytes(role: WorkerRole, mmproj: Path | None) -> int:
2664 """Estimated bytes the engine allocates without printing a buffer line.
2666 A vision projector's weights: llama.cpp allocates them in clip's own loader,
2667 which prints a size but not the "buffer size = N MiB" shape the readback
2668 reads, so the report is short by exactly this and the self-check would warn
2669 on a load that was sized correctly.
2670 """
2671 if role is not WorkerRole.VISION or mmproj is None:
2672 return 0
2673 try:
2674 return mmproj.stat().st_size
2675 except OSError:
2676 return 0
2679def _chat_no_mmap(weights_bytes: int, *, on_network_fs: bool = False) -> bool:
2680 """Whether the chat server should malloc its weights instead of mmapping them.
2682 Local disk mmaps: lazy page-fault paging gives a faster first token on a cold
2683 cache -- the common desktop first launch -- matching mmap-by-default engines.
2684 The no-mmap buffered full read only wins on an already-hot cache and it
2685 pessimizes cold start, so it is not worth defaulting on for local disk. A
2686 network filesystem still prefers the buffered read whenever the host copy
2687 fits, because mmap page faults served over the wire can wedge the loader in
2688 uninterruptible I/O (see ``_NO_MMAP_NETWORK_RAM_FRACTION``).
2689 """
2690 if not on_network_fs:
2691 return False
2692 return weights_bytes <= model_cache.total_system_memory() * _NO_MMAP_NETWORK_RAM_FRACTION
2695def _device_names(devices: tuple[FleetDevice, ...]) -> tuple[str, ...]:
2696 """``--device`` names for *devices*, empty when the backend pins through env.
2698 Vulkan and SYCL, because neither one's environment variable speaks the space
2699 the probe enumerated. Vulkan's indexes the raw loader enumeration while the
2700 names come from the engine's filtered list, so the two disagree wherever ggml
2701 drops or merges a device. SYCL's is not an index list at all but a selector
2702 over a backend runtime, so a device the engine calls ``SYCL1`` need not be
2703 Level Zero ordinal 1: OpenCL devices interleave, discarded devices shift the
2704 numbering, and multi-tile cards appear as sub-devices.
2706 ``--device`` sidesteps both by naming devices exactly as ``--list-devices``
2707 printed them, which is where these indices were read from. CUDA and ROCm
2708 keep composing their variables, which do share the probe's space.
2709 """
2710 if not devices or devices[0].backend not in _NAME_PINNED_BACKENDS:
2711 return ()
2712 if any(d.from_loader for d in devices):
2713 # These indices are raw loader ordinals, and --device speaks the engine's
2714 # own post-filter naming, so Vulkan1 here can name Vulkan0 there or
2715 # nothing at all. Sizing against them is still worth doing; pinning by
2716 # them is not. Left unpinned, ggml applies its own device selection,
2717 # which is the filtering lilbee is trying to agree with in the first
2718 # place. The env pin is not the answer either: it takes raw ordinals but
2719 # switches off the type filter, the support check and the dedup with them.
2720 return ()
2721 return tuple(f"{d.backend}{d.index}" for d in devices)
2724def _unified_memory_budget(devices: list[FleetDevice]) -> int | None:
2725 """Shared-RAM placement budget (free RAM minus the OS floor), or ``None``.
2727 ``None`` once any device has memory of its own, since dedicated VRAM is the
2728 constraint there rather than system RAM. A host whose only devices are
2729 integrated, and a host with no devices at all, both stay inside the system
2730 budget: their GPU memory is the system's memory.
2731 """
2732 # Only a device with memory of its own lifts the system-RAM constraint. An
2733 # integrated GPU or an Apple Silicon Mac reports a slice of the same RAM the
2734 # OS is using, so treating its total as headroom over-commits the machine by
2735 # roughly the whole system footprint.
2736 if any(not device.unified for device in devices):
2737 return None
2738 return _capped_by_device_memory(
2739 max(0, _plan_free_system_memory() - _system_memory_floor()), devices
2740 )
2743def _unified_admission_budget(devices: list[FleetDevice]) -> int | None:
2744 """Shared-RAM pool a role set is *admitted* against, or ``None`` if dedicated.
2746 Total installed RAM minus the OS floor, not what happens to be free. Sizing
2747 asks a different question and keeps using free RAM: how much context can be
2748 backed right now. Admission asks whether the machine can host this fleet at
2749 all, and the plan defines the whole intended residency, so charging it
2750 against a live figure refuses a 600 MB model on a box that is merely busy at
2751 the moment, which is what happened. The GPU path already charges total
2752 capacity for exactly this reason.
2753 """
2754 if _unified_memory_budget(devices) is None:
2755 return None
2756 return _capped_by_device_memory(
2757 max(0, model_cache.total_system_memory() - _system_memory_floor()), devices
2758 )
2761def _system_memory_floor() -> int:
2762 """RAM held back for the OS when placing against system memory.
2764 ``cfg.system_memory_reserve_gb``, still capped at a quarter of installed RAM:
2765 a fixed reserve leaves a 7-8 GB host with no budget at all and refuses even
2766 tiny models, so the proportional cap holds however the reserve is set.
2767 """
2768 from lilbee.core.config import cfg
2770 total = model_cache.total_system_memory()
2771 return min(int(cfg.system_memory_reserve_gb * 1024**3), total // _SYSTEM_MEMORY_FLOOR_DIVISOR)
2774def _capped_by_device_memory(budget: int, devices: Sequence[FleetDevice]) -> int:
2775 """*budget*, never above what the devices can address between them.
2777 A shared-memory device still has a ceiling of its own: an integrated GPU
2778 addresses a fixed aperture of system RAM, and Metal will not allocate past
2779 ``recommendedMaxWorkingSetSize``. Both report that ceiling as their total, so
2780 a host budget derived from installed RAM promises memory the devices cannot
2781 reach. Unchanged where the engine reports no device, since the fleet is then
2782 running on the CPU and the host figure is the true one.
2783 """
2784 if not devices:
2785 return budget
2786 return min(budget, sum(d.total_bytes for d in devices))
2789def _device_capacity(devices: list[FleetDevice], charge_against_free: bool) -> dict[int, int]:
2790 """Per-device memory placement may charge against, keyed by device index.
2792 A card's total is what it holds, not what is going spare. A compositor, a
2793 browser, or a training job sitting on VRAM is invisible in the total, and the
2794 usable fraction placement applies covers fragmentation and driver overhead
2795 rather than other tenants, so a plan fits on paper and OOMs at load.
2797 Free bytes answer that, but only where they mean "everyone else's residency":
2798 that is the clean-box snapshot, taken after stale servers are reaped and
2799 before anything is built. Read live on a warm box they also exclude the
2800 fleet's own models, and since a plan always describes the complete intended
2801 residency, charging them there would count the fleet against itself and
2802 report a running plan as unplaceable. Those callers keep the total.
2804 Placement applies its usable fraction to whatever this returns, so a card
2805 with a tenant keeps a proportional margin rather than being packed to its
2806 last free byte, where fragmentation is worst.
2807 """
2808 packable = _packable_devices(devices)
2809 if not charge_against_free:
2810 return {d.index: d.total_bytes for d in packable}
2811 return {d.index: min(d.total_bytes, d.free_bytes) for d in packable}
2814def _packable_devices(devices: list[FleetDevice]) -> list[FleetDevice]:
2815 """The devices bin-packing may charge against.
2817 An integrated GPU's memory is the host's. Packing it beside a dedicated card
2818 promises the same RAM twice, once to its own budget and once to everything
2819 else on the machine, and its heap is often the larger number, so the packer
2820 prefers it: a 32 GiB shared heap outbids a 24 GiB card that actually has the
2821 memory. Where a dedicated device exists it is the one to serve from, and the
2822 integrated one is left to the shared-memory budget.
2824 A host with nothing but integrated devices keeps them. There is nothing else
2825 to serve from, and that path is governed by the system budget rather than by
2826 per-device packing.
2827 """
2828 dedicated = [d for d in devices if not d.unified]
2829 return dedicated or devices
2832def _resolve_placement(
2833 placement: PlacementSpec | None,
2834 inputs: list[ModelPlacementInput],
2835 model_refs: dict[WorkerRole, str],
2836 devices: list[FleetDevice],
2837 *,
2838 unified_budget: int | None,
2839 charge_against_free: bool = False,
2840) -> Placement:
2841 """Resolve a Placement from the manual spec when set, else the auto planner."""
2842 estimate_peak = _peak_estimator(model_refs)
2843 capacity = _device_capacity(devices, charge_against_free)
2844 if placement is not None:
2845 return placement_from_spec(
2846 placement,
2847 tuple(model_refs),
2848 capacity,
2849 estimate_peak=estimate_peak,
2850 )
2851 # The chat split's card count is decided against the snapshot's free VRAM (what the
2852 # launch sizes its context against) so placement and launch agree. A split needs
2853 # >=2 GPUs, so skip the chat model's gguf read entirely below that.
2854 chat_ctx_fit, chat_ctx_target = (
2855 _chat_split_ctx_objective(model_refs) if len(capacity) >= _MIN_SPLIT_GPUS else (None, 0)
2856 )
2857 return plan_placement(
2858 inputs,
2859 [(idx, budget) for idx, budget in capacity.items()],
2860 estimate_peak=estimate_peak,
2861 unified_budget=unified_budget,
2862 chat_ctx_fit=chat_ctx_fit,
2863 chat_ctx_target=chat_ctx_target,
2864 free_headroom={d.index: d.free_bytes for d in devices},
2865 )
2868def _placement_or_auto(
2869 placement: PlacementSpec | None,
2870 inputs: list[ModelPlacementInput],
2871 model_refs: dict[WorkerRole, str],
2872 devices: list[FleetDevice],
2873 *,
2874 unified_budget: int | None,
2875 charge_against_free: bool = False,
2876) -> tuple[Placement, bool]:
2877 """Resolve a saved spec, falling back to auto when it no longer fits the hardware.
2879 Returns the placement and whether the spec was the one applied. Hardware moves
2880 under a saved placement: a card is removed, a driver stops enumerating a GPU, a
2881 container starts without one. Refusing to plan there takes chat, embed and
2882 ingest down over a pin set on hardware the host no longer has, so the fleet
2883 degrades to automatic placement and logs why. An interactive apply still fails
2884 loud (:func:`lilbee.app.placement.set_placement`), where the pin is what the
2885 caller just asked for and a silent substitution would be the surprise.
2886 """
2887 if placement is None:
2888 return _resolve_placement(
2889 None,
2890 inputs,
2891 model_refs,
2892 devices,
2893 unified_budget=unified_budget,
2894 charge_against_free=charge_against_free,
2895 ), False
2896 try:
2897 return _resolve_placement(
2898 placement,
2899 inputs,
2900 model_refs,
2901 devices,
2902 unified_budget=unified_budget,
2903 charge_against_free=charge_against_free,
2904 ), True
2905 except PlacementError as exc:
2906 log.warning(
2907 "The saved GPU placement does not fit this hardware (%s); using automatic "
2908 "placement instead. Set a new placement to replace it.",
2909 exc,
2910 )
2911 return _resolve_placement(
2912 None,
2913 inputs,
2914 model_refs,
2915 devices,
2916 unified_budget=unified_budget,
2917 charge_against_free=charge_against_free,
2918 ), False
2921@dataclass(frozen=True)
2922class ResolvedPlacement:
2923 """Devices + resolved instance plans + model refs for the placement view."""
2925 devices: tuple[FleetDevice, ...]
2926 instances: tuple[InstancePlan, ...]
2927 unplaceable_roles: tuple[WorkerRole, ...]
2928 model_refs: dict[WorkerRole, str]
2929 # Roles placed anyway despite not fitting, with the shortfall in bytes. The
2930 # planner has always known this and only logged it, so a surface showed a
2931 # tight role as comfortably placed.
2932 tight_roles: dict[WorkerRole, int] = field(default_factory=dict)
2933 co_tenants: frozenset[WorkerRole] = frozenset()
2934 # False when a spec was given but did not fit the hardware, so these instances
2935 # are the auto planner's and a surface must not present them as the manual plan.
2936 spec_applied: bool = True
2937 # Roles configured but skipped because their model isn't installed (role -> ref).
2938 # Distinct from unplaceable_roles (installed but won't fit); lets a surface show
2939 # "not downloaded" instead of an empty table on a fresh install.
2940 skipped_not_installed: dict[WorkerRole, str] = field(default_factory=dict)
2941 # The backend the engine selected, which no surface may infer from ``devices``:
2942 # an empty list is a CPU host and a failed probe alike.
2943 engine_backend: EngineBackend = EngineBackend.UNKNOWN
2946def resolve_placement_plan(
2947 placement: PlacementSpec | None, *, fall_back_to_auto: bool = False
2948) -> ResolvedPlacement:
2949 """Probe devices and resolve the auto-or-manual placement, without launching.
2951 ``fall_back_to_auto`` reads *placement* as a saved setting rather than a
2952 request: one that no longer fits the hardware resolves to the auto plan with
2953 ``spec_applied`` False instead of raising.
2955 The reported backend comes from the same reading as the device list, not from
2956 :func:`engine_backend`. This path deliberately reports live hardware rather
2957 than the plan snapshot, so borrowing the snapshot's backend would pair a live
2958 device list with a backend read at fleet build.
2959 """
2960 binary = resolve_llama_server()
2961 reading = _read_device_cache.get(binary)
2962 devices = reading.devices
2963 unified_budget = _unified_memory_budget(devices)
2964 inputs, model_refs, _, skipped_not_installed = _server_model_inputs(
2965 None, unified_budget=unified_budget, total_vram=sum(d.total_bytes for d in devices)
2966 )
2967 admission_budget = _unified_admission_budget(devices)
2968 if fall_back_to_auto:
2969 resolved, spec_applied = _placement_or_auto(
2970 placement, inputs, model_refs, devices, unified_budget=admission_budget
2971 )
2972 else:
2973 resolved = _resolve_placement(
2974 placement, inputs, model_refs, devices, unified_budget=admission_budget
2975 )
2976 spec_applied = placement is not None
2977 return ResolvedPlacement(
2978 devices=tuple(devices),
2979 instances=resolved.instances,
2980 unplaceable_roles=resolved.unplaceable_roles,
2981 model_refs=model_refs,
2982 co_tenants=resolved.co_tenants,
2983 skipped_not_installed=skipped_not_installed,
2984 spec_applied=spec_applied,
2985 tight_roles=dict(resolved.tight_roles),
2986 engine_backend=reading.backend,
2987 )
2990@dataclass(frozen=True)
2991class FleetPlan:
2992 """The servers to start, and the roles that share one swap group."""
2994 launches: tuple[InstanceLaunch, ...]
2995 co_tenants: frozenset[WorkerRole] = frozenset()
2996 # Configured roles left unplaced because their model isn't installed (role ->
2997 # ref), so the warm path can fail a not-installed chat with a named reason
2998 # instead of spinning the warm line forever.
2999 skipped_not_installed: dict[WorkerRole, str] = field(default_factory=dict)
3000 # Launches refused for a window below the minimum grounded prompt
3001 # (role -> user-facing reason with the numbers).
3002 skipped_unusable_ctx: dict[WorkerRole, str] = field(default_factory=dict)
3005def _log_placement_findings(placement: Placement, model_refs: dict[WorkerRole, str]) -> None:
3006 """Warn about placements that exceed the memory budget.
3008 Shared-memory roles that fit nowhere get no server (loading them would OOM the
3009 host). GPU roles are never refused: one whose estimate exceeds the free VRAM
3010 still loads on demand, with a warning carrying the shortfall.
3011 """
3012 for role in placement.unplaceable_roles:
3013 log.warning(
3014 "%s model %s does not fit available memory and will not be served; "
3015 "free up memory or use a smaller model.",
3016 role.value,
3017 model_refs[role],
3018 )
3019 for role, shortfall in placement.tight_roles.items():
3020 log.warning(
3021 "Memory is tight for the %s model %s: it is estimated to need %.1f GiB more "
3022 "GPU memory than is available. It will still load on demand, keeping the "
3023 "layers that fit on the GPU and the rest in system memory; if it runs "
3024 "slowly, free up GPU memory or use a smaller model.",
3025 role.value,
3026 model_refs[role],
3027 # A sub-0.05 GiB shortfall would render as "0.0 GiB more".
3028 max(shortfall / 1024**3, 0.1),
3029 )
3030 if placement.co_tenants:
3031 log.info(
3032 "%s share GPU memory and load on demand; only one is resident at a time.",
3033 ", ".join(sorted(role.value for role in placement.co_tenants)),
3034 )
3037def _gib_text(num_bytes: int) -> str:
3038 """*num_bytes* in GiB with at most one decimal and no trailing zero."""
3039 return f"{round(num_bytes / 1024**3, 1):g}"
3042def _other_processes_clause(launch: InstanceLaunch, devices: list[FleetDevice]) -> str:
3043 """A sentence naming the memory other processes hold on the cards, or ``""``.
3045 Empty when there is no device, or when others hold no more than the model's
3046 own weights (then the weights are the cause, not the neighbours).
3047 """
3048 total = sum(d.total_bytes for d in devices)
3049 held = total - sum(d.free_bytes for d in devices)
3050 if not devices or held <= launch.weights_bytes:
3051 return ""
3052 owner = "the cards'" if len(devices) > 1 else "the card's"
3053 return f" Other processes hold {_gib_text(held)} GiB of {owner} {_gib_text(total)} GiB."
3056def _unusable_chat_ctx_reason(launch: InstanceLaunch, devices: list[FleetDevice]) -> str | None:
3057 """Reason to refuse a chat launch whose window cannot hold a grounded prompt.
3059 ``None`` for non-chat roles, for a window that holds the minimum grounded
3060 prompt, and for user knobs that ask for a smaller one (a ``num_ctx`` pin,
3061 a sub-minimum ``num_ctx_max`` / ``chat_n_ctx_target``). *devices* are the
3062 cards the launch runs on; they let the reason name memory that other
3063 processes hold.
3064 """
3065 from lilbee.core.config import cfg
3067 if launch.role is not WorkerRole.CHAT or cfg.num_ctx is not None:
3068 return None
3069 needed = engine_params.min_usable_chat_ctx()
3070 # User knobs capping the window below the minimum are honored (the num_ctx
3071 # pin bypasses above).
3072 asked = min(cfg.chat_n_ctx_target, cfg.num_ctx_max or cfg.chat_n_ctx_target)
3073 if asked < needed:
3074 return None
3075 if launch.ctx >= needed:
3076 return None
3077 if devices and launch.weights_bytes > sum(d.free_bytes for d in devices):
3078 opening = (
3079 f"The chat model {launch.model} does not fit: its weights are larger than the "
3080 f"free memory, which backs only a {launch.ctx}-token context"
3081 )
3082 else:
3083 opening = (
3084 f"The chat model {launch.model} loads, but the memory left after its weights "
3085 f"backs only a {launch.ctx}-token context"
3086 )
3087 return (
3088 f"{opening}, and a grounded answer needs about "
3089 f"{needed} tokens (system prompt, a retrieved source, the question, and room "
3090 "for the answer), so it will not be served."
3091 f"{_other_processes_clause(launch, devices)}"
3092 " Use a smaller model or a smaller quant, or set num_ctx to force a larger window."
3093 )
3096def plan_launches(
3097 roles: tuple[WorkerRole, ...] | None,
3098 binary: Path,
3099 by_index: dict[int, FleetDevice],
3100 devices: list[FleetDevice],
3101) -> FleetPlan:
3102 """Plan placement for *roles* (``None`` = all configured) and build their launches."""
3103 with _one_engine_per_pass():
3104 return _planned_launches(roles, binary, by_index, devices)
3107def _planned_launches(
3108 roles: tuple[WorkerRole, ...] | None,
3109 binary: Path,
3110 by_index: dict[int, FleetDevice],
3111 devices: list[FleetDevice],
3112) -> FleetPlan:
3113 """The placement and launches for *roles*, sized and placed against one snapshot."""
3114 from lilbee.core.config import cfg
3116 unified_budget = _unified_memory_budget(devices)
3117 # Only the clean-box snapshot's free bytes mean "what other tenants hold";
3118 # a live probe is also missing the fleet's own residency.
3119 clean_box = _current_plan_probe() is not None
3120 inputs, model_refs, reservation, skipped_not_installed = _server_model_inputs(
3121 roles,
3122 unified_budget=unified_budget,
3123 device_count=len(devices),
3124 total_vram=sum(d.total_bytes for d in devices),
3125 )
3126 spec = PlacementSpec.from_json(cfg.placement) if cfg.placement else None
3127 placement, _spec_applied = _placement_or_auto(
3128 spec,
3129 inputs,
3130 model_refs,
3131 devices,
3132 unified_budget=_unified_admission_budget(devices),
3133 charge_against_free=clean_box,
3134 )
3135 _log_placement_findings(placement, model_refs)
3136 reserved_by_device = _non_chat_reservation(placement.instances, inputs, placement.co_tenants)
3137 reserved_after_chat = _non_chat_reservation(
3138 [inst for inst in placement.instances if inst.role not in _LOADED_BEFORE_CHAT],
3139 inputs,
3140 placement.co_tenants,
3141 )
3142 charged = {inp.role: inp.est_vram_bytes for inp in inputs}
3143 launches: list[InstanceLaunch] = []
3144 skipped_unusable_ctx: dict[WorkerRole, str] = {}
3145 for plan in placement.instances:
3146 launch = _launch_for(
3147 plan,
3148 model_refs[plan.role],
3149 binary,
3150 by_index,
3151 unified_budget=unified_budget,
3152 chat_reservation=reservation,
3153 reserved_by_device=reserved_by_device,
3154 reserved_after_chat=reserved_after_chat,
3155 est_vram_bytes=charged.get(plan.role, 0),
3156 )
3157 reason = _unusable_chat_ctx_reason(
3158 launch, [by_index[i] for i in plan.devices if i in by_index] if clean_box else []
3159 )
3160 if reason is not None:
3161 skipped_unusable_ctx[launch.role] = reason
3162 log.warning(reason)
3163 continue
3164 launches.append(launch)
3165 return FleetPlan(
3166 launches=tuple(launches),
3167 co_tenants=placement.co_tenants,
3168 skipped_not_installed=skipped_not_installed,
3169 skipped_unusable_ctx=skipped_unusable_ctx,
3170 )
3173class _VisionRequestGrant:
3174 """Whether a request that turned OCR on has asked for the vision role this process."""
3176 def __init__(self) -> None:
3177 self._lock = threading.Lock()
3178 self._granted = False
3180 @property
3181 def granted(self) -> bool:
3182 with self._lock:
3183 return self._granted
3185 def set(self, granted: bool) -> None:
3186 with self._lock:
3187 self._granted = granted
3190_vision_request_grant = _VisionRequestGrant()
3193def grant_vision_on_request() -> None:
3194 """Plan the vision role even with ``ocr`` off, for a request that turned OCR on."""
3195 _vision_request_grant.set(True)
3198def revoke_vision_on_request() -> None:
3199 """Plan the vision role from the OCR setting alone again."""
3200 _vision_request_grant.set(False)
3203def vision_role_wanted(ref: str) -> bool:
3204 """Whether a plan serves vision on *ref*: the OCR setting uses it, or a request asked."""
3205 from lilbee.core.config import cfg
3207 chosen = OcrBackendUsed.chosen(cfg.ocr, ref)
3208 return chosen is OcrBackendUsed.VISION or _vision_request_grant.granted
3211def _launched_roles() -> tuple[WorkerRole, ...]:
3212 """The roles a launch plan may serve: vision only while OCR can reach it."""
3213 from lilbee.core.config import cfg
3215 return tuple(
3216 role
3217 for role in ROLE_REGISTRY
3218 if role is not WorkerRole.VISION or vision_role_wanted(str(cfg.vision_model))
3219 )
3222def plan_all_launches() -> FleetPlan:
3223 """Probe devices and plan launches for the configured roles.
3225 Writes no environment: the probe child carries the Vulkan loader safety and any
3226 ``cfg.gpu_devices`` pin in its own env, and the plan reads the pin from there.
3227 """
3228 binary = resolve_llama_server()
3229 with _one_engine_per_pass():
3230 devices = _plan_devices(binary)
3231 by_index = {d.index: d for d in devices}
3232 return plan_launches(_launched_roles(), binary, by_index, devices)