Coverage for src/lilbee/providers/engine_params.py: 100%
106 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""Engine-neutral parameter helpers: model-path resolution, context and GPU-layer
2sizing, and chat-option translation.
4These derive launch/generation parameters from cfg + a model's GGUF metadata
5without loading the model, so both the local llama-server engine and any other
6provider compute them identically. No native binding.
7"""
9from __future__ import annotations
11import logging
12from dataclasses import dataclass
13from pathlib import Path
14from typing import TYPE_CHECKING, Any
16if TYPE_CHECKING:
17 from lilbee.modelhub.registry import ModelRegistry
18from lilbee.core.config import DEFAULT_NUM_CTX, cfg
19from lilbee.core.config.enums import KV_CACHE_TYPE_BYTES, KvCacheType
20from lilbee.providers.base import (
21 CONTEXT_WINDOW_MARGIN_TOKENS,
22 GENERATION_RESERVE_TOKENS,
23 ProviderError,
24 ProviderErrorKind,
25 estimate_budget_tokens,
26 normalize_generation_options,
27)
28from lilbee.providers.gguf_meta import read_gguf_metadata, train_ctx_from_meta
29from lilbee.providers.model_cache import (
30 compute_dynamic_ctx,
31 get_available_memory,
32 kv_bytes_per_token,
33)
35log = logging.getLogger(__name__)
37EMBED_FALLBACK_CTX = 2048
38"""Context used for embed/rerank when a GGUF reports junk (e.g. context_length=0)."""
40# Sized above chunk_size so BOS re-added on re-tokenization doesn't overflow a full-chunk input.
41_EMBED_CTX_MARGIN = 8
44def resolve_embed_ctx(meta: dict[str, str] | None, model_path: Path) -> int:
45 """Embed/rerank context: worst-case chunk tokenization, capped by trained context.
47 ``chunk_size`` is token-denominated but the chunker enforces a CHARACTER
48 budget (``chunk_size * CHARS_PER_TOKEN``). A BPE token is at least one
49 character, so that char budget is also the provable token ceiling for any
50 chunk the chunker can emit: size the context to it and embed-time
51 truncation becomes impossible. Token-dense text (numeric tables, dense
52 identifiers) otherwise reaches ~2x chunk_size tokens against a 1x cap and
53 silently loses its tail at embed time."""
54 from lilbee.data.extract.chunk import CHARS_PER_TOKEN
56 train_ctx = train_ctx_from_meta(meta, fallback=EMBED_FALLBACK_CTX, model_path=model_path)
57 return min(train_ctx, cfg.chunk_size * CHARS_PER_TOKEN + _EMBED_CTX_MARGIN)
60_LLM_RERANK_HEADROOM = 512
61"""Tokens reserved above chunk_size for an LLM reranker's query, prompt, and 1-token answer."""
64def resolve_llm_rerank_ctx(meta: dict[str, str] | None, model_path: Path) -> int:
65 """LLM-reranker context: a query+candidate pair, capped by the model's trained context."""
66 train_ctx = train_ctx_from_meta(meta, fallback=EMBED_FALLBACK_CTX, model_path=model_path)
67 return min(train_ctx, cfg.chunk_size + _LLM_RERANK_HEADROOM)
70_VISION_FALLBACK_N_CTX = 4096
71"""Context for a vision load when the GGUF reports no usable context_length."""
73_VISION_PAGE_CTX_CAP = 32768
74"""Per-page ceiling on a vision OCR server's context: covers a single high-res page's
75image tokens plus prompt, while keeping a long-context VLM placeable beside a chat giant."""
77N_GPU_LAYERS_AUTO = -1
78"""llama.cpp's "fit as many layers as the device holds" value for n_gpu_layers.
80The engine measures free VRAM at load and picks the count itself, spilling the
81rest to system memory. That is a better answer than any number lilbee can
82compute ahead of time, because it is taken on the real device after every other
83tenant, so the planner passes this rather than a layer count of its own.
84"""
85# llama.cpp's "offload nothing"; the user's CPU-only opt-out rather than a budget.
86_N_GPU_LAYERS_NONE = 0
89def chat_options_to_kwargs(options: dict[str, Any] | None) -> dict[str, Any]:
90 """Translate user-facing chat options into generation kwargs.
92 The output keys (``temperature``/``top_p``/``top_k``/``seed``/``max_tokens``/
93 ``repeat_penalty``) are accepted by llama-server's OpenAI body. ``top_k`` is
94 kept (local llama.cpp honors it), unlike the SDK/API translator which drops it.
95 ``think`` becomes ``chat_template_kwargs.enable_thinking``, which thinking
96 templates honor and others ignore.
97 """
98 kwargs = normalize_generation_options(options)
99 think = kwargs.pop("think", None)
100 if think is not None:
101 kwargs["chat_template_kwargs"] = {"enable_thinking": think}
102 return kwargs
105def resolve_model_path(model: str, registry: ModelRegistry | None = None) -> Path:
106 """Resolve a model name to a .gguf file path.
108 Resolution order: (1) registry (canonical source for installed models),
109 (2) an absolute path to an existing file. Pass *registry* to resolve without
110 reaching for ``get_services()`` (callers running inside its construction).
111 """
112 if not model:
113 raise ProviderError(
114 "No model is configured for this role. Pick one from the catalog "
115 "or run 'lilbee model pull <model>'.",
116 provider="llama-server",
117 kind=ProviderErrorKind.NOT_FOUND,
118 )
119 if registry is None:
120 # call-time import: keeps the app-layer container off this module's import graph
121 from lilbee.app.services import get_services
123 registry = get_services().registry
124 try:
125 return registry.resolve(model)
126 except (KeyError, ValueError):
127 pass
129 candidate = Path(model)
130 if candidate.is_absolute():
131 if candidate.exists():
132 return candidate
133 raise ProviderError(
134 f"Model file not found: {model}",
135 provider="llama-server",
136 kind=ProviderErrorKind.NOT_FOUND,
137 )
139 raise ProviderError(
140 f"Model {model!r} is not installed. Run 'lilbee model pull {model}' to download it.",
141 provider="llama-server",
142 kind=ProviderErrorKind.NOT_FOUND,
143 )
146def chat_kv_elem_bytes() -> tuple[float, float]:
147 """Per-element (K, V) byte costs of the KV cache a chat launch allocates.
149 Reads the same flags the launch passes
150 (:func:`lilbee.providers.fleet.planning.chat_cache_type_flags`): K carries
151 ``cfg.kv_cache_type``, while V carries it only when flash attention is
152 certain to be on, because llama.cpp refuses a quantized V cache without it
153 and the launch then leaves V at f16. Budgeting from the launch flags keeps
154 the granted window in step with the cache the engine actually allocates.
155 """
156 # call-time import: planning imports this module at load
157 from lilbee.providers.fleet.planning import chat_cache_type_flags
159 def elem_bytes(flag: str | None) -> float:
160 return KV_CACHE_TYPE_BYTES[KvCacheType(flag) if flag else KvCacheType.F16]
162 k_flag, v_flag = chat_cache_type_flags()
163 return elem_bytes(k_flag), elem_bytes(v_flag)
166def chat_ctx_ceiling(meta: dict[str, str] | None, model_path: Path) -> int:
167 """Hard upper bound on a chat per-slot n_ctx: trained context, capped by ``cfg.num_ctx_max``."""
168 training_ctx = train_ctx_from_meta(meta, fallback=DEFAULT_NUM_CTX, model_path=model_path)
169 if cfg.num_ctx_max is not None:
170 return min(training_ctx, cfg.num_ctx_max)
171 return training_ctx
174@dataclass(frozen=True)
175class ChatFit:
176 """The GPU offload and per-slot window one chat launch runs with.
178 The pair is inseparable: the window was sized against the memory this
179 offload leaves free, so serving it at any other offload overruns the card.
180 """
182 gpu_layers: int
183 ctx: int
186def resolve_chat_fit(
187 model_path: Path, meta: dict[str, str] | None, *, available_bytes: int | None = None
188) -> ChatFit:
189 """Pick a single-GPU offload and n_ctx aiming for ``cfg.chat_n_ctx_target``,
190 clamped to model + host.
192 A gguf-parser fit answers first
193 (:func:`lilbee.providers.fleet.planning.fit_chat_ctx`), because it prices the
194 cache each layer of this architecture holds, and it may leave layers in
195 system memory to free the KV room a usable window needs. Header math takes
196 over when the estimator cannot answer; it charges every layer as dense
197 attention over the whole window at the configured offload, which
198 under-grants linear-attention, sliding-window and MLA models. Either way the
199 window stops at the smallest of the trained context, ``cfg.num_ctx_max`` and
200 the target.
202 A multi-GPU tensor-split chat is sized separately by the fleet against its
203 per-device headroom (see :func:`lilbee.providers.fleet.ctx.fit_split_ctx`).
204 ``available_bytes`` overrides the live host-memory read, and every caller
205 that is sizing a real launch passes it: the fleet and the surfaces that
206 mirror it hand over
207 :func:`lilbee.providers.fleet.planning.plan_sizing_budget`, which reports the
208 memory of the GPU that will run the model and holds a clean-box snapshot so
209 a reload sizes ctx like the boot did.
210 """
211 # call-time import: planning imports this module at load
212 from lilbee.providers.fleet.planning import fit_chat_ctx
214 training_ctx = train_ctx_from_meta(meta, fallback=DEFAULT_NUM_CTX, model_path=model_path)
215 ceiling = cfg.num_ctx_max if cfg.num_ctx_max is not None else training_ctx
216 if available_bytes is None:
217 available_bytes = get_available_memory(cfg.gpu_memory_fraction)
218 upper = min(training_ctx, ceiling, cfg.chat_n_ctx_target)
220 try:
221 return fit_chat_ctx(model_path, meta, available_bytes=available_bytes, ctx_ceiling=upper)
222 except (ProviderError, OSError, ValueError):
223 log.debug("gguf-parser ctx fit failed for %s, using header math", model_path, exc_info=True)
224 return ChatFit(
225 resolve_n_gpu_layers(embedding=False),
226 _header_math_chat_ctx(
227 model_path,
228 meta,
229 available_bytes=available_bytes,
230 training_ctx=training_ctx,
231 ceiling=ceiling,
232 ),
233 )
236def _header_math_chat_ctx(
237 model_path: Path,
238 meta: dict[str, str] | None,
239 *,
240 available_bytes: int,
241 training_ctx: int,
242 ceiling: int,
243) -> int:
244 """Window from GGUF header arithmetic: weights plus a dense-attention cache."""
245 try:
246 model_bytes = model_path.stat().st_size
247 kv_per_tok = kv_bytes_per_token(meta, *chat_kv_elem_bytes())
248 return compute_dynamic_ctx(
249 model_bytes=model_bytes,
250 available_bytes=available_bytes,
251 training_ctx=training_ctx,
252 kv_bytes_per_tok=kv_per_tok,
253 ceiling=ceiling,
254 target=cfg.chat_n_ctx_target,
255 )
256 except (OSError, ValueError):
257 log.debug("dynamic ctx sizing failed for %s, using static cap", model_path, exc_info=True)
258 return min(training_ctx, cfg.chat_n_ctx_target)
261def resolve_chat_ctx(
262 model_path: Path, meta: dict[str, str] | None, *, available_bytes: int | None = None
263) -> int:
264 """The window half of :func:`resolve_chat_fit`, for callers that only serve
265 or advertise the context."""
266 return resolve_chat_fit(model_path, meta, available_bytes=available_bytes).ctx
269# Tokens the minimum grounded prompt allows for the question plus the context
270# template's framing, beyond the system prompt and one retrieved source.
271_GROUNDED_QUESTION_TOKENS = 128
274def min_usable_chat_ctx() -> int:
275 """Smallest chat window that serves one grounded answer: the system prompt,
276 one retrieved source, the question, and the generation reserve plus margin."""
277 return (
278 estimate_budget_tokens(cfg.rag_system_prompt)
279 + cfg.chunk_size
280 + _GROUNDED_QUESTION_TOKENS
281 + GENERATION_RESERVE_TOKENS
282 + CONTEXT_WINDOW_MARGIN_TOKENS
283 )
286def resolve_n_gpu_layers(*, embedding: bool) -> int:
287 """Resolve ``cfg.n_gpu_layers`` (None=all) to llama.cpp's offload integer.
289 Zero is honoured for every role. It is not a layer budget but the way a user
290 says "run this on the CPU", and the search roles used to take the
291 full-offload sentinel before the setting was read, so embed, rerank and
292 vision kept loading onto the GPU that had just been excluded.
294 Any other value is a chat-shaped budget and says nothing useful about a small
295 embedding model, which still offloads fully.
296 """
297 if cfg.n_gpu_layers == _N_GPU_LAYERS_NONE:
298 return _N_GPU_LAYERS_NONE
299 if embedding or cfg.n_gpu_layers is None:
300 return N_GPU_LAYERS_AUTO
301 return cfg.n_gpu_layers
304def resolve_vision_ctx(model_path: Path) -> int:
305 """Pick n_ctx for a vision OCR load: the model's training context, capped per page.
307 Uses the model's ``<arch>.context_length`` (not the chat-tuned ``cfg.num_ctx``: a
308 vision pass packs image-token embeddings plus the prompt, and a small chat ctx
309 truncates OCR output) but caps it at ``_VISION_PAGE_CTX_CAP``. OCR processes one page
310 per request, so a single page never exceeds the cap, yet a long-context VLM's full
311 context would otherwise estimate too large to place alongside a chat giant.
312 """
313 try:
314 meta = read_gguf_metadata(model_path)
315 except Exception:
316 log.debug("read_gguf_metadata failed for vision %s", model_path, exc_info=True)
317 meta = None
318 train_ctx = train_ctx_from_meta(meta, fallback=_VISION_FALLBACK_N_CTX, model_path=model_path)
319 return min(train_ctx, _VISION_PAGE_CTX_CAP)