Coverage for src/lilbee/providers/fleet/planning.py: 100%
928 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +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 KV_CACHE_TYPE_BYTES, 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 (
28 engine_build_id,
29 llama_server_runtime_env,
30 resolve_llama_server,
31)
32from lilbee.providers.fleet.devices import (
33 VULKAN_BACKEND,
34 FleetDevice,
35 host_lacks_nvlink,
36 probe_devices,
37 visible_env,
38)
39from lilbee.providers.fleet.launch import InstanceLaunch
40from lilbee.providers.fleet.placement import (
41 InstancePlan,
42 ModelPlacementInput,
43 PeakEstimator,
44 Placement,
45 SplitCtxFitter,
46 placement_from_spec,
47 plan_placement,
48)
49from lilbee.providers.fleet.placement_spec import PlacementError, PlacementSpec
50from lilbee.providers.fleet.readback import supports_memory_readback
51from lilbee.providers.fleet.replicas import resolve_replica_count
52from lilbee.providers.fleet.vram import estimate_instance_footprint, usable_vram_fraction
53from lilbee.providers.model_cache import free_system_memory, total_system_memory
54from lilbee.providers.model_ref import parse_model_ref
55from lilbee.providers.roles import ROLE_REGISTRY, RerankMode, WorkerRole
57log = logging.getLogger(__name__)
59if TYPE_CHECKING:
60 from collections.abc import Callable, Iterable, Mapping, Sequence
62# Fleet-only concurrency: continuous-batching slots (--parallel) per server.
63_CHAT_SLOTS = 4
64# Stand-ins in the launch log for a host whose devices could not be read.
65_UNKNOWN_BACKEND = "unknown"
66_NO_DEVICES = "none"
69def log_engine_launch(launch: InstanceLaunch, *, owner_pid: int | None = None) -> None:
70 """Log the binary, build, backend and devices serving *launch*.
72 *owner_pid* is the engine's owner when this process adopted it rather than
73 spawned it.
74 """
75 if owner_pid is None:
76 log.info("Launched %s serving %s on %s", launch.model_id, launch.model, _engine_id(launch))
77 return
78 log.info(
79 "Adopted %s serving %s from engine pid %d on %s",
80 launch.model_id,
81 launch.model,
82 owner_pid,
83 _engine_id(launch),
84 )
87def _engine_id(launch: InstanceLaunch) -> str:
88 """*launch*'s binary with the engine build, backend, and probed devices."""
89 devices = probed_devices()
90 backend = next((device.backend for device in devices), _UNKNOWN_BACKEND)
91 names = ", ".join(f"{d.backend}{d.index}: {d.name}" for d in devices) or _NO_DEVICES
92 return f"{launch.binary} (build {engine_build_id()}, backend {backend}, devices: {names})"
95def warn_when_chat_downsized(launch: InstanceLaunch) -> None:
96 """Log when a chat engine's granted shape ends below the requested one.
98 Runs at engine adoption so the warning lands on the serving process's
99 own log, where the operator of that engine reads it.
100 """
101 below_ctx = launch.built_ctx_target > 0 and launch.ctx < launch.built_ctx_target
102 below_slots = launch.slots < _CHAT_SLOTS
103 if not (below_ctx or below_slots):
104 return
105 # A pre-field launch record carries built_ctx_target=0; show the granted
106 # window as the requested one rather than "x 0 context". The remedy names
107 # both target knobs because the recorded target does not say which of
108 # num_ctx (pin) or chat_n_ctx_target produced it.
109 requested_ctx = launch.built_ctx_target or launch.ctx
110 log.warning(
111 "Chat engine downsized: serving %d slot(s) x %d context "
112 "(requested: up to %d slots x %d context). Requests beyond %d at once "
113 "wait for a free slot, so parallel agents can look stalled. "
114 "A smaller model or a lower context target (num_ctx or "
115 "chat_n_ctx_target) frees room for more slots.",
116 launch.slots,
117 launch.ctx,
118 _CHAT_SLOTS,
119 requested_ctx,
120 launch.slots,
121 )
124# Slots the PLACEMENT estimate reserves KV for on a tensor-split chat: one full
125# window, the minimum any split must hold. The launch may serve more than this
126# (see _resolve_split_chat_slots) when the cards measurably have room for several
127# full windows, so this is a planning floor, not the served slot count.
128_SPLIT_CHAT_SLOTS = 1
129# Floor context the PLACEMENT estimate reserves KV for, so a large model is never
130# single-carded into a KV corner too small for real use (a 17GB model on a 24GB
131# card leaves ~no KV room -> n_ctx collapses to a few hundred tokens). Sizing the
132# placement reserve against this floor forces a tensor-split when one card cannot
133# hold weights + a usable context; the served ctx is then grown by resolve_chat_ctx
134# (single) / fit_split_ctx (split) toward the cards' real headroom, with each
135# sequence capped at the working-context target. A split may then serve several
136# such sequences, so the served total can exceed the single-window reserve; the
137# per-device headroom test in fit_split_ctx is what bounds it. A fit still at
138# its floor after the forced split is refused by plan_launches
139# (_unusable_chat_ctx_reason).
140_MIN_USABLE_CHAT_CTX = 8192
141# Offload-search upper bound meaning "no reduced offload is worth probing": the
142# search runs 0..upper, so a negative bound skips it.
143_NO_OFFLOAD_PROBE = -1
144# Embed and cross-encoder rerank serve one request at a time. Raising it was
145# tried and measured worse: on 8xA40 with an 8B Q8 embedder and one ~100-token
146# passage per request, --parallel 1 gave 133 docs/sec at 81% SM while
147# --parallel 8 gave 100 at 63%. At one slot the card is already busy, so there is
148# no stall for slot-batching to reclaim and the extra slots only add
149# continuous-batching and KV-fragmentation overhead. Batch on the request side
150# (embed_batch_sequences) instead.
151_AUX_SLOTS = 1
152# A tensor-split needs at least this many GPUs; below it the chat context objective
153# (a gguf read) is pointless because the model can only single-card or stay unplaced.
154_MIN_SPLIT_GPUS = 2
155# Pooled single-slot search roles (embed/cross-encoder rerank) whose whole input
156# batches in one pass; derived from the role registry.
157_EMBED_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.pooled)
158# Roles whose loaders offload every layer regardless of cfg.n_gpu_layers; only
159# chat honors cfg.n_gpu_layers.
160_ALL_LAYER_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.offload_all_layers)
161_FLASH_ON = "on"
162_FLASH_OFF = "off"
163_FLASH_AUTO = "auto"
164# llama-server's documented way to say "offload nothing": --device none.
165_NO_DEVICE = "none"
166# Backends pinned by the name the engine printed rather than through an env var,
167# because their variables index a different space than --list-devices reports.
168_NAME_PINNED_BACKENDS = frozenset({VULKAN_BACKEND, "SYCL"})
169# Backends whose flash-attention coverage in llama.cpp is complete enough to ask
170# for it outright. Vulkan and SYCL are behind CUDA's and have been incomplete on
171# Intel's mesa driver, so those are left to the engine's own auto, which enables
172# flash attention only where the backend really supports it.
173_TRUSTED_FLASH_BACKENDS = frozenset({"CUDA", "ROCm", "HIP", "MTL", "Metal"})
174# Roles to which flash attention applies; embed/rerank run without it.
175_FLASH_ROLES = tuple(role for role, info in ROLE_REGISTRY.items() if info.flash_attn)
178# Cap vision's own KV footprint at this fraction of usable VRAM when sizing its
179# batching slots, leaving room for the weights and any co-located role.
180_VISION_VRAM_FRACTION = 0.5
182# Cap an LLM reranker's footprint at this fraction of usable VRAM when sizing its
183# slots; its per-slot ctx is tiny, so a normal GPU fits the full fan-out and a
184# small one steps down toward 1.
185_LLM_RERANK_VRAM_FRACTION = 0.5
187# RAM kept free for the OS when placing against system memory (no discrete GPU):
188# a quarter of total RAM, capped at 4 GiB. A fixed 4 GiB floor leaves a small
189# host (7-8 GB) with no budget at all, refusing to serve even tiny models.
190_SYSTEM_MEMORY_FLOOR_DIVISOR = 4
191# A GPU driver still initializing at boot answers with no devices. Ask again
192# before letting that decide the daemon's whole run; two extra probes cost a
193# couple of seconds only on a host that has a card the engine could not see.
194_PROBE_RETRIES = 2
195_PROBE_RETRY_DELAY_S = 1.0
197# A network filesystem makes mmap dangerous (page faults served over the wire can
198# wedge the loader in uninterruptible I/O), so the chat server loads its weights
199# into a malloc'd host copy (--no-mmap) whenever that copy fits in this fraction
200# of total system RAM. Local disk keeps mmap: its lazy paging gives a faster first
201# token on a cold cache -- the common desktop first launch -- and --no-mmap's
202# buffered full read only wins on an already-hot cache (#474: 33s vs 43s for a
203# 112GB model on 3 GPUs) while pessimizing cold start. Keyed on TOTAL memory
204# (stable), not free (fluctuates), so replans do not flap the launch argv. The
205# exact ceiling is tuned on a network-volume host.
206_NO_MMAP_NETWORK_RAM_FRACTION = 0.85
208# llama.cpp split-GGUF shard naming ("%s-%05d-of-%05d.gguf"); the cold-load
209# timeout must scale with the SUM of the shards, not the first file alone.
210_SPLIT_GGUF_NAME = re.compile(r"^(?P<prefix>.+)-(?P<index>\d{5})-of-(?P<total>\d{5})\.gguf$")
213def _weights_bytes(model_path: Path) -> int:
214 """Total weights size on disk; a split GGUF sums every sibling shard."""
215 match = _SPLIT_GGUF_NAME.fullmatch(model_path.name)
216 if match is None:
217 return model_path.stat().st_size
218 return sum(
219 sibling.stat().st_size
220 for sibling in model_path.parent.iterdir()
221 if _is_sibling_shard(sibling.name, match)
222 )
225def _is_sibling_shard(name: str, match: re.Match[str]) -> bool:
226 """Whether *name* is a shard of the same split GGUF as *match*."""
227 shard = _SPLIT_GGUF_NAME.fullmatch(name)
228 return (
229 shard is not None
230 and shard["prefix"] == match["prefix"]
231 and shard["total"] == match["total"]
232 )
235def _slots_for(
236 role: WorkerRole,
237 model_path: Path,
238 ctx: int,
239 *,
240 mmproj_path: Path | None = None,
241 unified_budget: int | None = None,
242 chat_reservation: int = 0,
243 rerank_mode: RerankMode | None = None,
244 device: FleetDevice | None = None,
245) -> int:
246 """Continuous-batching slots (--parallel) for a role's server.
248 Chat batches concurrent turns; vision batches concurrent OCR pages since a
249 one-page decode underutilizes the GPU; an LLM reranker batches its per-candidate
250 chat requests; embed and cross-encoder rerank are single-slot (their batching is
251 request-side). The memory-aware roles drop toward 1 on a small or shared host
252 instead of overcommitting. ``unified_budget`` caps sizing against free system RAM
253 with no discrete GPU; ``chat_reservation`` is the search-role footprint held back
254 from chat; ``device`` is the card the role was placed on, whose memory the
255 budget comes from once placement has chosen one.
256 """
257 if role is WorkerRole.CHAT:
258 return _resolve_chat_slots(
259 model_path,
260 ctx,
261 mmproj_path=mmproj_path,
262 unified_budget=unified_budget,
263 chat_reservation=chat_reservation,
264 device=device,
265 )
266 if role is WorkerRole.VISION:
267 return _resolve_vision_slots(
268 model_path, ctx, mmproj_path=mmproj_path, unified_budget=unified_budget, device=device
269 )
270 if role is WorkerRole.RERANK and rerank_mode is RerankMode.LLM:
271 return _resolve_llm_rerank_slots(
272 model_path, ctx, unified_budget=unified_budget, device=device
273 )
274 return _AUX_SLOTS
277def _resolve_split_chat_slots(fit_fn: Callable[[int], int]) -> tuple[int, int]:
278 """Largest split-chat slot count whose sequences each keep the full window.
280 ``fit_fn(n)`` is the per-slot context that fits when serving ``n`` sequences
281 (``fit_split_ctx``, capped at the working target and verified against real
282 per-card headroom). More slots divide the KV, so a split whose cards hold
283 several full windows can serve that many agents concurrently instead of one.
284 Returns ``(slots, per_slot_ctx)``, falling to one slot when only one full
285 window fits (or the fit degenerated to the floor), which preserves the
286 max-context single-sequence behaviour on a tight card.
288 Found by bisection rather than a scan because every ``fit_fn`` call is a
289 complete binary search whose probes each shell out to gguf-parser, and the
290 whole thing runs while this process holds the cross-process build lock that
291 every other lilbee start waits on without a deadline. A descending scan paid
292 for all of ``_CHAT_SLOTS - 1`` searches in exactly the tight-card case where
293 none of them fit. Bisection is sound here because the fit is non-increasing
294 in the slot count: more sequences divide the same headroom, so once a count
295 fails no larger one can succeed.
296 """
297 full = fit_fn(1)
298 if full <= model_cache._DYNAMIC_CTX_FLOOR:
299 return 1, full
300 low, high = 1, _CHAT_SLOTS
301 while low < high:
302 mid = (low + high + 1) // 2
303 if fit_fn(mid) >= full:
304 low = mid
305 else:
306 high = mid - 1
307 return low, full
310def _resolve_chat_slots(
311 model_path: Path,
312 ctx: int,
313 *,
314 mmproj_path: Path | None = None,
315 unified_budget: int | None = None,
316 chat_reservation: int = 0,
317 device: FleetDevice | None = None,
318) -> int:
319 """Largest chat slot count (<= ``_CHAT_SLOTS``) whose footprint fits the budget
320 after reserving the search roles; steps to 1 when none fit.
322 The budget is the whole serve budget less the search roles' measured
323 footprint. ``cfg.gpu_memory_fraction`` is already the margin held back from
324 the card, and the room for co-located embed/rerank is ``chat_reservation``,
325 which is what those servers were sized at rather than a flat share. A second
326 fraction on top charged that room twice and left a fifth of the card unused
327 on a box with no search roles at all.
328 """
329 budget = _slot_budget(unified_budget, device) - chat_reservation
330 return _fit_slots(
331 _CHAT_SLOTS,
332 WorkerRole.CHAT,
333 model_path,
334 ctx,
335 mmproj_path=mmproj_path,
336 unified=unified_budget is not None,
337 budget=budget,
338 )
341def _resolve_vision_slots(
342 model_path: Path,
343 ctx: int,
344 *,
345 mmproj_path: Path | None = None,
346 unified_budget: int | None = None,
347 device: FleetDevice | None = None,
348) -> int:
349 """Largest OCR batching slot count (<= ``cfg.vision_ocr_concurrency``) that fits
350 the memory budget; 1 when the ceiling is 1 or nothing larger fits."""
351 from lilbee.core.config import cfg
353 ceiling = max(1, cfg.vision_ocr_concurrency)
354 if ceiling == 1:
355 return 1
356 return _fit_slots(
357 ceiling,
358 WorkerRole.VISION,
359 model_path,
360 ctx,
361 mmproj_path=mmproj_path,
362 unified=unified_budget is not None,
363 budget=_slot_budget(unified_budget, device, vram_fraction=_VISION_VRAM_FRACTION),
364 )
367def _resolve_llm_rerank_slots(
368 model_path: Path,
369 ctx: int,
370 *,
371 unified_budget: int | None = None,
372 device: FleetDevice | None = None,
373) -> int:
374 """Largest LLM-reranker slot count (<= ``LLM_RERANK_CONCURRENCY``) that fits the
375 memory budget; 1 when nothing larger fits. Matches the client's request fan-out."""
376 return _fit_slots(
377 LLM_RERANK_CONCURRENCY,
378 WorkerRole.RERANK,
379 model_path,
380 ctx,
381 mmproj_path=None,
382 unified=unified_budget is not None,
383 budget=_slot_budget(unified_budget, device, vram_fraction=_LLM_RERANK_VRAM_FRACTION),
384 rerank_mode=RerankMode.LLM,
385 )
388def _slot_budget(
389 unified_budget: int | None,
390 device: FleetDevice | None = None,
391 *,
392 vram_fraction: float = 1.0,
393) -> int:
394 """Memory budget for slot sizing: the usable memory on *device* (the fleet's
395 smallest when placement has not chosen one yet), capped by ``unified_budget``
396 (free system RAM) when there is no discrete GPU so the count steps down to fit
397 free memory instead of overcommitting.
399 *vram_fraction* takes less than the whole for a role that shares its card by
400 design. Chat takes all of it and subtracts what the search roles were sized
401 at, which is the same room stated as a measurement rather than a share."""
402 budget = int(plan_sizing_budget(device) * vram_fraction)
403 if unified_budget is not None:
404 budget = min(budget, unified_budget)
405 return budget
408def _fit_slots(
409 ceiling: int,
410 role: WorkerRole,
411 model_path: Path,
412 ctx: int,
413 *,
414 mmproj_path: Path | None,
415 unified: bool,
416 budget: int,
417 rerank_mode: RerankMode | None = None,
418) -> int:
419 """Largest slot count in ``1..ceiling`` whose instance footprint fits *budget*;
420 1 when none larger fit."""
421 from lilbee.providers.base import ProviderError
423 for slots in range(ceiling, 1, -1):
424 try:
425 est = estimate_instance_footprint(
426 model_path,
427 ctx=ctx,
428 slots=slots,
429 gpu_layers=_role_gpu_layers(role),
430 flash_attn=_role_flash(role, rerank_mode),
431 kv_cache_type=_role_kv_cache_type(role),
432 kv_cache_type_v=_role_kv_cache_type_v(role),
433 mmproj_path=mmproj_path,
434 expert_offload=_role_expert_offload(model_path),
435 )
436 except (ProviderError, OSError):
437 # An unsizable model runs a single slot; the load decides the rest.
438 return 1
439 if est.footprint(unified=unified) <= budget:
440 return slots
441 return 1
444def fit_chat_ctx(
445 model_path: Path,
446 meta: dict[str, str] | None,
447 *,
448 available_bytes: int,
449 ctx_ceiling: int,
450) -> engine_params.ChatFit:
451 """The offload and chat window gguf-parser says *available_bytes* backs.
453 Sizes one slot at the configured offload, as :func:`_slots_for` then grows
454 the slot count against what the window leaves. A window that cannot hold a
455 grounded prompt buys KV room by leaving layers in system memory, and the
456 largest offload whose window reaches ``min_usable_chat_ctx`` wins; the
457 alternative for those models is no service at all
458 (:func:`_unusable_chat_ctx_reason`). Raises when the estimator cannot
459 answer, which sends :func:`engine_params.resolve_chat_fit` to its
460 header-math fallback.
461 """
463 def fit_at(gpu_layers: int) -> int:
464 return fleet_ctx.fit_single_ctx(
465 model_path,
466 meta=meta,
467 slots=1,
468 available_bytes=available_bytes,
469 gpu_layers=gpu_layers,
470 flash_attn=_role_flash(WorkerRole.CHAT),
471 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
472 kv_cache_type_v=_role_kv_cache_type_v(WorkerRole.CHAT),
473 unified=plan_sizing_is_unified(),
474 ctx_ceiling=ctx_ceiling,
475 expert_offload=_role_expert_offload(model_path),
476 )
478 configured = _role_gpu_layers(WorkerRole.CHAT)
479 fit = engine_params.ChatFit(configured, fit_at(configured))
480 needed = engine_params.min_usable_chat_ctx()
481 if fit.ctx >= needed or not _chat_offload_is_tradable(
482 model_path,
483 meta=meta,
484 available_bytes=available_bytes,
485 ctx_ceiling=ctx_ceiling,
486 needed=needed,
487 ):
488 return fit
489 traded = _traded_chat_fit(
490 fit_at, upper=_chat_offload_probe_ceiling(meta, configured), needed=needed
491 )
492 return traded or fit
495# Architectures whose head scores every position it proposes. Their compute
496# buffer grows with the window, and the estimator prices it by the physical
497# batch instead, so the estimate falls further behind the longer the window
498# gets. The trade maximises exactly that window, which turns a fixed
499# under-estimate into an unbounded one: measured on a 1 GiB card, the trade
500# granted gemma-4-26B-A4B-eagle3 59648 tokens against a real 1413 MiB.
501_NON_TRADABLE_ARCHITECTURES = frozenset({"eagle3"})
504def _chat_offload_is_tradable(
505 model_path: Path,
506 *,
507 meta: dict[str, str] | None,
508 available_bytes: int,
509 ctx_ceiling: int,
510 needed: int,
511) -> bool:
512 """Whether a too-small chat window may buy KV room by moving layers off the card.
514 Only while the weights alone fit *available_bytes*. Past that the layers the
515 trade leaves behind are served from system memory, and a model that answers
516 from host RAM is the unusable load :func:`_unusable_chat_ctx_reason` refuses
517 on purpose. A window the user capped below *needed* is their choice, and the
518 refusal already honours it, so nothing is bought by trading for it either.
520 Never for a speculator head. The estimate the search reads is not sound over
521 the window for those, so the search optimises a number that grows wrong.
522 """
523 if (meta or {}).get("architecture") in _NON_TRADABLE_ARCHITECTURES:
524 return False
525 return ctx_ceiling >= needed and _weights_bytes(model_path) <= available_bytes
528def _chat_offload_probe_ceiling(meta: dict[str, str] | None, configured: int) -> int:
529 """Highest layer count the offload search probes, below *configured*.
531 ``_NO_OFFLOAD_PROBE`` when the architecture does not report a layer count,
532 which leaves no grid to search.
533 """
534 try:
535 layers = int((meta or {})["block_count"])
536 except (KeyError, ValueError):
537 return _NO_OFFLOAD_PROBE
538 if configured == engine_params.N_GPU_LAYERS_AUTO:
539 return layers
540 return min(configured - 1, layers)
543def _traded_chat_fit(
544 fit_at: Callable[[int], int], *, upper: int, needed: int
545) -> engine_params.ChatFit | None:
546 """Largest offload at or below *upper* whose window reaches *needed*, or ``None``.
548 Binary search: the estimate charges less VRAM the fewer layers the card
549 holds, so the fitted window is monotone in the offload.
550 """
551 lo, hi, best = 0, upper, None
552 while lo <= hi:
553 mid = (lo + hi) // 2
554 ctx = fit_at(mid)
555 if ctx >= needed:
556 best, lo = engine_params.ChatFit(mid, ctx), mid + 1
557 else:
558 hi = mid - 1
559 return best
562def _role_ctx(
563 role: WorkerRole,
564 model_path: Path,
565 meta: dict[str, str] | None,
566 device: FleetDevice | None = None,
567) -> int:
568 """Per-slot context for a role, derived as the in-process loader does.
570 Embed/rerank use the embedding model's training context; vision uses the
571 vision loader's training-context picker; chat honors ``cfg.num_ctx`` then
572 falls back to the single-GPU dynamic chat-ctx picker, sized against *device*
573 once placement has chosen one. A tensor-split chat is sized against its
574 per-device headroom instead (see :func:`fit_split_ctx`).
575 """
576 from lilbee.core.config import cfg
578 if role is WorkerRole.EMBED:
579 return engine_params.resolve_embed_ctx(meta, model_path)
580 if role is WorkerRole.RERANK:
581 if _rerank_mode_for(meta) is RerankMode.LLM:
582 return engine_params.resolve_llm_rerank_ctx(meta, model_path)
583 return engine_params.resolve_embed_ctx(meta, model_path)
584 if role is WorkerRole.VISION:
585 return engine_params.resolve_vision_ctx(model_path)
586 if cfg.num_ctx is not None:
587 return _pinned_chat_ctx(model_path, meta)
588 return engine_params.resolve_chat_ctx(
589 model_path, meta, available_bytes=plan_sizing_budget(device)
590 )
593def _pinned_chat_ctx(model_path: Path, meta: dict[str, str] | None) -> int:
594 """``cfg.num_ctx``, clamped to what the model was trained for.
596 Every unpinned resolver already clamps, and both docstrings here claimed the
597 pin did too. It did not, so a pin past the trained window was passed straight
598 to the engine, which clamps it silently and serves a different number than
599 every budget was sized for.
601 Only against a window that is actually known. A GGUF whose header cannot be
602 read falls back to a default that is a guess, and contradicting an explicit
603 pin with a guess would break the hosts where the header is the thing that is
604 broken.
605 """
606 from lilbee.core.config import cfg
608 pinned = cfg.num_ctx
609 assert pinned is not None # noqa: S101 - callers check; this documents the contract
610 ceiling = _known_chat_ceiling(model_path, meta)
611 if ceiling is None or pinned <= ceiling:
612 return pinned
613 log.warning(
614 "num_ctx is set to %d but %s was trained for %d, so %d is what will be served. "
615 "Lower num_ctx to stop planning against a window this model does not have.",
616 pinned,
617 model_path.name,
618 ceiling,
619 ceiling,
620 )
621 return ceiling
624def _known_chat_ceiling(model_path: Path, meta: dict[str, str] | None) -> int | None:
625 """The largest chat window this model is known to support, or ``None``.
627 ``None`` when the GGUF header gave no usable context length and the user set
628 no ``cfg.num_ctx_max``: there is then no measured ceiling, only a default.
629 """
630 from lilbee.core.config import cfg
631 from lilbee.providers.gguf_meta import train_ctx_from_meta
633 sentinel = -1
634 trained = train_ctx_from_meta(meta, fallback=sentinel, model_path=model_path)
635 known = [value for value in (trained, cfg.num_ctx_max) if value is not None and value > 0]
636 return min(known) if known else None
639def _rerank_mode_for(meta: dict[str, str] | None) -> RerankMode:
640 """Resolve the RERANK serving mode from cfg + the reranker GGUF arch."""
641 from lilbee.core.config import cfg
643 arch = meta.get("architecture") if meta else None
644 return resolve_rerank_mode(cfg.reranker_type, arch)
647def _role_rerank_mode(role: WorkerRole, meta: dict[str, str] | None) -> RerankMode | None:
648 """The RERANK serving mode for *role*, or ``None`` for every other role."""
649 return _rerank_mode_for(meta) if role is WorkerRole.RERANK else None
652def _server_spec(
653 role: WorkerRole, rerank_mode: RerankMode | None, meta: dict[str, str] | None
654) -> RoleServerSpec:
655 """The llama-server spec for a launch: rerank mode, decoder-aware embed pooling,
656 or the role default. EMBED forces ``--pooling last`` for decoder-only archs."""
657 if rerank_mode is not None:
658 return rerank_spec(rerank_mode)
659 if role is WorkerRole.EMBED:
660 return embed_spec(meta)
661 return ROLE_SPECS[role]
664def _pooled_batch_size(role: WorkerRole, rerank_mode: RerankMode | None, ctx: int) -> int | None:
665 """The ``--batch-size``/``--ubatch-size`` the launch raises for pooled
666 embed/cross-encoder rerank (the full context), or ``None`` for other roles."""
667 if role in _EMBED_ROLES and rerank_mode is not RerankMode.LLM:
668 return ctx
669 return None
672def _role_gpu_layers(role: WorkerRole) -> int:
673 """GPU-layer offload: chat honors ``cfg.n_gpu_layers``, others offload all layers."""
675 return engine_params.resolve_n_gpu_layers(embedding=role in _ALL_LAYER_ROLES)
678def _flash_enabled() -> bool:
679 """Flash attention is on unless ``cfg.flash_attention`` is explicitly ``False``."""
680 from lilbee.core.config import cfg
682 return cfg.flash_attention is not False
685def probed_devices() -> tuple[FleetDevice, ...]:
686 """Devices the engine enumerated, empty when they could not be read.
688 Prefers the plan snapshot so a whole planning pass answers consistently, and
689 falls back to the short-TTL read cache rather than a fresh probe.
690 """
691 probe = _plan_probe_store.get()
692 if probe is not None:
693 return probe.devices
694 try:
695 return tuple(_read_device_cache.get(resolve_llama_server()))
696 except (ProviderError, OSError):
697 return ()
700def _fleet_backend() -> str | None:
701 """The engine backend this host plans onto, or ``None`` when unknown."""
702 return next((device.backend for device in probed_devices()), None)
705def _flash_attention_is_trusted() -> bool:
706 """Whether to ask for flash attention outright rather than let the engine decide.
708 Unknown backends answer yes, which keeps every host that works today on the
709 argv it has now; only the backends known to lag get the engine's own auto.
710 """
711 backend = _fleet_backend()
712 return backend is None or backend in _TRUSTED_FLASH_BACKENDS
715def flash_attn_flag() -> str:
716 """``--flash-attn`` argv value for chat and vision."""
717 if not _flash_enabled():
718 return _FLASH_OFF
719 return _FLASH_ON if _flash_attention_is_trusted() else _FLASH_AUTO
722def _role_launches_with_flash(role: WorkerRole, rerank_mode: RerankMode | None = None) -> bool:
723 """Whether the launch asks the engine for flash attention on *role*.
725 The one place that answers this. The registry marks RERANK as a non-flash
726 role because a cross-encoder pools in one batch, but an LLM reranker is
727 generative and launches exactly like chat, so the mode decides there.
728 """
729 if role is WorkerRole.RERANK:
730 return rerank_mode is RerankMode.LLM
731 return role in _FLASH_ROLES
734def _role_flash(role: WorkerRole, rerank_mode: RerankMode | None = None) -> bool:
735 """Whether the estimate may assume flash attention for *role*.
737 The launch's own answer, narrowed to a definite ``on``. Under ``auto`` the
738 engine decides at load time, and assuming it would size the KV cache below
739 what the launch may need.
740 """
741 return _role_launches_with_flash(role, rerank_mode) and flash_attn_flag() == _FLASH_ON
744def _role_kv_cache_type(role: WorkerRole) -> KvCacheType:
745 """Chat honors ``cfg.kv_cache_type``; embed/rerank/vision run f16 KV."""
746 from lilbee.core.config import cfg
748 return cfg.kv_cache_type if role is WorkerRole.CHAT else KvCacheType.F16
751def _replica_count(role: WorkerRole, device_count: int) -> int:
752 """Requested data-parallel instances for *role* via the shared resolver."""
753 return resolve_replica_count(role, device_count)
756def _role_kv_cache_type_v(role: WorkerRole) -> KvCacheType:
757 """The V cache type for *role*: the configured one only when flash attention is on.
759 llama.cpp refuses a quantized V cache without flash attention ("V cache
760 quantization requires flash_attn") and the server never starts, while a
761 quantized K cache needs nothing. So V follows the setting only where flash
762 attention is certain, and is f16 under ``auto`` or ``off``. That costs memory
763 rather than a launch, and the estimate moves with it.
764 """
765 from lilbee.core.config.enums import KvCacheType
767 configured = _role_kv_cache_type(role)
768 return configured if flash_attn_flag() == _FLASH_ON else KvCacheType.F16
771def chat_cache_type_flags() -> tuple[str | None, str | None]:
772 """``(--cache-type-k, --cache-type-v)`` for chat; ``None`` leaves the f16 default."""
773 from lilbee.core.config.enums import KvCacheType
775 def flag(kind: KvCacheType) -> str | None:
776 return None if kind is KvCacheType.F16 else kind.value
778 return flag(_role_kv_cache_type(WorkerRole.CHAT)), flag(_role_kv_cache_type_v(WorkerRole.CHAT))
781def _vision_mmproj(model_ref: str) -> Path | None:
782 """Resolve a vision model's mmproj sidecar, or ``None`` if absent."""
783 from lilbee.providers.base import ProviderError
784 from lilbee.providers.gguf_meta import find_mmproj_for_model
786 try:
787 return find_mmproj_for_model(engine_params.resolve_model_path(model_ref))
788 except (ProviderError, OSError, ValueError, KeyError):
789 return None
792def _estimate_role(
793 role: WorkerRole,
794 model_ref: str,
795 *,
796 slots: int | None = None,
797 unified_budget: int | None = None,
798 chat_reservation: int = 0,
799 device_count: int = 0,
800) -> ModelPlacementInput:
801 """Estimate one role-model's footprint via gguf-parser (+ mmproj for vision).
803 ``slots`` defaults to the role's resolved batching slots (chat and vision are
804 memory-aware); ``chat_reservation`` shrinks chat to leave room for the search
805 roles; ``device_count`` resolves an auto (0) replica knob to one per GPU.
806 Charges the unified footprint with no discrete GPU, else the VRAM one.
807 """
808 from lilbee.providers.gguf_meta import read_gguf_metadata
810 path = engine_params.resolve_model_path(model_ref)
811 mmproj = _vision_mmproj(model_ref) if role is WorkerRole.VISION else None
812 meta = read_gguf_metadata(path)
813 # Size the single-instance footprint against the placement reserve (a usable
814 # KV floor for chat), so a model that fits weights-only on one card but not
815 # weights + a usable context falls through to a tensor-split instead of being
816 # single-carded into a tiny n_ctx. Non-chat roles keep their launch ctx.
817 ctx = _placement_estimate_ctx(role, path, meta)
818 rerank_mode = _role_rerank_mode(role, meta)
819 if slots is None:
820 slots = _slots_for(
821 role,
822 path,
823 ctx,
824 mmproj_path=mmproj,
825 unified_budget=unified_budget,
826 chat_reservation=chat_reservation,
827 rerank_mode=rerank_mode,
828 )
829 est = estimate_instance_footprint(
830 path,
831 ctx=ctx,
832 slots=slots,
833 gpu_layers=_role_gpu_layers(role),
834 flash_attn=_role_flash(role, rerank_mode),
835 kv_cache_type=_role_kv_cache_type(role),
836 kv_cache_type_v=_role_kv_cache_type_v(role),
837 mmproj_path=mmproj,
838 batch_size=_pooled_batch_size(role, rerank_mode, ctx),
839 expert_offload=_role_expert_offload(path),
840 )
841 fp = est.footprint(unified=unified_budget is not None)
842 if role is WorkerRole.CHAT and unified_budget is None:
843 fp = _chat_serve_budget_footprint(fp)
844 return ModelPlacementInput(
845 role=role,
846 est_vram_bytes=fp,
847 replicas=_replica_count(role, device_count),
848 est_ram_bytes=est.ram_bytes,
849 )
852def _chat_serve_budget_footprint(footprint: int) -> int:
853 """Charge a chat instance against the serve budget, not the placement headroom.
855 The planner fits instances within ``cfg.usable_vram_fraction`` of a card, but a
856 single-card chat then sizes its KV cache against the smaller
857 ``cfg.gpu_memory_fraction`` budget (``resolve_chat_ctx``). A model that fills a
858 card at 0.9 leaves no room for KV at 0.75 and collapses to a few hundred tokens,
859 so scale its placement footprint by the budget ratio: it then needs a
860 tensor-split (pooling VRAM across cards) whenever single-carding it would starve
861 its context. Small models are unaffected -- they fit the serve budget with KV
862 room to spare.
863 """
864 from lilbee.core.config import cfg
866 # Never below 1.0. The ratio only compensates while the serve budget is the
867 # smaller of the two; a gpu_memory_fraction raised past the usable fraction
868 # inverts it, and the same line that exists to charge chat more starts
869 # charging it less than the model takes.
870 return int(footprint * max(1.0, usable_vram_fraction() / cfg.gpu_memory_fraction))
873def _placement_estimate_ctx(role: WorkerRole, model_path: Path, meta: dict[str, str] | None) -> int:
874 """Per-slot context the placement estimate sizes a role against.
876 For chat this reserves KV for a usable floor (``_MIN_USABLE_CHAT_CTX``, or the
877 user's ``cfg.num_ctx`` pin), capped by the model's trained ceiling -- not the
878 single-GPU dynamic ctx (which shrinks to fit one card and then confirms a
879 single-card placement) nor the full trained ceiling (which over-reserves). A
880 model that cannot hold weights + this floor on one card is tensor-split.
881 """
882 from lilbee.core.config import cfg
884 if role is WorkerRole.CHAT:
885 if cfg.num_ctx is not None:
886 return _pinned_chat_ctx(model_path, meta)
887 return apply_ctx_downshift(
888 role,
889 min(
890 engine_params.chat_ctx_ceiling(meta, model_path),
891 max(cfg.chat_n_ctx_target, _MIN_USABLE_CHAT_CTX),
892 ),
893 )
894 return apply_ctx_downshift(role, _role_ctx(role, model_path, meta))
897def _placement_estimate_slots(role: WorkerRole, meta: dict[str, str] | None) -> int:
898 """The slot count the placement estimate reserves KV for.
900 A tensor-split chat reserves one full-context sequence here: a conservative
901 floor for the card-count decision. The launch then fills the placed cards'
902 real headroom with as many full-context slots as fit (``_resolve_split_chat_slots``),
903 never exceeding what those cards hold, so a larger launch count can't OOM.
904 """
905 from lilbee.core.config import cfg
907 if role is WorkerRole.CHAT:
908 return _SPLIT_CHAT_SLOTS
909 if role is WorkerRole.VISION:
910 return max(1, cfg.vision_ocr_concurrency)
911 if role is WorkerRole.RERANK and _rerank_mode_for(meta) is RerankMode.LLM:
912 return LLM_RERANK_CONCURRENCY
913 return _AUX_SLOTS
916def _peak_estimator(model_refs: dict[WorkerRole, str]) -> PeakEstimator:
917 """Per-device VRAM-vector estimator for the planner, bound to the configured models.
919 Estimates each role at its launch ceiling (ctx x slots) with the candidate
920 tensor-split ratio, so the planner reserves enough cards for the busiest one.
921 """
922 from lilbee.providers.gguf_meta import read_gguf_metadata
924 def estimate_peak(role: WorkerRole, ratio: tuple[int, ...]) -> tuple[int, ...]:
925 path = engine_params.resolve_model_path(model_refs[role])
926 meta = read_gguf_metadata(path)
927 mmproj = _vision_mmproj(model_refs[role]) if role is WorkerRole.VISION else None
928 slots = _placement_estimate_slots(role, meta)
929 ctx = _placement_estimate_ctx(role, path, meta)
930 rerank_mode = _role_rerank_mode(role, meta)
931 est = estimate_instance_footprint(
932 path,
933 ctx=ctx,
934 slots=slots,
935 gpu_layers=_role_gpu_layers(role),
936 flash_attn=_role_flash(role, rerank_mode),
937 kv_cache_type=_role_kv_cache_type(role),
938 kv_cache_type_v=_role_kv_cache_type_v(role),
939 mmproj_path=mmproj,
940 tensor_split=ratio,
941 batch_size=_pooled_batch_size(role, rerank_mode, ctx),
942 expert_offload=_role_expert_offload(path),
943 )
944 return est.per_device_vram
946 return estimate_peak
949def _chat_split_ctx_objective(
950 model_refs: dict[WorkerRole, str],
951) -> tuple[SplitCtxFitter | None, int]:
952 """The chat split's context fitter and target, or ``(None, 0)`` with no chat model.
954 The fitter sizes a candidate shard's served context exactly as the launch does
955 (:func:`fit_split_ctx`), so the planner widens chat onto idle cards only when a
956 tighter shard would starve KV below the target. See docs/architecture.md.
957 """
958 if WorkerRole.CHAT not in model_refs:
959 return None, 0
960 from lilbee.providers.gguf_meta import read_gguf_metadata
962 path = engine_params.resolve_model_path(model_refs[WorkerRole.CHAT])
963 meta = read_gguf_metadata(path)
964 target = _placement_estimate_ctx(WorkerRole.CHAT, path, meta)
966 def fit(ratio: tuple[int, ...], per_device_free_bytes: Sequence[int]) -> int:
967 return fleet_ctx.fit_split_ctx(
968 path,
969 meta=meta,
970 slots=_SPLIT_CHAT_SLOTS,
971 ratio=ratio,
972 per_device_free_bytes=per_device_free_bytes,
973 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
974 flash_attn=_role_flash(WorkerRole.CHAT),
975 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
976 kv_cache_type_v=_role_kv_cache_type_v(WorkerRole.CHAT),
977 ctx_ceiling=target,
978 expert_offload=_role_expert_offload(path),
979 )
981 return fit, target
984def _search_reservation(inputs: dict[WorkerRole, ModelPlacementInput]) -> int:
985 """Total footprint of the placed search roles (all replicas), held back ahead
986 of chat."""
987 return sum(
988 inputs[role].est_vram_bytes * inputs[role].replicas
989 for role in _EMBED_ROLES
990 if role in inputs
991 )
994def _role_weights_bytes(role: WorkerRole, ref: str) -> int:
995 """The model's weight bytes on disk (plus the mmproj for vision): a
996 ground-truth lower bound on residency. 0 when the file cannot be resolved."""
997 from lilbee.providers.base import ProviderError
999 try:
1000 size = _weights_bytes(engine_params.resolve_model_path(ref))
1001 if role is WorkerRole.VISION:
1002 mmproj = _vision_mmproj(ref)
1003 if mmproj is not None:
1004 size += int(mmproj.stat().st_size)
1005 except (ProviderError, OSError):
1006 return 0
1007 return size
1010def _is_moe(meta: dict[str, str] | None) -> bool:
1011 """Whether the GGUF declares routed experts, so its experts can be offloaded."""
1012 count = (meta or {}).get("expert_count")
1013 try:
1014 return int(count) > 0 if count is not None else False
1015 except ValueError:
1016 return False
1019def expert_offload_all(meta: dict[str, str] | None) -> bool:
1020 """Whether to keep every layer's experts in system memory; MoE models only."""
1021 from lilbee.core.config import cfg
1023 return bool(cfg.cpu_moe) and _is_moe(meta)
1026def expert_offload_layers(meta: dict[str, str] | None) -> int | None:
1027 """How many layers' experts to keep in system memory, or None for no split.
1029 A non-positive ``n_cpu_moe`` offloads nothing (it would emit a no-op
1030 ``--n-cpu-moe 0``), so it reads as unset.
1031 """
1032 from lilbee.core.config import cfg
1034 if cfg.n_cpu_moe is None or cfg.n_cpu_moe < 1 or not _is_moe(meta):
1035 return None
1036 return cfg.n_cpu_moe
1039def _role_expert_offload(model_path: Path) -> tuple[str, ...]:
1040 """Expert patterns the launch will offload, for sizing the same way it runs.
1042 Reads the GGUF (cached) rather than taking metadata as an argument so every
1043 estimate site charges the same tensors the launch moves off the GPU.
1044 """
1045 from lilbee.providers.fleet.adapters import expert_offload_patterns
1046 from lilbee.providers.gguf_meta import read_gguf_metadata
1048 meta = read_gguf_metadata(model_path)
1049 return expert_offload_patterns(
1050 cpu_moe=expert_offload_all(meta), n_cpu_moe=expert_offload_layers(meta)
1051 )
1054def _expert_offload_configured() -> bool:
1055 """Whether the user asked for expert offload that would actually take effect.
1057 A non-positive ``n_cpu_moe`` offloads nothing, so it does not count.
1058 """
1059 from lilbee.core.config import cfg
1061 return bool(cfg.cpu_moe) or (cfg.n_cpu_moe is not None and cfg.n_cpu_moe >= 1)
1064def _weights_exceed_everything(size: int, *, total_vram: int, total_ram: int) -> bool:
1065 """True when a model's weights fit neither the GPUs nor system memory.
1067 File size is ground truth, not an estimate, so this bound cannot repeat the
1068 false-refusal class: no estimator error makes a 40 GiB file fit a 1 GiB box.
1069 Past both pools there is nowhere for a layer to go and no launch can win, so
1070 saying so beats a load that thrashes and then dies.
1071 """
1072 ceiling = total_vram + total_ram
1073 return ceiling > 0 and size > ceiling
1076def _weights_exceed_hardware(size: int, total_vram: int, *, is_moe: bool) -> bool:
1077 """True when this model cannot be served on this machine at all.
1079 Exceeding VRAM alone is not that. The engine chooses how many layers fit and
1080 keeps the rest in system memory, so a model larger than every card is a
1081 partial offload and lilbee's job is to launch it and say what will happen.
1082 Refusing there meant the fit never ran and the role was skipped, which left
1083 the user hand-tuning n_gpu_layers to get back what the engine does by itself.
1085 What still refuses is a model past VRAM and system memory together, where no
1086 arrangement of layers exists. A user-set n_gpu_layers or expert offload keeps
1087 standing the bound down entirely, since the user has said where the weights
1088 should go.
1089 """
1090 from lilbee.core.config import cfg
1092 if cfg.n_gpu_layers is not None:
1093 return False
1094 if is_moe and _expert_offload_configured():
1095 return False
1096 return _weights_exceed_everything(
1097 size, total_vram=total_vram, total_ram=model_cache.total_system_memory()
1098 )
1101def _vision_without_mmproj(role: WorkerRole, ref: str) -> bool:
1102 """True (with a warning) for a configured vision model whose mmproj is missing.
1104 The skip would silently disable OCR; the warning names the cause and the fix.
1105 """
1106 if role is not WorkerRole.VISION or _vision_mmproj(ref) is not None:
1107 return False
1108 log.warning(
1109 "Vision model %s has no mmproj (CLIP projector); OCR is disabled. "
1110 "Re-run 'lilbee model pull %s' to fetch the projector.",
1111 ref,
1112 ref,
1113 )
1114 return True
1117def _estimate_or_fallback(
1118 role: WorkerRole,
1119 ref: str,
1120 *,
1121 unified_budget: int | None,
1122 chat_reservation: int,
1123 device_count: int,
1124 total_vram: int,
1125 skipped_not_installed: dict[WorkerRole, str],
1126 host_committed: int = 0,
1127) -> ModelPlacementInput | None:
1128 """Size *role* for placement, degrading rather than refusing.
1130 A missing model is skipped and recorded; a sizing failure on an installed
1131 model falls back to the analytic floor; weights alone exceeding the physical
1132 VRAM refuse with a plain message (ground truth, not an estimate).
1133 """
1134 from lilbee.providers.base import ProviderError, ProviderErrorKind
1136 try:
1137 estimate = _estimate_role(
1138 role,
1139 ref,
1140 unified_budget=unified_budget,
1141 chat_reservation=chat_reservation,
1142 device_count=device_count,
1143 )
1144 except (ProviderError, OSError) as exc:
1145 if isinstance(exc, ProviderError) and exc.kind is ProviderErrorKind.NOT_FOUND:
1146 log.warning("Skipping %s server: model %r is not installed.", role.value, ref)
1147 skipped_not_installed[role] = ref
1148 return None
1149 return _sizing_failure_fallback(
1150 role,
1151 ref,
1152 exc,
1153 device_count=device_count,
1154 total_vram=total_vram,
1155 host_committed=host_committed,
1156 )
1157 return _admit_estimate(
1158 _floor_implausible_estimate(estimate, role, ref),
1159 role,
1160 ref,
1161 total_vram=total_vram,
1162 ram_bytes=estimate.est_ram_bytes,
1163 host_committed=host_committed,
1164 )
1167def _admit_estimate(
1168 estimate: ModelPlacementInput,
1169 role: WorkerRole,
1170 ref: str,
1171 *,
1172 total_vram: int,
1173 ram_bytes: int,
1174 host_committed: int = 0,
1175) -> ModelPlacementInput | None:
1176 """*estimate*, or ``None`` when this model cannot load on this machine.
1178 Two hardware bounds, one per kind of memory: the weights must fit the GPUs
1179 unless something offloads, and whatever offloading puts in system memory must
1180 fit the system.
1181 """
1182 weights = _role_weights_bytes(role, ref)
1183 if _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref)):
1184 _warn_weights_exceed(role, ref, weights, total_vram)
1185 return None
1186 if total_vram > 0 and weights > total_vram:
1187 _warn_weights_spill(role, ref, weights, total_vram)
1188 if _host_memory_refuses(role, ref, ram_bytes, host_committed):
1189 return None
1190 return estimate
1193def _analytic_footprint_floor(
1194 weights: int, *, role: WorkerRole, meta: dict[str, str] | None, ctx: int, slots: int
1195) -> int:
1196 """The least this instance can occupy: weights, its KV cache, and overhead.
1198 Used when the estimator cannot answer. Charging weight bytes alone was a
1199 knowing under-charge: the engine allocates a KV cache sized by context and
1200 slot count, plus compute buffers, and omitting all of it lets placement fit a
1201 model that cannot fit. The comment said the load would decide, and it did, by
1202 running out of memory.
1204 A floor rather than an estimate. It is derived from the header the same way
1205 the in-process sizing path derives it, and it is deliberately the smallest
1206 defensible number, because refusing a model that would have fit is its own
1207 failure. Without a readable header the per-token fallback still applies:
1208 zero is the one answer that is certainly wrong.
1209 """
1211 kv_bytes = (
1212 model_cache.kv_bytes_per_token(
1213 meta,
1214 KV_CACHE_TYPE_BYTES[_role_kv_cache_type(role)],
1215 KV_CACHE_TYPE_BYTES[_role_kv_cache_type_v(role)],
1216 )
1217 * ctx
1218 * max(slots, 1)
1219 )
1220 overhead = int(weights * model_cache._BUFFER_OVERHEAD_FRACTION)
1221 return weights + kv_bytes + overhead
1224def _estimate_is_implausible(*, estimated: int, floor: int) -> bool:
1225 """Whether *estimated* describes a load that cannot exist.
1227 Below the bytes the card must hold there is no arrangement of memory that
1228 serves the model. A floor of zero means nothing could be computed to compare
1229 against, and a guess is not grounds to discard the only measurement there is.
1230 """
1231 return floor > 0 and 0 < estimated < floor
1234def _floor_implausible_estimate(
1235 estimate: ModelPlacementInput, role: WorkerRole, ref: str
1236) -> ModelPlacementInput:
1237 """*estimate*, or the model's weight bytes when it reports less than those.
1239 The bound is the weights alone, not the analytic footprint. That figure
1240 sizes a KV cache as though every layer ran dense attention over the whole
1241 window, which linear-attention, sliding-window and MLA models do not, so it
1242 sits far above what they hold. Weights are architecture-independent, which
1243 makes an estimate under them the one answer the planner can call impossible
1244 without modelling attention.
1246 Offload lifts the bound: the layers in system memory are weight bytes the
1247 card never holds.
1248 """
1249 if _cpu_offload_in_play():
1250 return estimate
1251 weights = _role_weights_bytes(role, ref)
1252 if not _estimate_is_implausible(estimated=estimate.est_vram_bytes, floor=weights):
1253 return estimate
1254 log.warning(
1255 "The estimator sized the %s model %s at %.1f GiB, below the %.1f GiB of "
1256 "weights the card has to hold. Charging the weights instead.",
1257 role.value,
1258 ref,
1259 estimate.est_vram_bytes / 1024**3,
1260 weights / 1024**3,
1261 )
1262 return replace(estimate, est_vram_bytes=weights)
1265def _sizing_failure_fallback(
1266 role: WorkerRole,
1267 ref: str,
1268 exc: Exception,
1269 *,
1270 device_count: int,
1271 total_vram: int,
1272 host_committed: int = 0,
1273) -> ModelPlacementInput | None:
1274 """Analytic-floor placement input for an installed model the estimator cannot
1275 size; ``None`` skips the role (the file is unresolvable, its weights alone
1276 exceed the hardware, or offloading it would exceed system memory).
1278 The host bound applies here too. Charging the whole floor to VRAM and
1279 skipping it let an unsizable model past a check every sized model faces."""
1280 weights = _role_weights_bytes(role, ref)
1281 if weights == 0:
1282 log.warning("Skipping %s server: could not size model %r (%s).", role.value, ref, exc)
1283 return None
1284 if _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref)):
1285 _warn_weights_exceed(role, ref, weights, total_vram)
1286 return None
1287 floor = _fallback_floor_for(role, ref, weights)
1288 log.warning(
1289 "Could not size the %s model %s (%s). Charging %.1f GiB, its weights plus the "
1290 "cache and buffers it will allocate, which is a floor rather than an estimate: "
1291 "the load may still need more.",
1292 role.value,
1293 ref,
1294 exc,
1295 floor / 1024**3,
1296 )
1297 if _host_memory_refuses(role, ref, floor, host_committed):
1298 return None
1299 return ModelPlacementInput(
1300 role=role, est_vram_bytes=floor, replicas=_replica_count(role, device_count)
1301 )
1304def _fallback_floor_for(role: WorkerRole, ref: str, weights: int) -> int:
1305 """:func:`_analytic_footprint_floor` for *role*, reading what metadata it can."""
1306 from lilbee.providers.base import ProviderError
1307 from lilbee.providers.gguf_meta import read_gguf_metadata
1309 try:
1310 path = engine_params.resolve_model_path(ref)
1311 meta = read_gguf_metadata(path)
1312 except (ProviderError, OSError, ValueError):
1313 meta = None
1314 path = None
1315 ctx = _placement_estimate_ctx(role, path, meta) if path is not None else _MIN_USABLE_CHAT_CTX
1316 return _analytic_footprint_floor(
1317 weights, role=role, meta=meta, ctx=ctx, slots=_placement_estimate_slots(role, meta)
1318 )
1321def _ref_is_moe(ref: str) -> bool:
1322 """Whether *ref*'s GGUF declares routed experts; False when it cannot be read."""
1323 from lilbee.providers.base import ProviderError
1324 from lilbee.providers.gguf_meta import read_gguf_metadata
1326 try:
1327 return _is_moe(read_gguf_metadata(engine_params.resolve_model_path(ref)))
1328 except (ProviderError, OSError):
1329 return False
1332def _cpu_offload_in_play() -> bool:
1333 """Whether this configuration puts any of a model's weights in system memory.
1335 Expert offload moves the experts, a partial ``n_gpu_layers`` moves whole
1336 layers, and zero moves the model. Without one of these the engine keeps
1337 everything on the card and the estimator's host figure describes memory
1338 nobody will allocate.
1339 """
1340 from lilbee.core.config import cfg
1342 return _expert_offload_configured() or cfg.n_gpu_layers is not None
1345def _host_bytes_must_be_resident(role: WorkerRole, ref: str) -> bool:
1346 """Whether *role*'s host bytes have to fit RAM rather than page in and out.
1348 The estimator's host figure counts mmap pages, and llama.cpp maps CPU-side
1349 weights over that mapping instead of allocating them, so with mmap they are
1350 evictable page cache: a model far larger than RAM streams from disk and
1351 serves, which is a practiced setup for a large mixture-of-experts. Only
1352 ``--no-mmap`` turns them into a buffered read that must be resident, and the
1353 single path that asks for it is a chat model on a network filesystem.
1355 Anything this cannot determine counts as mappable, because a false refusal
1356 here has no override and costs the user a model that would have run.
1357 """
1358 if role is not WorkerRole.CHAT:
1359 return False
1360 try:
1361 path = engine_params.resolve_model_path(ref)
1362 except (ProviderError, OSError, ValueError):
1363 return False
1364 if not is_network_path(path):
1365 return False
1366 return _chat_no_mmap(_role_weights_bytes(role, ref), on_network_fs=True)
1369def _host_committed(admitted: Mapping[WorkerRole, ModelPlacementInput]) -> int:
1370 """System-memory bytes the roles already admitted to this plan will hold."""
1371 return sum(inp.est_ram_bytes for inp in admitted.values())
1374def _host_memory_refuses(role: WorkerRole, ref: str, ram_bytes: int, committed: int) -> bool:
1375 """Whether *role*'s system-memory half is too big for this machine to load.
1377 Charged only when something actually offloads, and only when the bytes must
1378 be resident: refusing a mapped model that would have streamed from disk is a
1379 false refusal with no override, which is worse than a slow load.
1381 Measured against the whole plan, not this role alone. Every role was
1382 previously compared to the entire machine on its own, so two roles that each
1383 fit and together do not were both admitted.
1384 """
1385 if not _cpu_offload_in_play() or ram_bytes <= 0:
1386 return False
1387 wanted = committed + ram_bytes
1388 total = total_system_memory()
1389 if total and wanted > total and _host_bytes_must_be_resident(role, ref):
1390 log.warning(
1391 "The %s model %s cannot load: this plan puts %.1f GiB in system memory, which "
1392 "cannot be paged out here, and the machine has %.1f GiB in total. Use a smaller "
1393 "model, or offload less.",
1394 role.value,
1395 ref,
1396 wanted / 1024**3,
1397 total / 1024**3,
1398 )
1399 return True
1400 free = free_system_memory()
1401 if free and wanted > free:
1402 log.warning(
1403 "Offloading the %s model %s brings this plan to %.1f GiB in system memory and "
1404 "only %.1f GiB is free. It will still load; close other programs if it swaps "
1405 "or runs slowly.",
1406 role.value,
1407 ref,
1408 wanted / 1024**3,
1409 free / 1024**3,
1410 )
1411 return False
1414def _warn_weights_exceed(role: WorkerRole, ref: str, weights: int, total_vram: int) -> None:
1415 log.warning(
1416 "The %s model %s cannot load: its weights are %.1f GiB and this machine has "
1417 "%.1f GiB of GPU memory and %.1f GiB of system memory, so there is nowhere "
1418 "for its layers to go. Use a smaller model or a smaller quantization.",
1419 role.value,
1420 ref,
1421 weights / 1024**3,
1422 total_vram / 1024**3,
1423 model_cache.total_system_memory() / 1024**3,
1424 )
1427def _warn_weights_spill(role: WorkerRole, ref: str, weights: int, total_vram: int) -> None:
1428 """Say that a model larger than the GPUs will run partly in system memory."""
1429 log.warning(
1430 "The %s model %s is %.1f GiB and this machine has %.1f GiB of GPU memory, so "
1431 "the engine will keep the layers that fit on the GPU and the rest in system "
1432 "memory. It will run, and it will be slower than a model that fits.",
1433 role.value,
1434 ref,
1435 weights / 1024**3,
1436 total_vram / 1024**3,
1437 )
1440def placeable_total_vram() -> int:
1441 """Physical VRAM across all cards, for the weights-exceed placeability bound.
1443 Physical total is box-state-independent (a running incumbent doesn't skew
1444 it), so it is safe to read without a clean box. Reuses the plan probe when
1445 one is captured; otherwise probes best-effort and returns ``0`` on failure,
1446 which disables only the weights-exceed filter (its own ``total > 0`` guard).
1447 """
1448 probe = _plan_probe_store.get()
1449 if probe is not None:
1450 return sum(d.total_bytes for d in probe.devices)
1451 from lilbee.providers.base import ProviderError
1452 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
1454 try:
1455 apply_fleet_gpu_env()
1456 return sum(d.total_bytes for d in resolve_devices(resolve_llama_server()))
1457 except (ProviderError, OSError):
1458 return 0
1461def role_model_placeable(role: WorkerRole, ref: str, total_vram: int) -> bool:
1462 """Whether a fresh plan would actually serve *role* on *ref*.
1464 Mirrors the planner's own drop conditions (SDK-routed role, vision without a
1465 projector, model not installed, weights exceeding physical VRAM) using the
1466 same primitives, so the acquisition ladder binds and replaces against what
1467 an engine can serve rather than the raw config. Without this a
1468 configured-but-unplaceable role keeps bind from ever matching a running
1469 engine and restarts the shared engine on every process start.
1470 """
1471 if parse_model_ref(ref).is_remote or _vision_without_mmproj(role, ref):
1472 return False
1473 weights = _role_weights_bytes(role, ref) # 0 when not installed / unresolvable
1474 if weights == 0:
1475 return False
1476 return not _weights_exceed_hardware(weights, total_vram, is_moe=_ref_is_moe(ref))
1479def _server_model_inputs(
1480 roles: tuple[WorkerRole, ...] | None = None,
1481 *,
1482 unified_budget: int | None = None,
1483 device_count: int = 0,
1484 total_vram: int = 0,
1485) -> tuple[list[ModelPlacementInput], dict[WorkerRole, str], int, dict[WorkerRole, str]]:
1486 """Build placement inputs for the configured server roles.
1488 The search and vision roles are estimated first; chat is then sized against the
1489 budget minus the search footprint (the ``reservation``) so a large chat cannot
1490 starve embed/rerank on a shared-memory host. ``device_count`` resolves an auto
1491 replica knob to one per GPU. When *roles* is given, only those are considered.
1492 Skips an unconfigured optional role, a vision model with no resolvable mmproj
1493 projector, a role whose model is not installed on disk (returned as
1494 ``skipped_not_installed`` so a surface can say so), and a model whose weight
1495 bytes alone exceed ``total_vram`` (physically unloadable under all-GPU layers).
1496 A model the estimator cannot size is enrolled at its file size instead of
1497 skipped, so the load, not the estimator, decides.
1498 """
1499 from lilbee.core.config import cfg
1501 inputs: dict[WorkerRole, ModelPlacementInput] = {}
1502 model_refs: dict[WorkerRole, str] = {}
1503 skipped_not_installed: dict[WorkerRole, str] = {}
1505 def consider(role: WorkerRole, *, chat_reservation: int = 0) -> None:
1506 if roles is not None and role not in roles:
1507 return
1508 # Any role may be "" (unconfigured) -> skipped, so that role has no
1509 # server and no not-installed complaint.
1510 ref = str(getattr(cfg, ROLE_REGISTRY[role].config_field))
1511 if not ref:
1512 return # unconfigured optional role -> no server
1513 if parse_model_ref(ref).is_remote:
1514 return # SDK-routed role: no local server to plan, not a missing install
1515 if _vision_without_mmproj(role, ref):
1516 return # no projector -> vision can't run on a server
1517 estimate = _estimate_or_fallback(
1518 role,
1519 ref,
1520 unified_budget=unified_budget,
1521 chat_reservation=chat_reservation,
1522 device_count=device_count,
1523 total_vram=total_vram,
1524 skipped_not_installed=skipped_not_installed,
1525 host_committed=_host_committed(inputs),
1526 )
1527 if estimate is None:
1528 return
1529 inputs[role] = estimate
1530 model_refs[role] = ref
1532 # Estimate every non-chat role first so the search footprint is known, then size
1533 # chat against the remainder. The reservation only applies on a shared-memory
1534 # host; discrete GPUs pin each role to its own VRAM and pack independently.
1535 for role in ROLE_REGISTRY:
1536 if role is not WorkerRole.CHAT:
1537 consider(role)
1538 reservation = _search_reservation(inputs) if unified_budget is not None else 0
1539 consider(WorkerRole.CHAT, chat_reservation=reservation)
1541 ordered = [inputs[role] for role in ROLE_REGISTRY if role in inputs]
1542 return ordered, model_refs, reservation, skipped_not_installed
1545def _non_chat_reservation(
1546 instances: Sequence[InstancePlan],
1547 inputs: Sequence[ModelPlacementInput],
1548 co_tenants: frozenset[WorkerRole] = frozenset(),
1549) -> dict[int, int]:
1550 """Per-device VRAM the non-chat role servers occupy, keyed by device index.
1552 A tensor-split chat shard must size its KV against the headroom left after the
1553 embed/rerank/vision servers on the same card, not the card's raw free VRAM, or
1554 it over-commits and OOMs at launch. Chat is excluded because it sizes its own
1555 weights. Chat's own swap-group siblings are excluded too: they are evicted while
1556 chat is resident, so their VRAM is chat's to use. That only holds when chat is
1557 itself a co-tenant; a co-tenant group that does not include chat runs behind its
1558 own swap process and can be resident beside chat, so it is charged normally.
1559 Non-chat roles are single-device, so each charges its full footprint (once per
1560 replica) to its card.
1561 """
1562 chat_siblings = co_tenants if WorkerRole.CHAT in co_tenants else frozenset()
1563 charge_by_role = {inp.role: inp.est_vram_bytes for inp in inputs}
1564 reserved: dict[int, int] = {}
1565 for inst in instances:
1566 if inst.role is WorkerRole.CHAT or inst.role in chat_siblings:
1567 continue
1568 charge = charge_by_role[inst.role]
1569 for device in inst.devices:
1570 reserved[device] = reserved.get(device, 0) + charge
1571 return reserved
1574def _charge_by_device(
1575 chosen: tuple[FleetDevice, ...], ratio: tuple[int, ...], total: int
1576) -> dict[str, int]:
1577 """What each of *chosen* was charged, keyed by the name the engine prints.
1579 A single-card instance carries the whole charge. A split carries it in the
1580 proportions it launches with, which is what the planner decided and therefore
1581 what the engine's own report should be compared against.
1582 """
1583 from lilbee.providers.fleet.readback import device_label
1585 if total <= 0 or not chosen:
1586 return {}
1587 if len(chosen) == 1:
1588 return {device_label(chosen[0]): total}
1589 weights = ratio if len(ratio) == len(chosen) else (1,) * len(chosen)
1590 denominator = sum(weights) or len(chosen)
1591 return {
1592 device_label(device): total * weight // denominator
1593 for device, weight in zip(chosen, weights, strict=True)
1594 }
1597def _launch_for(
1598 plan: InstancePlan,
1599 model_ref: str,
1600 binary: Path,
1601 by_index: dict[int, FleetDevice],
1602 *,
1603 unified_budget: int | None = None,
1604 chat_reservation: int = 0,
1605 reserved_by_device: dict[int, int] | None = None,
1606 est_vram_bytes: int = 0,
1607 model_path: Path | None = None,
1608) -> InstanceLaunch:
1609 """Build the launch spec (argv + device-pinning env) for one planned instance."""
1610 from lilbee.providers.gguf_meta import read_gguf_metadata
1612 # The self-check holds a downloaded file rather than a configured reference,
1613 # so it hands the path over instead of asking for one to be resolved.
1614 model_path = model_path or engine_params.resolve_model_path(model_ref)
1615 weights_bytes = _weights_bytes(model_path)
1616 meta = read_gguf_metadata(model_path)
1617 from lilbee.core.config import cfg
1619 chosen = tuple(by_index[i] for i in plan.devices)
1620 # ctx and slots are sized against the card this role landed on, not the fleet's
1621 # smallest, which is all the pre-placement estimate had to go on. A role spread
1622 # over several cards has no one budget; the split chat, the only such role today,
1623 # sizes against per-device headroom below.
1624 placed_device = chosen[0] if len(chosen) == 1 else None
1625 is_chat = plan.role is WorkerRole.CHAT
1626 is_vision = plan.role is WorkerRole.VISION
1627 mmproj = _vision_mmproj(model_ref) if is_vision else None
1628 chat_on_network_fs = is_chat and is_network_path(model_path)
1629 if chat_on_network_fs and not _chat_no_mmap(weights_bytes, on_network_fs=True):
1630 log.warning(
1631 "Chat model %s is served from a network filesystem and is too large to load "
1632 "into host RAM; mmap over the network can stall the load in uninterruptible "
1633 "I/O. Stage it on local disk for a reliable load.",
1634 model_ref,
1635 )
1636 # A tensor-split chat serves one full-context sequence sized against the busiest
1637 # card's headroom. A cfg.num_ctx pin overrides the fit (handled by _role_ctx).
1638 multi_card_chat = is_chat and len(chosen) > 1
1639 split_chat = multi_card_chat and cfg.num_ctx is None
1640 if multi_card_chat and host_lacks_nvlink():
1641 log.warning(
1642 "Chat model %s is tensor-split across GPUs %s on a host without NVLink; "
1643 "generation is PCIe all-reduce bound and can be very slow. A model that fits "
1644 "on fewer cards will generate faster.",
1645 model_ref,
1646 list(plan.devices),
1647 )
1648 split_slots = _SPLIT_CHAT_SLOTS
1649 # A single-card chat is the one launch whose offload its own window fit can
1650 # lower; the split and pinned paths never run that fit.
1651 fit_offload = is_chat and not multi_card_chat and cfg.num_ctx is None
1652 n_gpu_layers = _role_gpu_layers(plan.role)
1653 if split_chat:
1654 reserved = reserved_by_device or {}
1655 # Headroom left after the embed/rerank servers on each shared card, not the
1656 # card's raw free VRAM, so the chat KV doesn't over-commit.
1657 per_device_free = [max(0, d.free_bytes - reserved.get(d.index, 0)) for d in chosen]
1659 def _split_fit(slots: int) -> int:
1660 return fleet_ctx.fit_split_ctx(
1661 model_path,
1662 meta=meta,
1663 slots=slots,
1664 ratio=plan.tensor_split,
1665 per_device_free_bytes=per_device_free,
1666 gpu_layers=_role_gpu_layers(WorkerRole.CHAT),
1667 flash_attn=_role_flash(WorkerRole.CHAT),
1668 kv_cache_type=_role_kv_cache_type(WorkerRole.CHAT),
1669 kv_cache_type_v=_role_kv_cache_type_v(WorkerRole.CHAT),
1670 ctx_ceiling=_placement_estimate_ctx(WorkerRole.CHAT, model_path, meta),
1671 expert_offload=_role_expert_offload(model_path),
1672 )
1674 split_slots, ctx = _resolve_split_chat_slots(_split_fit)
1675 elif fit_offload:
1676 # One call answers both halves. The window was sized against the memory
1677 # this offload leaves free, so a second call for the offload could pair a
1678 # window with an offload that never fitted it.
1679 chat_fit = engine_params.resolve_chat_fit(
1680 model_path, meta, available_bytes=plan_sizing_budget(placed_device)
1681 )
1682 n_gpu_layers = chat_fit.gpu_layers
1683 ctx = apply_ctx_downshift(plan.role, chat_fit.ctx)
1684 else:
1685 # Downshifted here and not only in the estimate: the role resolvers are
1686 # pure functions of model and config, so without this the retry after a
1687 # load OOM re-emits a byte-identical argv and dies the same way. The
1688 # split branch above already inherits it through its ctx_ceiling.
1689 ctx = apply_ctx_downshift(plan.role, _role_ctx(plan.role, model_path, meta, placed_device))
1690 rerank_mode = _role_rerank_mode(plan.role, meta)
1691 is_llm_rerank = rerank_mode is RerankMode.LLM
1692 # A multi-card chat runs as many full-context slots as its cards' KV headroom
1693 # holds (split_slots, one when a num_ctx pin skips the fit); other roles size
1694 # --parallel against the budget the same way the estimator did.
1695 slots = (
1696 split_slots
1697 if multi_card_chat
1698 else _slots_for(
1699 plan.role,
1700 model_path,
1701 ctx,
1702 mmproj_path=mmproj,
1703 unified_budget=unified_budget,
1704 chat_reservation=chat_reservation,
1705 rerank_mode=rerank_mode,
1706 device=placed_device,
1707 )
1708 )
1709 spec = _server_spec(plan.role, rerank_mode, meta)
1710 # Cross-encoder embed/rerank pools the whole input in one batch; an LLM reranker
1711 # is generative and uses the default batching plus flash attention.
1712 cross_encoder_pooled = plan.role in _EMBED_ROLES and not is_llm_rerank
1713 cache_type_k, cache_type_v = chat_cache_type_flags() if is_chat else (None, None)
1714 argv = build_server_argv(
1715 binary=binary,
1716 spec=spec,
1717 model_path=model_path,
1718 devices=plan.devices,
1719 n_gpu_layers=n_gpu_layers,
1720 slots=slots,
1721 ctx_per_slot=ctx,
1722 tensor_split=plan.tensor_split,
1723 mmproj=mmproj,
1724 flash_attn=flash_attn_flag() if _role_launches_with_flash(plan.role, rerank_mode) else None,
1725 cache_type_k=cache_type_k,
1726 cache_type_v=cache_type_v,
1727 batch_size=_pooled_batch_size(plan.role, rerank_mode, ctx),
1728 no_mmap=is_chat and _chat_no_mmap(weights_bytes, on_network_fs=chat_on_network_fs),
1729 cpu_moe=expert_offload_all(meta),
1730 n_cpu_moe=expert_offload_layers(meta),
1731 device_names=_device_names(chosen) or _cpu_pin_when_every_device_was_refused(),
1732 memory_endpoint=supports_memory_readback(binary),
1733 )
1734 return InstanceLaunch(
1735 role=plan.role,
1736 argv=argv,
1737 env_overrides={**visible_env(chosen), **llama_server_runtime_env()},
1738 model=model_ref,
1739 # token_cap drives cross-encoder/embed input truncation; the LLM rerank path
1740 # doesn't truncate (it relies on the per-slot ctx headroom), so leave it None.
1741 token_cap=max(1, ctx - engine_params._EMBED_CTX_MARGIN) if cross_encoder_pooled else None,
1742 # Weights size scales the cold-load ready timeout (larger model = longer).
1743 weights_bytes=weights_bytes,
1744 # Slots is the chat concurrency the gate admits; ctx is what a client fits to.
1745 slots=slots,
1746 ctx=ctx,
1747 built_ctx_target=(
1748 (cfg.num_ctx if cfg.num_ctx is not None else cfg.chat_n_ctx_target) if is_chat else 0
1749 ),
1750 built_slots_target=(
1751 max(1, cfg.vision_ocr_concurrency) if plan.role is WorkerRole.VISION else 0
1752 ),
1753 replica=plan.replica,
1754 rerank_mode=rerank_mode,
1755 # What placement charged this instance, for the post-launch check against
1756 # the engine's own report of what it really allocated.
1757 est_vram_bytes=est_vram_bytes,
1758 est_vram_by_device=_charge_by_device(chosen, plan.tensor_split, est_vram_bytes),
1759 est_unreported_bytes=_unreported_bytes(plan.role, mmproj),
1760 )
1763def build_single_role_launch(role: WorkerRole, model_path: Path) -> InstanceLaunch:
1764 """The launch the fleet would build for *role* serving *model_path*, alone.
1766 One construction path. The self-check used to assemble its own beside this
1767 one and the two disagreed on slot count, on the context that follows from it,
1768 on device pinning and on the tensor split, so a green check proved nothing
1769 about the launch serving actually performs, and a red one could be a
1770 configuration serving would never have chosen.
1772 Placement is the planner's, on the devices the plan snapshot holds, so the
1773 check runs on the card the role would really land on.
1774 """
1775 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
1776 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
1778 apply_fleet_gpu_env()
1779 binary = resolve_llama_server()
1780 apply_cuda_runtime_env(binary)
1781 devices = _plan_devices(binary)
1782 by_index = {d.index: d for d in devices}
1783 # The whole machine, since nothing else is resident during a self-check.
1784 placed = (min(by_index),) if by_index else ()
1785 plan = InstancePlan(role=role, devices=placed)
1786 return _launch_for(
1787 plan,
1788 str(model_path),
1789 binary,
1790 by_index,
1791 unified_budget=_unified_memory_budget(devices),
1792 model_path=model_path,
1793 )
1796def resolve_devices(binary: Path) -> list[FleetDevice]:
1797 """Enumerate devices in the binary's index space, or the Vulkan VRAM probe."""
1798 return _resolve_devices_and_refusal(binary)[0]
1801# The visibility variable each vendor's runtime reads, named in the warning so
1802# the reader checks the one that applies to the card they actually have.
1803_VENDOR_VISIBILITY_HINT = {
1804 "NVIDIA": "CUDA_VISIBLE_DEVICES",
1805 "AMD": "ROCR_VISIBLE_DEVICES / HIP_VISIBLE_DEVICES",
1806 "Intel": "ONEAPI_DEVICE_SELECTOR",
1807}
1810def _warn_gpu_present_but_unenumerated(binary: Path) -> None:
1811 """Say so when the host has a GPU the engine did not list.
1813 Previously asked only whether an NVIDIA card was present, so an AMD or Intel
1814 host whose engine enumerated nothing produced the identical symptom, a fleet
1815 quietly planned for CPU, and said nothing. The vendor lookup is the same one
1816 the Vulkan ICD rules use, and it works on Windows as well as Linux.
1817 """
1818 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
1819 from lilbee.providers.fleet.gpu_select import PCIVendorID
1821 present = installed_gpu_vendor_ids()
1822 names = sorted(
1823 v.name.title() if v.name == "INTEL" else v.name for v in PCIVendorID if v in present
1824 )
1825 if not names:
1826 return
1827 hints = sorted(
1828 {_VENDOR_VISIBILITY_HINT[name] for name in names if name in _VENDOR_VISIBILITY_HINT}
1829 )
1830 log.warning(
1831 "This host has a %s GPU but the engine's device probe (%s --list-devices) "
1832 "reported none; placement is falling back to shared-memory mode with unpinned "
1833 "GPUs. Check the GPU driver, %s, and that this llama-server build supports "
1834 "that GPU.",
1835 " and ".join(names),
1836 binary,
1837 " / ".join(hints) if hints else "the vendor's visibility variable",
1838 )
1841def _resolve_devices_and_refusal(binary: Path) -> tuple[list[FleetDevice], bool]:
1842 """:func:`resolve_devices`, plus whether every GPU the engine listed was refused.
1844 One function because both answers come from one ``--list-devices`` run, and
1845 that run costs a subprocess against a driver that may be wedged. Asking twice
1846 would pay it twice.
1848 The binary's ``--list-devices`` is authoritative, including when it lists
1849 nothing: it prints every non-CPU device it can use, so an empty list means
1850 the engine has no usable GPU rather than that we failed to look. The Vulkan
1851 VRAM probe is consulted only when the binary produced no output at all, and
1852 it reports the same index space. A
1853 probe that times out raises instead (a wedged GPU driver); falling through
1854 to the in-process Vulkan probe there could hang this thread unkillably.
1855 """
1856 from lilbee.providers.fleet.cuda_runtime import assert_cuda_devices_usable
1857 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
1858 from lilbee.providers.fleet.gpu_select import enumerate_gpu_vram
1859 from lilbee.providers.fleet.rocm_runtime import assert_rocm_devices_usable
1861 probe = probe_devices(binary)
1862 # An engine that answered but listed no device while a GPU is physically
1863 # present may be hitting a transient init error (the card momentarily held by
1864 # another process, e.g. an embedder served alongside -- the "ggml_cuda_init:
1865 # initialization error" symptom). Re-probe before treating the empty list as
1866 # fatal; a persistently empty list still hits the fail-loud asserts below.
1867 for _ in range(_DEVICE_PROBE_EMPTY_RETRIES):
1868 if probe.devices or not probe.spoke_protocol or not installed_gpu_vendor_ids():
1869 break
1870 time.sleep(_DEVICE_PROBE_EMPTY_RETRY_DELAY_S)
1871 probe = probe_devices(binary)
1872 devices = probe.devices
1873 # A GPU build that links a runtime it cannot serve the host's GPU with must
1874 # fail loud, not silently fall back to CPU (the Vulkan VRAM probe below
1875 # would mask it). Only when the engine actually answered, though: a binary
1876 # that does not support --list-devices enumerated nothing because it was
1877 # never asked, and accusing its driver of failing would be wrong and fatal.
1878 if probe.spoke_protocol:
1879 assert_cuda_devices_usable(binary, devices, probe.output)
1880 assert_rocm_devices_usable(binary, devices, probe.output)
1881 if not devices and probe.spoke_protocol:
1882 _warn_gpu_present_but_unenumerated(binary)
1883 if not devices and not probe.spoke_protocol:
1884 # Only when the binary never answered the question. An engine that ran
1885 # and listed nothing is reporting a fact, not a gap: believing the host
1886 # loader instead invents devices the engine cannot see. A CPU-only build
1887 # on a desktop with mesa is the clearest case, and the cost is not merely
1888 # a wrong device list. The fleet is planned onto GPUs, the pins are
1889 # no-ops, the shared-RAM guard is off because devices looked non-empty,
1890 # and every role then loads its full weights into system RAM while
1891 # running on the CPU anyway.
1892 #
1893 # Keyed on the exit code and the header rather than on there being no
1894 # output at all: the probe merges stderr into stdout, so a build that
1895 # predates --list-devices prints usage text and would otherwise be read
1896 # as an authoritative "no GPUs here".
1897 from lilbee.providers.fleet.gpu_select import integrated_vulkan_indices
1899 integrated = integrated_vulkan_indices()
1900 devices = [
1901 FleetDevice(
1902 VULKAN_BACKEND, idx, "", vram, free, unified=idx in integrated, from_loader=True
1903 )
1904 for idx, vram, free in (enumerate_gpu_vram() or [])
1905 ]
1906 if devices:
1907 log.warning(
1908 "The engine's device probe returned nothing, so placement is using "
1909 "the host's Vulkan loader instead and found %d device(s). If the "
1910 "engine has no Vulkan backend it will run on CPU regardless; set %s "
1911 "to override the engine location if that is wrong.",
1912 len(devices),
1913 "LILBEE_ENGINE_DIR",
1914 )
1915 return devices, probe.refused_all
1918_DEVICE_PROBE_TTL_S = 2.0
1919# A failed probe is cached much longer than a good one: each retry against a
1920# wedged GPU driver costs a full probe timeout, so a per-poll retry would stall
1921# every placement read for a minute at a time.
1922_DEVICE_PROBE_FAILURE_TTL_S = 60.0
1923# An engine that lists no device on a GPU host may be hitting a transient GPU-init
1924# error (the card momentarily held by another process); re-probe before treating
1925# the empty list as fatal.
1926_DEVICE_PROBE_EMPTY_RETRIES = 3
1927_DEVICE_PROBE_EMPTY_RETRY_DELAY_S = 0.5
1930class _ReadDeviceCache:
1931 """Short-TTL device-probe cache for the read/view path.
1933 Not a ``cachetools.TTLCache``: it caches the *failure* too, under its own
1934 longer TTL, and re-raises it. A memoizing cache stores return values only,
1935 so a failing probe would re-spawn the subprocess on every placement read.
1937 Inspecting placement (GET placement/gpus, preview, ``placement show``)
1938 resolves devices on every call, which spawns a ``llama-server --list-devices``
1939 subprocess; a brief TTL collapses a burst of reads onto one probe. A probe
1940 failure is cached too (with its own TTL) and re-raised to every read in the
1941 window. The launch path is never served from here -- it sizes against the
1942 clean-box plan snapshot below (captured after stale-server reaping).
1943 """
1945 def __init__(self, ttl_s: float, failure_ttl_s: float) -> None:
1946 self._ttl_s = ttl_s
1947 self._failure_ttl_s = failure_ttl_s
1948 self._lock = threading.Lock()
1949 self._at: float | None = None
1950 self._devices: list[FleetDevice] | None = None
1951 self._failure: ProviderError | None = None
1953 def get(self, binary: Path) -> list[FleetDevice]:
1954 with self._lock:
1955 ttl = self._ttl_s if self._failure is None else self._failure_ttl_s
1956 fresh = self._at is not None and time.monotonic() - self._at < ttl
1957 if fresh and self._failure is not None:
1958 raise self._failure
1959 if self._devices is None or not fresh:
1960 self._at = time.monotonic()
1961 try:
1962 self._devices = resolve_devices(binary)
1963 except ProviderError as exc:
1964 self._devices = None
1965 self._failure = exc
1966 raise
1967 self._failure = None
1968 return self._devices
1970 def clear(self) -> None:
1971 with self._lock:
1972 self._at = None
1973 self._devices = None
1974 self._failure = None
1977_read_device_cache = _ReadDeviceCache(_DEVICE_PROBE_TTL_S, _DEVICE_PROBE_FAILURE_TTL_S)
1980def clear_read_device_cache() -> None:
1981 """Drop the read-path device probe cache (e.g. after the fleet is reconfigured).
1983 Also drops what the host's Vulkan loader told us about device types, which is
1984 otherwise held for the process lifetime and would survive a driver reload or
1985 an eGPU being plugged in.
1986 """
1987 from lilbee.providers.fleet.gpu_select import (
1988 integrated_vulkan_indices,
1989 vulkan_device_types_by_name,
1990 )
1992 _read_device_cache.clear()
1993 vulkan_device_types_by_name.cache_clear()
1994 integrated_vulkan_indices.cache_clear()
1997@dataclass(frozen=True)
1998class _PlanProbe:
1999 """Clean-box memory snapshot every plan is sized against.
2001 Captured once, right after stale-server reaping and before the first build,
2002 when nothing lilbee owns is loaded. Reloads re-plan against this same
2003 snapshot instead of re-probing: a live probe under a loaded fleet reports
2004 our own residency as unavailable, which would shrink chat context and slot
2005 counts, widen splits, and (on a unified-memory host) evict roles outright.
2006 Launches stay a pure function of config + hardware + this snapshot, so the
2007 reload diff restarts only real changes. Cleared on full fleet teardown so
2008 the next boot probes the clean box afresh.
2009 """
2011 devices: tuple[FleetDevice, ...]
2012 # What one role may size its ctx and slots against, already scaled by
2013 # cfg.gpu_memory_fraction. System memory only on a host with no GPU.
2014 sizing_budget: int
2015 free_system: int
2016 # The engine listed GPUs and lilbee rejected all of them, so the plan is
2017 # CPU-shaped while the engine would still choose one of those devices.
2018 engine_devices_all_refused: bool = False
2021class _PlanProbeStore:
2022 """Holds the captured plan snapshot; a single instance below (no bare global)."""
2024 def __init__(self) -> None:
2025 self._lock = threading.Lock()
2026 self._probe: _PlanProbe | None = None
2028 def set(self, probe: _PlanProbe) -> None:
2029 with self._lock:
2030 self._probe = probe
2032 def get(self) -> _PlanProbe | None:
2033 with self._lock:
2034 return self._probe
2036 def clear(self) -> None:
2037 with self._lock:
2038 self._probe = None
2041_plan_probe_store = _PlanProbeStore()
2044# Where the ladder stops. Sized for chat, below which the answers are too short
2045# to be useful, so a role that still will not load here has a real problem the
2046# planner cannot size its way out of and the failure should surface. Roles whose
2047# window already sits under it (a small embedding context) are left alone rather
2048# than raised to meet it, so for them the ladder is a no-op and the failure
2049# surfaces after the one retry.
2050MIN_DOWNSHIFT_CTX = 4096
2053class _CtxDownshiftStore:
2054 """How many halvings each role's auto context has taken after a load OOM.
2056 An estimate that was too optimistic is only recoverable if the retry asks
2057 for something different. Halving the auto context does that, and keeping the
2058 count here rather than in the launch means the whole plan is re-predicted
2059 against the smaller number, including the placement it implies.
2060 """
2062 def __init__(self) -> None:
2063 self._lock = threading.Lock()
2064 self._steps: dict[WorkerRole, int] = {}
2065 # The last unshifted context each role was sized from, recorded as it is
2066 # applied. Deciding whether another halving would change anything needs
2067 # the number being halved, and this is the only place that sees it.
2068 self._base: dict[WorkerRole, int] = {}
2070 def steps(self, role: WorkerRole) -> int:
2071 with self._lock:
2072 return self._steps.get(role, 0)
2074 def note_base(self, role: WorkerRole, ctx: int) -> None:
2075 with self._lock:
2076 self._base[role] = ctx
2078 def base(self, role: WorkerRole) -> int | None:
2079 with self._lock:
2080 return self._base.get(role)
2082 def step(self, role: WorkerRole) -> int:
2083 with self._lock:
2084 taken = self._steps.get(role, 0) + 1
2085 self._steps[role] = taken
2086 return taken
2088 def clear(self, role: WorkerRole | None = None) -> None:
2089 with self._lock:
2090 if role is None:
2091 self._steps.clear()
2092 self._base.clear()
2093 return
2094 self._steps.pop(role, None)
2095 self._base.pop(role, None)
2098_ctx_downshift_store = _CtxDownshiftStore()
2101def apply_ctx_downshift(role: WorkerRole, ctx: int) -> int:
2102 """*ctx* halved once per downshift step recorded for *role*, floored.
2104 Never more than *ctx*. The floor is a stopping point, not a target: applied
2105 to a context already below it (a small embedding window, a model trained for
2106 2048 tokens) a bare floor would hand back a larger number, and the retry
2107 after a load OOM would ask for more memory than the launch that just ran out
2108 of it. Such a role simply has nothing to give back, and its failure surfaces
2109 after the one retry instead.
2111 A user's ``cfg.num_ctx`` pin is returned untouched: serving a window smaller
2112 than the one that was asked for, without being asked, is worse than failing
2113 to load and saying so.
2114 """
2115 from lilbee.core.config import cfg
2117 if role is WorkerRole.CHAT and cfg.num_ctx is not None:
2118 return ctx
2119 _ctx_downshift_store.note_base(role, ctx)
2120 return _shifted(ctx, _ctx_downshift_store.steps(role))
2123def _shifted(ctx: int, steps: int) -> int:
2124 """*ctx* halved *steps* times, never below the floor and never above *ctx*."""
2125 return min(ctx, max(MIN_DOWNSHIFT_CTX, ctx >> steps)) if steps else ctx
2128def record_ctx_downshift(role: WorkerRole) -> bool:
2129 """Take one downshift step for *role*; False when there is none left to take.
2131 False means the retry would ask for the same thing again, so the caller must
2132 surface the load failure instead of respawning an identical launch.
2133 """
2134 from lilbee.core.config import cfg
2136 if role is WorkerRole.CHAT and cfg.num_ctx is not None:
2137 return False
2138 base = _ctx_downshift_store.base(role)
2139 if base is None:
2140 # Nothing has been sized for this role yet, so there is no number to
2141 # decide against. Allow one step rather than trusting that a plan always
2142 # runs first: an unbounded grant here would let a caller that never
2143 # sizes anything loop forever.
2144 if _ctx_downshift_store.steps(role):
2145 return False
2146 _ctx_downshift_store.step(role)
2147 return True
2148 steps = _ctx_downshift_store.steps(role)
2149 if _shifted(base, steps + 1) == _shifted(base, steps):
2150 return False
2151 _ctx_downshift_store.step(role)
2152 return True
2155def clear_ctx_downshift(role: WorkerRole | None = None) -> None:
2156 """Forget *role*'s recorded downshift, or every role's, back to full size.
2158 Called when a role's engine reports ready, which is proof the reduced plan
2159 loaded: keeping the reduction after that would carry a shrunken window into
2160 a machine that has since freed memory, or into a smaller model the user
2161 switched to, and would then refuse on its first failure with a budget it
2162 had already spent.
2163 """
2164 _ctx_downshift_store.clear(role)
2167def _probe_engine_devices() -> tuple[list[FleetDevice], bool]:
2168 """Apply the fleet GPU/CUDA env, resolve the binary, and enumerate devices.
2170 This is the wedge point: a missing binary raises NOT_FOUND, and a CUDA build
2171 that cannot init a GPU (a broken-runtime host) raises loud from resolve_devices
2172 rather than silently degrading. Device enumeration reads no residency, so it is
2173 safe to run while an incumbent engine is still up.
2174 """
2175 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
2176 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
2178 apply_fleet_gpu_env()
2179 binary = resolve_llama_server()
2180 apply_cuda_runtime_env(binary)
2181 devices, refused = _resolve_devices_and_refusal(binary)
2182 if devices:
2183 return devices, refused
2184 return _reprobe_while_a_gpu_is_installed(binary, refused)
2187def _reprobe_while_a_gpu_is_installed(
2188 binary: Path, refused: bool
2189) -> tuple[list[FleetDevice], bool]:
2190 """Ask again when the host has a GPU the engine did not list.
2192 The plan snapshot is taken once, on a clean box, and is not retaken until a
2193 full teardown, so an empty first answer decides the whole run. A GPU driver
2194 that is still initializing when the daemon starts, which is ordinary under
2195 systemd or right after a container gains a device, would leave a GPU host
2196 serving on CPU until someone noticed and restarted it.
2198 Only where a card is actually installed. A host with no GPU answers empty
2199 every time and must not pay a retry for it on every start.
2200 """
2201 from lilbee.providers.fleet.gpu_hardware import installed_gpu_vendor_ids
2203 if not installed_gpu_vendor_ids():
2204 return [], refused
2205 for attempt in range(1, _PROBE_RETRIES + 1):
2206 log.info(
2207 "The engine listed no GPU on a host that has one; asking again in %.1fs "
2208 "(attempt %d of %d) in case the driver is still initializing.",
2209 _PROBE_RETRY_DELAY_S,
2210 attempt,
2211 _PROBE_RETRIES,
2212 )
2213 time.sleep(_PROBE_RETRY_DELAY_S)
2214 clear_read_device_cache()
2215 devices, refused = _resolve_devices_and_refusal(binary)
2216 if devices:
2217 return devices, refused
2218 return [], refused
2221def assert_engine_probeable() -> None:
2222 """Raise if the engine cannot be probed; capture no snapshot.
2224 A build precondition that must run BEFORE stopping a replaceable incumbent:
2225 it surfaces a wedged GPU probe or an unusable CUDA runtime without taking the
2226 residency-dependent memory snapshot (that belongs on the clean box, after the
2227 stop, in capture_plan_probe). resolve_devices caches within its TTL, so the
2228 follow-up capture reuses this enumeration rather than re-probing the hardware.
2229 """
2230 _probe_engine_devices()
2233def capture_plan_probe() -> None:
2234 """Snapshot devices and memory for planning; call only on a clean box."""
2235 devices, refused_all = _probe_engine_devices()
2236 _plan_probe_store.set(
2237 _PlanProbe(
2238 devices=tuple(devices),
2239 sizing_budget=_device_sizing_budget(devices),
2240 free_system=model_cache.free_system_memory(),
2241 engine_devices_all_refused=refused_all,
2242 )
2243 )
2246def _structural(devices: Iterable[FleetDevice]) -> tuple[FleetDevice, ...]:
2247 """*devices* with the volatile free reading zeroed, so equality is structural.
2249 This probe runs while the fleet is resident, so its free figures are deflated
2250 by the very models a reload is about to stop; comparing or keying on them
2251 would read every reload as a hardware change.
2252 """
2253 return tuple(replace(d, free_bytes=0) for d in devices)
2256def refresh_plan_devices() -> None:
2257 """Re-read which devices exist, keeping the clean-box memory figures.
2259 The snapshot is captured once and only a full teardown clears it, so an eGPU
2260 unplugged, a driver reset, or a VM hot-remove left the fleet pinning a device
2261 that is no longer there and every rebuild replanning onto it.
2263 Only the structural half is restated. The memory figures are what make a
2264 reload plan the way the boot did, and re-taking them while the fleet is
2265 resident would charge it against itself, which is the whole reason the
2266 snapshot exists. So a card that survives the refresh keeps the snapshot's
2267 free figure; only a genuinely new card contributes a fresh one.
2269 A probe that cannot run leaves the snapshot alone: the last known device list
2270 is a better answer than none, and the loud paths for an unreachable engine
2271 live in the build, not here.
2272 """
2273 probe = _plan_probe_store.get()
2274 if probe is None:
2275 return
2276 clear_read_device_cache()
2277 try:
2278 devices, refused_all = _probe_engine_devices()
2279 except (ProviderError, OSError) as exc:
2280 log.debug("Device rediscovery could not run, keeping the previous list: %s", exc)
2281 return
2282 if _structural(devices) == _structural(probe.devices):
2283 return
2284 log.info(
2285 "The set of GPUs changed since this fleet was planned (%d device(s) now, %d before); "
2286 "replanning against the ones that are here.",
2287 len(devices),
2288 len(probe.devices),
2289 )
2290 snapshot_free = {replace(d, free_bytes=0): d.free_bytes for d in probe.devices}
2291 merged = tuple(
2292 replace(d, free_bytes=snapshot_free.get(replace(d, free_bytes=0), d.free_bytes))
2293 for d in devices
2294 )
2295 _plan_probe_store.set(
2296 _PlanProbe(
2297 devices=merged,
2298 sizing_budget=_device_sizing_budget(merged),
2299 free_system=probe.free_system,
2300 engine_devices_all_refused=refused_all,
2301 )
2302 )
2305def clear_plan_probe() -> None:
2306 """Drop the plan snapshot (full fleet teardown); the next build re-captures."""
2307 _plan_probe_store.clear()
2310def _cpu_pin_when_every_device_was_refused() -> tuple[str, ...]:
2311 """``("none",)`` when the engine offered GPUs that lilbee refused, else empty.
2313 Dropping a device from lilbee's view does not stop the engine using it. With
2314 no pin at all, ggml applies its own selection, and its fallback takes the
2315 first non-CPU adapter, which is exactly the paravirtual device just refused;
2316 with every layer offloaded by default the model then runs on it while
2317 placement budgeted against system RAM. Naming no device keeps the engine on
2318 the CPU the plan was shaped for.
2319 """
2320 probe = _plan_probe_store.get()
2321 if probe is None or not probe.engine_devices_all_refused:
2322 return ()
2323 log.warning(
2324 "The engine listed GPU devices that lilbee will not plan onto, so it is being "
2325 "run on the CPU. Serving from one of them would be slower than the CPU or fail "
2326 "outright, and placement has been sized for system RAM."
2327 )
2328 return (_NO_DEVICE,)
2331def _plan_devices(binary: Path) -> list[FleetDevice]:
2332 """Devices the plan paths size against: the snapshot, else a live probe."""
2333 probe = _plan_probe_store.get()
2334 return list(probe.devices) if probe is not None else resolve_devices(binary)
2337def plan_sizing_budget(device: FleetDevice | None = None) -> int:
2338 """Usable memory for ctx/slot sizing: *device*'s own, else the snapshot, else live."""
2339 from lilbee.core.config import cfg
2341 if device is not None:
2342 return int(device.total_bytes * cfg.gpu_memory_fraction)
2343 probe = _plan_probe_store.get()
2344 if probe is not None:
2345 return probe.sizing_budget
2346 return _device_sizing_budget(_live_sizing_devices())
2349def plan_sizing_is_unified() -> bool:
2350 """Whether ctx sizing charges the shared-memory footprint rather than VRAM.
2352 True when no device has memory of its own -- an integrated GPU, an Apple
2353 Silicon Mac, or no GPU at all -- because there a model's would-be VRAM and
2354 its host bytes are the same memory, and charging only the VRAM half
2355 under-reports the load by everything it maps. The same test decides the pool
2356 placement charges against (:func:`_unified_memory_budget`).
2357 """
2358 probe = _plan_probe_store.get()
2359 devices = list(probe.devices) if probe is not None else _live_sizing_devices()
2360 return all(dev.unified for dev in devices)
2363def _device_sizing_budget(devices: Sequence[FleetDevice]) -> int:
2364 """Memory one role may size its ctx and slots against, in bytes.
2366 Read from the engine's own device report, which ran under the environment the
2367 servers will run under and states each device's memory whatever the backend.
2368 A host-memory read answers with system RAM on every host without an NVIDIA
2369 card, which gave a 24 GiB AMD card a budget the size of the machine, and on
2370 Apple Silicon it ignores that Metal will not allocate past
2371 ``recommendedMaxWorkingSetSize``, which is the figure the probe carries.
2373 The smallest device, since this is asked before placement has picked one;
2374 :func:`_launch_for` re-sizes against the card the role actually landed on.
2375 System memory only when the engine reports no device at all, where the fleet
2376 runs on the CPU and system memory is the budget.
2377 """
2378 from lilbee.core.config import cfg
2380 if devices:
2381 return int(min(d.total_bytes for d in devices) * cfg.gpu_memory_fraction)
2382 return int(model_cache.total_system_memory() * cfg.gpu_memory_fraction)
2385def _live_sizing_devices() -> list[FleetDevice]:
2386 """Devices to size against with no plan snapshot; empty when none can be read."""
2387 try:
2388 return _read_device_cache.get(resolve_llama_server())
2389 except (ProviderError, OSError):
2390 return []
2393def _plan_free_system_memory() -> int:
2394 """Free system RAM for the unified-memory budget: the snapshot, else live."""
2395 probe = _plan_probe_store.get()
2396 return probe.free_system if probe is not None else model_cache.free_system_memory()
2399def _unreported_bytes(role: WorkerRole, mmproj: Path | None) -> int:
2400 """Estimated bytes the engine allocates without printing a buffer line.
2402 A vision projector's weights: llama.cpp allocates them in clip's own loader,
2403 which prints a size but not the "buffer size = N MiB" shape the readback
2404 reads, so the report is short by exactly this and the self-check would warn
2405 on a load that was sized correctly.
2406 """
2407 if role is not WorkerRole.VISION or mmproj is None:
2408 return 0
2409 try:
2410 return mmproj.stat().st_size
2411 except OSError:
2412 return 0
2415def _chat_no_mmap(weights_bytes: int, *, on_network_fs: bool = False) -> bool:
2416 """Whether the chat server should malloc its weights instead of mmapping them.
2418 Local disk mmaps: lazy page-fault paging gives a faster first token on a cold
2419 cache -- the common desktop first launch -- matching mmap-by-default engines.
2420 ``--no-mmap``'s buffered full read only wins on an already-hot cache and it
2421 pessimizes cold start, so it is not worth defaulting on for local disk. A
2422 network filesystem still prefers the buffered read whenever the host copy
2423 fits, because mmap page faults served over the wire can wedge the loader in
2424 uninterruptible I/O (see ``_NO_MMAP_NETWORK_RAM_FRACTION``).
2425 """
2426 if not on_network_fs:
2427 return False
2428 return weights_bytes <= model_cache.total_system_memory() * _NO_MMAP_NETWORK_RAM_FRACTION
2431def _device_names(devices: tuple[FleetDevice, ...]) -> tuple[str, ...]:
2432 """``--device`` names for *devices*, empty when the backend pins through env.
2434 Vulkan and SYCL, because neither one's environment variable speaks the space
2435 the probe enumerated. Vulkan's indexes the raw loader enumeration while the
2436 names come from the engine's filtered list, so the two disagree wherever ggml
2437 drops or merges a device. SYCL's is not an index list at all but a selector
2438 over a backend runtime, so a device the engine calls ``SYCL1`` need not be
2439 Level Zero ordinal 1: OpenCL devices interleave, discarded devices shift the
2440 numbering, and multi-tile cards appear as sub-devices.
2442 ``--device`` sidesteps both by naming devices exactly as ``--list-devices``
2443 printed them, which is where these indices were read from. CUDA and ROCm
2444 keep composing their variables, which do share the probe's space.
2445 """
2446 if not devices or devices[0].backend not in _NAME_PINNED_BACKENDS:
2447 return ()
2448 if any(d.from_loader for d in devices):
2449 # These indices are raw loader ordinals, and --device speaks the engine's
2450 # own post-filter naming, so Vulkan1 here can name Vulkan0 there or
2451 # nothing at all. Sizing against them is still worth doing; pinning by
2452 # them is not. Left unpinned, ggml applies its own device selection,
2453 # which is the filtering lilbee is trying to agree with in the first
2454 # place. The env pin is not the answer either: it takes raw ordinals but
2455 # switches off the type filter, the support check and the dedup with them.
2456 return ()
2457 return tuple(f"{d.backend}{d.index}" for d in devices)
2460def _unified_memory_budget(devices: list[FleetDevice]) -> int | None:
2461 """Shared-RAM placement budget (free RAM minus the OS floor), or ``None``.
2463 ``None`` once any device has memory of its own, since dedicated VRAM is the
2464 constraint there rather than system RAM. A host whose only devices are
2465 integrated, and a host with no devices at all, both stay inside the system
2466 budget: their GPU memory is the system's memory.
2467 """
2468 # Only a device with memory of its own lifts the system-RAM constraint. An
2469 # integrated GPU or an Apple Silicon Mac reports a slice of the same RAM the
2470 # OS is using, so treating its total as headroom over-commits the machine by
2471 # roughly the whole system footprint.
2472 if any(not device.unified for device in devices):
2473 return None
2474 return _capped_by_device_memory(
2475 max(0, _plan_free_system_memory() - _system_memory_floor()), devices
2476 )
2479def _unified_admission_budget(devices: list[FleetDevice]) -> int | None:
2480 """Shared-RAM pool a role set is *admitted* against, or ``None`` if dedicated.
2482 Total installed RAM minus the OS floor, not what happens to be free. Sizing
2483 asks a different question and keeps using free RAM: how much context can be
2484 backed right now. Admission asks whether the machine can host this fleet at
2485 all, and the plan defines the whole intended residency, so charging it
2486 against a live figure refuses a 600 MB model on a box that is merely busy at
2487 the moment, which is what happened. The GPU path already charges total
2488 capacity for exactly this reason.
2489 """
2490 if _unified_memory_budget(devices) is None:
2491 return None
2492 return _capped_by_device_memory(
2493 max(0, model_cache.total_system_memory() - _system_memory_floor()), devices
2494 )
2497def _system_memory_floor() -> int:
2498 """RAM held back for the OS when placing against system memory.
2500 ``cfg.system_memory_reserve_gb``, still capped at a quarter of installed RAM:
2501 a fixed reserve leaves a 7-8 GB host with no budget at all and refuses even
2502 tiny models, so the proportional cap holds however the reserve is set.
2503 """
2504 from lilbee.core.config import cfg
2506 total = model_cache.total_system_memory()
2507 return min(int(cfg.system_memory_reserve_gb * 1024**3), total // _SYSTEM_MEMORY_FLOOR_DIVISOR)
2510def _capped_by_device_memory(budget: int, devices: Sequence[FleetDevice]) -> int:
2511 """*budget*, never above what the devices can address between them.
2513 A shared-memory device still has a ceiling of its own: an integrated GPU
2514 addresses a fixed aperture of system RAM, and Metal will not allocate past
2515 ``recommendedMaxWorkingSetSize``. Both report that ceiling as their total, so
2516 a host budget derived from installed RAM promises memory the devices cannot
2517 reach. Unchanged where the engine reports no device, since the fleet is then
2518 running on the CPU and the host figure is the true one.
2519 """
2520 if not devices:
2521 return budget
2522 return min(budget, sum(d.total_bytes for d in devices))
2525def _device_capacity(devices: list[FleetDevice], charge_against_free: bool) -> dict[int, int]:
2526 """Per-device memory placement may charge against, keyed by device index.
2528 A card's total is what it holds, not what is going spare. A compositor, a
2529 browser, or a training job sitting on VRAM is invisible in the total, and the
2530 usable fraction placement applies covers fragmentation and driver overhead
2531 rather than other tenants, so a plan fits on paper and OOMs at load.
2533 Free bytes answer that, but only where they mean "everyone else's residency":
2534 that is the clean-box snapshot, taken after stale servers are reaped and
2535 before anything is built. Read live on a warm box they also exclude the
2536 fleet's own models, and since a plan always describes the complete intended
2537 residency, charging them there would count the fleet against itself and
2538 report a running plan as unplaceable. Those callers keep the total.
2540 Placement applies its usable fraction to whatever this returns, so a card
2541 with a tenant keeps a proportional margin rather than being packed to its
2542 last free byte, where fragmentation is worst.
2543 """
2544 packable = _packable_devices(devices)
2545 if not charge_against_free:
2546 return {d.index: d.total_bytes for d in packable}
2547 return {d.index: min(d.total_bytes, d.free_bytes) for d in packable}
2550def _packable_devices(devices: list[FleetDevice]) -> list[FleetDevice]:
2551 """The devices bin-packing may charge against.
2553 An integrated GPU's memory is the host's. Packing it beside a dedicated card
2554 promises the same RAM twice, once to its own budget and once to everything
2555 else on the machine, and its heap is often the larger number, so the packer
2556 prefers it: a 32 GiB shared heap outbids a 24 GiB card that actually has the
2557 memory. Where a dedicated device exists it is the one to serve from, and the
2558 integrated one is left to the shared-memory budget.
2560 A host with nothing but integrated devices keeps them. There is nothing else
2561 to serve from, and that path is governed by the system budget rather than by
2562 per-device packing.
2563 """
2564 dedicated = [d for d in devices if not d.unified]
2565 return dedicated or devices
2568def _resolve_placement(
2569 placement: PlacementSpec | None,
2570 inputs: list[ModelPlacementInput],
2571 model_refs: dict[WorkerRole, str],
2572 devices: list[FleetDevice],
2573 *,
2574 unified_budget: int | None,
2575 charge_against_free: bool = False,
2576) -> Placement:
2577 """Resolve a Placement from the manual spec when set, else the auto planner."""
2578 estimate_peak = _peak_estimator(model_refs)
2579 capacity = _device_capacity(devices, charge_against_free)
2580 if placement is not None:
2581 return placement_from_spec(
2582 placement,
2583 tuple(model_refs),
2584 capacity,
2585 estimate_peak=estimate_peak,
2586 )
2587 # The chat split's card count is decided against the snapshot's free VRAM (what the
2588 # launch sizes its context against) so placement and launch agree. A split needs
2589 # >=2 GPUs, so skip the chat model's gguf read entirely below that.
2590 chat_ctx_fit, chat_ctx_target = (
2591 _chat_split_ctx_objective(model_refs) if len(capacity) >= _MIN_SPLIT_GPUS else (None, 0)
2592 )
2593 return plan_placement(
2594 inputs,
2595 [(idx, budget) for idx, budget in capacity.items()],
2596 estimate_peak=estimate_peak,
2597 unified_budget=unified_budget,
2598 chat_ctx_fit=chat_ctx_fit,
2599 chat_ctx_target=chat_ctx_target,
2600 free_headroom={d.index: d.free_bytes for d in devices},
2601 )
2604def _placement_or_auto(
2605 placement: PlacementSpec | None,
2606 inputs: list[ModelPlacementInput],
2607 model_refs: dict[WorkerRole, str],
2608 devices: list[FleetDevice],
2609 *,
2610 unified_budget: int | None,
2611 charge_against_free: bool = False,
2612) -> tuple[Placement, bool]:
2613 """Resolve a saved spec, falling back to auto when it no longer fits the hardware.
2615 Returns the placement and whether the spec was the one applied. Hardware moves
2616 under a saved placement: a card is removed, a driver stops enumerating a GPU, a
2617 container starts without one. Refusing to plan there takes chat, embed and
2618 ingest down over a pin set on hardware the host no longer has, so the fleet
2619 degrades to automatic placement and logs why. An interactive apply still fails
2620 loud (:func:`lilbee.app.placement.set_placement`), where the pin is what the
2621 caller just asked for and a silent substitution would be the surprise.
2622 """
2623 if placement is None:
2624 return _resolve_placement(
2625 None,
2626 inputs,
2627 model_refs,
2628 devices,
2629 unified_budget=unified_budget,
2630 charge_against_free=charge_against_free,
2631 ), False
2632 try:
2633 return _resolve_placement(
2634 placement,
2635 inputs,
2636 model_refs,
2637 devices,
2638 unified_budget=unified_budget,
2639 charge_against_free=charge_against_free,
2640 ), True
2641 except PlacementError as exc:
2642 log.warning(
2643 "The saved GPU placement does not fit this hardware (%s); using automatic "
2644 "placement instead. Set a new placement to replace it.",
2645 exc,
2646 )
2647 return _resolve_placement(
2648 None,
2649 inputs,
2650 model_refs,
2651 devices,
2652 unified_budget=unified_budget,
2653 charge_against_free=charge_against_free,
2654 ), False
2657@dataclass(frozen=True)
2658class ResolvedPlacement:
2659 """Devices + resolved instance plans + model refs for the placement view."""
2661 devices: tuple[FleetDevice, ...]
2662 instances: tuple[InstancePlan, ...]
2663 unplaceable_roles: tuple[WorkerRole, ...]
2664 model_refs: dict[WorkerRole, str]
2665 # Roles placed anyway despite not fitting, with the shortfall in bytes. The
2666 # planner has always known this and only logged it, so a surface showed a
2667 # tight role as comfortably placed.
2668 tight_roles: dict[WorkerRole, int] = field(default_factory=dict)
2669 co_tenants: frozenset[WorkerRole] = frozenset()
2670 # False when a spec was given but did not fit the hardware, so these instances
2671 # are the auto planner's and a surface must not present them as the manual plan.
2672 spec_applied: bool = True
2673 # Roles configured but skipped because their model isn't installed (role -> ref).
2674 # Distinct from unplaceable_roles (installed but won't fit); lets a surface show
2675 # "not downloaded" instead of an empty table on a fresh install.
2676 skipped_not_installed: dict[WorkerRole, str] = field(default_factory=dict)
2679def resolve_placement_plan(
2680 placement: PlacementSpec | None, *, fall_back_to_auto: bool = False
2681) -> ResolvedPlacement:
2682 """Probe devices and resolve the auto-or-manual placement, without launching.
2684 ``fall_back_to_auto`` reads *placement* as a saved setting rather than a
2685 request: one that no longer fits the hardware resolves to the auto plan with
2686 ``spec_applied`` False instead of raising.
2687 """
2688 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
2689 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
2691 apply_fleet_gpu_env()
2692 binary = resolve_llama_server()
2693 apply_cuda_runtime_env(binary)
2694 devices = _read_device_cache.get(binary)
2695 unified_budget = _unified_memory_budget(devices)
2696 inputs, model_refs, _, skipped_not_installed = _server_model_inputs(
2697 None, unified_budget=unified_budget, total_vram=sum(d.total_bytes for d in devices)
2698 )
2699 admission_budget = _unified_admission_budget(devices)
2700 if fall_back_to_auto:
2701 resolved, spec_applied = _placement_or_auto(
2702 placement, inputs, model_refs, devices, unified_budget=admission_budget
2703 )
2704 else:
2705 resolved = _resolve_placement(
2706 placement, inputs, model_refs, devices, unified_budget=admission_budget
2707 )
2708 spec_applied = placement is not None
2709 return ResolvedPlacement(
2710 devices=tuple(devices),
2711 instances=resolved.instances,
2712 unplaceable_roles=resolved.unplaceable_roles,
2713 model_refs=model_refs,
2714 co_tenants=resolved.co_tenants,
2715 skipped_not_installed=skipped_not_installed,
2716 spec_applied=spec_applied,
2717 tight_roles=dict(resolved.tight_roles),
2718 )
2721@dataclass(frozen=True)
2722class FleetPlan:
2723 """The servers to start, and the roles that share one swap group."""
2725 launches: tuple[InstanceLaunch, ...]
2726 co_tenants: frozenset[WorkerRole] = frozenset()
2727 # Configured roles left unplaced because their model isn't installed (role ->
2728 # ref), so the warm path can fail a not-installed chat with a named reason
2729 # instead of spinning the warm line forever.
2730 skipped_not_installed: dict[WorkerRole, str] = field(default_factory=dict)
2731 # Launches refused for a window below the minimum grounded prompt
2732 # (role -> user-facing reason with the numbers).
2733 skipped_unusable_ctx: dict[WorkerRole, str] = field(default_factory=dict)
2736def _log_placement_findings(placement: Placement, model_refs: dict[WorkerRole, str]) -> None:
2737 """Warn about placements that exceed the memory budget.
2739 Shared-memory roles that fit nowhere get no server (loading them would OOM the
2740 host). GPU roles are never refused: one whose estimate exceeds the free VRAM
2741 still loads on demand, with a warning carrying the shortfall.
2742 """
2743 for role in placement.unplaceable_roles:
2744 log.warning(
2745 "%s model %s does not fit available memory and will not be served; "
2746 "free up memory or use a smaller model.",
2747 role.value,
2748 model_refs[role],
2749 )
2750 for role, shortfall in placement.tight_roles.items():
2751 log.warning(
2752 "Memory is tight for the %s model %s: it is estimated to need %.1f GiB more "
2753 "GPU memory than is available. It will still load on demand, keeping the "
2754 "layers that fit on the GPU and the rest in system memory; if it runs "
2755 "slowly, free up GPU memory or use a smaller model.",
2756 role.value,
2757 model_refs[role],
2758 # A sub-0.05 GiB shortfall would render as "0.0 GiB more".
2759 max(shortfall / 1024**3, 0.1),
2760 )
2761 if placement.co_tenants:
2762 log.info(
2763 "%s share GPU memory and load on demand; only one is resident at a time.",
2764 ", ".join(sorted(role.value for role in placement.co_tenants)),
2765 )
2768def _unusable_chat_ctx_reason(launch: InstanceLaunch) -> str | None:
2769 """Reason to refuse a chat launch whose window cannot hold a grounded prompt.
2771 ``None`` for non-chat roles, for a window that holds the minimum grounded
2772 prompt, and for user knobs that ask for a smaller one (a ``num_ctx`` pin,
2773 a sub-minimum ``num_ctx_max`` / ``chat_n_ctx_target``).
2774 """
2775 from lilbee.core.config import cfg
2777 if launch.role is not WorkerRole.CHAT or cfg.num_ctx is not None:
2778 return None
2779 needed = engine_params.min_usable_chat_ctx()
2780 # User knobs capping the window below the minimum are honored (the num_ctx
2781 # pin bypasses above).
2782 asked = min(cfg.chat_n_ctx_target, cfg.num_ctx_max or cfg.chat_n_ctx_target)
2783 if asked < needed:
2784 return None
2785 if launch.ctx >= needed:
2786 return None
2787 return (
2788 f"The chat model {launch.model} loads, but the memory left after its weights "
2789 f"backs only a {launch.ctx}-token context, and a grounded answer needs about "
2790 f"{needed} tokens (system prompt, a retrieved source, the question, and room "
2791 "for the answer), so it will not be served. Use a smaller model or a smaller "
2792 "quant, or set num_ctx to force a larger window."
2793 )
2796def plan_launches(
2797 roles: tuple[WorkerRole, ...] | None,
2798 binary: Path,
2799 by_index: dict[int, FleetDevice],
2800 devices: list[FleetDevice],
2801) -> FleetPlan:
2802 """Plan placement for *roles* (``None`` = all configured) and build their launches."""
2803 from lilbee.core.config import cfg
2805 unified_budget = _unified_memory_budget(devices)
2806 inputs, model_refs, reservation, skipped_not_installed = _server_model_inputs(
2807 roles,
2808 unified_budget=unified_budget,
2809 device_count=len(devices),
2810 total_vram=sum(d.total_bytes for d in devices),
2811 )
2812 spec = PlacementSpec.from_json(cfg.placement) if cfg.placement else None
2813 placement, _spec_applied = _placement_or_auto(
2814 spec,
2815 inputs,
2816 model_refs,
2817 devices,
2818 unified_budget=_unified_admission_budget(devices),
2819 # Only the clean-box snapshot's free bytes mean "what other tenants hold";
2820 # a live probe here would also be missing the fleet's own residency.
2821 charge_against_free=_plan_probe_store.get() is not None,
2822 )
2823 _log_placement_findings(placement, model_refs)
2824 reserved_by_device = _non_chat_reservation(placement.instances, inputs, placement.co_tenants)
2825 charged = {inp.role: inp.est_vram_bytes for inp in inputs}
2826 launches: list[InstanceLaunch] = []
2827 skipped_unusable_ctx: dict[WorkerRole, str] = {}
2828 for plan in placement.instances:
2829 launch = _launch_for(
2830 plan,
2831 model_refs[plan.role],
2832 binary,
2833 by_index,
2834 unified_budget=unified_budget,
2835 chat_reservation=reservation,
2836 reserved_by_device=reserved_by_device,
2837 est_vram_bytes=charged.get(plan.role, 0),
2838 )
2839 reason = _unusable_chat_ctx_reason(launch)
2840 if reason is not None:
2841 skipped_unusable_ctx[launch.role] = reason
2842 log.warning(reason)
2843 continue
2844 launches.append(launch)
2845 return FleetPlan(
2846 launches=tuple(launches),
2847 co_tenants=placement.co_tenants,
2848 skipped_not_installed=skipped_not_installed,
2849 skipped_unusable_ctx=skipped_unusable_ctx,
2850 )
2853def plan_all_launches() -> FleetPlan:
2854 """Apply GPU env, probe devices, and plan launches for every configured role.
2856 Disables crash-prone Vulkan layers / dual-vendor ICDs and applies any
2857 ``cfg.gpu_devices`` pin before the probe and plan (both inherit the env).
2858 """
2859 from lilbee.providers.fleet.cuda_runtime import apply_cuda_runtime_env
2860 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
2862 apply_fleet_gpu_env()
2863 binary = resolve_llama_server()
2864 # Put the CUDA-runtime wheels on the process path so the device probe sees the
2865 # same runtime the servers will, before resolve_devices enumerates GPUs.
2866 apply_cuda_runtime_env()
2867 devices = _plan_devices(binary)
2868 by_index = {d.index: d for d in devices}
2869 return plan_launches(None, binary, by_index, devices)