Coverage for src/lilbee/providers/fleet/readback.py: 100%
141 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""What the engine actually allocated, read back from its own report.
3Compares the planner's estimate against reality per device, and warns naming the
4role, the estimate and the reality when they diverge.
6Two sources, picked per engine binary. An engine that advertises ``--memory``
7in its ``--help`` (llama.cpp PR 26130; the bundled engine is built from that
8fork) serves ``GET /memory``: per-device rows carrying model, context, compute
9and mmproj bytes under the same device names ``--device`` takes. That is the
10preferred source -- a promised JSON shape, the vision projector reported on its
11real device, and no trace-level log required.
13An engine without the flag falls back to its startup log, which is then the
14only place these numbers exist: ``llama_model_size`` is the whole model's
15weights and nothing per device, and a stock server's HTTP surface carries none
16of it (``/props`` is metadata, ``/metrics`` is token counters; both checked
17against a running server).
19The log format, from upstream source:
21 src/llama-model.cpp "%s: %12s model buffer size = %8.2f MiB"
22 src/llama-kv-cache.cpp "%s: %10s KV buffer size = %8.2f MiB"
23 src/llama-context.cpp "%s: %10s compute buffer size = %8.2f MiB"
25These are format strings, not a promised interface. A bundled-engine version bump
26can break this: re-capture the fixture and confirm :func:`parse_device_buffers`
27still finds every line. A build that stops matching is reported, not swallowed
28(see :func:`check_launch`).
29"""
31from __future__ import annotations
33import functools
34import logging
35import re
36import subprocess
37from pathlib import Path
39from lilbee.providers.fleet.devices import FleetDevice
40from lilbee.providers.fleet.vram import usable_vram_fraction
41from lilbee.providers.roles import WorkerRole
43log = logging.getLogger(__name__)
45MIB = 1024 * 1024
47# The build the checked-in fixture was captured from, named in the drift warning
48# so a report says what to compare with. Tracks the fixture, not the shipped
49# engine. A landmark, not a gate: drift is detected by the parse coming back
50# empty on a load that finished, not by a version comparison.
51VERIFIED_ENGINE_BUILD = "9310 (e2ef8fe42)"
53# "load_tensors: MTL0_Mapped model buffer size = 82.41 MiB", and its siblings.
54#
55# Matches the shape rather than a list of buffer kinds, so LoRA, RS and the
56# DeepSeek V4 state buffer parse without being enumerated here.
57#
58# The "= N MiB" is load-bearing: it excludes the three lines carrying these words
59# that are not allocations -- the self-check pair ("compute buffer size is N MiB,
60# matches expectation" / "... does not match expectation") and ggml-opencl's
61# "buffer size reduced from A to B". None uses "=".
62#
63# The device label is whatever the backend calls itself: CUDA0, MTL0, Vulkan1,
64# CPU. A timestamp and level prefix the line under --log-file, so the match is
65# not anchored to the start.
66_BUFFER_RE = re.compile(r"\S+:\s+(?P<device>\S+)\s+.*?buffer size\s*=\s*(?P<mib>[\d.]+)\s*MiB")
67# The engine names an mmapped weight buffer "<device>_Mapped" beside the same
68# device's other buffers. Same memory, so the suffix is folded away rather than
69# splitting one card's total across two keys.
70_MAPPED_SUFFIX = "_Mapped"
71# ggml names a row-split buffer "<backend>_Split", one allocation shared by every
72# card in the split rather than a device of its own.
73_SPLIT_SUFFIX = "_Split"
74# Host memory rather than a GPU, in three shapes: the CPU backend's own
75# buffers (CPU, CPU_Mapped, CPU_REPACK, and AMX, which is a CPU extension),
76# every GPU backend's pinned-host allocator, named "<backend>_Host" by
77# ggml-cuda, -sycl, -vulkan, -cann and -hip alike (observed as CUDA_Host,
78# Vulkan_Host, ROCm_Host), and the single "Host" row GET /memory aggregates
79# them into. None of it occupies VRAM; charging it to a card reports a phantom
80# overrun on every partially offloaded model.
81_HOST_PREFIXES = ("CPU", "AMX", "Host")
82_HOST_SUFFIX = "_Host"
85def _is_host_device(device: str) -> bool:
86 """Whether *device* names host memory rather than a GPU."""
87 return device.startswith(_HOST_PREFIXES) or device.endswith(_HOST_SUFFIX)
90def parse_device_buffers(text: str) -> dict[str, int]:
91 """Bytes the engine reported allocating, per device label, from *text*.
93 Sums the model, KV, compute and output buffers, which is the same total the
94 estimate predicts. Empty when the text carries no buffer report: a load that
95 failed before allocating, a log that has rotated past it, or an engine whose
96 verbosity is below the level that prints it.
97 """
98 totals: dict[str, int] = {}
99 for match in _BUFFER_RE.finditer(text):
100 device = match.group("device")
101 if device.endswith(_SPLIT_SUFFIX):
102 # A row-split buffer is spread across every card in the split, so it
103 # belongs to no single one. Keeping it would invent a device that the
104 # per-device comparison then reports as an unplanned allocation.
105 continue
106 device = device.removesuffix(_MAPPED_SUFFIX)
107 totals[device] = totals.get(device, 0) + int(float(match.group("mib")) * MIB)
108 return totals
111# Printed once the weights are in and slots are being wired up, so its presence
112# separates "the report is not written yet" from "this engine writes none here".
113#
114# Matches only the word the engine has kept: "initializing slots" through b9665,
115# "initializing, n_slots = N" from b9829. This gate arms the format-drift
116# warning, so pinning either exact phrase would silence the warning on the other.
117# The buffer lines held identical across all three builds; only the prose moved.
118_LOAD_FINISHED_RE = re.compile(r"load_model:\s+initializing\b")
119# "common_params_print_info: build 9310 (e2ef8fe42) with AppleClang ...", the
120# engine's own first line. Carried into the format-drift warning so the report
121# names the exact build to re-verify against.
122_BUILD_RE = re.compile(r"build\s+(?P<build>\d+)\s+\((?P<commit>[0-9a-f]+)\)")
125def engine_build(text: str) -> str:
126 """The engine build the log was written by, or empty when it does not say."""
127 match = _BUILD_RE.search(text)
128 return f"{match.group('build')} ({match.group('commit')})" if match else ""
131def load_finished(text: str) -> bool:
132 """Whether the engine got far enough to have reported its buffers."""
133 return _LOAD_FINISHED_RE.search(text) is not None
136def device_footprint(text: str) -> int:
137 """Total GPU bytes the engine reported, host buffers excluded."""
138 return sum(
139 size for device, size in parse_device_buffers(text).items() if not _is_host_device(device)
140 )
143def device_label(device: FleetDevice) -> str:
144 """The name the engine prints for *device*, and the join between the two sides.
146 ``ggml_backend_dev_name`` produces ``CUDA0`` / ``MTL0`` / ``Vulkan1``, which is
147 the same token ``--device`` and ``--tensor-split`` take and the same one the
148 buffer report is keyed by. Joining on it keeps the check out of the index-space
149 ambiguity that ``FleetDevice.from_loader`` exists to mark.
150 """
151 return f"{device.backend}{device.index}"
154def report_divergence(
155 role: WorkerRole,
156 model: str,
157 estimated_bytes: int,
158 actual_bytes: int,
159 *,
160 tolerance: float,
161) -> bool:
162 """Warn when the engine's real footprint diverges materially from the estimate.
164 Returns whether a warning was emitted, so a caller can record that this
165 instance has already been checked and not repeat it on every request.
167 Both directions are worth saying. An under-estimate is how a plan that fit on
168 paper OOMs, and it is the one that ends in a failed load. A large
169 over-estimate is quieter but costs capacity: it is why a role gets fewer
170 slots, a narrower context, or a split it did not need.
171 """
172 if estimated_bytes <= 0 or actual_bytes <= 0:
173 return False
174 ratio = actual_bytes / estimated_bytes
175 limit = min(tolerance, _absorbable_overrun()) if ratio > 1.0 else tolerance
176 if abs(ratio - 1.0) <= limit:
177 return False
178 log.warning(
179 "The %s model %s allocated %.1f GiB of GPU memory but was planned for %.1f GiB "
180 "(%+.0f%%). Placement decisions for this model were made on the smaller figure; "
181 "if it fails to load or runs slowly, that gap is why.",
182 role.value,
183 model,
184 actual_bytes / 1024**3,
185 estimated_bytes / 1024**3,
186 (ratio - 1.0) * 100,
187 )
188 return True
191# The engine's own log, one per instance, beside the swap process's log. Named
192# by model id so a role's replicas do not overwrite each other.
193_ENGINE_LOG_TEMPLATE = "engine-{model_id}.log"
194# Env the engine reads for its log destination and threshold. Set through the
195# environment rather than argv: the launch is planned before the data directory
196# holding these logs is chosen, and neither affects sizing.
197#
198# Both spellings, because the engine renamed them. common/arg.cpp registers
199# LLAMA_ARG_LOG_FILE / LLAMA_ARG_LOG_VERBOSITY on master; builds around 9310 read
200# LLAMA_LOG_FILE / LLAMA_LOG_VERBOSITY. Verified against both. An unread variable
201# costs nothing; picking one produced no log at all on half the builds in use.
202ENV_LOG_FILE = "LLAMA_LOG_FILE"
203ENV_LOG_VERBOSITY = "LLAMA_LOG_VERBOSITY"
204ENV_ARG_LOG_FILE = "LLAMA_ARG_LOG_FILE"
205ENV_ARG_LOG_VERBOSITY = "LLAMA_ARG_LOG_VERBOSITY"
206# Level 4 ("trace") is where the per-device buffer report appears. Measured
207# against the bundled engine: the default 3 omits it entirely, and 5 adds a
208# per-layer and per-slot flood for the same six lines.
209LOAD_REPORT_VERBOSITY = "4"
212def engine_log_path(log_dir: Path, model_id: str) -> Path:
213 """Where the engine serving *model_id* writes its own log."""
214 return log_dir / _ENGINE_LOG_TEMPLATE.format(model_id=model_id)
217def engine_log_env(log_dir: Path, model_id: str) -> dict[str, str]:
218 """Environment that makes the engine report what it allocated, and where."""
219 path = str(engine_log_path(log_dir, model_id))
220 return {
221 ENV_LOG_FILE: path,
222 ENV_ARG_LOG_FILE: path,
223 ENV_LOG_VERBOSITY: LOAD_REPORT_VERBOSITY,
224 ENV_ARG_LOG_VERBOSITY: LOAD_REPORT_VERBOSITY,
225 }
228def check_launch(
229 log_dir: Path,
230 model_id: str,
231 role: WorkerRole,
232 model: str,
233 estimated_bytes: int,
234 est_by_device: dict[str, int] | None = None,
235 unreported_bytes: int = 0,
236) -> bool:
237 """Compare the engine's own report for *model_id* against the estimate.
239 Checked per device when *est_by_device* says what each card was planned for,
240 because per device is the only dimension the planner decides in: a split is a
241 ratio, a placement is a card, and a shortfall is recorded against a role on a
242 card. Two cards planned 50/50 that land 80/20 sum to exactly the planned
243 total, so a scalar comparison sees nothing while card 0 is the one that runs
244 out. Falls back to the total for a model the estimator could only size as one
245 number.
247 Three outcomes, and the third is the one that matters. The engine has no API
248 for any of this: /props carries no memory keys and /metrics is token
249 counters, both checked against a running server, so its log is the only
250 place these numbers exist. That makes this the one part of the fleet whose
251 input is a format nobody promises to keep.
253 So a load that finished without a readable report is reported, not swallowed.
254 Left silent it would look exactly like a correct estimate, and the check
255 would quietly become decoration the first time llama.cpp renames a line or
256 renumbers its verbosity levels. Loud, it names itself as the thing to fix.
257 """
258 try:
259 text = engine_log_path(log_dir, model_id).read_text(encoding="utf-8", errors="replace")
260 except OSError:
261 # No log at all. Usually the engine simply has not written one yet, so
262 # this is silent by default. It is also exactly what a wrong environment
263 # variable name looks like, which is how an earlier spelling went
264 # unnoticed: the check returned False forever and read as "estimate fine".
265 # report_missing_log is how a caller that knows the engine is up says so.
266 return False
267 per_device = {
268 label: size
269 for label, size in parse_device_buffers(text).items()
270 if not _is_host_device(label)
271 }
272 actual = sum(per_device.values())
273 if actual > 0 and est_by_device:
274 return _report_per_device(
275 role, model, _without_unreported(est_by_device, unreported_bytes), per_device
276 )
277 if actual <= 0:
278 if load_finished(text):
279 log.warning(
280 "The %s engine (build %s) finished loading but reported no memory usage where "
281 "lilbee reads it, so its estimate could not be checked. The engine's log format "
282 "or verbosity levels have most likely changed since build %s, which lilbee's "
283 "parser was written against; placement estimates are unverified until it is "
284 "updated to match.",
285 role.value,
286 engine_build(text) or "unknown",
287 VERIFIED_ENGINE_BUILD,
288 )
289 return False
290 return report_divergence(
291 role, model, estimated_bytes - unreported_bytes, actual, tolerance=_TOLERANCE
292 )
295def _without_unreported(est_by_device: dict[str, int], unreported: int) -> dict[str, int]:
296 """*est_by_device* less the bytes the engine allocates without reporting them.
298 Charged to the busiest device, which is where the planner put them: a vision
299 projector loads on the main GPU rather than across a split. Comparing the
300 full estimate against a report that structurally cannot contain these bytes
301 warns on every correctly sized vision load.
302 """
303 if unreported <= 0 or not est_by_device:
304 return est_by_device
305 main = max(est_by_device, key=lambda label: est_by_device[label])
306 adjusted = dict(est_by_device)
307 adjusted[main] = max(0, adjusted[main] - unreported)
308 return adjusted
311# How far the engine may land from the estimate before it is worth saying. Wide
312# enough that the estimator's normal error is quiet, narrow enough to catch the
313# whole-slot and whole-cache mistakes this exists to surface.
314_TOLERANCE = 0.25
316# Share of the card's remaining margin an overrun may eat before it is worth
317# saying, leaving the operator a gap between the warning and the overflow.
318_MARGIN_WARN_FRACTION = 0.75
321def _absorbable_overrun() -> float:
322 """How far past its estimate a load may land while the card still holds it.
324 Placement packs a card up to ``cfg.usable_vram_fraction`` and a single-card
325 chat sizes its cache against ``cfg.gpu_memory_fraction``; the binding one is
326 whichever leaves less room, since either can be raised past the other. A load
327 filling that share overflows once it exceeds its estimate by the remainder,
328 so the warning has to come before that, which makes the threshold a function
329 of the margin rather than a constant. At the stock 0.9 usable fraction the
330 room is 11%, well inside the flat 25% this replaces.
331 """
332 from lilbee.core.config import cfg
334 committed = max(cfg.gpu_memory_fraction, usable_vram_fraction())
335 return max(0.0, 1.0 / committed - 1.0) * _MARGIN_WARN_FRACTION
338def _report_per_device(
339 role: WorkerRole,
340 model: str,
341 estimated: dict[str, int],
342 actual: dict[str, int],
343) -> bool:
344 """Warn about the card that diverged worst, naming both figures.
346 One warning rather than one per card: the operator needs to know the plan did
347 not hold and which card to look at, and a split that skews puts every card out
348 at once by construction.
349 """
350 worst_label, worst_gap, worst_over = "", 0.0, False
351 for label in set(estimated) | set(actual):
352 planned, landed = estimated.get(label, 0), actual.get(label, 0)
353 gap = abs(landed - planned) / planned if planned else float(landed)
354 over = landed > planned
355 # An overrun outranks an equal shortfall: a card holding more than it was
356 # planned for is the one that fails to load, while its partner holding
357 # less is only the symptom of the same skew.
358 if (over, gap) > (worst_over, worst_gap):
359 worst_label, worst_gap, worst_over = label, gap, over
360 limit = min(_TOLERANCE, _absorbable_overrun()) if worst_over else _TOLERANCE
361 if not worst_label or (estimated.get(worst_label) and worst_gap <= limit):
362 return False
363 log.warning(
364 "The %s model %s did not land where it was planned: %s holds %.1f GiB but was "
365 "planned for %.1f GiB. Placement, the tensor split and the context were all "
366 "decided per card, so a total that looks right can still overrun one of them.",
367 role.value,
368 model,
369 worst_label,
370 actual.get(worst_label, 0) / 1024**3,
371 estimated.get(worst_label, 0) / 1024**3,
372 )
373 return True
376# llama-server flag (llama.cpp PR 26130) that both enables GET /memory and
377# marks, via --help, an engine that has it. Launch argv and the swap config key
378# the readback mode off its presence.
379MEMORY_FLAG = "--memory"
381# The flag as --help advertises it. The lookahead keeps historic longer flags
382# (--memory-f32) from reading as support for the endpoint.
383_HELP_MEMORY_RE = re.compile(r"--memory(?!\S)")
384# --help exits from arg parsing, before any backend or model load, so this is
385# generous; a binary that cannot even print usage in this time is not one the
386# fleet should silently trust either way.
387_HELP_PROBE_TIMEOUT_S = 15.0
390@functools.lru_cache(maxsize=8)
391def supports_memory_readback(binary: Path) -> bool:
392 """Whether *binary* serves ``GET /memory`` when launched with ``--memory``.
394 Read from the binary's own ``--help``, once per path: a stock llama-server
395 exits with "unknown argument" when handed a flag it lacks, so the launch
396 must never carry it on advertisement it did not see. Any failure to run or
397 to answer reads as unsupported, which degrades to the log path rather than
398 a failed launch.
399 """
400 try:
401 done = subprocess.run( # noqa: S603 - argv is the resolved engine binary and a literal
402 [str(binary), "--help"],
403 capture_output=True,
404 encoding="utf-8",
405 errors="replace",
406 timeout=_HELP_PROBE_TIMEOUT_S,
407 check=False,
408 )
409 except (OSError, subprocess.SubprocessError):
410 return False
411 return bool(_HELP_MEMORY_RE.search(f"{done.stdout}\n{done.stderr}"))
414# The per-device byte fields of one GET /memory row. Summed because the total
415# they form is the same one the estimate predicts and the log path sums from
416# its buffer lines; mmproj is the projector the log path could never see.
417_MEMORY_ROW_FIELDS = ("model", "context", "compute", "mmproj")
420def parse_memory_rows(payload: object) -> dict[str, int]:
421 """Bytes the engine reported allocating per device, from a /memory payload.
423 Empty for anything that is not the endpoint's ``{"data": [rows]}`` shape,
424 and per row for junk within it: this crosses a process boundary, and the
425 caller says "unverified" for an empty parse rather than crashing readiness.
426 """
427 rows = payload.get("data") if isinstance(payload, dict) else None
428 if not isinstance(rows, list):
429 return {}
430 totals: dict[str, int] = {}
431 for row in rows:
432 if not isinstance(row, dict):
433 continue
434 name = row.get("name")
435 if not isinstance(name, str) or not name:
436 continue
437 size = 0
438 for field in _MEMORY_ROW_FIELDS:
439 value = row.get(field, 0)
440 if isinstance(value, (int, float)) and not isinstance(value, bool):
441 size += int(value)
442 totals[name] = totals.get(name, 0) + size
443 return totals
446def check_memory_report(
447 role: WorkerRole,
448 model: str,
449 estimated_bytes: int,
450 est_by_device: dict[str, int] | None,
451 payload: object,
452) -> bool:
453 """Compare a ``GET /memory`` payload against the estimate; the API-mode twin
454 of :func:`check_launch`.
456 No ``unreported_bytes`` adjustment exists here on purpose: the endpoint
457 reports the vision projector per device in its ``mmproj`` field, so the
458 quantity that adjustment approximates is in the rows themselves and the
459 comparison is exact where the log path had to guess.
461 An engine that took the flag but produced no usable rows is reported, not
462 swallowed, for the same reason a finished load with an unreadable log is:
463 silent, it looks exactly like a correct estimate.
464 """
465 per_device = {
466 label: size
467 for label, size in parse_memory_rows(payload).items()
468 if not _is_host_device(label)
469 }
470 if not per_device:
471 log.warning(
472 "The %s engine answered /memory without any per-device rows, so its "
473 "estimate is unverified. The endpoint's response shape has most likely "
474 "changed since the build lilbee's reader was written against.",
475 role.value,
476 )
477 return True
478 if est_by_device:
479 return _report_per_device(role, model, est_by_device, per_device)
480 return report_divergence(
481 role, model, estimated_bytes, sum(per_device.values()), tolerance=_TOLERANCE
482 )
485def report_missing_log(log_dir: Path, model_id: str, role: WorkerRole) -> bool:
486 """Warn when a ready engine wrote no log where lilbee told it to.
488 Separate from :func:`check_launch` because only the caller knows the engine
489 finished loading; an absent file before that is ordinary. After it, the file
490 should exist, and its absence means the engine never accepted the settings
491 that produce it. That is a silent no-op rather than a wrong answer, which is
492 the harder kind to notice, so it is stated.
493 """
494 if engine_log_path(log_dir, model_id).exists():
495 return False
496 log.warning(
497 "The %s engine is running but wrote no log to %s, so its memory use could not "
498 "be checked against the estimate. The engine build most likely does not read "
499 "the variables lilbee sets to ask for one (%s or %s); placement estimates are "
500 "unverified until that is updated.",
501 role.value,
502 engine_log_path(log_dir, model_id),
503 ENV_LOG_FILE,
504 ENV_ARG_LOG_FILE,
505 )
506 return True