Coverage for src/lilbee/providers/gguf_meta.py: 100%
143 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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 struct
15import threading
16from pathlib import Path
17from typing import cast
19from gguf import GGUFReader, GGUFValueType
21from lilbee.catalog.header_probe import GGUF_ARCH_KEY, gguf_scalar_str
22from lilbee.providers.base import ProviderError
24log = logging.getLogger(__name__)
26_HF_BLOBS_DIR_NAME = "blobs"
27_HF_SNAPSHOTS_DIR_NAME = "snapshots"
28_CLIP_PROJECTOR_TYPE_KEY = "clip.projector_type"
29_DEFAULT_ARCH = "llama"
30_CHAT_TEMPLATE_KEY = "tokenizer.chat_template"
31_FILE_TYPE_KEY = "general.file_type"
32_NAME_KEY = "general.name"
34# Arch-prefixed metadata key suffix -> the lilbee field name it maps to. The
35# prefix is the GGUF's general.architecture value (e.g. "qwen3.context_length").
36_ARCH_FIELD_SUFFIXES: dict[str, str] = {
37 "context_length": "context_length",
38 "embedding_length": "embedding_length",
39 "block_count": "block_count",
40 "attention.head_count_kv": "head_count_kv",
41 "attention.head_count": "head_count",
42 "attention.key_length": "key_length",
43 "attention.value_length": "value_length",
44 # Embedding pooling the model was trained for; absent on most non-embedders.
45 "pooling_type": "pooling_type",
46 # Routed expert count; present only on MoE models, whose experts offload.
47 "expert_count": "expert_count",
48}
51def train_ctx_from_meta(
52 meta: dict[str, str] | None,
53 *,
54 fallback: int,
55 model_path: Path,
56) -> int:
57 """Resolve ``<arch>.context_length`` from GGUF metadata, clamping junk to ``fallback``.
59 Some published GGUFs (nomic-embed, certain Qwen3 and vision builds)
60 report ``context_length=0`` in their headers. Passing zero into
61 ``Llama(n_ctx=...)`` cascades into ``n_batch=0`` / ``n_ubatch=0``,
62 which trips ggml's Vulkan dispatch into undefined behaviour and
63 surfaces as STATUS_HEAP_CORRUPTION on Windows. Unparseable values
64 and non-positive integers both route to ``fallback``.
65 """
66 if not meta:
67 return fallback
68 raw = meta.get("context_length", str(fallback))
69 try:
70 value = int(raw)
71 except (TypeError, ValueError):
72 log.warning(
73 "GGUF %s has unparseable context_length=%r; using %d",
74 model_path.name,
75 raw,
76 fallback,
77 )
78 return fallback
79 if value <= 0:
80 log.warning(
81 "GGUF %s reports context_length=%d; using %d to avoid n_batch=0 crash",
82 model_path.name,
83 value,
84 fallback,
85 )
86 return fallback
87 return value
90_METADATA_CACHE: dict[tuple[str, int, int], dict[str, str] | None] = {}
91_METADATA_CACHE_LOCK = threading.Lock()
93# Bump when the extracted field set changes, so old entries are ignored rather
94# than served stale.
95_DISK_CACHE_VERSION = 1
96_DISK_CACHE_DIRNAME = "gguf-meta"
97# Distinguishes "no entry on disk" from a cached "this file has no metadata".
98_DISK_MISS = object()
101def _disk_cache_file(key: tuple[str, int, int]) -> Path | None:
102 """Where *key*'s metadata is cached on disk, or None if no state dir works."""
103 from lilbee.core.system import default_cache_dir
105 digest = hashlib.sha256(
106 "\0".join(str(part) for part in (_DISK_CACHE_VERSION, *key)).encode()
107 ).hexdigest()
108 try:
109 # A cache dir, not the state dir: losing this costs a re-parse, never a
110 # lost handle on a running engine.
111 return default_cache_dir() / _DISK_CACHE_DIRNAME / f"{digest}.json"
112 except OSError: # pragma: no cover - unwritable/undiscoverable state dir
113 return None
116def _disk_cache_load(key: tuple[str, int, int]) -> object:
117 """The cached metadata for *key*, or ``_DISK_MISS`` when not usable."""
118 path = _disk_cache_file(key)
119 if path is None:
120 return _DISK_MISS
121 try:
122 payload = json.loads(path.read_text(encoding="utf-8"))
123 except (OSError, ValueError):
124 return _DISK_MISS
125 if not isinstance(payload, dict) or "metadata" not in payload:
126 return _DISK_MISS
127 meta = payload["metadata"]
128 if meta is None:
129 return None
130 if not isinstance(meta, dict) or not all(
131 isinstance(k, str) and isinstance(v, str) for k, v in meta.items()
132 ):
133 return _DISK_MISS
134 return meta
137def _disk_cache_store(key: tuple[str, int, int], result: dict[str, str] | None) -> None:
138 """Persist *result* for *key*; best effort, a failure just costs a re-parse."""
139 path = _disk_cache_file(key)
140 if path is None:
141 return
142 with contextlib.suppress(OSError, TypeError, ValueError):
143 path.parent.mkdir(parents=True, exist_ok=True)
144 # Write-then-rename so a concurrent reader never sees a half-written file.
145 tmp = path.with_suffix(f".{os.getpid()}.tmp")
146 tmp.write_text(json.dumps({"metadata": result}), encoding="utf-8")
147 tmp.replace(path)
150def read_gguf_metadata(model_path: Path) -> dict[str, str] | None:
151 """Read header metadata from a GGUF file with the ``gguf`` parser.
153 Cached by ``(path, mtime, size)`` in memory and on disk. ``GGUFReader`` parses
154 the whole header -- every tensor descriptor and the large tokenizer arrays --
155 to hand back a dozen scalar fields, which measured at ~60s for a 2.6GB model
156 on a spinning-rust-era CPU. Planning reads the same model several times per
157 fleet build, so the in-memory cache collapses those to one parse; the on-disk
158 cache carries it across processes, which is what stops a relaunch paying that
159 minute again while an already-warm engine sits idle waiting to be adopted.
160 Keying on mtime and size means an edited or replaced file re-reads. Returns a
161 copy so callers can't mutate the shared entry.
162 """
163 try:
164 stat = model_path.stat()
165 key: tuple[str, int, int] | None = (str(model_path), stat.st_mtime_ns, stat.st_size)
166 except OSError:
167 key = None
168 if key is not None:
169 with _METADATA_CACHE_LOCK:
170 if key in _METADATA_CACHE:
171 cached = _METADATA_CACHE[key]
172 return dict(cached) if cached is not None else None
173 from_disk = _disk_cache_load(key)
174 if from_disk is not _DISK_MISS:
175 entry = cast("dict[str, str] | None", from_disk)
176 with _METADATA_CACHE_LOCK:
177 _METADATA_CACHE[key] = entry
178 return dict(entry) if entry is not None else None
179 result = _read_gguf_metadata_uncached(model_path)
180 if key is not None:
181 with _METADATA_CACHE_LOCK:
182 _METADATA_CACHE[key] = result
183 _disk_cache_store(key, result)
184 return dict(result) if result is not None else None
187def _read_gguf_metadata_uncached(model_path: Path) -> dict[str, str] | None:
188 try:
189 reader = GGUFReader(str(model_path))
190 fields = reader.fields
191 except (ValueError, KeyError, IndexError, struct.error, OSError, UnicodeDecodeError) as exc:
192 # A truncated or corrupt GGUF header surfaces as a parser error. Report
193 # "no readable metadata" (None, an outcome callers already handle) rather
194 # than letting a raw parse error abort the whole fleet build.
195 log.warning("Could not parse GGUF metadata from %s: %s", model_path, exc)
196 return None
197 result: dict[str, str] = {}
199 arch = gguf_scalar_str(fields.get(GGUF_ARCH_KEY))
200 if arch is not None:
201 result["architecture"] = arch
202 arch = arch or _DEFAULT_ARCH
204 for suffix, out_key in _ARCH_FIELD_SUFFIXES.items():
205 value = gguf_scalar_str(fields.get(f"{arch}.{suffix}"))
206 if value is not None:
207 result[out_key] = value
208 for raw_key, out_key in (
209 (_CHAT_TEMPLATE_KEY, "chat_template"),
210 (_FILE_TYPE_KEY, "file_type"),
211 (_NAME_KEY, "name"),
212 ):
213 value = gguf_scalar_str(fields.get(raw_key))
214 if value is not None:
215 result[out_key] = value
216 return result or None
219def _find_mmproj_in_hf_snapshots(model_dir: Path) -> Path | None:
220 """Walk an HF-cache ``blobs/`` dir up to its sibling ``snapshots/`` tree."""
221 if model_dir.name != _HF_BLOBS_DIR_NAME:
222 return None
223 snapshots_dir = model_dir.parent / _HF_SNAPSHOTS_DIR_NAME
224 if not snapshots_dir.is_dir():
225 return None
226 for snapshot in snapshots_dir.iterdir():
227 candidates = sorted(snapshot.glob("*mmproj*.gguf"))
228 if candidates:
229 return candidates[0]
230 return None
233def _find_mmproj_in_flat_dir(model_dir: Path) -> Path | None:
234 """Glob ``*mmproj*.gguf`` siblings of a model GGUF (sideloaded layout)."""
235 candidates = sorted(model_dir.glob("*mmproj*.gguf"))
236 return candidates[0] if candidates else None
239def find_mmproj_for_model(model_path: Path) -> Path:
240 """Find the mmproj (CLIP projection) file for a vision model.
242 Resolution order: (1) the HuggingFace-cache ``snapshots/`` sibling of
243 ``blobs/``, (2) same-directory glob for flat sideloaded layouts. Both look
244 beside the model file itself, so a projector is never borrowed from another
245 repo. Raises ``ProviderError`` if neither finds a file.
246 """
247 found = _find_mmproj_in_hf_snapshots(model_path.parent) or _find_mmproj_in_flat_dir(
248 model_path.parent
249 )
250 if found is not None:
251 return found
253 raise ProviderError(
254 f"No mmproj (CLIP projection) file found for vision model {model_path.name}. "
255 f"Download the mmproj file to {model_path.parent} or re-download the vision "
256 "model through the catalog to get both files.",
257 provider="llama-server",
258 )
261def read_mmproj_projector_type(mmproj_path: Path) -> str | None:
262 """Read ``clip.projector_type`` from a GGUF mmproj without loading the model."""
263 try:
264 reader = GGUFReader(str(mmproj_path))
265 field = reader.get_field(_CLIP_PROJECTOR_TYPE_KEY)
266 except Exception:
267 log.debug("Failed to read mmproj metadata from %s", mmproj_path, exc_info=True)
268 return None
269 if field is None or not field.types or field.types[-1] != GGUFValueType.STRING:
270 return None
271 return bytes(field.parts[field.data[0]]).decode("utf-8", errors="replace")