Coverage for src/lilbee/providers/fleet/adapters.py: 100%
99 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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
11import logging
12from dataclasses import dataclass, replace
13from enum import StrEnum
14from pathlib import Path
16from lilbee.core.config.enums import RerankerType
17from lilbee.providers.roles import RerankMode, WorkerRole
18from lilbee.runtime.cpu import engine_thread_count
20log = logging.getLogger(__name__)
22_HOST = "127.0.0.1"
23# llama-server batch flags; gguf-parser accepts the same names, so vram.py shares these.
24FLAG_BATCH_SIZE = "--batch-size"
25FLAG_UBATCH_SIZE = "--ubatch-size"
28@dataclass(frozen=True)
29class RoleServerSpec:
30 """How one role maps onto a llama-server instance."""
32 role: WorkerRole
33 endpoint_path: str
34 extra_args: tuple[str, ...]
35 server_capable: bool
38# Every role runs on the fleet by mirroring the in-process primitive over HTTP:
39# rerank uses rank-pooling embeddings (--pooling rank -> /v1/embeddings, with the
40# same query</s></s>candidate pairing as in-process), NOT the template-dependent
41# /v1/rerank; vision uses the chat endpoint with an --mmproj projector. This keeps
42# the in-process robustness without depending on a model's embedded rerank template.
43ROLE_SPECS: dict[WorkerRole, RoleServerSpec] = {
44 WorkerRole.CHAT: RoleServerSpec(
45 role=WorkerRole.CHAT,
46 endpoint_path="/v1/chat/completions",
47 # --jinja renders the model's own chat template and parses native
48 # tool-call syntax into structured message.tool_calls.
49 # --reasoning-format deepseek makes the server parse every model's native
50 # reasoning dialect (<think>, gpt-oss harmony, ...) into reasoning_content;
51 # the chat client re-inlines it as <think> so downstream parsing stays
52 # format-agnostic and control tokens never leak into answers.
53 # --no-prefill-assistant keeps a trailing assistant message a finished turn:
54 # by default the server continues its text instead of answering, and rejects
55 # two trailing assistant messages with a 400. Agents compacting their history
56 # send both shapes, and OpenAI's API accepts them.
57 extra_args=("--jinja", "--reasoning-format", "deepseek", "--no-prefill-assistant"),
58 server_capable=True,
59 ),
60 WorkerRole.EMBED: RoleServerSpec(
61 role=WorkerRole.EMBED,
62 endpoint_path="/v1/embeddings",
63 extra_args=("--embeddings",),
64 server_capable=True,
65 ),
66 WorkerRole.RERANK: RoleServerSpec(
67 role=WorkerRole.RERANK,
68 endpoint_path="/v1/embeddings",
69 extra_args=("--embeddings", "--pooling", "rank"),
70 server_capable=True,
71 ),
72 WorkerRole.VISION: RoleServerSpec(
73 role=WorkerRole.VISION,
74 endpoint_path="/v1/chat/completions",
75 extra_args=(),
76 server_capable=True,
77 ),
78}
80# Decoder-only architectures. Served generatively as rerankers (Qwen3-Reranker,
81# mxbai-rerank-v2), and their embeddings use last-token (EOS) pooling, not the
82# encoder default of mean/cls.
83_DECODER_ARCHS: frozenset[str] = frozenset(
84 {"qwen2", "qwen3", "llama", "mistral", "gemma", "gemma2", "gemma3", "phi3"}
85)
87LLM_RERANK_SPEC = RoleServerSpec(
88 role=WorkerRole.RERANK,
89 endpoint_path="/v1/chat/completions",
90 extra_args=("--jinja",),
91 server_capable=True,
92)
94_RERANK_MODE_SPECS: dict[RerankMode, RoleServerSpec] = {
95 RerankMode.CROSS_ENCODER: ROLE_SPECS[WorkerRole.RERANK],
96 RerankMode.LLM: LLM_RERANK_SPEC,
97}
99# An LLM reranker scores one chat request per candidate; this is both the client's
100# per-rerank request fan-out and the server's --parallel slot ceiling, so the
101# server can decode concurrently instead of serializing the fan-out.
102LLM_RERANK_CONCURRENCY = 8
104_FLAG_MAIN_GPU = "--main-gpu"
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(
183 flash_attn: str | None, cache_type_k: str | None, cache_type_v: str | None
184) -> list[str]:
185 """Flash-attention and KV cache flags; each stays absent to keep the engine default."""
186 args: list[str] = []
187 if flash_attn is not None:
188 args += ["--flash-attn", flash_attn]
189 if cache_type_k is not None:
190 args += ["--cache-type-k", cache_type_k]
191 if cache_type_v is not None:
192 args += ["--cache-type-v", cache_type_v]
193 return args
196def _main_gpu_args(devices: tuple[int, ...]) -> list[str]:
197 """``--main-gpu`` for this instance, empty when it does not apply.
199 The index is into this instance's own device list, which is the space its
200 visibility pin exposes to the engine, and it is what the setting's help text
201 already describes. llama.cpp ignores the flag with a single device, so it is
202 only emitted where it changes something: which card holds the model under
203 split-mode none, and the intermediate results and KV under split-mode row.
205 An index past the end is refused and said out loud. It was previously not
206 emitted at all, so a user could set it, watch it save, and get nothing.
207 """
208 from lilbee.core.config import cfg
210 main_gpu = cfg.main_gpu
211 if main_gpu is None or len(devices) <= 1:
212 return []
213 if not 0 <= main_gpu < len(devices):
214 log.warning(
215 "Ignoring main_gpu=%d: this server runs on %d device(s), so the index has "
216 "to be between 0 and %d. It counts within the cards this role was placed "
217 "on, not across every card in the machine.",
218 main_gpu,
219 len(devices),
220 len(devices) - 1,
221 )
222 return []
223 return [_FLAG_MAIN_GPU, str(main_gpu)]
226def _thread_args() -> list[str]:
227 """``--threads`` when a CPU cap applies to this process, else empty.
229 llama.cpp counts host cores and sees neither a cgroup quota nor an affinity
230 mask, so a container-bound engine oversubscribes its quota badly. Only the
231 generation flag is emitted: ``--threads-batch`` defaults to ``--threads``.
232 """
233 threads = engine_thread_count()
234 return [] if threads is None else [_FLAG_THREADS, str(threads)]
237def build_server_argv(
238 *,
239 binary: Path,
240 spec: RoleServerSpec,
241 model_path: Path,
242 devices: tuple[int, ...],
243 device_names: tuple[str, ...] = (),
244 n_gpu_layers: int,
245 slots: int,
246 ctx_per_slot: int,
247 tensor_split: tuple[int, ...] = (),
248 mmproj: Path | None = None,
249 flash_attn: str | None = None,
250 cache_type_k: str | None = None,
251 cache_type_v: str | None = None,
252 batch_size: int | None = None,
253 no_mmap: bool = False,
254 cpu_moe: bool = False,
255 n_cpu_moe: int | None = None,
256) -> list[str]:
257 """Assemble the llama-server command line for one instance, minus ``--port``.
259 ``--ctx-size`` is the per-slot context times the slot count, since
260 llama-server divides total context across parallel slots. ``n_cpu_moe``
261 wins over ``cpu_moe``; the pair would offload the same tensors twice.
263 """
264 argv = [
265 str(binary),
266 "--model",
267 str(model_path),
268 "--host",
269 _HOST,
270 "--n-gpu-layers",
271 str(n_gpu_layers),
272 "--parallel",
273 str(slots),
274 "--cont-batching",
275 "--ctx-size",
276 str(ctx_per_slot * slots),
277 ]
278 argv += _main_gpu_args(devices)
279 argv += _thread_args()
280 argv += _attention_args(flash_attn, cache_type_k, cache_type_v)
281 if batch_size is not None:
282 argv += [FLAG_BATCH_SIZE, str(batch_size), FLAG_UBATCH_SIZE, str(batch_size)]
283 if mmproj is not None: # vision: the CLIP/mtmd projector sidecar
284 argv += ["--mmproj", str(mmproj)]
285 if device_names:
286 # Names as --list-devices prints them, which is the space they were parsed
287 # from. The Vulkan visible-devices variable takes RAW loader indices
288 # instead, so re-emitting parsed ordinals into it silently changes index
289 # space, and setting it also turns off ggml's own type filter, support
290 # check and same-UUID dedup. This keeps all of that active.
291 argv += ["--device", ",".join(device_names)]
292 if len(devices) > 1 and tensor_split:
293 # Only a ratio the planner actually chose. Inventing an even one for a
294 # group that did not choose disables the engine's own fit: it aborts the
295 # fit pass when tensor_split is user-set, and a negative n_gpu_layers
296 # then means every layer. The one group that arrives without a ratio is
297 # the tight placement, whose whole promise is that the engine keeps what
298 # fits and spills the rest, so an invented split turns that promise into
299 # a load-time out-of-memory.
300 argv += ["--tensor-split", ",".join(str(r) for r in tensor_split)]
301 if no_mmap:
302 argv += ["--no-mmap"]
303 if n_cpu_moe is not None:
304 argv += ["--n-cpu-moe", str(n_cpu_moe)]
305 elif cpu_moe:
306 argv += ["--cpu-moe"]
307 argv += list(spec.extra_args)
308 return argv