Coverage for src/lilbee/providers/fleet/readback.py: 100%
158 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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 logging
34import re
35from pathlib import Path
37from lilbee.core.health_warnings import HealthWarning, WarningCode
38from lilbee.providers.fleet.binary import engine_help, help_lists_option
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 divergence_warning(
155 role: WorkerRole,
156 model: str,
157 estimated_bytes: int,
158 actual_bytes: int,
159 *,
160 tolerance: float,
161) -> HealthWarning | None:
162 """The placement warning when the engine's footprint diverges from the estimate."""
163 if estimated_bytes <= 0 or actual_bytes <= 0:
164 return None
165 ratio = actual_bytes / estimated_bytes
166 limit = min(tolerance, _absorbable_overrun()) if ratio > 1.0 else tolerance
167 if abs(ratio - 1.0) <= limit:
168 return None
169 return HealthWarning(
170 code=WarningCode.PLACEMENT_DIVERGED,
171 message=(
172 f"The {role.value} model {model} allocated {actual_bytes / 1024**3:.1f} GiB "
173 f"of GPU memory but was planned for {estimated_bytes / 1024**3:.1f} GiB "
174 f"({(ratio - 1.0) * 100:+.0f}%). Placement decisions for this model were "
175 f"made on the smaller figure; if it fails to load or runs slowly, that "
176 f"gap is why."
177 ),
178 remedy="Free up GPU memory or use a smaller model." if ratio > 1.0 else None,
179 )
182def report_divergence(
183 role: WorkerRole,
184 model: str,
185 estimated_bytes: int,
186 actual_bytes: int,
187 *,
188 tolerance: float,
189) -> HealthWarning | None:
190 """Warn when the engine's real footprint diverges materially from the estimate.
192 Returns the warning when one was emitted, so a caller can surface it and not
193 repeat the check on every request.
195 Both directions are worth saying. An under-estimate is how a plan that fit on
196 paper OOMs, and it is the one that ends in a failed load. A large
197 over-estimate is quieter but costs capacity: it is why a role gets fewer
198 slots, a narrower context, or a split it did not need.
199 """
200 warning = divergence_warning(role, model, estimated_bytes, actual_bytes, tolerance=tolerance)
201 if warning is not None:
202 log.warning("%s", warning.message)
203 return warning
206# The engine's own log, one per instance, beside the swap process's log. Named
207# by model id so a role's replicas do not overwrite each other.
208_ENGINE_LOG_TEMPLATE = "engine-{model_id}.log"
209# Env the engine reads for its log destination and threshold. Set through the
210# environment rather than argv: the launch is planned before the data directory
211# holding these logs is chosen, and neither affects sizing.
212#
213# Both spellings, because the engine renamed them. common/arg.cpp registers
214# LLAMA_ARG_LOG_FILE / LLAMA_ARG_LOG_VERBOSITY on master; builds around 9310 read
215# LLAMA_LOG_FILE / LLAMA_LOG_VERBOSITY. Verified against both. An unread variable
216# costs nothing; picking one produced no log at all on half the builds in use.
217ENV_LOG_FILE = "LLAMA_LOG_FILE"
218ENV_LOG_VERBOSITY = "LLAMA_LOG_VERBOSITY"
219ENV_ARG_LOG_FILE = "LLAMA_ARG_LOG_FILE"
220ENV_ARG_LOG_VERBOSITY = "LLAMA_ARG_LOG_VERBOSITY"
221# Level 4 ("trace") is where the per-device buffer report appears. Measured
222# against the bundled engine: the default 3 omits it entirely, and 5 adds a
223# per-layer and per-slot flood for the same six lines.
224LOAD_REPORT_VERBOSITY = "4"
227def engine_log_path(log_dir: Path, model_id: str) -> Path:
228 """Where the engine serving *model_id* writes its own log."""
229 return log_dir / _ENGINE_LOG_TEMPLATE.format(model_id=model_id)
232def writes_engine_log(argv: list[str]) -> bool:
233 """Whether a launch with *argv* writes the engine log file.
235 A launch carrying ``--memory`` serves its numbers over HTTP and gets no log
236 env, so no file exists for it.
237 """
238 return MEMORY_FLAG not in argv
241def engine_log_env(log_dir: Path, model_id: str) -> dict[str, str]:
242 """Environment that makes the engine report what it allocated, and where."""
243 path = str(engine_log_path(log_dir, model_id))
244 return {
245 ENV_LOG_FILE: path,
246 ENV_ARG_LOG_FILE: path,
247 ENV_LOG_VERBOSITY: LOAD_REPORT_VERBOSITY,
248 ENV_ARG_LOG_VERBOSITY: LOAD_REPORT_VERBOSITY,
249 }
252def check_launch(
253 log_dir: Path,
254 model_id: str,
255 role: WorkerRole,
256 model: str,
257 estimated_bytes: int,
258 est_by_device: dict[str, int] | None = None,
259 unreported_bytes: int = 0,
260 est_host_bytes: int = 0,
261) -> HealthWarning | None:
262 """Compare the engine's own report for *model_id* against the estimate.
264 Checked per device when *est_by_device* says what each card was planned for,
265 because per device is the only dimension the planner decides in: a split is a
266 ratio, a placement is a card, and a shortfall is recorded against a role on a
267 card. Two cards planned 50/50 that land 80/20 sum to exactly the planned
268 total, so a scalar comparison sees nothing while card 0 is the one that runs
269 out. Falls back to the total for a model the estimator could only size as one
270 number.
272 Three outcomes, and the third is the one that matters. The engine has no API
273 for any of this: /props carries no memory keys and /metrics is token
274 counters, both checked against a running server, so its log is the only
275 place these numbers exist. That makes this the one part of the fleet whose
276 input is a format nobody promises to keep.
278 So a load that finished without a readable report is reported, not swallowed.
279 Left silent it would look exactly like a correct estimate, and the check
280 would quietly become decoration the first time llama.cpp renames a line or
281 renumbers its verbosity levels. Loud, it names itself as the thing to fix.
282 """
283 try:
284 text = engine_log_path(log_dir, model_id).read_text(encoding="utf-8", errors="replace")
285 except OSError:
286 # No log at all. Usually the engine simply has not written one yet, so
287 # this is silent by default. It is also exactly what a wrong environment
288 # variable name looks like, which is how an earlier spelling went
289 # unnoticed: the check returned None forever and read as "estimate fine".
290 # report_missing_log is how a caller that knows the engine is up says so.
291 return None
292 buffers = parse_device_buffers(text)
293 report_host_spill(role, model, _host_bytes(buffers), est_host_bytes)
294 per_device = {label: size for label, size in buffers.items() if not _is_host_device(label)}
295 actual = sum(per_device.values())
296 if actual > 0 and est_by_device:
297 return _report_per_device(
298 role, model, _without_unreported(est_by_device, unreported_bytes), per_device
299 )
300 if actual <= 0:
301 if load_finished(text):
302 log.warning(
303 "The %s engine (build %s) finished loading but reported no memory usage where "
304 "lilbee reads it, so its estimate could not be checked. The engine's log format "
305 "or verbosity levels have most likely changed since build %s, which lilbee's "
306 "parser was written against; placement estimates are unverified until it is "
307 "updated to match.",
308 role.value,
309 engine_build(text) or "unknown",
310 VERIFIED_ENGINE_BUILD,
311 )
312 return None
313 return report_divergence(
314 role, model, estimated_bytes - unreported_bytes, actual, tolerance=_TOLERANCE
315 )
318def _host_bytes(per_device: dict[str, int]) -> int:
319 """The bytes of *per_device* that sit in host memory."""
320 return sum(size for label, size in per_device.items() if _is_host_device(label))
323def report_host_spill(role: WorkerRole, model: str, host_bytes: int, est_host_bytes: int) -> bool:
324 """Warn when the engine's fit left more of *model* in system memory than estimated.
326 *est_host_bytes* is what the host holds with every layer on the card; 0
327 means the engine's fit does not run for this launch, and nothing is said.
328 """
329 if est_host_bytes <= 0 or host_bytes <= est_host_bytes * (1.0 + _TOLERANCE):
330 return False
331 log.warning(
332 "The %s model %s does not fit in GPU memory: the engine keeps %.1f GiB of it in "
333 "system memory, where it runs slower. Free GPU memory or use a smaller model.",
334 role.value,
335 model,
336 host_bytes / 1024**3,
337 )
338 return True
341def _without_unreported(est_by_device: dict[str, int], unreported: int) -> dict[str, int]:
342 """*est_by_device* less the bytes the engine allocates without reporting them.
344 Charged to the busiest device, which is where the planner put them: a vision
345 projector loads on the main GPU rather than across a split. Comparing the
346 full estimate against a report that structurally cannot contain these bytes
347 warns on every correctly sized vision load.
348 """
349 if unreported <= 0 or not est_by_device:
350 return est_by_device
351 main = max(est_by_device, key=lambda label: est_by_device[label])
352 adjusted = dict(est_by_device)
353 adjusted[main] = max(0, adjusted[main] - unreported)
354 return adjusted
357# How far the engine may land from the estimate before it is worth saying. Wide
358# enough that the estimator's normal error is quiet, narrow enough to catch the
359# whole-slot and whole-cache mistakes this exists to surface.
360_TOLERANCE = 0.25
362# Share of the card's remaining margin an overrun may eat before it is worth
363# saying, leaving the operator a gap between the warning and the overflow.
364_MARGIN_WARN_FRACTION = 0.75
367def _absorbable_overrun() -> float:
368 """How far past its estimate a load may land while the card still holds it.
370 Placement packs a card up to ``cfg.usable_vram_fraction`` and a single-card
371 chat sizes its cache against ``cfg.gpu_memory_fraction``; the binding one is
372 whichever leaves less room, since either can be raised past the other. A load
373 filling that share overflows once it exceeds its estimate by the remainder,
374 so the warning has to come before that, which makes the threshold a function
375 of the margin rather than a constant. At the stock 0.9 usable fraction the
376 room is 11%, well inside the flat 25% this replaces.
377 """
378 from lilbee.core.config import cfg
380 committed = max(cfg.gpu_memory_fraction, usable_vram_fraction())
381 return max(0.0, 1.0 / committed - 1.0) * _MARGIN_WARN_FRACTION
384def _per_device_warning(
385 role: WorkerRole,
386 model: str,
387 estimated: dict[str, int],
388 actual: dict[str, int],
389) -> HealthWarning | None:
390 """The placement warning for the card that diverged worst, naming both figures."""
391 worst_label, worst_gap, worst_over = "", 0.0, False
392 for label in set(estimated) | set(actual):
393 planned, landed = estimated.get(label, 0), actual.get(label, 0)
394 gap = abs(landed - planned) / planned if planned else float(landed)
395 over = landed > planned
396 # An overrun outranks an equal shortfall: a card holding more than it was
397 # planned for is the one that fails to load, while its partner holding
398 # less is only the symptom of the same skew.
399 if (over, gap) > (worst_over, worst_gap):
400 worst_label, worst_gap, worst_over = label, gap, over
401 limit = min(_TOLERANCE, _absorbable_overrun()) if worst_over else _TOLERANCE
402 if not worst_label or (estimated.get(worst_label) and worst_gap <= limit):
403 return None
404 return HealthWarning(
405 code=WarningCode.PLACEMENT_DIVERGED,
406 message=(
407 f"The {role.value} model {model} did not land where it was planned: "
408 f"{worst_label} holds {actual.get(worst_label, 0) / 1024**3:.1f} GiB but was "
409 f"planned for {estimated.get(worst_label, 0) / 1024**3:.1f} GiB. Placement, "
410 f"the tensor split and the context were all decided per card, so a total "
411 f"that looks right can still overrun one of them."
412 ),
413 remedy="Free up GPU memory or use a smaller model." if worst_over else None,
414 )
417def _report_per_device(
418 role: WorkerRole,
419 model: str,
420 estimated: dict[str, int],
421 actual: dict[str, int],
422) -> HealthWarning | None:
423 """Warn about the card that diverged worst, naming both figures.
425 One warning rather than one per card: the operator needs to know the plan did
426 not hold and which card to look at, and a split that skews puts every card out
427 at once by construction.
428 """
429 warning = _per_device_warning(role, model, estimated, actual)
430 if warning is not None:
431 log.warning("%s", warning.message)
432 return warning
435# llama-server flag (llama.cpp PR 26130) that both enables GET /memory and
436# marks, via --help, an engine that has it. Launch argv and the swap config key
437# the readback mode off its presence.
438MEMORY_FLAG = "--memory"
440# The flag as --help advertises it. The lookahead keeps historic longer flags
441# (--memory-f32) from reading as support for the endpoint.
442_HELP_MEMORY_RE = re.compile(r"--memory(?!\S)")
445def supports_memory_readback(binary: Path) -> bool:
446 """Whether *binary* serves ``GET /memory`` when launched with ``--memory``.
448 Read from the binary's own ``--help``: a stock llama-server exits with
449 "unknown argument" when handed a flag it lacks, so the launch must never
450 carry it on advertisement it did not see. A binary that cannot run or answer
451 reads as unsupported, which degrades to the log path rather than a failed launch.
452 """
453 return bool(_HELP_MEMORY_RE.search(engine_help(binary)))
456def engine_advertises(binary: Path, flag: str) -> bool:
457 """Whether an option line of *binary*'s ``--help`` names *flag*."""
458 return help_lists_option(engine_help(binary), flag)
461# The per-device byte fields of one GET /memory row. Summed because the total
462# they form is the same one the estimate predicts and the log path sums from
463# its buffer lines; mmproj is the projector the log path could never see.
464_MEMORY_ROW_FIELDS = ("model", "context", "compute", "mmproj")
467def parse_memory_rows(payload: object) -> dict[str, int]:
468 """Bytes the engine reported allocating per device, from a /memory payload.
470 Empty for anything that is not the endpoint's ``{"data": [rows]}`` shape,
471 and per row for junk within it: this crosses a process boundary, and the
472 caller says "unverified" for an empty parse rather than crashing readiness.
473 """
474 rows = payload.get("data") if isinstance(payload, dict) else None
475 if not isinstance(rows, list):
476 return {}
477 totals: dict[str, int] = {}
478 for row in rows:
479 if not isinstance(row, dict):
480 continue
481 name = row.get("name")
482 if not isinstance(name, str) or not name:
483 continue
484 size = 0
485 for field in _MEMORY_ROW_FIELDS:
486 value = row.get(field, 0)
487 if isinstance(value, (int, float)) and not isinstance(value, bool):
488 size += int(value)
489 totals[name] = totals.get(name, 0) + size
490 return totals
493def check_memory_report(
494 role: WorkerRole,
495 model: str,
496 estimated_bytes: int,
497 est_by_device: dict[str, int] | None,
498 payload: object,
499 est_host_bytes: int = 0,
500) -> HealthWarning | None:
501 """Compare a ``GET /memory`` payload against the estimate; the API-mode twin
502 of :func:`check_launch`.
504 No ``unreported_bytes`` adjustment exists here on purpose: the endpoint
505 reports the vision projector per device in its ``mmproj`` field, so the
506 quantity that adjustment approximates is in the rows themselves and the
507 comparison is exact where the log path had to guess.
509 An engine that took the flag but produced no usable rows is reported, not
510 swallowed, for the same reason a finished load with an unreadable log is:
511 silent, it looks exactly like a correct estimate.
512 """
513 rows = parse_memory_rows(payload)
514 report_host_spill(role, model, _host_bytes(rows), est_host_bytes)
515 per_device = {label: size for label, size in rows.items() if not _is_host_device(label)}
516 if not per_device:
517 log.warning(
518 "The %s engine answered /memory without any per-device rows, so its "
519 "estimate is unverified. The endpoint's response shape has most likely "
520 "changed since the build lilbee's reader was written against.",
521 role.value,
522 )
523 return None
524 if est_by_device:
525 return _report_per_device(role, model, est_by_device, per_device)
526 return report_divergence(
527 role, model, estimated_bytes, sum(per_device.values()), tolerance=_TOLERANCE
528 )
531def report_missing_log(log_dir: Path, model_id: str, role: WorkerRole) -> bool:
532 """Warn when a ready engine wrote no log where lilbee told it to.
534 Separate from :func:`check_launch` because only the caller knows the engine
535 finished loading; an absent file before that is ordinary. After it, the file
536 should exist, and its absence means the engine never accepted the settings
537 that produce it. That is a silent no-op rather than a wrong answer, which is
538 the harder kind to notice, so it is stated.
539 """
540 if engine_log_path(log_dir, model_id).exists():
541 return False
542 log.warning(
543 "The %s engine is running but wrote no log to %s, so its memory use could not "
544 "be checked against the estimate. The engine build most likely does not read "
545 "the variables lilbee sets to ask for one (%s or %s); placement estimates are "
546 "unverified until that is updated.",
547 role.value,
548 engine_log_path(log_dir, model_id),
549 ENV_LOG_FILE,
550 ENV_ARG_LOG_FILE,
551 )
552 return True