Coverage for src/lilbee/providers/gguf_meta.py: 100%
162 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-04 17:08 +0000
1"""GGUF metadata helpers: header reads, mmproj sidecar lookup, projector type.
3Reads GGUF headers with the standalone ``gguf`` parser (no native binding), so
4metadata is available to any provider without loading a model into the engine.
5"""
7from __future__ import annotations
9import contextlib
10import hashlib
11import json
12import logging
13import os
14import subprocess
15import threading
16from pathlib import Path
17from typing import NamedTuple, cast
19from lilbee.catalog.header_probe import GGUF_ARCH_KEY
20from lilbee.providers.base import ProviderError
22log = logging.getLogger(__name__)
24_HF_BLOBS_DIR_NAME = "blobs"
25_HF_SNAPSHOTS_DIR_NAME = "snapshots"
26_CLIP_PROJECTOR_TYPE_KEY = "clip.projector_type"
27_DEFAULT_ARCH = "llama"
28_CHAT_TEMPLATE_KEY = "tokenizer.chat_template"
29_FILE_TYPE_KEY = "general.file_type"
30_NAME_KEY = "general.name"
32# Arch-prefixed metadata key suffix -> the lilbee field name it maps to. The
33# prefix is the GGUF's general.architecture value (e.g. "qwen3.context_length").
34_ARCH_FIELD_SUFFIXES: dict[str, str] = {
35 "context_length": "context_length",
36 "embedding_length": "embedding_length",
37 "block_count": "block_count",
38 "attention.head_count_kv": "head_count_kv",
39 "attention.head_count": "head_count",
40 "attention.key_length": "key_length",
41 "attention.value_length": "value_length",
42 # Embedding pooling the model was trained for; absent on most non-embedders.
43 "pooling_type": "pooling_type",
44 # Routed expert count; present only on MoE models, whose experts offload.
45 "expert_count": "expert_count",
46}
49def train_ctx_from_meta(
50 meta: dict[str, str] | None,
51 *,
52 fallback: int,
53 model_path: Path,
54) -> int:
55 """Resolve ``<arch>.context_length`` from GGUF metadata, clamping junk to ``fallback``.
57 Some published GGUFs (nomic-embed, certain Qwen3 and vision builds)
58 report ``context_length=0`` in their headers. Passing zero into
59 ``Llama(n_ctx=...)`` cascades into ``n_batch=0`` / ``n_ubatch=0``,
60 which trips ggml's Vulkan dispatch into undefined behaviour and
61 surfaces as STATUS_HEAP_CORRUPTION on Windows. Unparseable values
62 and non-positive integers both route to ``fallback``.
63 """
64 if not meta:
65 return fallback
66 raw = meta.get("context_length", str(fallback))
67 try:
68 value = int(raw)
69 except (TypeError, ValueError):
70 log.warning(
71 "GGUF %s has unparseable context_length=%r; using %d",
72 model_path.name,
73 raw,
74 fallback,
75 )
76 return fallback
77 if value <= 0:
78 log.warning(
79 "GGUF %s reports context_length=%d; using %d to avoid n_batch=0 crash",
80 model_path.name,
81 value,
82 fallback,
83 )
84 return fallback
85 return value
88_METADATA_CACHE: dict[tuple[str, int, int], dict[str, str] | None] = {}
89_METADATA_CACHE_LOCK = threading.Lock()
91# Bump when the extracted field set changes, or when the reader does, so old
92# entries are ignored rather than served stale. Version 2 reads through
93# gguf-parser: entries written by version 1 recorded "no metadata" for files
94# whose tensor type gguf-py did not know, and the key is (path, mtime, size),
95# so an unchanged file would keep serving that answer after the upgrade.
96_DISK_CACHE_VERSION = 2
97_DISK_CACHE_DIRNAME = "gguf-meta"
98# Distinguishes "no entry on disk" from a cached "this file has no metadata".
99_DISK_MISS = object()
102def _disk_cache_file(key: tuple[str, int, int]) -> Path | None:
103 """Where *key*'s metadata is cached on disk, or None if no state dir works."""
104 from lilbee.core.system import default_cache_dir
106 digest = hashlib.sha256(
107 "\0".join(str(part) for part in (_DISK_CACHE_VERSION, *key)).encode()
108 ).hexdigest()
109 try:
110 # A cache dir, not the state dir: losing this costs a re-parse, never a
111 # lost handle on a running engine.
112 return default_cache_dir() / _DISK_CACHE_DIRNAME / f"{digest}.json"
113 except OSError: # pragma: no cover - unwritable/undiscoverable state dir
114 return None
117def _disk_cache_load(key: tuple[str, int, int]) -> object:
118 """The cached metadata for *key*, or ``_DISK_MISS`` when not usable."""
119 path = _disk_cache_file(key)
120 if path is None:
121 return _DISK_MISS
122 try:
123 payload = json.loads(path.read_text(encoding="utf-8"))
124 except (OSError, ValueError):
125 return _DISK_MISS
126 if not isinstance(payload, dict) or "metadata" not in payload:
127 return _DISK_MISS
128 meta = payload["metadata"]
129 if meta is None:
130 return None
131 if not isinstance(meta, dict) or not all(
132 isinstance(k, str) and isinstance(v, str) for k, v in meta.items()
133 ):
134 return _DISK_MISS
135 return meta
138def _disk_cache_store(key: tuple[str, int, int], result: dict[str, str] | None) -> None:
139 """Persist *result* for *key*; best effort, a failure just costs a re-parse."""
140 path = _disk_cache_file(key)
141 if path is None:
142 return
143 with contextlib.suppress(OSError, TypeError, ValueError):
144 path.parent.mkdir(parents=True, exist_ok=True)
145 # Write-then-rename so a concurrent reader never sees a half-written file.
146 tmp = path.with_suffix(f".{os.getpid()}.tmp")
147 tmp.write_text(json.dumps({"metadata": result}), encoding="utf-8")
148 tmp.replace(path)
151def read_gguf_metadata(model_path: Path) -> dict[str, str] | None:
152 """Read header metadata from a GGUF file with the engine's own parser.
154 Cached by ``(path, mtime, size)`` in memory and on disk. The read itself is
155 cheap now that it goes through gguf-parser, which reports an array as a type
156 and a length rather than its contents; what the cache saves is the process
157 spawn, and planning reads the same model several times per fleet build.
158 Keying on mtime and size means an edited or replaced file re-reads. Returns a
159 copy so callers can't mutate the shared entry.
160 """
161 try:
162 stat = model_path.stat()
163 key: tuple[str, int, int] | None = (str(model_path), stat.st_mtime_ns, stat.st_size)
164 except OSError:
165 key = None
166 if key is not None:
167 with _METADATA_CACHE_LOCK:
168 if key in _METADATA_CACHE:
169 cached = _METADATA_CACHE[key]
170 return dict(cached) if cached is not None else None
171 from_disk = _disk_cache_load(key)
172 if from_disk is not _DISK_MISS:
173 entry = cast("dict[str, str] | None", from_disk)
174 with _METADATA_CACHE_LOCK:
175 _METADATA_CACHE[key] = entry
176 return dict(entry) if entry is not None else None
177 result = _read_gguf_metadata_uncached(model_path)
178 if key is not None:
179 with _METADATA_CACHE_LOCK:
180 _METADATA_CACHE[key] = result
181 _disk_cache_store(key, result)
182 return dict(result) if result is not None else None
185# Array values report as ``{type, len, startOffset}`` rather than their contents,
186# so a 151k-token vocabulary costs nothing to skip. Only scalars are read here.
187_PARSER_ARRAY_VALUE_TYPE = 9
188# GGUF string value type, in both readers' numbering.
189_GGUF_STRING_VALUE_TYPE = 8
190_PARSER_TIMEOUT_SECONDS = 60.0
191_PARSER_KILL_WAIT_SECONDS = 5.0
192_PARSER_LABEL = "gguf-parser"
195class _Scalar(NamedTuple):
196 """One scalar metadata value, and whether the file typed it as a string.
198 A caller that wants text out of a field the file wrote as an integer is
199 reading a malformed file, so the type travels with the value rather than
200 being inferred from how the digits look.
201 """
203 text: str
204 is_string: bool
207def _kv_via_parser(model_path: Path) -> dict[str, _Scalar] | None:
208 """Every scalar metadata key in *model_path*, read by the engine's own parser.
210 Returns ``None`` when there is no parser to ask.
212 gguf-parser is built from the pin that builds llama-server, so the tensor
213 types it accepts are the ones the engine can decode. gguf-py carries its own
214 table and trails llama.cpp: a Q2_0 file (tensor type 42) that the engine
215 loads makes ``GGUFReader`` raise while it builds tensor descriptors nothing
216 here reads, and every field in the file becomes unreadable with it.
217 """
218 from lilbee.providers.fleet.binary import resolve_gguf_parser
219 from lilbee.providers.fleet.proc import run_bounded
221 try:
222 parser = resolve_gguf_parser()
223 except Exception as exc:
224 log.debug("No gguf-parser to read %s with: %s", model_path, exc)
225 return None
226 try:
227 # merge_stderr stays off: the caller parses stdout as JSON.
228 out, code = run_bounded(
229 [str(parser), "--path", str(model_path), "--raw"],
230 timeout_s=_PARSER_TIMEOUT_SECONDS,
231 kill_wait_s=_PARSER_KILL_WAIT_SECONDS,
232 label=_PARSER_LABEL,
233 )
234 except (OSError, subprocess.SubprocessError) as exc:
235 log.debug("gguf-parser did not run against %s: %s", model_path, exc)
236 return None
237 if code != 0:
238 # The engine's own reader rejected the file. Report "no readable
239 # metadata" rather than asking gguf-py for a second opinion the engine
240 # would not honour anyway.
241 log.warning("Could not parse GGUF metadata from %s: gguf-parser exit %d", model_path, code)
242 return {}
243 try:
244 entries = json.loads(out)["header"]["metadataKV"]
245 except (ValueError, KeyError, TypeError) as exc:
246 log.warning("gguf-parser returned unreadable output for %s: %s", model_path, exc)
247 return {}
248 return {
249 e["key"]: _Scalar(str(e["value"]), e.get("valueType") == _GGUF_STRING_VALUE_TYPE)
250 for e in entries
251 if e.get("valueType") != _PARSER_ARRAY_VALUE_TYPE and e.get("value") is not None
252 }
255def _scalars(model_path: Path) -> dict[str, _Scalar]:
256 """Scalar metadata for *model_path*, or empty when there is no parser to ask.
258 An install without the engine extra has no gguf-parser, and also no
259 llama-server to load a model with, so the fields here have nothing left to
260 size. Callers already treat absent metadata as a fallback case.
261 """
262 return _kv_via_parser(model_path) or {}
265def _read_gguf_metadata_uncached(model_path: Path) -> dict[str, str] | None:
266 kv = _scalars(model_path)
267 result: dict[str, str] = {}
269 arch_field = kv.get(GGUF_ARCH_KEY)
270 if arch_field is not None:
271 result["architecture"] = arch_field.text
272 arch = arch_field.text if arch_field is not None else _DEFAULT_ARCH
274 for suffix, out_key in _ARCH_FIELD_SUFFIXES.items():
275 value = kv.get(f"{arch}.{suffix}")
276 if value is not None:
277 result[out_key] = value.text
278 for raw_key, out_key in (
279 (_CHAT_TEMPLATE_KEY, "chat_template"),
280 (_FILE_TYPE_KEY, "file_type"),
281 (_NAME_KEY, "name"),
282 ):
283 value = kv.get(raw_key)
284 if value is not None:
285 result[out_key] = value.text
286 return result or None
289def _find_mmproj_in_hf_snapshots(model_dir: Path) -> Path | None:
290 """Walk an HF-cache ``blobs/`` dir up to its sibling ``snapshots/`` tree."""
291 if model_dir.name != _HF_BLOBS_DIR_NAME:
292 return None
293 snapshots_dir = model_dir.parent / _HF_SNAPSHOTS_DIR_NAME
294 if not snapshots_dir.is_dir():
295 return None
296 for snapshot in snapshots_dir.iterdir():
297 candidates = sorted(snapshot.glob("*mmproj*.gguf"))
298 if candidates:
299 return candidates[0]
300 return None
303def _find_mmproj_in_flat_dir(model_dir: Path) -> Path | None:
304 """Glob ``*mmproj*.gguf`` siblings of a model GGUF (sideloaded layout)."""
305 candidates = sorted(model_dir.glob("*mmproj*.gguf"))
306 return candidates[0] if candidates else None
309def find_mmproj_for_model(model_path: Path) -> Path:
310 """Find the mmproj (CLIP projection) file for a vision model.
312 Resolution order: (1) the HuggingFace-cache ``snapshots/`` sibling of
313 ``blobs/``, (2) same-directory glob for flat sideloaded layouts. Both look
314 beside the model file itself, so a projector is never borrowed from another
315 repo. Raises ``ProviderError`` if neither finds a file.
316 """
317 found = _find_mmproj_in_hf_snapshots(model_path.parent) or _find_mmproj_in_flat_dir(
318 model_path.parent
319 )
320 if found is not None:
321 return found
323 raise ProviderError(
324 f"No mmproj (CLIP projection) file found for vision model {model_path.name}. "
325 f"Download the mmproj file to {model_path.parent} or re-download the vision "
326 "model through the catalog to get both files.",
327 provider="llama-server",
328 )
331def read_mmproj_projector_type(mmproj_path: Path) -> str | None:
332 """Read ``clip.projector_type`` from a GGUF mmproj without loading the model."""
333 value = _scalars(mmproj_path).get(_CLIP_PROJECTOR_TYPE_KEY)
334 return value.text if value is not None and value.is_string else None