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