Coverage for src/lilbee/providers/gguf_meta.py: 100%

1"""GGUF metadata helpers: header reads, mmproj sidecar lookup, projector type. 

2 

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""" 

6 

7from __future__ import annotations 

8 

9import contextlib 

10import hashlib 

11import json 

12import logging 

13import os 

14import subprocess 

15import threading 

16from pathlib import Path 

17from typing import NamedTuple, cast 

18 

19from lilbee.catalog.header_probe import GGUF_ARCH_KEY 

20from lilbee.providers.base import ProviderError 

21 

22log = logging.getLogger(__name__) 

23 

24_HF_BLOBS_DIR_NAME = "blobs" 

25_HF_SHARD_LEN = 2 

26_HF_SNAPSHOTS_DIR_NAME = "snapshots" 

27_CLIP_PROJECTOR_TYPE_KEY = "clip.projector_type" 

28_DEFAULT_ARCH = "llama" 

29_CHAT_TEMPLATE_KEY = "tokenizer.chat_template" 

30_FILE_TYPE_KEY = "general.file_type" 

31_NAME_KEY = "general.name" 

32 

33# Arch-prefixed metadata key suffix -> the lilbee field name it maps to. The 

34# prefix is the GGUF's general.architecture value (e.g. "qwen3.context_length"). 

35_ARCH_FIELD_SUFFIXES: dict[str, str] = { 

36 "context_length": "context_length", 

37 "embedding_length": "embedding_length", 

38 "block_count": "block_count", 

39 "attention.head_count_kv": "head_count_kv", 

40 "attention.head_count": "head_count", 

41 "attention.key_length": "key_length", 

42 "attention.value_length": "value_length", 

43 # Embedding pooling the model was trained for; absent on most non-embedders. 

44 "pooling_type": "pooling_type", 

45 # Routed expert count; present only on MoE models, whose experts offload. 

46 "expert_count": "expert_count", 

47} 

48 

49 

50def train_ctx_from_meta( 

51 meta: dict[str, str] | None, 

52 *, 

53 fallback: int, 

54 model_path: Path, 

55) -> int: 

56 """Resolve ``<arch>.context_length`` from GGUF metadata, clamping junk to ``fallback``. 

57 

58 Some published GGUFs (nomic-embed, certain Qwen3 and vision builds) 

59 report ``context_length=0`` in their headers. Passing zero into 

60 ``Llama(n_ctx=...)`` cascades into ``n_batch=0`` / ``n_ubatch=0``, 

61 which trips ggml's Vulkan dispatch into undefined behaviour and 

62 surfaces as STATUS_HEAP_CORRUPTION on Windows. Unparseable values 

63 and non-positive integers both route to ``fallback``. 

64 """ 

65 if not meta: 

66 return fallback 

67 raw = meta.get("context_length", str(fallback)) 

68 try: 

69 value = int(raw) 

70 except (TypeError, ValueError): 

71 log.warning( 

72 "GGUF %s has unparseable context_length=%r; using %d", 

73 model_path.name, 

74 raw, 

75 fallback, 

76 ) 

77 return fallback 

78 if value <= 0: 

79 log.warning( 

80 "GGUF %s reports context_length=%d; using %d to avoid n_batch=0 crash", 

81 model_path.name, 

82 value, 

83 fallback, 

84 ) 

85 return fallback 

86 return value 

87 

88 

89_METADATA_CACHE: dict[tuple[str, int, int], dict[str, str] | None] = {} 

90_METADATA_CACHE_LOCK = threading.Lock() 

91 

92# Bump when the extracted field set changes, or when the reader does, so old 

93# entries are ignored rather than served stale. Version 2 reads through 

94# gguf-parser: entries written by version 1 recorded "no metadata" for files 

95# whose tensor type gguf-py did not know, and the key is (path, mtime, size), 

96# so an unchanged file would keep serving that answer after the upgrade. 

97_DISK_CACHE_VERSION = 2 

98_DISK_CACHE_DIRNAME = "gguf-meta" 

99# Distinguishes "no entry on disk" from a cached "this file has no metadata". 

100_DISK_MISS = object() 

101 

102 

103def _disk_cache_file(key: tuple[str, int, int]) -> Path | None: 

104 """Where *key*'s metadata is cached on disk, or None if no state dir works.""" 

105 from lilbee.core.system import default_cache_dir 

106 

107 digest = hashlib.sha256( 

108 "\0".join(str(part) for part in (_DISK_CACHE_VERSION, *key)).encode() 

109 ).hexdigest() 

110 try: 

111 # A cache dir, not the state dir: losing this costs a re-parse, never a 

112 # lost handle on a running engine. 

113 return default_cache_dir() / _DISK_CACHE_DIRNAME / f"{digest}.json" 

114 except OSError: # pragma: no cover - unwritable/undiscoverable state dir 

115 return None 

116 

117 

118def _disk_cache_load(key: tuple[str, int, int]) -> object: 

119 """The cached metadata for *key*, or ``_DISK_MISS`` when not usable.""" 

120 path = _disk_cache_file(key) 

121 if path is None: 

122 return _DISK_MISS 

123 try: 

124 payload = json.loads(path.read_text(encoding="utf-8")) 

125 except (OSError, ValueError): 

126 return _DISK_MISS 

127 if not isinstance(payload, dict) or "metadata" not in payload: 

128 return _DISK_MISS 

129 meta = payload["metadata"] 

130 if meta is None: 

131 return None 

132 if not isinstance(meta, dict) or not all( 

133 isinstance(k, str) and isinstance(v, str) for k, v in meta.items() 

134 ): 

135 return _DISK_MISS 

136 return meta 

137 

138 

139def _disk_cache_store(key: tuple[str, int, int], result: dict[str, str] | None) -> None: 

140 """Persist *result* for *key*; best effort, a failure just costs a re-parse.""" 

141 path = _disk_cache_file(key) 

142 if path is None: 

143 return 

144 with contextlib.suppress(OSError, TypeError, ValueError): 

145 path.parent.mkdir(parents=True, exist_ok=True) 

146 # Write-then-rename so a concurrent reader never sees a half-written file. 

147 tmp = path.with_suffix(f".{os.getpid()}.tmp") 

148 tmp.write_text(json.dumps({"metadata": result}), encoding="utf-8") 

149 tmp.replace(path) 

150 

151 

152def read_gguf_metadata(model_path: Path) -> dict[str, str] | None: 

153 """Read header metadata from a GGUF file with the engine's own parser. 

154 

155 Cached by ``(path, mtime, size)`` in memory and on disk. The read itself is 

156 cheap now that it goes through gguf-parser, which reports an array as a type 

157 and a length rather than its contents; what the cache saves is the process 

158 spawn, and planning reads the same model several times per fleet build. 

159 Keying on mtime and size means an edited or replaced file re-reads. Returns a 

160 copy so callers can't mutate the shared entry. 

161 """ 

162 try: 

163 stat = model_path.stat() 

164 key: tuple[str, int, int] | None = (str(model_path), stat.st_mtime_ns, stat.st_size) 

165 except OSError: 

166 key = None 

167 if key is not None: 

168 with _METADATA_CACHE_LOCK: 

169 if key in _METADATA_CACHE: 

170 cached = _METADATA_CACHE[key] 

171 return dict(cached) if cached is not None else None 

172 from_disk = _disk_cache_load(key) 

173 if from_disk is not _DISK_MISS: 

174 entry = cast("dict[str, str] | None", from_disk) 

175 with _METADATA_CACHE_LOCK: 

176 _METADATA_CACHE[key] = entry 

177 return dict(entry) if entry is not None else None 

178 result = _read_gguf_metadata_uncached(model_path) 

179 if key is not None: 

180 with _METADATA_CACHE_LOCK: 

181 _METADATA_CACHE[key] = result 

182 _disk_cache_store(key, result) 

183 return dict(result) if result is not None else None 

184 

185 

186# Array values report as ``{type, len, startOffset}`` rather than their contents, 

187# so a 151k-token vocabulary costs nothing to skip. Only scalars are read here. 

188_PARSER_ARRAY_VALUE_TYPE = 9 

189# GGUF string value type, in both readers' numbering. 

190_GGUF_STRING_VALUE_TYPE = 8 

191_PARSER_TIMEOUT_SECONDS = 60.0 

192_PARSER_KILL_WAIT_SECONDS = 5.0 

193_PARSER_LABEL = "gguf-parser" 

194 

195 

196class _Scalar(NamedTuple): 

197 """One scalar metadata value, and whether the file typed it as a string. 

198 

199 A caller that wants text out of a field the file wrote as an integer is 

200 reading a malformed file, so the type travels with the value rather than 

201 being inferred from how the digits look. 

202 """ 

203 

204 text: str 

205 is_string: bool 

206 

207 

208def _kv_via_parser(model_path: Path) -> dict[str, _Scalar] | None: 

209 """Every scalar metadata key in *model_path*, read by the engine's own parser. 

210 

211 Returns ``None`` when there is no parser to ask. 

212 

213 gguf-parser is built from the pin that builds llama-server, so the tensor 

214 types it accepts are the ones the engine can decode. gguf-py carries its own 

215 table and trails llama.cpp: a Q2_0 file (tensor type 42) that the engine 

216 loads makes ``GGUFReader`` raise while it builds tensor descriptors nothing 

217 here reads, and every field in the file becomes unreadable with it. 

218 """ 

219 from lilbee.providers.fleet.binary import resolve_gguf_parser 

220 from lilbee.providers.fleet.proc import run_bounded 

221 

222 try: 

223 parser = resolve_gguf_parser() 

224 except Exception as exc: 

225 log.debug("No gguf-parser to read %s with: %s", model_path, exc) 

226 return None 

227 try: 

228 # merge_stderr stays off: the caller parses stdout as JSON. 

229 out, code = run_bounded( 

230 [str(parser), "--path", str(model_path), "--raw"], 

231 timeout_s=_PARSER_TIMEOUT_SECONDS, 

232 kill_wait_s=_PARSER_KILL_WAIT_SECONDS, 

233 label=_PARSER_LABEL, 

234 ) 

235 except (OSError, subprocess.SubprocessError) as exc: 

236 log.debug("gguf-parser did not run against %s: %s", model_path, exc) 

237 return None 

238 if code != 0: 

239 # The engine's own reader rejected the file. Report "no readable 

240 # metadata" rather than asking gguf-py for a second opinion the engine 

241 # would not honour anyway. 

242 log.warning("Could not parse GGUF metadata from %s: gguf-parser exit %d", model_path, code) 

243 return {} 

244 try: 

245 entries = json.loads(out)["header"]["metadataKV"] 

246 except (ValueError, KeyError, TypeError) as exc: 

247 log.warning("gguf-parser returned unreadable output for %s: %s", model_path, exc) 

248 return {} 

249 return { 

250 e["key"]: _Scalar(str(e["value"]), e.get("valueType") == _GGUF_STRING_VALUE_TYPE) 

251 for e in entries 

252 if e.get("valueType") != _PARSER_ARRAY_VALUE_TYPE and e.get("value") is not None 

253 } 

254 

255 

256def _scalars(model_path: Path) -> dict[str, _Scalar]: 

257 """Scalar metadata for *model_path*, or empty when there is no parser to ask. 

258 

259 An install without the engine extra has no gguf-parser, and also no 

260 llama-server to load a model with, so the fields here have nothing left to 

261 size. Callers already treat absent metadata as a fallback case. 

262 """ 

263 return _kv_via_parser(model_path) or {} 

264 

265 

266def _read_gguf_metadata_uncached(model_path: Path) -> dict[str, str] | None: 

267 kv = _scalars(model_path) 

268 result: dict[str, str] = {} 

269 

270 arch_field = kv.get(GGUF_ARCH_KEY) 

271 if arch_field is not None: 

272 result["architecture"] = arch_field.text 

273 arch = arch_field.text if arch_field is not None else _DEFAULT_ARCH 

274 

275 for suffix, out_key in _ARCH_FIELD_SUFFIXES.items(): 

276 value = kv.get(f"{arch}.{suffix}") 

277 if value is not None: 

278 result[out_key] = value.text 

279 for raw_key, out_key in ( 

280 (_CHAT_TEMPLATE_KEY, "chat_template"), 

281 (_FILE_TYPE_KEY, "file_type"), 

282 (_NAME_KEY, "name"), 

283 ): 

284 value = kv.get(raw_key) 

285 if value is not None: 

286 result[out_key] = value.text 

287 return result or None 

288 

289 

290def _find_mmproj_in_hf_snapshots(model_dir: Path) -> Path | None: 

291 """Walk an HF-cache ``blobs/`` dir up to its sibling ``snapshots/`` tree. 

292 

293 The blob sits in ``blobs/`` directly, or in a two-character shard dir under 

294 it (``blobs/a0/``) with newer huggingface_hub releases. 

295 """ 

296 if model_dir.name == _HF_BLOBS_DIR_NAME: 

297 cache_dir = model_dir.parent 

298 elif model_dir.parent.name == _HF_BLOBS_DIR_NAME and len(model_dir.name) == _HF_SHARD_LEN: 

299 cache_dir = model_dir.parent.parent 

300 else: 

301 return None 

302 snapshots_dir = cache_dir / _HF_SNAPSHOTS_DIR_NAME 

303 if not snapshots_dir.is_dir(): 

304 return None 

305 for snapshot in snapshots_dir.iterdir(): 

306 candidates = sorted(snapshot.glob("*mmproj*.gguf")) 

307 if candidates: 

308 return candidates[0] 

309 return None 

310 

311 

312def _find_mmproj_in_flat_dir(model_dir: Path) -> Path | None: 

313 """Glob ``*mmproj*.gguf`` siblings of a model GGUF (sideloaded layout).""" 

314 candidates = sorted(model_dir.glob("*mmproj*.gguf")) 

315 return candidates[0] if candidates else None 

316 

317 

318def find_mmproj_for_model(model_path: Path) -> Path: 

319 """Find the mmproj (CLIP projection) file for a vision model. 

320 

321 Resolution order: (1) the HuggingFace-cache ``snapshots/`` sibling of 

322 ``blobs/``, (2) same-directory glob for flat sideloaded layouts. Both look 

323 beside the model file itself, so a projector is never borrowed from another 

324 repo. Raises ``ProviderError`` if neither finds a file. 

325 """ 

326 found = _find_mmproj_in_hf_snapshots(model_path.parent) or _find_mmproj_in_flat_dir( 

327 model_path.parent 

328 ) 

329 if found is not None: 

330 return found 

331 

332 raise ProviderError( 

333 f"No mmproj (CLIP projection) file found for vision model {model_path.name}. " 

334 f"Download the mmproj file to {model_path.parent} or re-download the vision " 

335 "model through the catalog to get both files.", 

336 provider="llama-server", 

337 ) 

338 

339 

340def read_mmproj_projector_type(mmproj_path: Path) -> str | None: 

341 """Read ``clip.projector_type`` from a GGUF mmproj without loading the model.""" 

342 value = _scalars(mmproj_path).get(_CLIP_PROJECTOR_TYPE_KEY) 

343 return value.text if value is not None and value.is_string else None