Coverage for src/lilbee/providers/fleet/planning.py: 100%
859 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Launch planning for the fleet: device probe, VRAM estimate, placement, argv."""
3from __future__ import annotations
5import logging
6import re
7import threading
8import time
9from dataclasses import dataclass, field, replace
10from pathlib import Path
11from typing import TYPE_CHECKING
13from lilbee.core.config.enums import KvCacheType
14from lilbee.core.system import is_network_path
15from lilbee.providers import engine_params, model_cache
16from lilbee.providers.base import ProviderError
17from lilbee.providers.fleet import ctx as fleet_ctx
18from lilbee.providers.fleet.adapters import (
19 LLM_RERANK_CONCURRENCY,
20 ROLE_SPECS,
21 RoleServerSpec,
22 build_server_argv,
23 embed_spec,
24 rerank_spec,
25 resolve_rerank_mode,
26)
27from lilbee.providers.fleet.binary import llama_server_runtime_env, resolve_llama_server
28from lilbee.providers.fleet.devices import (
29 VULKAN_BACKEND,
30 FleetDevice,
31 host_lacks_nvlink,
32 probe_devices,
33 visible_env,
34)
35from lilbee.providers.fleet.launch import InstanceLaunch
36from lilbee.providers.fleet.placement import (
37 InstancePlan,
38 ModelPlacementInput,
39 PeakEstimator,
40 Placement,
41 SplitCtxFitter,
42 placement_from_spec,
43 plan_placement,
44)
45from lilbee.providers.fleet.placement_spec import PlacementError, PlacementSpec
46from lilbee.providers.fleet.replicas import resolve_replica_count
47from lilbee.providers.fleet.vram import estimate_instance_footprint, usable_vram_fraction
48from lilbee.providers.model_cache import free_system_memory, total_system_memory
49from lilbee.providers.model_ref import parse_model_ref
50from lilbee.providers.roles import ROLE_REGISTRY, RerankMode, WorkerRole
52log = logging.getLogger(__name__)
54if TYPE_CHECKING:
55 from collections.abc import Callable, Mapping, Sequence
57# Fleet-only concurrency: continuous-batching slots (--parallel) per server.
58_CHAT_SLOTS = 4
59# Slots the PLACEMENT estimate reserves KV for on a tensor-split chat: one full
60# window, the minimum any split must hold. The launch may serve more than this
61# (see _resolve_split_chat_slots) when the cards measurably have room for several
62# full windows, so this is a planning floor, not the served slot count.
63_SPLIT_CHAT_SLOTS = 1
64# Floor context the PLACEMENT estimate reserves KV for, so a large model is never
65# single-carded into a KV corner too small for real use (a 17GB model on a 24GB
66# card leaves ~no KV room -> n_ctx collapses to a few hundred tokens). Sizing the
67# placement reserve against this floor forces a tensor-split when one card cannot
68# hold weights + a usable context; the served ctx is then grown by resolve_chat_ctx
69# (single) / fit_split_ctx (split) toward the cards' real headroom, with each
70# sequence capped at the working-context target. A split may then serve several
71# such sequences, so the served total can exceed the single-window reserve; the
72# per-device headroom test in fit_split_ctx is what bounds it. A fit still at
73# its floor after the forced split is refused by plan_launches
74# (_unusable_chat_ctx_reason).
75_MIN_USABLE_CHAT_CTX = 8192
76# Embed and cross-encoder rerank serve one request at a time. Raising it was
77# tried and measured worse: on 8xA40 with an 8B Q8 embedder and one ~100-token
78# passage per request, --parallel 1 gave 133 docs/sec at 81% SM while
79# --parallel 8 gave 100 at 63%. At one slot the card is already busy, so there is
80# no stall for slot-batching to reclaim and the extra slots only add
81# continuous-batching and KV-fragmentation overhead. Batch on the request side
82# (embed_batch_sequences) instead.
83_AUX_SLOTS = 1
84# A tensor-split needs at least this many GPUs; below it the chat context objective
85# (a gguf read) is pointless because the model can only single-card or stay unplaced.
86_MIN_SPLIT_GPUS = 2
87# Pooled single-slot search roles (embed/cross-encoder rerank) whose whole input
88# batches in one pass; derived from the role registry.
89_EMBED_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.pooled)
90# Roles whose loaders offload every layer regardless of cfg.n_gpu_layers; only
91# chat honors cfg.n_gpu_layers.
92_ALL_LAYER_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.offload_all_layers)
93_FLASH_ON = "on"
94_FLASH_OFF = "off"
95_FLASH_AUTO = "auto"
96# llama-server's documented way to say "offload nothing": --device none.
97_NO_DEVICE = "none"
98# Backends pinned by the name the engine printed rather than through an env var,
99# because their variables index a different space than --list-devices reports.
100_NAME_PINNED_BACKENDS = frozenset({VULKAN_BACKEND, "SYCL"})
101# Backends whose flash-attention coverage in llama.cpp is complete enough to ask
102# for it outright. Vulkan and SYCL are behind CUDA's and have been incomplete on
103# Intel's mesa driver, so those are left to the engine's own auto, which enables
104# flash attention only where the backend really supports it.
105_TRUSTED_FLASH_BACKENDS = frozenset({"CUDA", "ROCm", "HIP", "MTL", "Metal"})
106# Roles to which flash attention applies; embed/rerank run without it.
107_FLASH_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.flash_attn)
110# Cap vision's own KV footprint at this fraction of usable VRAM when sizing its
111# batching slots, leaving room for the weights and any co-located role.
112_VISION_VRAM_FRACTION = 0.5
114# Cap chat's footprint at this fraction of usable VRAM when sizing its batching
115# slots, reserving room for the co-located embed/rerank servers and the decode
116# compute buffers (which the flat overhead term only partly covers).
117_CHAT_VRAM_FRACTION = 0.8
119# Cap an LLM reranker's footprint at this fraction of usable VRAM when sizing its
120# slots; its per-slot ctx is tiny, so a normal GPU fits the full fan-out and a
121# small one steps down toward 1.
122_LLM_RERANK_VRAM_FRACTION = 0.5
124# RAM kept free for the OS when placing against system memory (no discrete GPU):
125# a quarter of total RAM, capped at 4 GiB. A fixed 4 GiB floor leaves a small
126# host (7-8 GB) with no budget at all, refusing to serve even tiny models.
127_SYSTEM_MEMORY_FLOOR_DIVISOR = 4
128# A GPU driver still initializing at boot answers with no devices. Ask again
129# before letting that decide the daemon's whole run; two extra probes cost a
130# couple of seconds only on a host that has a card the engine could not see.
131_PROBE_RETRIES = 2
132_PROBE_RETRY_DELAY_S = 1.0
134# A network filesystem makes mmap dangerous (page faults served over the wire can
135# wedge the loader in uninterruptible I/O), so the chat server loads its weights
136# into a malloc'd host copy (--no-mmap) whenever that copy fits in this fraction
137# of total system RAM. Local disk keeps mmap: its lazy paging gives a faster first
138# token on a cold cache -- the common desktop first launch -- and --no-mmap's
139# buffered full read only wins on an already-hot cache (#474: 33s vs 43s for a
140# 112GB model on 3 GPUs) while pessimizing cold start. Keyed on TOTAL memory
141# (stable), not free (fluctuates), so replans do not flap the launch argv. The
142# exact ceiling is tuned on a network-volume host.
143_NO_MMAP_NETWORK_RAM_FRACTION = 0.85
145# llama.cpp split-GGUF shard naming ("%s-%05d-of-%05d.gguf"); the cold-load
146# timeout must scale with the SUM of the shards, not the first file alone.
147_SPLIT_GGUF_NAME = re.compile(r"^(?P<prefix>.+)-(?P<index>\d{5})-of-(?P<total>\d{5})\.gguf$")
150def _weights_bytes(model_path: Path) -> int:
151 """Total weights size on disk; a split GGUF sums every sibling shard."""
152 match = _SPLIT_GGUF_NAME.fullmatch(model_path.name)
153 if match is None:
154 return model_path.stat().st_size
155 return sum(
156 sibling.stat().st_size
157 for sibling in model_path.parent.iterdir()
158 if _is_sibling_shard(sibling.name, match)
159 )
162def _is_sibling_shard(name: str, match: re.Match[str]) -> bool:
163 """Whether *name* is a shard of the same split GGUF as *match*."""
164 shard = _SPLIT_GGUF_NAME.fullmatch(name)
165 return (
166 shard is not None
167 and shard["prefix"] == match["prefix"]
168 and shard["total"] == match["total"]
169 )
172def _slots_for(
173 role: WorkerRole,
174 model_path: Path,
175 ctx: int,
176 *,
177 mmproj_path: Path | None = None,
178 unified_budget: int | None = None,
179 chat_reservation: int = 0,
180 rerank_mode: RerankMode | None = None,
181 device: FleetDevice | None = None,
182) -> int:
183 """Continuous-batching slots (--parallel) for a role's server.
185 Chat batches concurrent turns; vision batches concurrent OCR pages since a
186 one-page decode underutilizes the GPU; an LLM reranker batches its per-candidate
187 chat requests; embed and cross-encoder rerank are single-slot (their batching is
188 request-side). The memory-aware roles drop toward 1 on a small or shared host
189 instead of overcommitting. ``unified_budget`` caps sizing against free system RAM
190 with no discrete GPU; ``chat_reservation`` is the search-role footprint held back
191 from chat; ``device`` is the card the role was placed on, whose memory the
192 budget comes from once placement has chosen one.
193 """
194 if role is WorkerRole.CHAT:
195 return _resolve_chat_slots(
196 model_path,
197 ctx,
198 mmproj_path=mmproj_path,
199 unified_budget=unified_budget,
200 chat_reservation=chat_reservation,
201 device=device,
202 )
203 if role is WorkerRole.VISION:
204 return _resolve_vision_slots(
205 model_path, ctx, mmproj_path=mmproj_path, unified_budget=unified_budget, device=device
206 )
207 if role is WorkerRole.RERANK and rerank_mode is RerankMode.LLM:
208 return _resolve_llm_rerank_slots(
209 model_path, ctx, unified_budget=unified_budget, device=device
210 )
211 return _AUX_SLOTS
214def _resolve_split_chat_slots(fit_fn: Callable[[int], int]) -> tuple[int, int]:
215 """Largest split-chat slot count whose sequences each keep the full window.
217 ``fit_fn(n)`` is the per-slot context that fits when serving ``n`` sequences
218 (``fit_split_ctx``, capped at the working target and verified against real
219 per-card headroom). More slots divide the KV, so a split whose cards hold
220 several full windows can serve that many agents concurrently instead of one.
221 Returns ``(slots, per_slot_ctx)``, falling to one slot when only one full
222 window fits (or the fit degenerated to the floor), which preserves the
223 max-context single-sequence behaviour on a tight card.
225 Found by bisection rather than a scan because every ``fit_fn`` call is a
226 complete binary search whose probes each shell out to gguf-parser, and the
227 whole thing runs while this process holds the cross-process build lock that
228 every other lilbee start waits on without a deadline. A descending scan paid
229 for all of ``_CHAT_SLOTS - 1`` searches in exactly the tight-card case where
230 none of them fit. Bisection is sound here because the fit is non-increasing
231 in the slot count: more sequences divide the same headroom, so once a count
232 fails no larger one can succeed.
233 """
234 full = fit_fn(1)
235 if full <= model_cache._DYNAMIC_CTX_FLOOR:
236 return 1, full
237 low, high = 1, _CHAT_SLOTS
238 while low < high:
239 mid = (low + high + 1) // 2
240 if fit_fn(mid) >= full:
241 low = mid
242 else:
243 high = mid - 1
244 return low, full
247def _resolve_chat_slots(
248 model_path: Path,
249 ctx: int,
250 *,
251 mmproj_path: Path | None = None,
252 unified_budget: int | None = None,
253 chat_reservation: int = 0,
254 device: FleetDevice | None = None,
255) -> int:
256 """Largest chat slot count (<= ``_CHAT_SLOTS``) whose footprint fits the budget
257 after reserving the search roles; steps to 1 when none fit."""
258 budget = _slot_budget(_CHAT_VRAM_FRACTION, unified_budget, device) - chat_reservation
259 return _fit_slots(
260 _CHAT_SLOTS,
261 WorkerRole.CHAT,
262 model_path,
263 ctx,
264 mmproj_path=mmproj_path,
265 unified=unified_budget is not None,
266 budget=budget,
267 )
270def _resolve_vision_slots(
271 model_path: Path,
272 ctx: int,
273 *,
274 mmproj_path: Path | None = None,
275 unified_budget: int | None = None,
276 device: FleetDevice | None = None,
277) -> int:
278 """Largest OCR batching slot count (<= ``cfg.vision_ocr_concurrency``) that fits
279 the memory budget; 1 when the ceiling is 1 or nothing larger fits."""
280 from lilbee.core.config import cfg
282 ceiling = max(1, cfg.vision_ocr_concurrency)
283 if ceiling == 1:
284 return 1
285 return _fit_slots(
286 ceiling,
287 WorkerRole.VISION,
288 model_path,
289 ctx,
290 mmproj_path=mmproj_path,
291 unified=unified_budget is not None,
292 budget=_slot_budget(_VISION_VRAM_FRACTION, unified_budget, device),
293 )
296def _resolve_llm_rerank_slots(
297 model_path: Path,
298 ctx: int,
299 *,
300 unified_budget: int | None = None,
301 device: FleetDevice | None = None,
302) -> int:
303 """Largest LLM-reranker slot count (<= ``LLM_RERANK_CONCURRENCY``) that fits the
304 memory budget; 1 when nothing larger fits. Matches the client's request fan-out."""
305 return _fit_slots(
306 LLM_RERANK_CONCURRENCY,
307 WorkerRole.RERANK,
308 model_path,
309 ctx,
310 mmproj_path=None,
311 unified=unified_budget is not None,
312 budget=_slot_budget(_LLM_RERANK_VRAM_FRACTION, unified_budget, device),
313 rerank_mode=RerankMode.LLM,
314 )
317def _slot_budget(
318 vram_fraction: float, unified_budget: int | None, device: FleetDevice | None = None
319) -> int:
320 """Memory budget for slot sizing: *vram_fraction* of the usable memory on *device*
321 (the fleet's smallest when placement has not chosen one yet), capped by
322 ``unified_budget`` (free system RAM) when there is no discrete GPU so the count
323 steps down to fit free memory instead of overcommitting."""
324 budget = int(plan_sizing_budget(device) * vram_fraction)
325 if unified_budget is not None:
326 budget = min(budget, unified_budget)
327 return budget
330def _fit_slots(
331 ceiling: int,
332 role: WorkerRole,
333 model_path: Path,
334 ctx: int,
335 *,
336 mmproj_path: Path | None,
337 unified: bool,
338 budget: int,
339 rerank_mode: RerankMode | None = None,
340) -> int:
341 """Largest slot count in ``1..ceiling`` whose instance footprint fits *budget*;
342 1 when none larger fit."""
343 from lilbee.providers.base import ProviderError
345 for slots in range(ceiling, 1, -1):
346 try:
347 est = estimate_instance_footprint(
348 model_path,
349 ctx=ctx,
350 slots=slots,
351 gpu_layers=_role_gpu_layers(role),
352 flash_attn=_role_flash(role, rerank_mode),
353 kv_cache_type=_role_kv_cache_type(role),
354 kv_cache_type_v=_role_kv_cache_type_v(role),
355 mmproj_path=mmproj_path,
356 expert_offload=_role_expert_offload(model_path),
357 )
358 except (ProviderError, OSError):
359 # An unsizable model runs a single slot; the load decides the rest.
360 return 1
361 if est.footprint(unified=unified) <= budget:
362 return slots
363 return 1
366def _role_ctx(
367 role: WorkerRole,
368 model_path: Path,
369 meta: dict[str, str] | None,
370 device: FleetDevice | None = None,
371) -> int:
372 """Per-slot context for a role, derived as the in-process loader does.
374 Embed/rerank use the embedding model's training context; vision uses the
375 vision loader's training-context picker; chat honors ``cfg.num_ctx`` then
376 falls back to the single-GPU dynamic chat-ctx picker, sized against *device*
377 once placement has chosen one. A tensor-split chat is sized against its
378 per-device headroom instead (see :func:`fit_split_ctx`).
379 """
380 from lilbee.core.config import cfg
382 if role is WorkerRole.EMBED:
383 return engine_params.resolve_embed_ctx(meta, model_path)
384 if role is WorkerRole.RERANK:
385 if _rerank_mode_for(meta) is RerankMode.LLM:
386 return engine_params.resolve_llm_rerank_ctx(meta, model_path)
387 return engine_params.resolve_embed_ctx(meta, model_path)
388 if role is WorkerRole.VISION:
389 return engine_params.resolve_vision_ctx(model_path)
390 if cfg.num_ctx is not None:
391 return _pinned_chat_ctx(model_path, meta)
392 return engine_params.resolve_chat_ctx(
393 model_path, meta, available_bytes=plan_sizing_budget(device)
394 )
397def _pinned_chat_ctx(model_path: Path, meta: dict[str, str] | None) -> int:
398 """``cfg.num_ctx``, clamped to what the model was trained for.
400 Every unpinned resolver already clamps, and both docstrings here claimed the
401 pin did too. It did not, so a pin past the trained window was passed straight
402 to the engine, which clamps it silently and serves a different number than
403 every budget was sized for.
405 Only against a window that is actually known. A GGUF whose header cannot be
406 read falls back to a default that is a guess, and contradicting an explicit
407 pin with a guess would break the hosts where the header is the thing that is
408 broken.
409 """
410 from lilbee.core.config import cfg
412 pinned = cfg.num_ctx
413 assert pinned is not None # noqa: S101 - callers check; this documents the contract
414 ceiling = _known_chat_ceiling(model_path, meta)
415 if ceiling is None or pinned <= ceiling:
416 return pinned
417 log.warning(
418 "num_ctx is set to %d but %s was trained for %d, so %d is what will be served. "
419 "Lower num_ctx to stop planning against a window this model does not have.",
420 pinned,
421 model_path.name,
422 ceiling,
423 ceiling,
424 )
425 return ceiling
428def _known_chat_ceiling(model_path: Path, meta: dict[str, str] | None) -> int | None:
429 """The largest chat window this model is known to support, or ``None``.
431 ``None`` when the GGUF header gave no usable context length and the user set
432 no ``cfg.num_ctx_max``: there is then no measured ceiling, only a default.
433 """
434 from lilbee.core.config import cfg
435 from lilbee.providers.gguf_meta import train_ctx_from_meta
437 sentinel = -1
438 trained = train_ctx_from_meta(meta, fallback=sentinel, model_path=model_path)
439 known = [value for value in (trained, cfg.num_ctx_max) if value is not None and value > 0]
440 return min(known) if known else None
443def _rerank_mode_for(meta: dict[str, str] | None) -> RerankMode:
444 """Resolve the RERANK serving mode from cfg + the reranker GGUF arch."""
445 from lilbee.core.config import cfg
447 arch = meta.get("architecture") if meta else None
448 return resolve_rerank_mode(cfg.reranker_type, arch)
451def _role_rerank_mode(role: WorkerRole, meta: dict[str, str] | None) -> RerankMode | None:
452 """The RERANK serving mode for *role*, or ``None`` for every other role."""
453 return _rerank_mode_for(meta) if role is WorkerRole.RERANK else None
456def _server_spec(
457 role: WorkerRole, rerank_mode: RerankMode | None, meta: dict[str, str] | None
458) -> RoleServerSpec:
459 """The llama-server spec for a launch: rerank mode, decoder-aware embed pooling,
460 or the role default. EMBED forces ``--pooling last`` for decoder-only archs."""
461 if rerank_mode is not None:
462 return rerank_spec(rerank_mode)
463 if role is WorkerRole.EMBED:
464 return embed_spec(meta)
465 return ROLE_SPECS[role]
468def _pooled_batch_size(role: WorkerRole, rerank_mode: RerankMode | None, ctx: int) -> int | None:
469 """The ``--batch-size``/``--ubatch-size`` the launch raises for pooled
470 embed/cross-encoder rerank (the full context), or ``None`` for other roles."""
471 if role in _EMBED_ROLES and rerank_mode is not RerankMode.LLM:
472 return ctx
473 return None
476def _role_gpu_layers(role: WorkerRole) -> int:
477 """GPU-layer offload: chat honors ``cfg.n_gpu_layers``, others offload all layers."""
479 return engine_params.resolve_n_gpu_layers(embedding=role in _ALL_LAYER_ROLES)
482def _flash_enabled() -> bool:
483 """Flash attention is on unless ``cfg.flash_attention`` is explicitly ``False``."""
484 from lilbee.core.config import cfg
486 return cfg.flash_attention is not False
489def _fleet_backend() -> str | None:
490 """The engine backend this host plans onto, or ``None`` when unknown.
492 Prefers the plan snapshot so a whole planning pass answers consistently, and
493 falls back to the short-TTL read cache rather than a fresh probe.
494 """
495 probe = _plan_probe_store.get()
496 if probe is not None:
497 return probe.devices[0].backend if probe.devices else None
498 try:
499 devices = _read_device_cache.get(resolve_llama_server())
500 except (ProviderError, OSError):
501 return None
502 return devices[0].backend if devices else None
505def _flash_attention_is_trusted() -> bool:
506 """Whether to ask for flash attention outright rather than let the engine decide.
508 Unknown backends answer yes, which keeps every host that works today on the
509 argv it has now; only the backends known to lag get the engine's own auto.
510 """
511 backend = _fleet_backend()
512 return backend is None or backend in _TRUSTED_FLASH_BACKENDS
515def flash_attn_flag() -> str:
516 """``--flash-attn`` argv value for chat and vision."""
517 if not _flash_enabled():
518 return _FLASH_OFF
519 return _FLASH_ON if _flash_attention_is_trusted() else _FLASH_AUTO
522def _role_launches_with_flash(role: WorkerRole, rerank_mode: RerankMode | None = None) -> bool:
523 """Whether the launch asks the engine for flash attention on *role*.
525 The one place that answers this. The registry marks RERANK as a non-flash
526 role because a cross-encoder pools in one batch, but an LLM reranker is
527 generative and launches exactly like chat, so the mode decides there.
528 """
529 if role is WorkerRole.RERANK:
530 return rerank_mode is RerankMode.LLM
531 return role in _FLASH_ROLES
534def _role_flash(role: WorkerRole, rerank_mode: RerankMode | None = None) -> bool:
535 """Whether the estimate may assume flash attention for *role*.
537 The launch's own answer, narrowed to a definite ``on``. Under ``auto`` the
538 engine decides at load time, and assuming it would size the KV cache below
539 what the launch may need.
540 """
541 return _role_launches_with_flash(role, rerank_mode) and flash_attn_flag() == _FLASH_ON
544def _role_kv_cache_type(role: WorkerRole) -> KvCacheType:
545 """Chat honors ``cfg.kv_cache_type``; embed/rerank/vision run f16 KV."""
546 from lilbee.core.config import cfg
548 return cfg.kv_cache_type if role is WorkerRole.CHAT else KvCacheType.F16
551def _replica_count(role: WorkerRole, device_count: int) -> int:
552 """Requested data-parallel instances for *role* via the shared resolver."""
553 return resolve_replica_count(role, device_count)
556def _role_kv_cache_type_v(role: WorkerRole) -> KvCacheType:
557 """The V cache type for *role*: the configured one only when flash attention is on.
559 llama.cpp refuses a quantized V cache without flash attention ("V cache
560 quantization requires flash_attn") and the server never starts, while a
561 quantized K cache needs nothing. So V follows the setting only where flash
562 attention is certain, and is f16 under ``auto`` or ``off``. That costs memory
563 rather than a launch, and the estimate moves with it.
564 """
565 from lilbee.core.config.enums import KvCacheType
567 configured = _role_kv_cache_type(role)
568 return configured if flash_attn_flag() == _FLASH_ON else KvCacheType.F16
571def chat_cache_type_flags() -> tuple[str | None, str | None]:
572 """``(--cache-type-k, --cache-type-v)`` for chat; ``None`` leaves the f16 default."""
573 from lilbee.core.config.enums import KvCacheType
575 def flag(kind: KvCacheType) -> str | None:
576 return None if kind is KvCacheType.F16 else kind.value
578 return flag(_role_kv_cache_type(WorkerRole.CHAT)), flag(_role_kv_cache_type_v(WorkerRole.CHAT))
581def _vision_mmproj(model_ref: str) -> Path | None:
582 """Resolve a vision model's mmproj sidecar, or ``None`` if absent."""
583 from lilbee.providers.base import ProviderError
584 from lilbee.providers.gguf_meta import find_mmproj_for_model
586 try:
587 return find_mmproj_for_model(engine_params.resolve_model_path(model_ref))
588 except (ProviderError, OSError, ValueError, KeyError):
589 return None
592def _estimate_role(
593 role: WorkerRole,
594 model_ref: str,
595 *,
596 slots: int | None = None,
597 unified_budget: int | None = None,
598 chat_reservation: int = 0,
599 device_count: int = 0,
600) -> ModelPlacementInput:
601 """Estimate one role-model's footprint via gguf-parser (+ mmproj for vision).
603 ``slots`` defaults to the role's resolved batching slots (chat and vision are
604 memory-aware); ``chat_reservation`` shrinks chat to leave room for the search
605 roles; ``device_count`` resolves an auto (0) replica knob to one per GPU.
606 Charges the unified footprint with no discrete GPU, else the VRAM one.
607 """
608 from lilbee.providers.gguf_meta import read_gguf_metadata
610 path = engine_params.resolve_model_path(model_ref)
611 mmproj = _vision_mmproj(model_ref) if role is WorkerRole.VISION else None
612 meta = read_gguf_metadata(path)
613 # Size the single-instance footprint against the placement reserve (a usable
614 # KV floor for chat), so a model that fits weights-only on one card but not
615 # weights + a usable context falls through to a tensor-split instead of being
616 # single-carded into a tiny n_ctx. Non-chat roles keep their launch ctx.
617 ctx = _placement_estimate_ctx(role, path, meta)
618 rerank_mode = _role_rerank_mode(role, meta)
619 if slots is None:
620 slots = _slots_for(
621 role,
622 path,
623 ctx,
624 mmproj_path=mmproj,
625 unified_budget=unified_budget,
626 chat_reservation=chat_reservation,
627 rerank_mode=rerank_mode,
628 )
629 est = estimate_instance_footprint(
630 path,
631 ctx=ctx,
632 slots=slots,
633 gpu_layers=_role_gpu_layers(role),
634 flash_attn=_role_flash(role, rerank_mode),
635 kv_cache_type=_role_kv_cache_type(role),
636 kv_cache_type_v=_role_kv_cache_type_v(role),
637 mmproj_path=mmproj,
638 batch_size=_pooled_batch_size(role, rerank_mode, ctx),
639 expert_offload=_role_expert_offload(path),
640 )
641 fp = est.footprint(unified=unified_budget is not None)
642 if role is WorkerRole.CHAT and unified_budget is None:
643 fp = _chat_serve_budget_footprint(fp)
644 return ModelPlacementInput(
645 role=role,
646 est_vram_bytes=fp,
647 replicas=_replica_count(role, device_count),
648 est_ram_bytes=est.ram_bytes,
649 )
652def _chat_serve_budget_footprint(footprint: int) -> int:
653 """Charge a chat instance against the serve budget, not the placement headroom.
655 The planner fits instances within ``cfg.usable_vram_fraction`` of a card, but a
656 single-card chat then sizes its KV cache against the smaller
657 ``cfg.gpu_memory_fraction`` budget (``resolve_chat_ctx``). A model that fills a
658 card at 0.9 leaves no room for KV at 0.75 and collapses to a few hundred tokens,
659 so scale its placement footprint by the budget ratio: it then needs a
660 tensor-split (pooling VRAM across cards) whenever single-carding it would starve
661 its context. Small models are unaffected -- they fit the serve budget with KV
662 room to spare.
663 """
664 from lilbee.core.config import cfg
666 # Never below 1.0. The ratio only compensates while the serve budget is the
667 # smaller of the two; a gpu_memory_fraction raised past the usable fraction
668 # inverts it, and the same line that exists to charge chat more starts
669 # charging it less than the model takes.
670 return int(footprint * max(1.0, usable_vram_fraction() / cfg.gpu_memory_fraction))
673def _placement_estimate_ctx(role: WorkerRole, model_path: Path, meta: dict[str, str] | None) -> int:
674 """Per-slot context the placement estimate sizes a role against.
676 For chat this reserves KV for a usable floor (``_MIN_USABLE_CHAT_CTX``, or the
677 user's ``cfg.num_ctx`` pin), capped by the model's trained ceiling -- not the
678 single-GPU dynamic ctx (which shrinks to fit one card and then confirms a
679 single-card placement) nor the full trained ceiling (which over-reserves). A
680 model that cannot hold weights + this floor on one card is tensor-split.
681 """
682 from lilbee.core.config import cfg
684 if role is WorkerRole.CHAT:
685 if cfg.num_ctx is not None:
686 return _pinned_chat_ctx(model_path, meta)
687 return apply_ctx_downshift(
688 role,
689 min(
690 engine_params.chat_ctx_ceiling(meta, model_path),
691 max(cfg.chat_n_ctx_target, _MIN_USABLE_CHAT_CTX),
692 ),
693 )
694 return apply_ctx_downshift(role, _role_ctx(role, model_path, meta))
697def _placement_estimate_slots(role: WorkerRole, meta: dict[str, str] | None) -> int:
698 """The slot count the placement estimate reserves KV for.
700 A tensor-split chat reserves one full-context sequence here: a conservative
701 floor for the card-count decision. The launch then fills the placed cards'
702 real headroom with as many full-context slots as fit (``_resolve_split_chat_slots``),
703 never exceeding what those cards hold, so a larger launch count can't OOM.
704 """
705 from lilbee.core.config import cfg
707 if role is WorkerRole.CHAT:
708 return _SPLIT_CHAT_SLOTS
709 if role is WorkerRole.VISION:
710 return max(1, cfg.vision_ocr_concurrency)
711 if role is WorkerRole.RERANK and _rerank_mode_for(meta) is RerankMode.LLM:
712 return LLM_RERANK_CONCURRENCY
713 return _AUX_SLOTS
716def _peak_estimator(model_refs: dict[WorkerRole, str]) -> PeakEstimator:
717 """Per-device VRAM-vector estimator for the planner, bound to the configured models.
719 Estimates each role at its launch ceiling (ctx x slots) with the candidate
720 tensor-split ratio, so the planner reserves enough cards for the busiest one.
721 """
722 from lilbee.providers.gguf_meta import read_gguf_metadata
724 def estimate_peak(role: WorkerRole, ratio: tuple[int, ...]) -> tuple[int, ...]:
725 path = engine_params.resolve_model_path(model_refs[role])
726 meta = read_gguf_metadata(path)
727 mmproj = _vision_mmproj(model_refs[role]) if role is WorkerRole.VISION else None
728 slots = _placement_estimate_slots(role, meta)
729 ctx = _placement_estimate_ctx(role, path, meta)
730 rerank_mode = _role_rerank_mode(role, meta)
731 est = estimate_instance_footprint(
732 path,
733 ctx=ctx,
734 slots=slots,
735 gpu_layers=_role_gpu_layers(role),
736 flash_attn=_role_flash(role, rerank_mode),
737 kv_cache_type=_role_kv_cache_type(role),
738 kv_cache_type_v=_role_kv_cache_type_v(role),
739 mmproj_path=mmproj,
740 tensor_split=ratio,
741 batch_size=_pooled_batch_size(role, rerank_mode, ctx),
742 expert_offload=_role_expert_offload(path),
743 )
744 return est.per_device_vram
746 return estimate_peak
749def _chat_split_ctx_objective(
750 model_refs: dict[WorkerRole, str],
751) -> tuple[SplitCtxFitter | None, int]:
752 """The chat split's context fitter and target, or ``(None, 0)`` with no chat model.
754 The fitter sizes a candidate shard's served context exactly as the launch does
755 (:func:`fit_split_ctx`), so the planner widens chat onto idle cards only when a
756 tighter shard would starve KV below the target. See docs/architecture.md.
757 """
758 if WorkerRole.CHAT not in model_refs:
759 return None, 0
760 from lilbee.providers.gguf_meta import read_gguf_metadata
762 path = engine_params.resolve_model_path(model_refs[WorkerRole.CHAT])
763 meta = read_gguf_metadata(path)
764 target = _placement_estimate_ctx(WorkerRole.CHAT, path, meta)
766 def fit(ratio: tuple[int, ...], per_device_free_bytes: Sequence[int]) -> int:
767 return fleet_ctx.fit_split_ctx(
768 path,
769 meta=meta,
770 slots=_SPLIT_CHAT_SLOTS,
771 ratio=ratio,
772 per_device_free_bytes=per_device_free_bytes,
773 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
774 flash_attn=_role_flash(WorkerRole.CHAT),
775 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
776 kv_cache_type_v=_role_kv_cache_type_v(WorkerRole.CHAT),
777 ctx_ceiling=target,
778 )
780 return fit, target
783def _search_reservation(inputs: dict[WorkerRole, ModelPlacementInput]) -> int:
784 """Total footprint of the placed search roles (all replicas), held back ahead
785 of chat."""
786 return sum(
787 inputs[role].est_vram_bytes * inputs[role].replicas
788 for role in _EMBED_ROLES
789 if role in inputs
790 )
793def _role_weights_bytes(role: WorkerRole, ref: str) -> int:
794 """The model's weight bytes on disk (plus the mmproj for vision): a
795 ground-truth lower bound on residency. 0 when the file cannot be resolved."""
796 from lilbee.providers.base import ProviderError
798 try:
799 size = _weights_bytes(engine_params.resolve_model_path(ref))
800 if role is WorkerRole.VISION:
801 mmproj = _vision_mmproj(ref)
802 if mmproj is not None:
803 size += int(mmproj.stat().st_size)
804 except (ProviderError, OSError):
805 return 0
806 return size
809def _is_moe(meta: dict[str, str] | None) -> bool:
810 """Whether the GGUF declares routed experts, so its experts can be offloaded."""
811 count = (meta or {}).get("expert_count")
812 try:
813 return int(count) > 0 if count is not None else False
814 except ValueError:
815 return False
818def expert_offload_all(meta: dict[str, str] | None) -> bool:
819 """Whether to keep every layer's experts in system memory; MoE models only."""
820 from lilbee.core.config import cfg
822 return bool(cfg.cpu_moe) and _is_moe(meta)
825def expert_offload_layers(meta: dict[str, str] | None) -> int | None:
826 """How many layers' experts to keep in system memory, or None for no split.
828 A non-positive ``n_cpu_moe`` offloads nothing (it would emit a no-op
829 ``--n-cpu-moe 0``), so it reads as unset.
830 """
831 from lilbee.core.config import cfg
833 if cfg.n_cpu_moe is None or cfg.n_cpu_moe < 1 or not _is_moe(meta):
834 return None
835 return cfg.n_cpu_moe
838def _role_expert_offload(model_path: Path) -> tuple[str, ...]:
839 """Expert patterns the launch will offload, for sizing the same way it runs.
841 Reads the GGUF (cached) rather than taking metadata as an argument so every
842 estimate site charges the same tensors the launch moves off the GPU.
843 """
844 from lilbee.providers.fleet.adapters import expert_offload_patterns
845 from lilbee.providers.gguf_meta import read_gguf_metadata
847 meta = read_gguf_metadata(model_path)
848 return expert_offload_patterns(
849 cpu_moe=expert_offload_all(meta), n_cpu_moe=expert_offload_layers(meta)
850 )
853def _expert_offload_configured() -> bool:
854 """Whether the user asked for expert offload that would actually take effect.
856 A non-positive ``n_cpu_moe`` offloads nothing, so it does not count.
857 """
858 from lilbee.core.config import cfg
860 return bool(cfg.cpu_moe) or (cfg.n_cpu_moe is not None and cfg.n_cpu_moe >= 1)
863def _weights_exceed_everything(size: int, *, total_vram: int, total_ram: int) -> bool:
864 """True when a model's weights fit neither the GPUs nor system memory.
866 File size is ground truth, not an estimate, so this bound cannot repeat the
867 false-refusal class: no estimator error makes a 40 GiB file fit a 1 GiB box.
868 Past both pools there is nowhere for a layer to go and no launch can win, so
869 saying so beats a load that thrashes and then dies.
870 """
871 ceiling = total_vram + total_ram
872 return ceiling > 0 and size > ceiling
875def _weights_exceed_hardware(size: int, total_vram: int, *, is_moe: bool) -> bool:
876 """True when this model cannot be served on this machine at all.
878 Exceeding VRAM alone is not that. The engine chooses how many layers fit and
879 keeps the rest in system memory, so a model larger than every card is a
880 partial offload and lilbee's job is to launch it and say what will happen.
881 Refusing there meant the fit never ran and the role was skipped, which left
882 the user hand-tuning n_gpu_layers to get back what the engine does by itself.
884 What still refuses is a model past VRAM and system memory together, where no
885 arrangement of layers exists. A user-set n_gpu_layers or expert offload keeps
886 standing the bound down entirely, since the user has said where the weights
887 should go.
888 """
889 from lilbee.core.config import cfg
891 if cfg.n_gpu_layers is not None:
892 return False
893 if is_moe and _expert_offload_configured():
894 return False
895 return _weights_exceed_everything(
896 size, total_vram=total_vram, total_ram=model_cache.total_system_memory()
897 )
900def _vision_without_mmproj(role: WorkerRole, ref: str) -> bool:
901 """True (with a warning) for a configured vision model whose mmproj is missing.
903 The skip would silently disable OCR; the warning names the cause and the fix.
904 """
905 if role is not WorkerRole.VISION or _vision_mmproj(ref) is not None:
906 return False
907 log.warning(
908 "Vision model %s has no mmproj (CLIP projector); OCR is disabled. "
909 "Re-run 'lilbee model pull %s' to fetch the projector.",
910 ref,
911 ref,
912 )
913 return True
916def _estimate_or_fallback(
917 role: WorkerRole,
918 ref: str,
919 *,
920 unified_budget: int | None,
921 chat_reservation: int,
922 device_count: int,
923 total_vram: int,
924 skipped_not_installed: dict[WorkerRole, str],
925 host_committed: int = 0,
926) -> ModelPlacementInput | None:
927 """Size *role* for placement, degrading rather than refusing.
929 A missing model is skipped and recorded; a sizing failure on an installed
930 model falls back to its weight bytes; weights alone exceeding the physical
931 VRAM refuse with a plain message (ground truth, not an estimate).
932 """
933 from lilbee.providers.base import ProviderError, ProviderErrorKind
935 try:
936 estimate = _estimate_role(
937 role,
938 ref,
939 unified_budget=unified_budget,
940 chat_reservation=chat_reservation,
941 device_count=device_count,
942 )
943 except (ProviderError, OSError) as exc:
944 if isinstance(exc, ProviderError) and exc.kind is ProviderErrorKind.NOT_FOUND:
945 log.warning("Skipping %s server: model %r is not installed.", role.value, ref)
946 skipped_not_installed[role] = ref
947 return None
948 return _sizing_failure_fallback(
949 role,
950 ref,
951 exc,
952 device_count=device_count,
953 total_vram=total_vram,
954 host_committed=host_committed,
955 )
956 return _admit_estimate(
957 _floor_implausible_estimate(estimate, role, ref),
958 role,
959 ref,
960 total_vram=total_vram,
961 ram_bytes=estimate.est_ram_bytes,
962 host_committed=host_committed,
963 )
966def _admit_estimate(
967 estimate: ModelPlacementInput,
968 role: WorkerRole,
969 ref: str,
970 *,
971 total_vram: int,
972 ram_bytes: int,
973 host_committed: int = 0,
974) -> ModelPlacementInput | None:
975 """*estimate*, or ``None`` when this model cannot load on this machine.
977 Two hardware bounds, one per kind of memory: the weights must fit the GPUs
978 unless something offloads, and whatever offloading puts in system memory must
979 fit the system.
980 """
981 weights = _role_weights_bytes(role, ref)
982 if _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref)):
983 _warn_weights_exceed(role, ref, weights, total_vram)
984 return None
985 if total_vram > 0 and weights > total_vram:
986 _warn_weights_spill(role, ref, weights, total_vram)
987 if _host_memory_refuses(role, ref, ram_bytes, host_committed):
988 return None
989 return estimate
992def _analytic_footprint_floor(
993 weights: int, *, meta: dict[str, str] | None, ctx: int, slots: int
994) -> int:
995 """The least this instance can occupy: weights, its KV cache, and overhead.
997 Used when the estimator cannot answer. Charging weight bytes alone was a
998 knowing under-charge: the engine allocates a KV cache sized by context and
999 slot count, plus compute buffers, and omitting all of it lets placement fit a
1000 model that cannot fit. The comment said the load would decide, and it did, by
1001 running out of memory.
1003 A floor rather than an estimate. It is derived from the header the same way
1004 the in-process sizing path derives it, and it is deliberately the smallest
1005 defensible number, because refusing a model that would have fit is its own
1006 failure. Without a readable header the per-token fallback still applies:
1007 zero is the one answer that is certainly wrong.
1008 """
1010 kv_bytes = (
1011 model_cache.kv_bytes_per_token(meta, engine_params._kv_elem_bytes_for_cfg())
1012 * ctx
1013 * max(slots, 1)
1014 )
1015 overhead = int(weights * model_cache._BUFFER_OVERHEAD_FRACTION)
1016 return weights + kv_bytes + overhead
1019def _estimate_is_implausible(*, estimated: int, floor: int) -> bool:
1020 """Whether *estimated* describes a load that cannot exist.
1022 Below the analytic floor, which is the model's own weight bytes plus the
1023 cache and buffers it was asked to hold, there is no arrangement of memory
1024 that serves it. A floor of zero means nothing could be computed to compare
1025 against, and a guess is not grounds to discard the only measurement there is.
1026 """
1027 return floor > 0 and 0 < estimated < floor
1030def _floor_implausible_estimate(
1031 estimate: ModelPlacementInput, role: WorkerRole, ref: str
1032) -> ModelPlacementInput:
1033 """*estimate*, or the analytic floor when the estimator returned less than one.
1035 The fallback floor otherwise fires only when the estimator cannot answer, so
1036 an answer that is well formed and impossible went straight through, and
1037 placement committed a card against a number the load then overran.
1038 """
1039 floor = _fallback_floor_for(role, ref, _role_weights_bytes(role, ref))
1040 if not _estimate_is_implausible(estimated=estimate.est_vram_bytes, floor=floor):
1041 return estimate
1042 log.warning(
1043 "The estimator sized the %s model %s at %.1f GiB, below the %.1f GiB its "
1044 "weights and cache alone need. Charging the floor instead; the estimate "
1045 "cannot be describing this load.",
1046 role.value,
1047 ref,
1048 estimate.est_vram_bytes / 1024**3,
1049 floor / 1024**3,
1050 )
1051 return replace(estimate, est_vram_bytes=floor)
1054def _sizing_failure_fallback(
1055 role: WorkerRole,
1056 ref: str,
1057 exc: Exception,
1058 *,
1059 device_count: int,
1060 total_vram: int,
1061 host_committed: int = 0,
1062) -> ModelPlacementInput | None:
1063 """Analytic-floor placement input for an installed model the estimator cannot
1064 size; ``None`` skips the role (the file is unresolvable, its weights alone
1065 exceed the hardware, or offloading it would exceed system memory).
1067 The host bound applies here too. Charging the whole floor to VRAM and
1068 skipping it let an unsizable model past a check every sized model faces."""
1069 weights = _role_weights_bytes(role, ref)
1070 if weights == 0:
1071 log.warning("Skipping %s server: could not size model %r (%s).", role.value, ref, exc)
1072 return None
1073 if _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref)):
1074 _warn_weights_exceed(role, ref, weights, total_vram)
1075 return None
1076 floor = _fallback_floor_for(role, ref, weights)
1077 log.warning(
1078 "Could not size the %s model %s (%s). Charging %.1f GiB, its weights plus the "
1079 "cache and buffers it will allocate, which is a floor rather than an estimate: "
1080 "the load may still need more.",
1081 role.value,
1082 ref,
1083 exc,
1084 floor / 1024**3,
1085 )
1086 if _host_memory_refuses(role, ref, floor, host_committed):
1087 return None
1088 return ModelPlacementInput(
1089 role=role, est_vram_bytes=floor, replicas=_replica_count(role, device_count)
1090 )
1093def _fallback_floor_for(role: WorkerRole, ref: str, weights: int) -> int:
1094 """:func:`_analytic_footprint_floor` for *role*, reading what metadata it can."""
1095 from lilbee.providers.base import ProviderError
1096 from lilbee.providers.gguf_meta import read_gguf_metadata
1098 try:
1099 path = engine_params.resolve_model_path(ref)
1100 meta = read_gguf_metadata(path)
1101 except (ProviderError, OSError, ValueError):
1102 meta = None
1103 path = None
1104 ctx = _placement_estimate_ctx(role, path, meta) if path is not None else _MIN_USABLE_CHAT_CTX
1105 return _analytic_footprint_floor(
1106 weights, meta=meta, ctx=ctx, slots=_placement_estimate_slots(role, meta)
1107 )
1110def _ref_is_moe(ref: str) -> bool:
1111 """Whether *ref*'s GGUF declares routed experts; False when it cannot be read."""
1112 from lilbee.providers.base import ProviderError
1113 from lilbee.providers.gguf_meta import read_gguf_metadata
1115 try:
1116 return _is_moe(read_gguf_metadata(engine_params.resolve_model_path(ref)))
1117 except (ProviderError, OSError):
1118 return False
1121def _cpu_offload_in_play() -> bool:
1122 """Whether this configuration puts any of a model's weights in system memory.
1124 Expert offload moves the experts, a partial ``n_gpu_layers`` moves whole
1125 layers, and zero moves the model. Without one of these the engine keeps
1126 everything on the card and the estimator's host figure describes memory
1127 nobody will allocate.
1128 """
1129 from lilbee.core.config import cfg
1131 return _expert_offload_configured() or cfg.n_gpu_layers is not None
1134def _host_bytes_must_be_resident(role: WorkerRole, ref: str) -> bool:
1135 """Whether *role*'s host bytes have to fit RAM rather than page in and out.
1137 The estimator's host figure counts mmap pages, and llama.cpp maps CPU-side
1138 weights over that mapping instead of allocating them, so with mmap they are
1139 evictable page cache: a model far larger than RAM streams from disk and
1140 serves, which is a practiced setup for a large mixture-of-experts. Only
1141 ``--no-mmap`` turns them into a buffered read that must be resident, and the
1142 single path that asks for it is a chat model on a network filesystem.
1144 Anything this cannot determine counts as mappable, because a false refusal
1145 here has no override and costs the user a model that would have run.
1146 """
1147 if role is not WorkerRole.CHAT:
1148 return False
1149 try:
1150 path = engine_params.resolve_model_path(ref)
1151 except (ProviderError, OSError, ValueError):
1152 return False
1153 if not is_network_path(path):
1154 return False
1155 return _chat_no_mmap(_role_weights_bytes(role, ref), on_network_fs=True)
1158def _host_committed(admitted: Mapping[WorkerRole, ModelPlacementInput]) -> int:
1159 """System-memory bytes the roles already admitted to this plan will hold."""
1160 return sum(inp.est_ram_bytes for inp in admitted.values())
1163def _host_memory_refuses(role: WorkerRole, ref: str, ram_bytes: int, committed: int) -> bool:
1164 """Whether *role*'s system-memory half is too big for this machine to load.
1166 Charged only when something actually offloads, and only when the bytes must
1167 be resident: refusing a mapped model that would have streamed from disk is a
1168 false refusal with no override, which is worse than a slow load.
1170 Measured against the whole plan, not this role alone. Every role was
1171 previously compared to the entire machine on its own, so two roles that each
1172 fit and together do not were both admitted.
1173 """
1174 if not _cpu_offload_in_play() or ram_bytes <= 0:
1175 return False
1176 wanted = committed + ram_bytes
1177 total = total_system_memory()
1178 if total and wanted > total and _host_bytes_must_be_resident(role, ref):
1179 log.warning(
1180 "The %s model %s cannot load: this plan puts %.1f GiB in system memory, which "
1181 "cannot be paged out here, and the machine has %.1f GiB in total. Use a smaller "
1182 "model, or offload less.",
1183 role.value,
1184 ref,
1185 wanted / 1024**3,
1186 total / 1024**3,
1187 )
1188 return True
1189 free = free_system_memory()
1190 if free and wanted > free:
1191 log.warning(
1192 "Offloading the %s model %s brings this plan to %.1f GiB in system memory and "
1193 "only %.1f GiB is free. It will still load; close other programs if it swaps "
1194 "or runs slowly.",
1195 role.value,
1196 ref,
1197 wanted / 1024**3,
1198 free / 1024**3,
1199 )
1200 return False
1203def _warn_weights_exceed(role: WorkerRole, ref: str, weights: int, total_vram: int) -> None:
1204 log.warning(
1205 "The %s model %s cannot load: its weights are %.1f GiB and this machine has "
1206 "%.1f GiB of GPU memory and %.1f GiB of system memory, so there is nowhere "
1207 "for its layers to go. Use a smaller model or a smaller quantization.",
1208 role.value,
1209 ref,
1210 weights / 1024**3,
1211 total_vram / 1024**3,
1212 model_cache.total_system_memory() / 1024**3,
1213 )
1216def _warn_weights_spill(role: WorkerRole, ref: str, weights: int, total_vram: int) -> None:
1217 """Say that a model larger than the GPUs will run partly in system memory."""
1218 log.warning(
1219 "The %s model %s is %.1f GiB and this machine has %.1f GiB of GPU memory, so "
1220 "the engine will keep the layers that fit on the GPU and the rest in system "
1221 "memory. It will run, and it will be slower than a model that fits.",
1222 role.value,
1223 ref,
1224 weights / 1024**3,
1225 total_vram / 1024**3,
1226 )
1229def placeable_total_vram() -> int:
1230 """Physical VRAM across all cards, for the weights-exceed placeability bound.
1232 Physical total is box-state-independent (a running incumbent doesn't skew
1233 it), so it is safe to read without a clean box. Reuses the plan probe when
1234 one is captured; otherwise probes best-effort and returns ``0`` on failure,
1235 which disables only the weights-exceed filter (its own ``total > 0`` guard).
1236 """
1237 probe = _plan_probe_store.get()
1238 if probe is not None:
1239 return sum(d.total_bytes for d in probe.devices)
1240 from lilbee.providers.base import ProviderError
1241 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
1243 try:
1244 apply_fleet_gpu_env()
1245 return sum(d.total_bytes for d in resolve_devices(resolve_llama_server()))
1246 except (ProviderError, OSError):
1247 return 0
1250def role_model_placeable(role: WorkerRole, ref: str, total_vram: int) -> bool:
1251 """Whether a fresh plan would actually serve *role* on *ref*.
1253 Mirrors the planner's own drop conditions (SDK-routed role, vision without a
1254 projector, model not installed, weights exceeding physical VRAM) using the
1255 same primitives, so the acquisition ladder binds and replaces against what
1256 an engine can serve rather than the raw config. Without this a
1257 configured-but-unplaceable role keeps bind from ever matching a running
1258 engine and restarts the shared engine on every process start.
1259 """
1260 if parse_model_ref(ref).is_remote or _vision_without_mmproj(role, ref):
1261 return False
1262 weights = _role_weights_bytes(role, ref) # 0 when not installed / unresolvable
1263 if weights == 0:
1264 return False
1265 return not _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref))
1268def _server_model_inputs(
1269 roles: tuple[WorkerRole, ...] | None = None,
1270 *,
1271 unified_budget: int | None = None,
1272 device_count: int = 0,
1273 total_vram: int = 0,
1274) -> tuple[list[ModelPlacementInput], dict[WorkerRole, str], int, dict[WorkerRole, str]]:
1275 """Build placement inputs for the configured server roles.
1277 The search and vision roles are estimated first; chat is then sized against the
1278 budget minus the search footprint (the ``reservation``) so a large chat cannot
1279 starve embed/rerank on a shared-memory host. ``device_count`` resolves an auto
1280 replica knob to one per GPU. When *roles* is given, only those are considered.
1281 Skips an unconfigured optional role, a vision model with no resolvable mmproj
1282 projector, a role whose model is not installed on disk (returned as
1283 ``skipped_not_installed`` so a surface can say so), and a model whose weight
1284 bytes alone exceed ``total_vram`` (physically unloadable under all-GPU layers).
1285 A model the estimator cannot size is enrolled at its file size instead of
1286 skipped, so the load, not the estimator, decides.
1287 """
1288 from lilbee.core.config import cfg
1290 inputs: dict[WorkerRole, ModelPlacementInput] = {}
1291 model_refs: dict[WorkerRole, str] = {}
1292 skipped_not_installed: dict[WorkerRole, str] = {}
1294 def consider(role: WorkerRole, *, chat_reservation: int = 0) -> None:
1295 if roles is not None and role not in roles:
1296 return
1297 # Any role may be "" (unconfigured) -> skipped, so that role has no
1298 # server and no not-installed complaint.
1299 ref = str(getattr(cfg, ROLE_REGISTRY[role].config_field))
1300 if not ref:
1301 return # unconfigured optional role -> no server
1302 if parse_model_ref(ref).is_remote:
1303 return # SDK-routed role: no local server to plan, not a missing install
1304 if _vision_without_mmproj(role, ref):
1305 return # no projector -> vision can't run on a server
1306 estimate = _estimate_or_fallback(
1307 role,
1308 ref,
1309 unified_budget=unified_budget,
1310 chat_reservation=chat_reservation,
1311 device_count=device_count,
1312 total_vram=total_vram,
1313 skipped_not_installed=skipped_not_installed,
1314 host_committed=_host_committed(inputs),
1315 )
1316 if estimate is None:
1317 return
1318 inputs[role] = estimate
1319 model_refs[role] = ref
1321 # Estimate every non-chat role first so the search footprint is known, then size
1322 # chat against the remainder. The reservation only applies on a shared-memory
1323 # host; discrete GPUs pin each role to its own VRAM and pack independently.
1324 for role in ROLE_REGISTRY:
1325 if role is not WorkerRole.CHAT:
1326 consider(role)
1327 reservation = _search_reservation(inputs) if unified_budget is not None else 0
1328 consider(WorkerRole.CHAT, chat_reservation=reservation)
1330 ordered = [inputs[role] for role in ROLE_REGISTRY if role in inputs]
1331 return ordered, model_refs, reservation, skipped_not_installed
1334def _non_chat_reservation(
1335 instances: Sequence[InstancePlan],
1336 inputs: Sequence[ModelPlacementInput],
1337 co_tenants: frozenset[WorkerRole] = frozenset(),
1338) -> dict[int, int]:
1339 """Per-device VRAM the non-chat role servers occupy, keyed by device index.
1341 A tensor-split chat shard must size its KV against the headroom left after the
1342 embed/rerank/vision servers on the same card, not the card's raw free VRAM, or
1343 it over-commits and OOMs at launch. Chat is excluded because it sizes its own
1344 weights. Chat's own swap-group siblings are excluded too: they are evicted while
1345 chat is resident, so their VRAM is chat's to use. That only holds when chat is
1346 itself a co-tenant; a co-tenant group that does not include chat runs behind its
1347 own swap process and can be resident beside chat, so it is charged normally.
1348 Non-chat roles are single-device, so each charges its full footprint (once per
1349 replica) to its card.
1350 """
1351 chat_siblings = co_tenants if WorkerRole.CHAT in co_tenants else frozenset()
1352 charge_by_role = {inp.role: inp.est_vram_bytes for inp in inputs}
1353 reserved: dict[int, int] = {}
1354 for inst in instances:
1355 if inst.role is WorkerRole.CHAT or inst.role in chat_siblings:
1356 continue
1357 charge = charge_by_role[inst.role]
1358 for device in inst.devices:
1359 reserved[device] = reserved.get(device, 0) + charge
1360 return reserved
1363def _charge_by_device(
1364 chosen: tuple[FleetDevice, ...], ratio: tuple[int, ...], total: int
1365) -> dict[str, int]:
1366 """What each of *chosen* was charged, keyed by the name the engine prints.
1368 A single-card instance carries the whole charge. A split carries it in the
1369 proportions it launches with, which is what the planner decided and therefore
1370 what the engine's own report should be compared against.
1371 """
1372 from lilbee.providers.fleet.readback import device_label
1374 if total <= 0 or not chosen:
1375 return {}
1376 if len(chosen) == 1:
1377 return {device_label(chosen[0]): total}
1378 weights = ratio if len(ratio) == len(chosen) else (1,) * len(chosen)
1379 denominator = sum(weights) or len(chosen)
1380 return {
1381 device_label(device): total * weight // denominator
1382 for device, weight in zip(chosen, weights, strict=True)
1383 }
1386def _launch_for(
1387 plan: InstancePlan,
1388 model_ref: str,
1389 binary: Path,
1390 by_index: dict[int, FleetDevice],
1391 *,
1392 unified_budget: int | None = None,
1393 chat_reservation: int = 0,
1394 reserved_by_device: dict[int, int] | None = None,
1395 est_vram_bytes: int = 0,
1396 model_path: Path | None = None,
1397) -> InstanceLaunch:
1398 """Build the launch spec (argv + device-pinning env) for one planned instance."""
1399 from lilbee.providers.gguf_meta import read_gguf_metadata
1401 # The self-check holds a downloaded file rather than a configured reference,
1402 # so it hands the path over instead of asking for one to be resolved.
1403 model_path = model_path or engine_params.resolve_model_path(model_ref)
1404 weights_bytes = _weights_bytes(model_path)
1405 meta = read_gguf_metadata(model_path)
1406 from lilbee.core.config import cfg
1408 chosen = tuple(by_index[i] for i in plan.devices)
1409 # ctx and slots are sized against the card this role landed on, not the fleet's
1410 # smallest, which is all the pre-placement estimate had to go on. A role spread
1411 # over several cards has no one budget; the split chat, the only such role today,
1412 # sizes against per-device headroom below.
1413 placed_device = chosen[0] if len(chosen) == 1 else None
1414 is_chat = plan.role is WorkerRole.CHAT
1415 is_vision = plan.role is WorkerRole.VISION
1416 mmproj = _vision_mmproj(model_ref) if is_vision else None
1417 chat_on_network_fs = is_chat and is_network_path(model_path)
1418 if chat_on_network_fs and not _chat_no_mmap(weights_bytes, on_network_fs=True):
1419 log.warning(
1420 "Chat model %s is served from a network filesystem and is too large to load "
1421 "into host RAM; mmap over the network can stall the load in uninterruptible "
1422 "I/O. Stage it on local disk for a reliable load.",
1423 model_ref,
1424 )
1425 # A tensor-split chat serves one full-context sequence sized against the busiest
1426 # card's headroom. A cfg.num_ctx pin overrides the fit (handled by _role_ctx).
1427 multi_card_chat = is_chat and len(chosen) > 1
1428 split_chat = multi_card_chat and cfg.num_ctx is None
1429 if multi_card_chat and host_lacks_nvlink():
1430 log.warning(
1431 "Chat model %s is tensor-split across GPUs %s on a host without NVLink; "
1432 "generation is PCIe all-reduce bound and can be very slow. A model that fits "
1433 "on fewer cards will generate faster.",
1434 model_ref,
1435 list(plan.devices),
1436 )
1437 split_slots = _SPLIT_CHAT_SLOTS
1438 if split_chat:
1439 reserved = reserved_by_device or {}
1440 # Headroom left after the embed/rerank servers on each shared card, not the
1441 # card's raw free VRAM, so the chat KV doesn't over-commit.
1442 per_device_free = [max(0, d.free_bytes - reserved.get(d.index, 0)) for d in chosen]
1444 def _split_fit(slots: int) -> int:
1445 return fleet_ctx.fit_split_ctx(
1446 model_path,
1447 meta=meta,
1448 slots=slots,
1449 ratio=plan.tensor_split,
1450 per_device_free_bytes=per_device_free,
1451 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
1452 flash_attn=_role_flash(WorkerRole.CHAT),
1453 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
1454 kv_cache_type_v=_role_kv_cache_type_v(WorkerRole.CHAT),
1455 ctx_ceiling=_placement_estimate_ctx(WorkerRole.CHAT, model_path, meta),
1456 )
1458 split_slots, ctx = _resolve_split_chat_slots(_split_fit)
1459 else:
1460 # Downshifted here and not only in the estimate: the role resolvers are
1461 # pure functions of model and config, so without this the retry after a
1462 # load OOM re-emits a byte-identical argv and dies the same way. The
1463 # split branch above already inherits it through its ctx_ceiling.
1464 ctx = apply_ctx_downshift(plan.role, _role_ctx(plan.role, model_path, meta, placed_device))
1465 rerank_mode = _role_rerank_mode(plan.role, meta)
1466 is_llm_rerank = rerank_mode is RerankMode.LLM
1467 # A multi-card chat runs as many full-context slots as its cards' KV headroom
1468 # holds (split_slots, one when a num_ctx pin skips the fit); other roles size
1469 # --parallel against the budget the same way the estimator did.
1470 slots = (
1471 split_slots
1472 if multi_card_chat
1473 else _slots_for(
1474 plan.role,
1475 model_path,
1476 ctx,
1477 mmproj_path=mmproj,
1478 unified_budget=unified_budget,
1479 chat_reservation=chat_reservation,
1480 rerank_mode=rerank_mode,
1481 device=placed_device,
1482 )
1483 )
1484 spec = _server_spec(plan.role, rerank_mode, meta)
1485 # Cross-encoder embed/rerank pools the whole input in one batch; an LLM reranker
1486 # is generative and uses the default batching plus flash attention.
1487 cross_encoder_pooled = plan.role in _EMBED_ROLES and not is_llm_rerank
1488 cache_type_k, cache_type_v = chat_cache_type_flags() if is_chat else (None, None)
1489 argv = build_server_argv(
1490 binary=binary,
1491 spec=spec,
1492 model_path=model_path,
1493 devices=plan.devices,
1494 n_gpu_layers=_role_gpu_layers(plan.role),
1495 slots=slots,
1496 ctx_per_slot=ctx,
1497 tensor_split=plan.tensor_split,
1498 mmproj=mmproj,
1499 flash_attn=flash_attn_flag() if _role_launches_with_flash(plan.role, rerank_mode) else None,
1500 cache_type_k=cache_type_k,
1501 cache_type_v=cache_type_v,
1502 batch_size=_pooled_batch_size(plan.role, rerank_mode, ctx),
1503 no_mmap=is_chat and _chat_no_mmap(weights_bytes, on_network_fs=chat_on_network_fs),
1504 cpu_moe=expert_offload_all(meta),
1505 n_cpu_moe=expert_offload_layers(meta),
1506 device_names=_device_names(chosen) or _cpu_pin_when_every_device_was_refused(),
1507 )
1508 return InstanceLaunch(
1509 role=plan.role,
1510 argv=argv,
1511 env_overrides={**visible_env(chosen), **llama_server_runtime_env()},
1512 model=model_ref,
1513 # token_cap drives cross-encoder/embed input truncation; the LLM rerank path
1514 # doesn't truncate (it relies on the per-slot ctx headroom), so leave it None.
1515 token_cap=max(1, ctx - engine_params._EMBED_CTX_MARGIN) if cross_encoder_pooled else None,
1516 # Weights size scales the cold-load ready timeout (larger model = longer).
1517 weights_bytes=weights_bytes,
1518 # Slots is the chat concurrency the gate admits; ctx is what a client fits to.
1519 slots=slots,
1520 ctx=ctx,
1521 built_ctx_target=(
1522 (cfg.num_ctx if cfg.num_ctx is not None else cfg.chat_n_ctx_target) if is_chat else 0
1523 ),
1524 replica=plan.replica,
1525 rerank_mode=rerank_mode,
1526 # What placement charged this instance, for the post-launch check against
1527 # the engine's own report of what it really allocated.
1528 est_vram_bytes=est_vram_bytes,
1529 est_vram_by_device=_charge_by_device(chosen, plan.tensor_split, est_vram_bytes),
1530 est_unreported_bytes=_unreported_bytes(plan.role, mmproj),
1531 )
1534def build_single_role_launch(role: WorkerRole, model_path: Path) -> InstanceLaunch:
1535 """The launch the fleet would build for *role* serving *model_path*, alone.
1537 One construction path. The self-check used to assemble its own beside this
1538 one and the two disagreed on slot count, on the context that follows from it,
1539 on device pinning and on the tensor split, so a green check proved nothing
1540 about the launch serving actually performs, and a red one could be a
1541 configuration serving would never have chosen.
1543 Placement is the planner's, on the devices the plan snapshot holds, so the
1544 check runs on the card the role would really land on.
1545 """
1546 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
1547 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
1549 apply_fleet_gpu_env()
1550 binary = resolve_llama_server()
1551 apply_cuda_runtime_env(binary)
1552 devices = _plan_devices(binary)
1553 by_index = {d.index: d for d in devices}
1554 # The whole machine, since nothing else is resident during a self-check.
1555 placed = (min(by_index),) if by_index else ()
1556 plan = InstancePlan(role=role, devices=placed)
1557 return _launch_for(
1558 plan,
1559 str(model_path),
1560 binary,
1561 by_index,
1562 unified_budget=_unified_memory_budget(devices),
1563 model_path=model_path,
1564 )
1567def resolve_devices(binary: Path) -> list[FleetDevice]:
1568 """Enumerate devices in the binary's index space, or the Vulkan VRAM probe."""
1569 return _resolve_devices_and_refusal(binary)[0]
1572# The visibility variable each vendor's runtime reads, named in the warning so
1573# the reader checks the one that applies to the card they actually have.
1574_VENDOR_VISIBILITY_HINT = {
1575 "NVIDIA": "CUDA_VISIBLE_DEVICES",
1576 "AMD": "ROCR_VISIBLE_DEVICES / HIP_VISIBLE_DEVICES",
1577 "Intel": "ONEAPI_DEVICE_SELECTOR",
1578}
1581def _warn_gpu_present_but_unenumerated(binary: Path) -> None:
1582 """Say so when the host has a GPU the engine did not list.
1584 Previously asked only whether an NVIDIA card was present, so an AMD or Intel
1585 host whose engine enumerated nothing produced the identical symptom, a fleet
1586 quietly planned for CPU, and said nothing. The vendor lookup is the same one
1587 the Vulkan ICD rules use, and it works on Windows as well as Linux.
1588 """
1589 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
1590 from lilbee.providers.fleet.gpu_select import PCIVendorID
1592 present = installed_gpu_vendor_ids()
1593 names = sorted(
1594 v.name.title() if v.name == "INTEL" else v.name for v in PCIVendorID if v in present
1595 )
1596 if not names:
1597 return
1598 hints = sorted(
1599 {_VENDOR_VISIBILITY_HINT[name] for name in names if name in _VENDOR_VISIBILITY_HINT}
1600 )
1601 log.warning(
1602 "This host has a %s GPU but the engine's device probe (%s --list-devices) "
1603 "reported none; placement is falling back to shared-memory mode with unpinned "
1604 "GPUs. Check the GPU driver, %s, and that this llama-server build supports "
1605 "that GPU.",
1606 " and ".join(names),
1607 binary,
1608 " / ".join(hints) if hints else "the vendor's visibility variable",
1609 )
1612def _resolve_devices_and_refusal(binary: Path) -> tuple[list[FleetDevice], bool]:
1613 """:func:`resolve_devices`, plus whether every GPU the engine listed was refused.
1615 One function because both answers come from one ``--list-devices`` run, and
1616 that run costs a subprocess against a driver that may be wedged. Asking twice
1617 would pay it twice.
1619 The binary's ``--list-devices`` is authoritative, including when it lists
1620 nothing: it prints every non-CPU device it can use, so an empty list means
1621 the engine has no usable GPU rather than that we failed to look. The Vulkan
1622 VRAM probe is consulted only when the binary produced no output at all, and
1623 it reports the same index space. A
1624 probe that times out raises instead (a wedged GPU driver); falling through
1625 to the in-process Vulkan probe there could hang this thread unkillably.
1626 """
1627 from lilbee.providers.fleet.cuda_runtime import assert_cuda_devices_usable
1628 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
1629 from lilbee.providers.fleet.gpu_select import enumerate_gpu_vram
1630 from lilbee.providers.fleet.rocm_runtime import assert_rocm_devices_usable
1632 probe = probe_devices(binary)
1633 # An engine that answered but listed no device while a GPU is physically
1634 # present may be hitting a transient init error (the card momentarily held by
1635 # another process, e.g. an embedder served alongside -- the "ggml_cuda_init:
1636 # initialization error" symptom). Re-probe before treating the empty list as
1637 # fatal; a persistently empty list still hits the fail-loud asserts below.
1638 for _ in range(_DEVICE_PROBE_EMPTY_RETRIES):
1639 if probe.devices or not probe.spoke_protocol or not installed_gpu_vendor_ids():
1640 break
1641 time.sleep(_DEVICE_PROBE_EMPTY_RETRY_DELAY_S)
1642 probe = probe_devices(binary)
1643 devices = probe.devices
1644 # A GPU build that links a runtime it cannot serve the host's GPU with must
1645 # fail loud, not silently fall back to CPU (the Vulkan VRAM probe below
1646 # would mask it). Only when the engine actually answered, though: a binary
1647 # that does not support --list-devices enumerated nothing because it was
1648 # never asked, and accusing its driver of failing would be wrong and fatal.
1649 if probe.spoke_protocol:
1650 assert_cuda_devices_usable(binary, devices, probe.output)
1651 assert_rocm_devices_usable(binary, devices, probe.output)
1652 if not devices and probe.spoke_protocol:
1653 _warn_gpu_present_but_unenumerated(binary)
1654 if not devices and not probe.spoke_protocol:
1655 # Only when the binary never answered the question. An engine that ran
1656 # and listed nothing is reporting a fact, not a gap: believing the host
1657 # loader instead invents devices the engine cannot see. A CPU-only build
1658 # on a desktop with mesa is the clearest case, and the cost is not merely
1659 # a wrong device list. The fleet is planned onto GPUs, the pins are
1660 # no-ops, the shared-RAM guard is off because devices looked non-empty,
1661 # and every role then loads its full weights into system RAM while
1662 # running on the CPU anyway.
1663 #
1664 # Keyed on the exit code and the header rather than on there being no
1665 # output at all: the probe merges stderr into stdout, so a build that
1666 # predates --list-devices prints usage text and would otherwise be read
1667 # as an authoritative "no GPUs here".
1668 from lilbee.providers.fleet.gpu_select import integrated_vulkan_indices
1670 integrated = integrated_vulkan_indices()
1671 devices = [
1672 FleetDevice(
1673 VULKAN_BACKEND, idx, "", vram, free, unified=idx in integrated, from_loader=True
1674 )
1675 for idx, vram, free in (enumerate_gpu_vram() or [])
1676 ]
1677 if devices:
1678 log.warning(
1679 "The engine's device probe returned nothing, so placement is using "
1680 "the host's Vulkan loader instead and found %d device(s). If the "
1681 "engine has no Vulkan backend it will run on CPU regardless; set %s "
1682 "to override the engine location if that is wrong.",
1683 len(devices),
1684 "LILBEE_ENGINE_DIR",
1685 )
1686 return devices, probe.refused_all
1689_DEVICE_PROBE_TTL_S = 2.0
1690# A failed probe is cached much longer than a good one: each retry against a
1691# wedged GPU driver costs a full probe timeout, so a per-poll retry would stall
1692# every placement read for a minute at a time.
1693_DEVICE_PROBE_FAILURE_TTL_S = 60.0
1694# An engine that lists no device on a GPU host may be hitting a transient GPU-init
1695# error (the card momentarily held by another process); re-probe before treating
1696# the empty list as fatal.
1697_DEVICE_PROBE_EMPTY_RETRIES = 3
1698_DEVICE_PROBE_EMPTY_RETRY_DELAY_S = 0.5
1701class _ReadDeviceCache:
1702 """Short-TTL device-probe cache for the read/view path.
1704 Not a ``cachetools.TTLCache``: it caches the *failure* too, under its own
1705 longer TTL, and re-raises it. A memoizing cache stores return values only,
1706 so a failing probe would re-spawn the subprocess on every placement read.
1708 Inspecting placement (GET placement/gpus, preview, ``placement show``)
1709 resolves devices on every call, which spawns a ``llama-server --list-devices``
1710 subprocess; a brief TTL collapses a burst of reads onto one probe. A probe
1711 failure is cached too (with its own TTL) and re-raised to every read in the
1712 window. The launch path is never served from here -- it sizes against the
1713 clean-box plan snapshot below (captured after stale-server reaping).
1714 """
1716 def __init__(self, ttl_s: float, failure_ttl_s: float) -> None:
1717 self._ttl_s = ttl_s
1718 self._failure_ttl_s = failure_ttl_s
1719 self._lock = threading.Lock()
1720 self._at: float | None = None
1721 self._devices: list[FleetDevice] | None = None
1722 self._failure: ProviderError | None = None
1724 def get(self, binary: Path) -> list[FleetDevice]:
1725 with self._lock:
1726 ttl = self._ttl_s if self._failure is None else self._failure_ttl_s
1727 fresh = self._at is not None and time.monotonic() - self._at < ttl
1728 if fresh and self._failure is not None:
1729 raise self._failure
1730 if self._devices is None or not fresh:
1731 self._at = time.monotonic()
1732 try:
1733 self._devices = resolve_devices(binary)
1734 except ProviderError as exc:
1735 self._devices = None
1736 self._failure = exc
1737 raise
1738 self._failure = None
1739 return self._devices
1741 def clear(self) -> None:
1742 with self._lock:
1743 self._at = None
1744 self._devices = None
1745 self._failure = None
1748_read_device_cache = _ReadDeviceCache(_DEVICE_PROBE_TTL_S, _DEVICE_PROBE_FAILURE_TTL_S)
1751def clear_read_device_cache() -> None:
1752 """Drop the read-path device probe cache (e.g. after the fleet is reconfigured).
1754 Also drops what the host's Vulkan loader told us about device types, which is
1755 otherwise held for the process lifetime and would survive a driver reload or
1756 an eGPU being plugged in.
1757 """
1758 from lilbee.providers.fleet.gpu_select import (
1759 integrated_vulkan_indices,
1760 vulkan_device_types_by_name,
1761 )
1763 _read_device_cache.clear()
1764 vulkan_device_types_by_name.cache_clear()
1765 integrated_vulkan_indices.cache_clear()
1768@dataclass(frozen=True)
1769class _PlanProbe:
1770 """Clean-box memory snapshot every plan is sized against.
1772 Captured once, right after stale-server reaping and before the first build,
1773 when nothing lilbee owns is loaded. Reloads re-plan against this same
1774 snapshot instead of re-probing: a live probe under a loaded fleet reports
1775 our own residency as unavailable, which would shrink chat context and slot
1776 counts, widen splits, and (on a unified-memory host) evict roles outright.
1777 Launches stay a pure function of config + hardware + this snapshot, so the
1778 reload diff restarts only real changes. Cleared on full fleet teardown so
1779 the next boot probes the clean box afresh.
1780 """
1782 devices: tuple[FleetDevice, ...]
1783 # What one role may size its ctx and slots against, already scaled by
1784 # cfg.gpu_memory_fraction. System memory only on a host with no GPU.
1785 sizing_budget: int
1786 free_system: int
1787 # The engine listed GPUs and lilbee rejected all of them, so the plan is
1788 # CPU-shaped while the engine would still choose one of those devices.
1789 engine_devices_all_refused: bool = False
1792class _PlanProbeStore:
1793 """Holds the captured plan snapshot; a single instance below (no bare global)."""
1795 def __init__(self) -> None:
1796 self._lock = threading.Lock()
1797 self._probe: _PlanProbe | None = None
1799 def set(self, probe: _PlanProbe) -> None:
1800 with self._lock:
1801 self._probe = probe
1803 def get(self) -> _PlanProbe | None:
1804 with self._lock:
1805 return self._probe
1807 def clear(self) -> None:
1808 with self._lock:
1809 self._probe = None
1812_plan_probe_store = _PlanProbeStore()
1815# Where the ladder stops. Sized for chat, below which the answers are too short
1816# to be useful, so a role that still will not load here has a real problem the
1817# planner cannot size its way out of and the failure should surface. Roles whose
1818# window already sits under it (a small embedding context) are left alone rather
1819# than raised to meet it, so for them the ladder is a no-op and the failure
1820# surfaces after the one retry.
1821MIN_DOWNSHIFT_CTX = 4096
1824class _CtxDownshiftStore:
1825 """How many halvings each role's auto context has taken after a load OOM.
1827 An estimate that was too optimistic is only recoverable if the retry asks
1828 for something different. Halving the auto context does that, and keeping the
1829 count here rather than in the launch means the whole plan is re-predicted
1830 against the smaller number, including the placement it implies.
1831 """
1833 def __init__(self) -> None:
1834 self._lock = threading.Lock()
1835 self._steps: dict[WorkerRole, int] = {}
1836 # The last unshifted context each role was sized from, recorded as it is
1837 # applied. Deciding whether another halving would change anything needs
1838 # the number being halved, and this is the only place that sees it.
1839 self._base: dict[WorkerRole, int] = {}
1841 def steps(self, role: WorkerRole) -> int:
1842 with self._lock:
1843 return self._steps.get(role, 0)
1845 def note_base(self, role: WorkerRole, ctx: int) -> None:
1846 with self._lock:
1847 self._base[role] = ctx
1849 def base(self, role: WorkerRole) -> int | None:
1850 with self._lock:
1851 return self._base.get(role)
1853 def step(self, role: WorkerRole) -> int:
1854 with self._lock:
1855 taken = self._steps.get(role, 0) + 1
1856 self._steps[role] = taken
1857 return taken
1859 def clear(self, role: WorkerRole | None = None) -> None:
1860 with self._lock:
1861 if role is None:
1862 self._steps.clear()
1863 self._base.clear()
1864 return
1865 self._steps.pop(role, None)
1866 self._base.pop(role, None)
1869_ctx_downshift_store = _CtxDownshiftStore()
1872def apply_ctx_downshift(role: WorkerRole, ctx: int) -> int:
1873 """*ctx* halved once per downshift step recorded for *role*, floored.
1875 Never more than *ctx*. The floor is a stopping point, not a target: applied
1876 to a context already below it (a small embedding window, a model trained for
1877 2048 tokens) a bare floor would hand back a larger number, and the retry
1878 after a load OOM would ask for more memory than the launch that just ran out
1879 of it. Such a role simply has nothing to give back, and its failure surfaces
1880 after the one retry instead.
1882 A user's ``cfg.num_ctx`` pin is returned untouched: serving a window smaller
1883 than the one that was asked for, without being asked, is worse than failing
1884 to load and saying so.
1885 """
1886 from lilbee.core.config import cfg
1888 if role is WorkerRole.CHAT and cfg.num_ctx is not None:
1889 return ctx
1890 _ctx_downshift_store.note_base(role, ctx)
1891 return _shifted(ctx, _ctx_downshift_store.steps(role))
1894def _shifted(ctx: int, steps: int) -> int:
1895 """*ctx* halved *steps* times, never below the floor and never above *ctx*."""
1896 return min(ctx, max(MIN_DOWNSHIFT_CTX, ctx >> steps)) if steps else ctx
1899def record_ctx_downshift(role: WorkerRole) -> bool:
1900 """Take one downshift step for *role*; False when there is none left to take.
1902 False means the retry would ask for the same thing again, so the caller must
1903 surface the load failure instead of respawning an identical launch.
1904 """
1905 from lilbee.core.config import cfg
1907 if role is WorkerRole.CHAT and cfg.num_ctx is not None:
1908 return False
1909 base = _ctx_downshift_store.base(role)
1910 if base is None:
1911 # Nothing has been sized for this role yet, so there is no number to
1912 # decide against. Allow one step rather than trusting that a plan always
1913 # runs first: an unbounded grant here would let a caller that never
1914 # sizes anything loop forever.
1915 if _ctx_downshift_store.steps(role):
1916 return False
1917 _ctx_downshift_store.step(role)
1918 return True
1919 steps = _ctx_downshift_store.steps(role)
1920 if _shifted(base, steps + 1) == _shifted(base, steps):
1921 return False
1922 _ctx_downshift_store.step(role)
1923 return True
1926def clear_ctx_downshift(role: WorkerRole | None = None) -> None:
1927 """Forget *role*'s recorded downshift, or every role's, back to full size.
1929 Called when a role's engine reports ready, which is proof the reduced plan
1930 loaded: keeping the reduction after that would carry a shrunken window into
1931 a machine that has since freed memory, or into a smaller model the user
1932 switched to, and would then refuse on its first failure with a budget it
1933 had already spent.
1934 """
1935 _ctx_downshift_store.clear(role)
1938def _probe_engine_devices() -> tuple[list[FleetDevice], bool]:
1939 """Apply the fleet GPU/CUDA env, resolve the binary, and enumerate devices.
1941 This is the wedge point: a missing binary raises NOT_FOUND, and a CUDA build
1942 that cannot init a GPU (a broken-runtime host) raises loud from resolve_devices
1943 rather than silently degrading. Device enumeration reads no residency, so it is
1944 safe to run while an incumbent engine is still up.
1945 """
1946 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
1947 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
1949 apply_fleet_gpu_env()
1950 binary = resolve_llama_server()
1951 apply_cuda_runtime_env(binary)
1952 devices, refused = _resolve_devices_and_refusal(binary)
1953 if devices:
1954 return devices, refused
1955 return _reprobe_while_a_gpu_is_installed(binary, refused)
1958def _reprobe_while_a_gpu_is_installed(
1959 binary: Path, refused: bool
1960) -> tuple[list[FleetDevice], bool]:
1961 """Ask again when the host has a GPU the engine did not list.
1963 The plan snapshot is taken once, on a clean box, and is not retaken until a
1964 full teardown, so an empty first answer decides the whole run. A GPU driver
1965 that is still initializing when the daemon starts, which is ordinary under
1966 systemd or right after a container gains a device, would leave a GPU host
1967 serving on CPU until someone noticed and restarted it.
1969 Only where a card is actually installed. A host with no GPU answers empty
1970 every time and must not pay a retry for it on every start.
1971 """
1972 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
1974 if not installed_gpu_vendor_ids():
1975 return [], refused
1976 for attempt in range(1, _PROBE_RETRIES + 1):
1977 log.info(
1978 "The engine listed no GPU on a host that has one; asking again in %.1fs "
1979 "(attempt %d of %d) in case the driver is still initializing.",
1980 _PROBE_RETRY_DELAY_S,
1981 attempt,
1982 _PROBE_RETRIES,
1983 )
1984 time.sleep(_PROBE_RETRY_DELAY_S)
1985 clear_read_device_cache()
1986 devices, refused = _resolve_devices_and_refusal(binary)
1987 if devices:
1988 return devices, refused
1989 return [], refused
1992def assert_engine_probeable() -> None:
1993 """Raise if the engine cannot be probed; capture no snapshot.
1995 A build precondition that must run BEFORE stopping a replaceable incumbent:
1996 it surfaces a wedged GPU probe or an unusable CUDA runtime without taking the
1997 residency-dependent memory snapshot (that belongs on the clean box, after the
1998 stop, in capture_plan_probe). resolve_devices caches within its TTL, so the
1999 follow-up capture reuses this enumeration rather than re-probing the hardware.
2000 """
2001 _probe_engine_devices()
2004def capture_plan_probe() -> None:
2005 """Snapshot devices and memory for planning; call only on a clean box."""
2006 devices, refused_all = _probe_engine_devices()
2007 _plan_probe_store.set(
2008 _PlanProbe(
2009 devices=tuple(devices),
2010 sizing_budget=_device_sizing_budget(devices),
2011 free_system=model_cache.free_system_memory(),
2012 engine_devices_all_refused=refused_all,
2013 )
2014 )
2017def refresh_plan_devices() -> None:
2018 """Re-read which devices exist, keeping the clean-box memory figures.
2020 The snapshot is captured once and only a full teardown clears it, so an eGPU
2021 unplugged, a driver reset, or a VM hot-remove left the fleet pinning a device
2022 that is no longer there and every rebuild replanning onto it.
2024 Only the structural half is restated. The memory figures are what make a
2025 reload plan the way the boot did, and re-taking them while the fleet is
2026 resident would charge it against itself, which is the whole reason the
2027 snapshot exists.
2029 A probe that cannot run leaves the snapshot alone: the last known device list
2030 is a better answer than none, and the loud paths for an unreachable engine
2031 live in the build, not here.
2032 """
2033 probe = _plan_probe_store.get()
2034 if probe is None:
2035 return
2036 clear_read_device_cache()
2037 try:
2038 devices, refused_all = _probe_engine_devices()
2039 except (ProviderError, OSError) as exc:
2040 log.debug("Device rediscovery could not run, keeping the previous list: %s", exc)
2041 return
2042 if tuple(devices) == probe.devices:
2043 return
2044 log.info(
2045 "The set of GPUs changed since this fleet was planned (%d device(s) now, %d before); "
2046 "replanning against the ones that are here.",
2047 len(devices),
2048 len(probe.devices),
2049 )
2050 _plan_probe_store.set(
2051 _PlanProbe(
2052 devices=tuple(devices),
2053 sizing_budget=_device_sizing_budget(devices),
2054 free_system=probe.free_system,
2055 engine_devices_all_refused=refused_all,
2056 )
2057 )
2060def clear_plan_probe() -> None:
2061 """Drop the plan snapshot (full fleet teardown); the next build re-captures."""
2062 _plan_probe_store.clear()
2065def _cpu_pin_when_every_device_was_refused() -> tuple[str, ...]:
2066 """``("none",)`` when the engine offered GPUs that lilbee refused, else empty.
2068 Dropping a device from lilbee's view does not stop the engine using it. With
2069 no pin at all, ggml applies its own selection, and its fallback takes the
2070 first non-CPU adapter, which is exactly the paravirtual device just refused;
2071 with every layer offloaded by default the model then runs on it while
2072 placement budgeted against system RAM. Naming no device keeps the engine on
2073 the CPU the plan was shaped for.
2074 """
2075 probe = _plan_probe_store.get()
2076 if probe is None or not probe.engine_devices_all_refused:
2077 return ()
2078 log.warning(
2079 "The engine listed GPU devices that lilbee will not plan onto, so it is being "
2080 "run on the CPU. Serving from one of them would be slower than the CPU or fail "
2081 "outright, and placement has been sized for system RAM."
2082 )
2083 return (_NO_DEVICE,)
2086def _plan_devices(binary: Path) -> list[FleetDevice]:
2087 """Devices the plan paths size against: the snapshot, else a live probe."""
2088 probe = _plan_probe_store.get()
2089 return list(probe.devices) if probe is not None else resolve_devices(binary)
2092def plan_sizing_budget(device: FleetDevice | None = None) -> int:
2093 """Usable memory for ctx/slot sizing: *device*'s own, else the snapshot, else live."""
2094 from lilbee.core.config import cfg
2096 if device is not None:
2097 return int(device.total_bytes * cfg.gpu_memory_fraction)
2098 probe = _plan_probe_store.get()
2099 if probe is not None:
2100 return probe.sizing_budget
2101 return _device_sizing_budget(_live_sizing_devices())
2104def _device_sizing_budget(devices: Sequence[FleetDevice]) -> int:
2105 """Memory one role may size its ctx and slots against, in bytes.
2107 Read from the engine's own device report, which ran under the environment the
2108 servers will run under and states each device's memory whatever the backend.
2109 A host-memory read answers with system RAM on every host without an NVIDIA
2110 card, which gave a 24 GiB AMD card a budget the size of the machine, and on
2111 Apple Silicon it ignores that Metal will not allocate past
2112 ``recommendedMaxWorkingSetSize``, which is the figure the probe carries.
2114 The smallest device, since this is asked before placement has picked one;
2115 :func:`_launch_for` re-sizes against the card the role actually landed on.
2116 System memory only when the engine reports no device at all, where the fleet
2117 runs on the CPU and system memory is the budget.
2118 """
2119 from lilbee.core.config import cfg
2121 if devices:
2122 return int(min(d.total_bytes for d in devices) * cfg.gpu_memory_fraction)
2123 return int(model_cache.total_system_memory() * cfg.gpu_memory_fraction)
2126def _live_sizing_devices() -> list[FleetDevice]:
2127 """Devices to size against with no plan snapshot; empty when none can be read."""
2128 try:
2129 return _read_device_cache.get(resolve_llama_server())
2130 except (ProviderError, OSError):
2131 return []
2134def _plan_free_system_memory() -> int:
2135 """Free system RAM for the unified-memory budget: the snapshot, else live."""
2136 probe = _plan_probe_store.get()
2137 return probe.free_system if probe is not None else model_cache.free_system_memory()
2140def _unreported_bytes(role: WorkerRole, mmproj: Path | None) -> int:
2141 """Estimated bytes the engine allocates without printing a buffer line.
2143 A vision projector's weights: llama.cpp allocates them in clip's own loader,
2144 which prints a size but not the "buffer size = N MiB" shape the readback
2145 reads, so the report is short by exactly this and the self-check would warn
2146 on a load that was sized correctly.
2147 """
2148 if role is not WorkerRole.VISION or mmproj is None:
2149 return 0
2150 try:
2151 return mmproj.stat().st_size
2152 except OSError:
2153 return 0
2156def _chat_no_mmap(weights_bytes: int, *, on_network_fs: bool = False) -> bool:
2157 """Whether the chat server should malloc its weights instead of mmapping them.
2159 Local disk mmaps: lazy page-fault paging gives a faster first token on a cold
2160 cache -- the common desktop first launch -- matching mmap-by-default engines.
2161 ``--no-mmap``'s buffered full read only wins on an already-hot cache and it
2162 pessimizes cold start, so it is not worth defaulting on for local disk. A
2163 network filesystem still prefers the buffered read whenever the host copy
2164 fits, because mmap page faults served over the wire can wedge the loader in
2165 uninterruptible I/O (see ``_NO_MMAP_NETWORK_RAM_FRACTION``).
2166 """
2167 if not on_network_fs:
2168 return False
2169 return weights_bytes <= model_cache.total_system_memory() * _NO_MMAP_NETWORK_RAM_FRACTION
2172def _device_names(devices: tuple[FleetDevice, ...]) -> tuple[str, ...]:
2173 """``--device`` names for *devices*, empty when the backend pins through env.
2175 Vulkan and SYCL, because neither one's environment variable speaks the space
2176 the probe enumerated. Vulkan's indexes the raw loader enumeration while the
2177 names come from the engine's filtered list, so the two disagree wherever ggml
2178 drops or merges a device. SYCL's is not an index list at all but a selector
2179 over a backend runtime, so a device the engine calls ``SYCL1`` need not be
2180 Level Zero ordinal 1: OpenCL devices interleave, discarded devices shift the
2181 numbering, and multi-tile cards appear as sub-devices.
2183 ``--device`` sidesteps both by naming devices exactly as ``--list-devices``
2184 printed them, which is where these indices were read from. CUDA and ROCm
2185 keep composing their variables, which do share the probe's space.
2186 """
2187 if not devices or devices[0].backend not in _NAME_PINNED_BACKENDS:
2188 return ()
2189 if any(d.from_loader for d in devices):
2190 # These indices are raw loader ordinals, and --device speaks the engine's
2191 # own post-filter naming, so Vulkan1 here can name Vulkan0 there or
2192 # nothing at all. Sizing against them is still worth doing; pinning by
2193 # them is not. Left unpinned, ggml applies its own device selection,
2194 # which is the filtering lilbee is trying to agree with in the first
2195 # place. The env pin is not the answer either: it takes raw ordinals but
2196 # switches off the type filter, the support check and the dedup with them.
2197 return ()
2198 return tuple(f"{d.backend}{d.index}" for d in devices)
2201def _unified_memory_budget(devices: list[FleetDevice]) -> int | None:
2202 """Shared-RAM placement budget (free RAM minus the OS floor), or ``None``.
2204 ``None`` once any device has memory of its own, since dedicated VRAM is the
2205 constraint there rather than system RAM. A host whose only devices are
2206 integrated, and a host with no devices at all, both stay inside the system
2207 budget: their GPU memory is the system's memory.
2208 """
2209 # Only a device with memory of its own lifts the system-RAM constraint. An
2210 # integrated GPU or an Apple Silicon Mac reports a slice of the same RAM the
2211 # OS is using, so treating its total as headroom over-commits the machine by
2212 # roughly the whole system footprint.
2213 if any(not device.unified for device in devices):
2214 return None
2215 return _capped_by_device_memory(
2216 max(0, _plan_free_system_memory() - _system_memory_floor()), devices
2217 )
2220def _unified_admission_budget(devices: list[FleetDevice]) -> int | None:
2221 """Shared-RAM pool a role set is *admitted* against, or ``None`` if dedicated.
2223 Total installed RAM minus the OS floor, not what happens to be free. Sizing
2224 asks a different question and keeps using free RAM: how much context can be
2225 backed right now. Admission asks whether the machine can host this fleet at
2226 all, and the plan defines the whole intended residency, so charging it
2227 against a live figure refuses a 600 MB model on a box that is merely busy at
2228 the moment, which is what happened. The GPU path already charges total
2229 capacity for exactly this reason.
2230 """
2231 if _unified_memory_budget(devices) is None:
2232 return None
2233 return _capped_by_device_memory(
2234 max(0, model_cache.total_system_memory() - _system_memory_floor()), devices
2235 )
2238def _system_memory_floor() -> int:
2239 """RAM held back for the OS when placing against system memory.
2241 ``cfg.system_memory_reserve_gb``, still capped at a quarter of installed RAM:
2242 a fixed reserve leaves a 7-8 GB host with no budget at all and refuses even
2243 tiny models, so the proportional cap holds however the reserve is set.
2244 """
2245 from lilbee.core.config import cfg
2247 total = model_cache.total_system_memory()
2248 return min(int(cfg.system_memory_reserve_gb * 1024**3), total // _SYSTEM_MEMORY_FLOOR_DIVISOR)
2251def _capped_by_device_memory(budget: int, devices: Sequence[FleetDevice]) -> int:
2252 """*budget*, never above what the devices can address between them.
2254 A shared-memory device still has a ceiling of its own: an integrated GPU
2255 addresses a fixed aperture of system RAM, and Metal will not allocate past
2256 ``recommendedMaxWorkingSetSize``. Both report that ceiling as their total, so
2257 a host budget derived from installed RAM promises memory the devices cannot
2258 reach. Unchanged where the engine reports no device, since the fleet is then
2259 running on the CPU and the host figure is the true one.
2260 """
2261 if not devices:
2262 return budget
2263 return min(budget, sum(d.total_bytes for d in devices))
2266def _device_capacity(devices: list[FleetDevice], charge_against_free: bool) -> dict[int, int]:
2267 """Per-device memory placement may charge against, keyed by device index.
2269 A card's total is what it holds, not what is going spare. A compositor, a
2270 browser, or a training job sitting on VRAM is invisible in the total, and the
2271 usable fraction placement applies covers fragmentation and driver overhead
2272 rather than other tenants, so a plan fits on paper and OOMs at load.
2274 Free bytes answer that, but only where they mean "everyone else's residency":
2275 that is the clean-box snapshot, taken after stale servers are reaped and
2276 before anything is built. Read live on a warm box they also exclude the
2277 fleet's own models, and since a plan always describes the complete intended
2278 residency, charging them there would count the fleet against itself and
2279 report a running plan as unplaceable. Those callers keep the total.
2281 Placement applies its usable fraction to whatever this returns, so a card
2282 with a tenant keeps a proportional margin rather than being packed to its
2283 last free byte, where fragmentation is worst.
2284 """
2285 packable = _packable_devices(devices)
2286 if not charge_against_free:
2287 return {d.index: d.total_bytes for d in packable}
2288 return {d.index: min(d.total_bytes, d.free_bytes) for d in packable}
2291def _packable_devices(devices: list[FleetDevice]) -> list[FleetDevice]:
2292 """The devices bin-packing may charge against.
2294 An integrated GPU's memory is the host's. Packing it beside a dedicated card
2295 promises the same RAM twice, once to its own budget and once to everything
2296 else on the machine, and its heap is often the larger number, so the packer
2297 prefers it: a 32 GiB shared heap outbids a 24 GiB card that actually has the
2298 memory. Where a dedicated device exists it is the one to serve from, and the
2299 integrated one is left to the shared-memory budget.
2301 A host with nothing but integrated devices keeps them. There is nothing else
2302 to serve from, and that path is governed by the system budget rather than by
2303 per-device packing.
2304 """
2305 dedicated = [d for d in devices if not d.unified]
2306 return dedicated or devices
2309def _resolve_placement(
2310 placement: PlacementSpec | None,
2311 inputs: list[ModelPlacementInput],
2312 model_refs: dict[WorkerRole, str],
2313 devices: list[FleetDevice],
2314 *,
2315 unified_budget: int | None,
2316 charge_against_free: bool = False,
2317) -> Placement:
2318 """Resolve a Placement from the manual spec when set, else the auto planner."""
2319 estimate_peak = _peak_estimator(model_refs)
2320 capacity = _device_capacity(devices, charge_against_free)
2321 if placement is not None:
2322 return placement_from_spec(
2323 placement,
2324 tuple(model_refs),
2325 capacity,
2326 estimate_peak=estimate_peak,
2327 )
2328 # The chat split's card count is decided against the snapshot's free VRAM (what the
2329 # launch sizes its context against) so placement and launch agree. A split needs
2330 # >=2 GPUs, so skip the chat model's gguf read entirely below that.
2331 chat_ctx_fit, chat_ctx_target = (
2332 _chat_split_ctx_objective(model_refs) if len(capacity) >= _MIN_SPLIT_GPUS else (None, 0)
2333 )
2334 return plan_placement(
2335 inputs,
2336 [(idx, budget) for idx, budget in capacity.items()],
2337 estimate_peak=estimate_peak,
2338 unified_budget=unified_budget,
2339 chat_ctx_fit=chat_ctx_fit,
2340 chat_ctx_target=chat_ctx_target,
2341 free_headroom={d.index: d.free_bytes for d in devices},
2342 )
2345def _placement_or_auto(
2346 placement: PlacementSpec | None,
2347 inputs: list[ModelPlacementInput],
2348 model_refs: dict[WorkerRole, str],
2349 devices: list[FleetDevice],
2350 *,
2351 unified_budget: int | None,
2352 charge_against_free: bool = False,
2353) -> tuple[Placement, bool]:
2354 """Resolve a saved spec, falling back to auto when it no longer fits the hardware.
2356 Returns the placement and whether the spec was the one applied. Hardware moves
2357 under a saved placement: a card is removed, a driver stops enumerating a GPU, a
2358 container starts without one. Refusing to plan there takes chat, embed and
2359 ingest down over a pin set on hardware the host no longer has, so the fleet
2360 degrades to automatic placement and logs why. An interactive apply still fails
2361 loud (:func:`lilbee.app.placement.set_placement`), where the pin is what the
2362 caller just asked for and a silent substitution would be the surprise.
2363 """
2364 if placement is None:
2365 return _resolve_placement(
2366 None,
2367 inputs,
2368 model_refs,
2369 devices,
2370 unified_budget=unified_budget,
2371 charge_against_free=charge_against_free,
2372 ), False
2373 try:
2374 return _resolve_placement(
2375 placement,
2376 inputs,
2377 model_refs,
2378 devices,
2379 unified_budget=unified_budget,
2380 charge_against_free=charge_against_free,
2381 ), True
2382 except PlacementError as exc:
2383 log.warning(
2384 "The saved GPU placement does not fit this hardware (%s); using automatic "
2385 "placement instead. Set a new placement to replace it.",
2386 exc,
2387 )
2388 return _resolve_placement(
2389 None,
2390 inputs,
2391 model_refs,
2392 devices,
2393 unified_budget=unified_budget,
2394 charge_against_free=charge_against_free,
2395 ), False
2398@dataclass(frozen=True)
2399class ResolvedPlacement:
2400 """Devices + resolved instance plans + model refs for the placement view."""
2402 devices: tuple[FleetDevice, ...]
2403 instances: tuple[InstancePlan, ...]
2404 unplaceable_roles: tuple[WorkerRole, ...]
2405 model_refs: dict[WorkerRole, str]
2406 # Roles placed anyway despite not fitting, with the shortfall in bytes. The
2407 # planner has always known this and only logged it, so a surface showed a
2408 # tight role as comfortably placed.
2409 tight_roles: dict[WorkerRole, int] = field(default_factory=dict)
2410 co_tenants: frozenset[WorkerRole] = frozenset()
2411 # False when a spec was given but did not fit the hardware, so these instances
2412 # are the auto planner's and a surface must not present them as the manual plan.
2413 spec_applied: bool = True
2414 # Roles configured but skipped because their model isn't installed (role -> ref).
2415 # Distinct from unplaceable_roles (installed but won't fit); lets a surface show
2416 # "not downloaded" instead of an empty table on a fresh install.
2417 skipped_not_installed: dict[WorkerRole, str] = field(default_factory=dict)
2420def resolve_placement_plan(
2421 placement: PlacementSpec | None, *, fall_back_to_auto: bool = False
2422) -> ResolvedPlacement:
2423 """Probe devices and resolve the auto-or-manual placement, without launching.
2425 ``fall_back_to_auto`` reads *placement* as a saved setting rather than a
2426 request: one that no longer fits the hardware resolves to the auto plan with
2427 ``spec_applied`` False instead of raising.
2428 """
2429 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
2430 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
2432 apply_fleet_gpu_env()
2433 binary = resolve_llama_server()
2434 apply_cuda_runtime_env(binary)
2435 devices = _read_device_cache.get(binary)
2436 unified_budget = _unified_memory_budget(devices)
2437 inputs, model_refs, _, skipped_not_installed = _server_model_inputs(
2438 None, unified_budget=unified_budget, total_vram=sum(d.total_bytes for d in devices)
2439 )
2440 admission_budget = _unified_admission_budget(devices)
2441 if fall_back_to_auto:
2442 resolved, spec_applied = _placement_or_auto(
2443 placement, inputs, model_refs, devices, unified_budget=admission_budget
2444 )
2445 else:
2446 resolved = _resolve_placement(
2447 placement, inputs, model_refs, devices, unified_budget=admission_budget
2448 )
2449 spec_applied = placement is not None
2450 return ResolvedPlacement(
2451 devices=tuple(devices),
2452 instances=resolved.instances,
2453 unplaceable_roles=resolved.unplaceable_roles,
2454 model_refs=model_refs,
2455 co_tenants=resolved.co_tenants,
2456 skipped_not_installed=skipped_not_installed,
2457 spec_applied=spec_applied,
2458 tight_roles=dict(resolved.tight_roles),
2459 )
2462@dataclass(frozen=True)
2463class FleetPlan:
2464 """The servers to start, and the roles that share one swap group."""
2466 launches: tuple[InstanceLaunch, ...]
2467 co_tenants: frozenset[WorkerRole] = frozenset()
2468 # Configured roles left unplaced because their model isn't installed (role ->
2469 # ref), so the warm path can fail a not-installed chat with a named reason
2470 # instead of spinning the warm line forever.
2471 skipped_not_installed: dict[WorkerRole, str] = field(default_factory=dict)
2472 # Launches refused for a window below the minimum grounded prompt
2473 # (role -> user-facing reason with the numbers).
2474 skipped_unusable_ctx: dict[WorkerRole, str] = field(default_factory=dict)
2477def _log_placement_findings(placement: Placement, model_refs: dict[WorkerRole, str]) -> None:
2478 """Warn about placements that exceed the memory budget.
2480 Shared-memory roles that fit nowhere get no server (loading them would OOM the
2481 host). GPU roles are never refused: one whose estimate exceeds the free VRAM
2482 still loads on demand, with a warning carrying the shortfall.
2483 """
2484 for role in placement.unplaceable_roles:
2485 log.warning(
2486 "%s model %s does not fit available memory and will not be served; "
2487 "free up memory or use a smaller model.",
2488 role.value,
2489 model_refs[role],
2490 )
2491 for role, shortfall in placement.tight_roles.items():
2492 log.warning(
2493 "Memory is tight for the %s model %s: it is estimated to need %.1f GiB more "
2494 "GPU memory than is available. It will still load on demand, keeping the "
2495 "layers that fit on the GPU and the rest in system memory; if it runs "
2496 "slowly, free up GPU memory or use a smaller model.",
2497 role.value,
2498 model_refs[role],
2499 # A sub-0.05 GiB shortfall would render as "0.0 GiB more".
2500 max(shortfall / 1024**3, 0.1),
2501 )
2502 if placement.co_tenants:
2503 log.info(
2504 "%s share GPU memory and load on demand; only one is resident at a time.",
2505 ", ".join(sorted(role.value for role in placement.co_tenants)),
2506 )
2509def _unusable_chat_ctx_reason(launch: InstanceLaunch) -> str | None:
2510 """Reason to refuse a chat launch whose window cannot hold a grounded prompt.
2512 ``None`` for non-chat roles, for a window that holds the minimum grounded
2513 prompt, and for user knobs that ask for a smaller one (a ``num_ctx`` pin,
2514 a sub-minimum ``num_ctx_max`` / ``chat_n_ctx_target``).
2515 """
2516 from lilbee.core.config import cfg
2518 if launch.role is not WorkerRole.CHAT or cfg.num_ctx is not None:
2519 return None
2520 needed = engine_params.min_usable_chat_ctx()
2521 # User knobs capping the window below the minimum are honored (the num_ctx
2522 # pin bypasses above).
2523 asked = min(cfg.chat_n_ctx_target, cfg.num_ctx_max or cfg.chat_n_ctx_target)
2524 if asked < needed:
2525 return None
2526 if launch.ctx >= needed:
2527 return None
2528 return (
2529 f"The chat model {launch.model} loads, but the memory left after its weights "
2530 f"backs only a {launch.ctx}-token context, and a grounded answer needs about "
2531 f"{needed} tokens (system prompt, a retrieved source, the question, and room "
2532 "for the answer), so it will not be served. Use a smaller model or a smaller "
2533 "quant, or set num_ctx to force a larger window."
2534 )
2537def plan_launches(
2538 roles: tuple[WorkerRole, ...] | None,
2539 binary: Path,
2540 by_index: dict[int, FleetDevice],
2541 devices: list[FleetDevice],
2542) -> FleetPlan:
2543 """Plan placement for *roles* (``None`` = all configured) and build their launches."""
2544 from lilbee.core.config import cfg
2546 unified_budget = _unified_memory_budget(devices)
2547 inputs, model_refs, reservation, skipped_not_installed = _server_model_inputs(
2548 roles,
2549 unified_budget=unified_budget,
2550 device_count=len(devices),
2551 total_vram=sum(d.total_bytes for d in devices),
2552 )
2553 spec = PlacementSpec.from_json(cfg.placement) if cfg.placement else None
2554 placement, _spec_applied = _placement_or_auto(
2555 spec,
2556 inputs,
2557 model_refs,
2558 devices,
2559 unified_budget=_unified_admission_budget(devices),
2560 # Only the clean-box snapshot's free bytes mean "what other tenants hold";
2561 # a live probe here would also be missing the fleet's own residency.
2562 charge_against_free=_plan_probe_store.get() is not None,
2563 )
2564 _log_placement_findings(placement, model_refs)
2565 reserved_by_device = _non_chat_reservation(placement.instances, inputs, placement.co_tenants)
2566 charged = {inp.role: inp.est_vram_bytes for inp in inputs}
2567 launches: list[InstanceLaunch] = []
2568 skipped_unusable_ctx: dict[WorkerRole, str] = {}
2569 for plan in placement.instances:
2570 launch = _launch_for(
2571 plan,
2572 model_refs[plan.role],
2573 binary,
2574 by_index,
2575 unified_budget=unified_budget,
2576 chat_reservation=reservation,
2577 reserved_by_device=reserved_by_device,
2578 est_vram_bytes=charged.get(plan.role, 0),
2579 )
2580 reason = _unusable_chat_ctx_reason(launch)
2581 if reason is not None:
2582 skipped_unusable_ctx[launch.role] = reason
2583 log.warning(reason)
2584 continue
2585 launches.append(launch)
2586 return FleetPlan(
2587 launches=tuple(launches),
2588 co_tenants=placement.co_tenants,
2589 skipped_not_installed=skipped_not_installed,
2590 skipped_unusable_ctx=skipped_unusable_ctx,
2591 )
2594def plan_all_launches() -> FleetPlan:
2595 """Apply GPU env, probe devices, and plan launches for every configured role.
2597 Disables crash-prone Vulkan layers / dual-vendor ICDs and applies any
2598 ``cfg.gpu_devices`` pin before the probe and plan (both inherit the env).
2599 """
2600 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
2601 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
2603 apply_fleet_gpu_env()
2604 binary = resolve_llama_server()
2605 # Put the CUDA-runtime wheels on the process path so the device probe sees the
2606 # same runtime the servers will, before resolve_devices enumerates GPUs.
2607 apply_cuda_runtime_env()
2608 devices = _plan_devices(binary)
2609 by_index = {d.index: d for d in devices}
2610 return plan_launches(None, binary, by_index, devices)