Coverage for src/lilbee/providers/fleet/vram.py: 100%
107 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""gguf-parser-backed, UMA-aware memory estimation for one llama-server instance.
3See docs/architecture.md (VRAM estimation).
4"""
6from __future__ import annotations
8import json
9import subprocess
10from dataclasses import dataclass
11from functools import lru_cache
12from pathlib import Path
14from lilbee.core.config.enums import KvCacheType
15from lilbee.providers.base import ProviderError, ProviderErrorKind
16from lilbee.providers.fleet.adapters import FLAG_BATCH_SIZE, FLAG_UBATCH_SIZE
17from lilbee.providers.fleet.binary import resolve_gguf_parser
18from lilbee.providers.fleet.proc import run_bounded
20# gguf-parser CLI flags (the batch flags are shared with the llama-server argv builder).
21_FLAG_PATH = "--path"
22_FLAG_CTX = "--ctx-size"
23_FLAG_PARALLEL = "--parallel"
24_FLAG_GPU_LAYERS = "--gpu-layers"
25_FLAG_CACHE_K = "--cache-type-k"
26_FLAG_CACHE_V = "--cache-type-v"
27_FLAG_MMPROJ = "--mmproj-path"
28_FLAG_FLASH = "--flash-attention"
29_FLAG_NO_FLASH = "--no-flash-attention"
30_FLAG_TENSOR_SPLIT = "--tensor-split"
31_FLAG_SPLIT_MODE = "--split-mode"
32_SPLIT_MODE_LAYER = "layer"
33_FLAG_JSON = "--json"
34_FLAG_OVERRIDE_TENSOR = "--override-tensor"
35_BUFFER_TYPE_CPU = "CPU"
37# gguf-parser JSON keys: the per-instance footprint lives under estimate.items
38# (v0.24.x) or estimate.memory (upstream's post-v0.24.1 rename of the same list).
39_KEY_ESTIMATE = "estimate"
40_KEY_ITEMS = "items"
41_KEY_MEMORY = "memory"
42_KEY_RAM = "ram"
43_KEY_VRAMS = "vrams"
44_KEY_UMA = "uma"
45_KEY_NONUMA = "nonuma"
47_PROVIDER = "llama-server"
48_PARSE_TIMEOUT_S = 60
49# Bounded wait for a timed-out parser to die before abandoning it (matches the
50# device probe): a parser wedged in driver I/O must not hang the warm-up thread.
51_PARSE_KILL_WAIT_S = 5.0
52# Sized for a whole plan on a wide box, not for one sweep. The ratio ladder and
53# the context bisection key separately (the key carries ctx), and slot fitting
54# adds more, so an eight-GPU chat split touches on the order of a hundred keys.
55# Too small and the winning candidate's keys are evicted before the launch reads
56# them back, which spawns gguf-parser again to recompute what was just measured.
57_CACHE_SIZE = 256
59# Mirrors vLLM's gpu_memory_utilization default: never charge a GPU past 90% of
60# its free VRAM, leaving headroom for allocator fragmentation and driver overhead.
61# The default for cfg.usable_vram_fraction, which is what callers should read.
62USABLE_VRAM_FRACTION = 0.9
65def usable_vram_fraction() -> float:
66 """Share of a card placement may charge.
68 Configurable because it decides admission rather than merely tuning it: at
69 the default, a host whose chat model lands just over the line is refused chat
70 with no way for its owner to say the card has the room.
71 """
72 from lilbee.core.config import cfg
74 return cfg.usable_vram_fraction
77@dataclass(frozen=True)
78class GgufVramEstimate:
79 """One instance's footprint from gguf-parser, for both memory models.
81 ``per_device_*`` carry the per-GPU breakdown gguf-parser returns once a
82 ``tensor_split`` is supplied. A tensor-split instance OOMs on its busiest card,
83 so the planner fits/charges ``peak_footprint`` (the max device), never the sum.
84 """
86 vram_bytes: int
87 """Discrete-GPU model: bytes resident in device VRAM (summed over devices)."""
88 ram_bytes: int
89 """Discrete-GPU model: host RAM bytes (mmap pages, compute buffers)."""
90 unified_bytes: int
91 """Unified-memory model: total resident footprint (RAM + would-be VRAM)."""
92 per_device_vram: tuple[int, ...] = ()
93 """Discrete-GPU VRAM per device (gguf-parser ``vrams[].nonuma``)."""
94 per_device_unified: tuple[int, ...] = ()
95 """Unified-memory footprint per device (``vrams[].uma``)."""
97 def footprint(self, *, unified: bool) -> int:
98 """Total bytes to charge against a shared budget for this memory model."""
99 return self.unified_bytes if unified else self.vram_bytes
101 def peak_footprint(self, *, unified: bool) -> int:
102 """The busiest single device's bytes -- what must fit on one GPU.
104 Falls back to the total when there is no per-device breakdown (a
105 single-device estimate, or an estimate run without a tensor split).
106 """
107 per_device = self.per_device_unified if unified else self.per_device_vram
108 return max(per_device) if per_device else self.footprint(unified=unified)
111def estimate_instance_footprint(
112 model_path: Path,
113 *,
114 ctx: int,
115 slots: int,
116 gpu_layers: int,
117 flash_attn: bool,
118 kv_cache_type: KvCacheType,
119 mmproj_path: Path | None = None,
120 tensor_split: tuple[int, ...] = (),
121 batch_size: int | None = None,
122 expert_offload: tuple[str, ...] = (),
123) -> GgufVramEstimate:
124 """gguf-parser's UMA-aware footprint for one llama-server instance.
126 Pass *tensor_split* (the per-device proportions a multi-GPU instance launches
127 with) so gguf-parser reports the real per-device breakdown; without it the
128 estimate is single-device and the per-GPU peak that actually OOMs is invisible.
129 Pass *batch_size* when the launch raises ``--batch-size``/``--ubatch-size``
130 (pooled embed/rerank), so the compute-buffer estimate matches the launch.
132 With an *mmproj_path* the discrete-GPU number is corrected: gguf-parser's
133 nonuma merge overcharges a multimodal projector by roughly 10 GiB of compute
134 buffer (v0.24.x and current main), while its unified-memory accounting stays
135 accurate, so the projector is re-charged from that side
136 (:func:`_corrected_projector_estimate`).
137 """
139 def run(mmproj: Path | None) -> GgufVramEstimate:
140 return _cached_footprint(
141 engine_build_identity(),
142 str(model_path),
143 model_path.stat().st_mtime_ns,
144 ctx,
145 slots,
146 gpu_layers,
147 flash_attn,
148 kv_cache_type.value,
149 str(mmproj) if mmproj is not None else None,
150 mmproj.stat().st_mtime_ns if mmproj is not None else 0,
151 tensor_split,
152 batch_size,
153 expert_offload,
154 )
156 if mmproj_path is None:
157 return run(None)
158 with_projector = run(mmproj_path)
159 return _corrected_projector_estimate(run(None), with_projector, mmproj_path.stat().st_size)
162def _corrected_projector_estimate(
163 base: GgufVramEstimate, with_projector: GgufVramEstimate, mmproj_bytes: int
164) -> GgufVramEstimate:
165 """Charge the projector at its unified-memory delta, floored at its weights.
167 The floor covers a mmap-shared projector whose uma delta hides weights that
168 still occupy VRAM once offloaded. The charge lands on the first device:
169 llama.cpp loads the projector on the main GPU, not across a split.
170 """
171 projector = max(with_projector.unified_bytes - base.unified_bytes, mmproj_bytes)
172 per_device_vram = tuple(
173 vram + (projector if i == 0 else 0) for i, vram in enumerate(base.per_device_vram)
174 )
175 return GgufVramEstimate(
176 vram_bytes=base.vram_bytes + projector,
177 ram_bytes=with_projector.ram_bytes,
178 unified_bytes=with_projector.unified_bytes,
179 per_device_vram=per_device_vram,
180 per_device_unified=with_projector.per_device_unified,
181 )
184def engine_build_identity() -> str:
185 """Which engine build these numbers describe.
187 Part of the memo key. The estimate prices what one particular llama-server
188 will allocate, and the key held the model, the sizing and the parser's own
189 arguments without a trace of that, so swapping the engine kept the previous
190 engine's answers.
191 """
192 from lilbee.providers.fleet.binary import engine_build_id
194 return engine_build_id()
197@lru_cache(maxsize=_CACHE_SIZE)
198def _cached_footprint(
199 _engine_id: str,
200 path_str: str,
201 _mtime_ns: int,
202 ctx: int,
203 slots: int,
204 gpu_layers: int,
205 flash_attn: bool,
206 kv_cache_type: str,
207 mmproj: str | None,
208 _mmproj_mtime_ns: int,
209 tensor_split: tuple[int, ...],
210 batch_size: int | None,
211 expert_offload: tuple[str, ...],
212) -> GgufVramEstimate:
213 """Memoised gguf-parser run keyed on engine + path + mtime + sizing.
215 The mtime and engine args participate in the cache key only; a re-pulled file
216 at the same path invalidates automatically because its mtime changes, and a
217 swapped engine invalidates because its build identity does.
218 """
219 argv = estimator_argv(
220 path_str,
221 ctx=ctx,
222 slots=slots,
223 gpu_layers=gpu_layers,
224 flash_attn=flash_attn,
225 kv_cache_type=kv_cache_type,
226 mmproj=mmproj,
227 tensor_split=tensor_split,
228 batch_size=batch_size,
229 expert_offload=expert_offload,
230 )
231 return _parse_estimate(_run_parser(argv, path_str), path_str)
234def estimator_argv(
235 path_str: str,
236 *,
237 ctx: int,
238 slots: int,
239 gpu_layers: int,
240 flash_attn: bool,
241 kv_cache_type: str,
242 mmproj: str | None,
243 tensor_split: tuple[int, ...],
244 batch_size: int | None,
245 expert_offload: tuple[str, ...] = (),
246) -> list[str]:
247 """The gguf-parser command line for one instance's sizing parameters.
249 ``ctx`` is the per-slot context, as it is for the launch argv, and
250 ``--ctx-size`` carries the total across slots. The parser's ``--parallel``
251 does not divide the context the way llama-server's does, so the multiply has
252 to reach the parser through ``--ctx-size`` or the attention cache is
253 under-reserved by the slot count.
255 ``--parallel`` is still passed, and it is not inert. A hybrid or recurrent
256 model holds a per-sequence recurrent state that scales with the sequence
257 count and not with the context, exactly as llama.cpp sizes it, so the parser
258 reads the slot count from this flag. Both flags are load-bearing, for
259 different halves of the memory.
260 """
261 argv = [
262 str(resolve_gguf_parser()),
263 _FLAG_PATH,
264 path_str,
265 _FLAG_CTX,
266 str(ctx * slots),
267 _FLAG_PARALLEL,
268 str(slots),
269 _FLAG_GPU_LAYERS,
270 str(gpu_layers),
271 _FLAG_CACHE_K,
272 kv_cache_type,
273 _FLAG_CACHE_V,
274 kv_cache_type,
275 _FLAG_FLASH if flash_attn else _FLAG_NO_FLASH,
276 _FLAG_JSON,
277 ]
278 if batch_size is not None:
279 # Pooled embed/rerank launch with --batch-size/--ubatch-size raised to the
280 # context; the default ubatch (512) would under-estimate their compute buffer.
281 argv += [FLAG_BATCH_SIZE, str(batch_size), FLAG_UBATCH_SIZE, str(batch_size)]
282 if tensor_split:
283 # The split proportions are gguf-parser's only signal for the device count,
284 # so it returns one ``vrams[]`` entry per GPU instead of a single total.
285 argv += [
286 _FLAG_TENSOR_SPLIT,
287 ",".join(str(p) for p in tensor_split),
288 _FLAG_SPLIT_MODE,
289 _SPLIT_MODE_LAYER,
290 ]
291 if expert_offload:
292 # Without these the estimate charges the GPU for experts the launch keeps
293 # in system memory, and the planner sizes slots against a footprint that
294 # never materializes.
295 argv += [
296 _FLAG_OVERRIDE_TENSOR,
297 ",".join(f"{pattern}={_BUFFER_TYPE_CPU}" for pattern in expert_offload),
298 ]
299 if mmproj is not None:
300 argv += [_FLAG_MMPROJ, mmproj]
301 return argv
304def _run_parser(argv: list[str], path_str: str) -> str:
305 """Run gguf-parser, returning its JSON stdout or a user-facing error."""
306 failed = ProviderError(
307 f"Could not size the model {path_str!r}: the memory estimator failed to run.",
308 provider=_PROVIDER,
309 kind=ProviderErrorKind.SERVER,
310 )
311 try:
312 stdout, returncode = run_bounded(
313 argv, timeout_s=_PARSE_TIMEOUT_S, kill_wait_s=_PARSE_KILL_WAIT_S, label="gguf-parser"
314 )
315 except (OSError, subprocess.SubprocessError) as exc:
316 raise failed from exc
317 if returncode != 0:
318 raise failed
319 return stdout
322def _parse_estimate(stdout: str, path_str: str) -> GgufVramEstimate:
323 """Parse gguf-parser JSON into a UMA-aware footprint.
325 Accepts both estimate payload keys: ``items`` (the pinned v0.24.x releases)
326 and ``memory`` (upstream renamed the key after v0.24.1), so an engine built
327 past the pin still sizes instead of failing every launch plan.
328 """
329 try:
330 estimate = json.loads(stdout)[_KEY_ESTIMATE]
331 item = (estimate.get(_KEY_ITEMS) or estimate[_KEY_MEMORY])[0]
332 ram = item[_KEY_RAM]
333 vrams = item[_KEY_VRAMS]
334 per_device_vram = tuple(int(v[_KEY_NONUMA]) for v in vrams)
335 per_device_unified = tuple(int(v[_KEY_UMA]) for v in vrams)
336 return GgufVramEstimate(
337 vram_bytes=sum(per_device_vram),
338 ram_bytes=int(ram[_KEY_NONUMA]),
339 unified_bytes=int(ram[_KEY_UMA]) + sum(per_device_unified),
340 per_device_vram=per_device_vram,
341 per_device_unified=per_device_unified,
342 )
343 except (ValueError, KeyError, IndexError, TypeError) as exc:
344 raise ProviderError(
345 f"Could not size the model {path_str!r}: unexpected estimator output "
346 f"({type(exc).__name__}: {exc}).",
347 provider=_PROVIDER,
348 kind=ProviderErrorKind.SERVER,
349 ) from exc