Coverage for src/lilbee/providers/fleet/adapters.py: 100%
109 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Per-role llama-server specs and the argv builder for the fleet.
3A data table (not per-role functions) keyed by ``WorkerRole``: each spec carries
4the OpenAI endpoint path, the role-specific server flags, and whether the role is
5viable on a server today. ``build_server_argv`` reads a spec plus placement data
6to assemble one llama-server command line.
7"""
9from __future__ import annotations
11from dataclasses import dataclass, replace
12from enum import StrEnum
13from pathlib import Path
15from lilbee.core.config.enums import RerankerType
16from lilbee.providers.fleet.binary import engine_help, help_lists_option
17from lilbee.providers.roles import RerankMode, WorkerRole
18from lilbee.runtime.cpu import engine_thread_count
20_HOST = "127.0.0.1"
21# llama-server batch flags; gguf-parser accepts the same names, so vram.py shares these.
22FLAG_BATCH_SIZE = "--batch-size"
23FLAG_UBATCH_SIZE = "--ubatch-size"
24FLAG_FIT_TARGET = "--fit-target"
25FLAG_KV_UNIFIED = "--kv-unified"
26FLAG_KV_UNIFIED_PER_SLOT = "--kv-unified-per-slot"
29@dataclass(frozen=True)
30class RoleServerSpec:
31 """How one role maps onto a llama-server instance."""
33 role: WorkerRole
34 endpoint_path: str
35 extra_args: tuple[str, ...]
36 server_capable: bool
39# Every role runs on the fleet by mirroring the in-process primitive over HTTP:
40# rerank uses rank-pooling embeddings (--pooling rank -> /v1/embeddings, with the
41# same query</s></s>candidate pairing as in-process), NOT the template-dependent
42# /v1/rerank; vision uses the chat endpoint with an --mmproj projector. This keeps
43# the in-process robustness without depending on a model's embedded rerank template.
44ROLE_SPECS: dict[WorkerRole, RoleServerSpec] = {
45 WorkerRole.CHAT: RoleServerSpec(
46 role=WorkerRole.CHAT,
47 endpoint_path="/v1/chat/completions",
48 # --jinja renders the model's own chat template and parses native
49 # tool-call syntax into structured message.tool_calls.
50 # --reasoning-format deepseek makes the server parse every model's native
51 # reasoning dialect (<think>, gpt-oss harmony, ...) into reasoning_content;
52 # the chat client re-inlines it as <think> so downstream parsing stays
53 # format-agnostic and control tokens never leak into answers.
54 # --no-prefill-assistant keeps a trailing assistant message a finished turn:
55 # by default the server continues its text instead of answering, and rejects
56 # two trailing assistant messages with a 400. Agents compacting their history
57 # send both shapes, and OpenAI's API accepts them.
58 extra_args=("--jinja", "--reasoning-format", "deepseek", "--no-prefill-assistant"),
59 server_capable=True,
60 ),
61 WorkerRole.EMBED: RoleServerSpec(
62 role=WorkerRole.EMBED,
63 endpoint_path="/v1/embeddings",
64 extra_args=("--embeddings",),
65 server_capable=True,
66 ),
67 WorkerRole.RERANK: RoleServerSpec(
68 role=WorkerRole.RERANK,
69 endpoint_path="/v1/embeddings",
70 extra_args=("--embeddings", "--pooling", "rank"),
71 server_capable=True,
72 ),
73 WorkerRole.VISION: RoleServerSpec(
74 role=WorkerRole.VISION,
75 endpoint_path="/v1/chat/completions",
76 extra_args=(),
77 server_capable=True,
78 ),
79}
81# Decoder-only architectures. Served generatively as rerankers (Qwen3-Reranker,
82# mxbai-rerank-v2), and their embeddings use last-token (EOS) pooling, not the
83# encoder default of mean/cls.
84_DECODER_ARCHS: frozenset[str] = frozenset(
85 {"qwen2", "qwen3", "llama", "mistral", "gemma", "gemma2", "gemma3", "phi3"}
86)
88LLM_RERANK_SPEC = RoleServerSpec(
89 role=WorkerRole.RERANK,
90 endpoint_path="/v1/chat/completions",
91 extra_args=("--jinja",),
92 server_capable=True,
93)
95_RERANK_MODE_SPECS: dict[RerankMode, RoleServerSpec] = {
96 RerankMode.CROSS_ENCODER: ROLE_SPECS[WorkerRole.RERANK],
97 RerankMode.LLM: LLM_RERANK_SPEC,
98}
100# An LLM reranker scores one chat request per candidate; this is both the client's
101# per-rerank request fan-out and the server's --parallel slot ceiling, so the
102# server can decode concurrently instead of serializing the fan-out.
103LLM_RERANK_CONCURRENCY = 8
105_FLAG_THREADS = "--threads"
108def resolve_rerank_mode(reranker_type: RerankerType, arch: str | None) -> RerankMode:
109 """Pick the reranker serving mode from the config setting and GGUF arch.
111 ``auto`` serves a known decoder arch generatively; encoder/unknown archs stay
112 cross-encoder. Explicit settings override the arch.
113 """
114 if reranker_type is RerankerType.LLM:
115 return RerankMode.LLM
116 if reranker_type is RerankerType.CROSS_ENCODER:
117 return RerankMode.CROSS_ENCODER
118 if arch in _DECODER_ARCHS:
119 return RerankMode.LLM
120 return RerankMode.CROSS_ENCODER
123def rerank_spec(mode: RerankMode) -> RoleServerSpec:
124 """The server spec for a RERANK launch given its resolved mode."""
125 return _RERANK_MODE_SPECS[mode]
128class PoolingType(StrEnum):
129 """A llama-server ``--pooling`` value (also the GGUF pooling_type enum names)."""
131 NONE = "none"
132 MEAN = "mean"
133 CLS = "cls"
134 LAST = "last"
135 RANK = "rank"
138# GGUF <arch>.pooling_type integer (as read_gguf_metadata returns it, a string) ->
139# the --pooling value. NONE (0) is omitted: a 0 on an embedder is the unset default,
140# so it falls through to the arch-based choice rather than per-token (non-)pooling.
141_GGUF_POOLING: dict[str, PoolingType] = {
142 "1": PoolingType.MEAN,
143 "2": PoolingType.CLS,
144 "3": PoolingType.LAST,
145 "4": PoolingType.RANK,
146}
149def embed_spec(meta: dict[str, str] | None) -> RoleServerSpec:
150 """The EMBED server spec with the model's pooling resolved for llama-server."""
151 # The GGUF's declared pooling wins; else a decoder-only embedder pools on its
152 # last/EOS token (llama-server would otherwise default to mean, which is wrong).
153 arch = meta.get("architecture") if meta else None
154 pooling_type = meta.get("pooling_type") if meta else None
155 pooling = _GGUF_POOLING.get(pooling_type or "")
156 if pooling is None and arch in _DECODER_ARCHS:
157 pooling = PoolingType.LAST
158 base = ROLE_SPECS[WorkerRole.EMBED]
159 if pooling is None:
160 return base
161 return replace(base, extra_args=(*base.extra_args, "--pooling", pooling.value))
164# The expert tensors --cpu-moe/--n-cpu-moe move to system memory, copied from
165# llama.cpp's LLM_FFN_EXPS_REGEX (common/common.h). The estimator is handed the
166# same patterns so its sizing matches what the launch actually offloads; they
167# must stay in step with upstream or the estimate silently drifts from reality.
168EXPERT_TENSOR_REGEX = r"\.ffn_(up|down|gate|gate_up)_(ch|)exps"
171def expert_offload_patterns(*, cpu_moe: bool, n_cpu_moe: int | None) -> tuple[str, ...]:
172 """Tensor-name patterns whose experts live in system memory, launch order.
174 Mirrors llama.cpp's expansion: ``--cpu-moe`` is one blanket pattern, while
175 ``--n-cpu-moe N`` is one per-block pattern for the first N blocks.
176 """
177 if n_cpu_moe is not None:
178 return tuple(rf"blk\.{i}{EXPERT_TENSOR_REGEX}" for i in range(n_cpu_moe))
179 return (EXPERT_TENSOR_REGEX,) if cpu_moe else ()
182def _attention_args(flash_attn: str | None, cache_type: str | None) -> list[str]:
183 """Flash-attention and KV cache flags; each stays absent to keep the engine default.
185 One *cache_type* goes to both K and V: llama-server refuses K != V on MLA models.
186 """
187 args: list[str] = []
188 if flash_attn is not None:
189 args += ["--flash-attn", flash_attn]
190 if cache_type is not None:
191 args += ["--cache-type-k", cache_type, "--cache-type-v", cache_type]
192 return args
195def _fit_and_kv_pool_args(
196 fit_target_mib: tuple[int, ...], kv_unified_per_slot: int | None
197) -> list[str]:
198 """``--fit-target`` per device and the shared KV pool flags, each only when given."""
199 args: list[str] = []
200 if fit_target_mib:
201 args += [FLAG_FIT_TARGET, ",".join(str(mib) for mib in fit_target_mib)]
202 if kv_unified_per_slot is not None:
203 args += [FLAG_KV_UNIFIED, FLAG_KV_UNIFIED_PER_SLOT, str(kv_unified_per_slot)]
204 return args
207def _thread_args() -> list[str]:
208 """``--threads`` when a CPU cap applies to this process, else empty.
210 llama.cpp counts host cores and sees neither a cgroup quota nor an affinity
211 mask, so a container-bound engine oversubscribes its quota badly. Only the
212 generation flag is emitted: ``--threads-batch`` defaults to ``--threads``.
213 """
214 threads = engine_thread_count()
215 return [] if threads is None else [_FLAG_THREADS, str(threads)]
218# The flags that load the weights into a host copy instead of mapping the file.
219# llama.cpp replaced --no-mmap with --load-mode, and each build refuses the
220# other's flag at startup, so the binary's own --help picks one.
221_FLAG_LOAD_MODE = "--load-mode"
222_LOAD_MODE_NO_MMAP = "none"
223_FLAG_NO_MMAP = "--no-mmap"
226def _no_mmap_args(binary: Path) -> list[str]:
227 """The no-mmap flag *binary* lists in its ``--help``, ``--load-mode`` first; else empty."""
228 help_text = engine_help(binary)
229 if help_lists_option(help_text, _FLAG_LOAD_MODE):
230 return [_FLAG_LOAD_MODE, _LOAD_MODE_NO_MMAP]
231 if help_lists_option(help_text, _FLAG_NO_MMAP):
232 return [_FLAG_NO_MMAP]
233 return []
236def build_server_argv(
237 *,
238 binary: Path,
239 spec: RoleServerSpec,
240 model_path: Path,
241 devices: tuple[int, ...],
242 device_names: tuple[str, ...] = (),
243 n_gpu_layers: int,
244 slots: int,
245 ctx_per_slot: int,
246 tensor_split: tuple[int, ...] = (),
247 mmproj: Path | None = None,
248 flash_attn: str | None = None,
249 cache_type: str | None = None,
250 batch_size: int | None = None,
251 no_mmap: bool = False,
252 cpu_moe: bool = False,
253 n_cpu_moe: int | None = None,
254 memory_endpoint: bool = False,
255 fit_target_mib: tuple[int, ...] = (),
256 kv_unified_per_slot: int | None = None,
257) -> list[str]:
258 """Assemble the llama-server command line for one instance, minus ``--port``.
260 ``--ctx-size`` is the per-slot context times the slot count, since
261 llama-server divides total context across parallel slots. ``n_cpu_moe``
262 wins over ``cpu_moe``; the pair would offload the same tensors twice.
263 ``memory_endpoint`` is set only for a binary whose ``--help`` advertises
264 ``--memory`` (readback.supports_memory_readback): a stock engine dies at
265 launch on the unknown flag. ``no_mmap`` sends the no-mmap flag the binary's
266 ``--help`` lists, or nothing when it lists neither. ``fit_target_mib`` is
267 the memory the engine's own fit leaves free on each device, in
268 ``--device`` order.
269 ``kv_unified_per_slot`` makes the slots share one KV pool, each slot capped
270 at that window.
271 """
272 argv = [
273 str(binary),
274 "--model",
275 str(model_path),
276 "--host",
277 _HOST,
278 "--n-gpu-layers",
279 str(n_gpu_layers),
280 "--parallel",
281 str(slots),
282 "--cont-batching",
283 "--ctx-size",
284 str(ctx_per_slot * slots),
285 ]
286 argv += _thread_args()
287 argv += _attention_args(flash_attn, cache_type)
288 if batch_size is not None:
289 argv += [FLAG_BATCH_SIZE, str(batch_size), FLAG_UBATCH_SIZE, str(batch_size)]
290 if mmproj is not None: # vision: the CLIP/mtmd projector sidecar
291 argv += ["--mmproj", str(mmproj)]
292 if device_names:
293 # Names as --list-devices prints them, which is the space they were parsed
294 # from. The Vulkan visible-devices variable takes RAW loader indices
295 # instead, so re-emitting parsed ordinals into it silently changes index
296 # space, and setting it also turns off ggml's own type filter, support
297 # check and same-UUID dedup. This keeps all of that active.
298 argv += ["--device", ",".join(device_names)]
299 if len(devices) > 1 and tensor_split:
300 # Only a ratio the planner actually chose. Inventing an even one for a
301 # group that did not choose disables the engine's own fit: it aborts the
302 # fit pass when tensor_split is user-set, and a negative n_gpu_layers
303 # then means every layer. The one group that arrives without a ratio is
304 # the tight placement, whose whole promise is that the engine keeps what
305 # fits and spills the rest, so an invented split turns that promise into
306 # a load-time out-of-memory.
307 argv += ["--tensor-split", ",".join(str(r) for r in tensor_split)]
308 if no_mmap:
309 argv += _no_mmap_args(binary)
310 if n_cpu_moe is not None:
311 argv += ["--n-cpu-moe", str(n_cpu_moe)]
312 elif cpu_moe:
313 argv += ["--cpu-moe"]
314 argv += _fit_and_kv_pool_args(fit_target_mib, kv_unified_per_slot)
315 if memory_endpoint:
316 from lilbee.providers.fleet.readback import MEMORY_FLAG
318 argv += [MEMORY_FLAG]
319 argv += list(spec.extra_args)
320 return argv