Coverage for src/lilbee/providers/fleet/vram.py: 100%
107 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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 kv_cache_type_v: KvCacheType | None = None,
120 mmproj_path: Path | None = None,
121 tensor_split: tuple[int, ...] = (),
122 batch_size: int | None = None,
123 expert_offload: tuple[str, ...] = (),
124) -> GgufVramEstimate:
125 """gguf-parser's UMA-aware footprint for one llama-server instance.
127 Pass *tensor_split* (the per-device proportions a multi-GPU instance launches
128 with) so gguf-parser reports the real per-device breakdown; without it the
129 estimate is single-device and the per-GPU peak that actually OOMs is invisible.
130 Pass *batch_size* when the launch raises ``--batch-size``/``--ubatch-size``
131 (pooled embed/rerank), so the compute-buffer estimate matches the launch.
133 With an *mmproj_path* the discrete-GPU number is corrected: gguf-parser's
134 nonuma merge overcharges a multimodal projector by roughly 10 GiB of compute
135 buffer (v0.24.x and current main), while its unified-memory accounting stays
136 accurate, so the projector is re-charged from that side
137 (:func:`_corrected_projector_estimate`).
138 """
140 def run(mmproj: Path | None) -> GgufVramEstimate:
141 return _cached_footprint(
142 engine_build_identity(),
143 str(model_path),
144 model_path.stat().st_mtime_ns,
145 ctx,
146 slots,
147 gpu_layers,
148 flash_attn,
149 kv_cache_type.value,
150 (kv_cache_type_v or kv_cache_type).value,
151 str(mmproj) if mmproj is not None else None,
152 mmproj.stat().st_mtime_ns if mmproj is not None else 0,
153 tensor_split,
154 batch_size,
155 expert_offload,
156 )
158 if mmproj_path is None:
159 return run(None)
160 with_projector = run(mmproj_path)
161 return _corrected_projector_estimate(run(None), with_projector, mmproj_path.stat().st_size)
164def _corrected_projector_estimate(
165 base: GgufVramEstimate, with_projector: GgufVramEstimate, mmproj_bytes: int
166) -> GgufVramEstimate:
167 """Charge the projector at its unified-memory delta, floored at its weights.
169 The floor covers a mmap-shared projector whose uma delta hides weights that
170 still occupy VRAM once offloaded. The charge lands on the first device:
171 llama.cpp loads the projector on the main GPU, not across a split.
172 """
173 projector = max(with_projector.unified_bytes - base.unified_bytes, mmproj_bytes)
174 per_device_vram = tuple(
175 vram + (projector if i == 0 else 0) for i, vram in enumerate(base.per_device_vram)
176 )
177 return GgufVramEstimate(
178 vram_bytes=base.vram_bytes + projector,
179 ram_bytes=with_projector.ram_bytes,
180 unified_bytes=with_projector.unified_bytes,
181 per_device_vram=per_device_vram,
182 per_device_unified=with_projector.per_device_unified,
183 )
186def engine_build_identity() -> str:
187 """Which engine build these numbers describe.
189 Part of the memo key. The estimate prices what one particular llama-server
190 will allocate, and the key held the model, the sizing and the parser's own
191 arguments without a trace of that, so swapping the engine kept the previous
192 engine's answers.
193 """
194 from lilbee.providers.fleet.binary import _engine_build_id
196 return _engine_build_id()
199@lru_cache(maxsize=_CACHE_SIZE)
200def _cached_footprint(
201 _engine_id: str,
202 path_str: str,
203 _mtime_ns: int,
204 ctx: int,
205 slots: int,
206 gpu_layers: int,
207 flash_attn: bool,
208 kv_cache_type: str,
209 kv_cache_type_v: str,
210 mmproj: str | None,
211 _mmproj_mtime_ns: int,
212 tensor_split: tuple[int, ...],
213 batch_size: int | None,
214 expert_offload: tuple[str, ...],
215) -> GgufVramEstimate:
216 """Memoised gguf-parser run keyed on engine + path + mtime + sizing.
218 The mtime and engine args participate in the cache key only; a re-pulled file
219 at the same path invalidates automatically because its mtime changes, and a
220 swapped engine invalidates because its build identity does.
221 """
222 argv = estimator_argv(
223 path_str,
224 ctx=ctx,
225 slots=slots,
226 gpu_layers=gpu_layers,
227 flash_attn=flash_attn,
228 kv_cache_type=kv_cache_type,
229 kv_cache_type_v=kv_cache_type_v,
230 mmproj=mmproj,
231 tensor_split=tensor_split,
232 batch_size=batch_size,
233 expert_offload=expert_offload,
234 )
235 return _parse_estimate(_run_parser(argv, path_str), path_str)
238def estimator_argv(
239 path_str: str,
240 *,
241 ctx: int,
242 slots: int,
243 gpu_layers: int,
244 flash_attn: bool,
245 kv_cache_type: str,
246 kv_cache_type_v: str,
247 mmproj: str | None,
248 tensor_split: tuple[int, ...],
249 batch_size: int | None,
250 expert_offload: tuple[str, ...] = (),
251) -> list[str]:
252 """The gguf-parser command line for one instance's sizing parameters.
254 ``ctx`` is the per-slot context, as it is for the launch argv, and
255 ``--ctx-size`` carries the total across slots. The parser's ``--parallel``
256 does not divide the context the way llama-server's does: the KV estimate is
257 identical for every value of it, so the multiply has to reach the parser
258 through ``--ctx-size`` or the cache is under-reserved by the slot count.
259 """
260 argv = [
261 str(resolve_gguf_parser()),
262 _FLAG_PATH,
263 path_str,
264 _FLAG_CTX,
265 str(ctx * slots),
266 _FLAG_PARALLEL,
267 str(slots),
268 _FLAG_GPU_LAYERS,
269 str(gpu_layers),
270 _FLAG_CACHE_K,
271 kv_cache_type,
272 _FLAG_CACHE_V,
273 kv_cache_type_v,
274 _FLAG_FLASH if flash_attn else _FLAG_NO_FLASH,
275 _FLAG_JSON,
276 ]
277 if batch_size is not None:
278 # Pooled embed/rerank launch with --batch-size/--ubatch-size raised to the
279 # context; the default ubatch (512) would under-estimate their compute buffer.
280 argv += [FLAG_BATCH_SIZE, str(batch_size), FLAG_UBATCH_SIZE, str(batch_size)]
281 if tensor_split:
282 # The split proportions are gguf-parser's only signal for the device count,
283 # so it returns one ``vrams[]`` entry per GPU instead of a single total.
284 argv += [
285 _FLAG_TENSOR_SPLIT,
286 ",".join(str(p) for p in tensor_split),
287 _FLAG_SPLIT_MODE,
288 _SPLIT_MODE_LAYER,
289 ]
290 if expert_offload:
291 # Without these the estimate charges the GPU for experts the launch keeps
292 # in system memory, and the planner sizes slots against a footprint that
293 # never materializes.
294 argv += [
295 _FLAG_OVERRIDE_TENSOR,
296 ",".join(f"{pattern}={_BUFFER_TYPE_CPU}" for pattern in expert_offload),
297 ]
298 if mmproj is not None:
299 argv += [_FLAG_MMPROJ, mmproj]
300 return argv
303def _run_parser(argv: list[str], path_str: str) -> str:
304 """Run gguf-parser, returning its JSON stdout or a user-facing error."""
305 failed = ProviderError(
306 f"Could not size the model {path_str!r}: the memory estimator failed to run.",
307 provider=_PROVIDER,
308 kind=ProviderErrorKind.SERVER,
309 )
310 try:
311 stdout, returncode = run_bounded(
312 argv, timeout_s=_PARSE_TIMEOUT_S, kill_wait_s=_PARSE_KILL_WAIT_S, label="gguf-parser"
313 )
314 except (OSError, subprocess.SubprocessError) as exc:
315 raise failed from exc
316 if returncode != 0:
317 raise failed
318 return stdout
321def _parse_estimate(stdout: str, path_str: str) -> GgufVramEstimate:
322 """Parse gguf-parser JSON into a UMA-aware footprint.
324 Accepts both estimate payload keys: ``items`` (the pinned v0.24.x releases)
325 and ``memory`` (upstream renamed the key after v0.24.1), so an engine built
326 past the pin still sizes instead of failing every launch plan.
327 """
328 try:
329 estimate = json.loads(stdout)[_KEY_ESTIMATE]
330 item = (estimate.get(_KEY_ITEMS) or estimate[_KEY_MEMORY])[0]
331 ram = item[_KEY_RAM]
332 vrams = item[_KEY_VRAMS]
333 per_device_vram = tuple(int(v[_KEY_NONUMA]) for v in vrams)
334 per_device_unified = tuple(int(v[_KEY_UMA]) for v in vrams)
335 return GgufVramEstimate(
336 vram_bytes=sum(per_device_vram),
337 ram_bytes=int(ram[_KEY_NONUMA]),
338 unified_bytes=int(ram[_KEY_UMA]) + sum(per_device_unified),
339 per_device_vram=per_device_vram,
340 per_device_unified=per_device_unified,
341 )
342 except (ValueError, KeyError, IndexError, TypeError) as exc:
343 raise ProviderError(
344 f"Could not size the model {path_str!r}: unexpected estimator output "
345 f"({type(exc).__name__}: {exc}).",
346 provider=_PROVIDER,
347 kind=ProviderErrorKind.SERVER,
348 ) from exc