Coverage for src/lilbee/modelhub/registry.py: 100%
368 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-31 21:55 +0000
1"""Manifest store keyed by ``(hf_repo, gguf_filename)`` over the HF cache.
3Canonical ref: ``<hf_repo>/<gguf_filename>``. Two quants of the same
4repo are two distinct installations. Manifests live at
5``manifests/<repo--repo>/<filename>.json``; blobs at
6``models--<repo--repo>/blobs/<sha>``.
7"""
9from __future__ import annotations
11import contextlib
12import hashlib
13import json
14import logging
15import os
16import re
17import shutil
18import tempfile
19from dataclasses import asdict, dataclass, field
20from datetime import UTC, datetime
21from pathlib import Path
22from typing import TYPE_CHECKING
24from lilbee.catalog.query import reclassify_by_name
25from lilbee.catalog.refs import (
26 GGUF_SUFFIX,
27 NATIVE_GGUF_REF_MIN_SLASHES,
28 format_native_gguf_ref,
29 is_bare_hf_repo,
30 split_shard_filenames,
31)
32from lilbee.catalog.types import ModelTask
33from lilbee.core.config.model import cfg
34from lilbee.core.security import validate_path_within
36if TYPE_CHECKING:
37 from lilbee.catalog.models import CatalogModel
39log = logging.getLogger(__name__)
41_HASH_CHUNK_SIZE = 8192 # bytes read per iteration when hashing
42_REPO_SEGMENT_RE = re.compile(r"^[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+$")
43# A GGUF filename, optionally under repo subdirectories (unsloth stores quants
44# in e.g. ``Q4_K_M/Model-...-00001-of-00003.gguf``). Path separators are allowed;
45# ``..`` and absolute paths are rejected in the validator to stay inside the repo.
46_FILENAME_RE = re.compile(r"^[a-zA-Z0-9._/-]+\.gguf$")
48REPO_DIR_SEPARATOR = "--"
51def _validate_hf_repo(hf_repo: str) -> str:
52 """Validate that a HuggingFace repo id has the form ``org/name``."""
53 if not hf_repo or not _REPO_SEGMENT_RE.match(hf_repo) or ".." in hf_repo:
54 raise ValueError(f"Invalid hf_repo: {hf_repo!r}")
55 return hf_repo
58def _validate_gguf_filename(filename: str) -> str:
59 """Validate a ``.gguf`` filename, allowing repo subdirectories but no traversal."""
60 if (
61 not filename
62 or not _FILENAME_RE.match(filename)
63 or ".." in filename
64 or filename.startswith("/")
65 ):
66 raise ValueError(f"Invalid gguf_filename: {filename!r}")
67 return filename
70_REF_SHAPE_HINT = "Use '<org>/<repo>/<filename>.gguf'."
73def parse_hf_ref(ref: str) -> tuple[str, str]:
74 """Split ``<org>/<repo>/<file>.gguf`` into ``(hf_repo, gguf_filename)``.
76 The repo is always the first two segments (``<org>/<repo>``); everything
77 after is the filename, which may include repo subdirectories (unsloth stores
78 quants under e.g. ``Q4_K_M/Model-...-00001-of-00003.gguf``).
79 """
80 if not ref.endswith(GGUF_SUFFIX) or ref.count("/") < NATIVE_GGUF_REF_MIN_SLASHES:
81 raise ValueError(f"Model ref {ref!r} is not a HuggingFace ref. {_REF_SHAPE_HINT}")
82 parts = ref.split("/")
83 hf_repo = "/".join(parts[:NATIVE_GGUF_REF_MIN_SLASHES])
84 gguf_filename = "/".join(parts[NATIVE_GGUF_REF_MIN_SLASHES:])
85 return _validate_hf_repo(hf_repo), _validate_gguf_filename(gguf_filename)
88def repo_to_dir(hf_repo: str) -> str:
89 """Encode an HF repo for use as a directory name (HF cache convention)."""
90 return hf_repo.replace("/", REPO_DIR_SEPARATOR)
93@dataclass
94class ModelManifest:
95 """One installed model's metadata. Identity: ``(hf_repo, gguf_filename)``."""
97 hf_repo: str
98 gguf_filename: str
99 size_bytes: int # primary (first-shard) blob size; validated against the blob on disk
100 task: ModelTask
101 downloaded_at: str # ISO 8601
102 blob: str | None = None # SHA-256 hex of the blob in the HF cache; None pre-install
103 # A split GGUF has further shard blobs beyond ``blob``. ``total_size_bytes``
104 # is the sum across every shard (None = single file, use ``size_bytes``);
105 # ``shard_blobs`` are the non-primary shard digests, so removal frees them all.
106 total_size_bytes: int | None = None
107 shard_blobs: list[str] = field(default_factory=list)
109 @property
110 def ref(self) -> str:
111 return format_native_gguf_ref(self.hf_repo, self.gguf_filename)
113 @property
114 def disk_size_bytes(self) -> int:
115 """Total bytes this model occupies on disk, across all shards."""
116 return self.total_size_bytes if self.total_size_bytes is not None else self.size_bytes
119def _copy_atomic(source_path: Path, blob_path: Path) -> None:
120 """Copy *source_path* to *blob_path* via a temp file + atomic rename.
122 A crash mid-copy leaves only the temp file, never a partial blob at
123 the final digest path that callers would treat as complete.
124 """
125 fd, tmp_name = tempfile.mkstemp(dir=str(blob_path.parent), suffix=".part")
126 tmp_path = Path(tmp_name)
127 try:
128 with os.fdopen(fd, "wb") as dst, source_path.open("rb") as src:
129 shutil.copyfileobj(src, dst)
130 os.replace(tmp_path, blob_path)
131 except BaseException:
132 tmp_path.unlink(missing_ok=True)
133 raise
136def _blob_size_matches(blob_file: Path, expected_size: int) -> bool:
137 """True iff *blob_file* exists and its byte size equals *expected_size*.
139 A blob shorter than the manifest's recorded size is a truncated /
140 interrupted download and must not count as installed.
141 """
142 try:
143 return blob_file.stat().st_size == expected_size
144 except OSError:
145 return False
148def _sha256_file(path: Path) -> str:
149 """Compute SHA-256 hex digest of a file."""
150 h = hashlib.sha256()
151 with path.open("rb") as f:
152 while True:
153 chunk = f.read(_HASH_CHUNK_SIZE)
154 if not chunk:
155 break
156 h.update(chunk)
157 return h.hexdigest()
160_SHA256_HEX = re.compile(r"[0-9a-f]{64}")
163def _blob_digest(source_path: Path) -> str:
164 """Digest for *source_path*, reusing the HF-cache blob name when possible.
166 huggingface_hub names cache blobs by their sha256, so a snapshot path that
167 resolves into a ``blobs/`` dir already carries its digest. Re-hashing a
168 multi-GB GGUF only to recompute that name is slow and, on network volumes,
169 I/O-fragile enough to fail registration outright. Plain files still hash.
170 """
171 real = source_path.resolve()
172 if real.parent.name == "blobs" and _SHA256_HEX.fullmatch(real.name):
173 return real.name
174 return _sha256_file(source_path)
177class ModelRegistry:
178 """Read/write manifests and resolve refs to blobs in the HF cache."""
180 def __init__(self, models_dir: Path) -> None:
181 self._root = models_dir
182 self._manifests_dir = models_dir / "manifests"
184 def _repo_cache_dir(self, hf_repo: str) -> Path:
185 """The HuggingFace cache directory for *hf_repo* under this registry root."""
186 return self._root / f"models--{repo_to_dir(hf_repo)}"
188 def resolve(self, ref: str) -> Path:
189 """Return the loadable GGUF path for *ref*; ``KeyError`` if not installed.
191 A single-file GGUF resolves to its content-hashed blob (the snapshot
192 file on a cache without symlinks); a split GGUF to its first shard's
193 snapshot symlink (so llama.cpp finds the sibling shards).
195 The canonical *ref* is ``<org>/<repo>/<file>.gguf`` resolved via the
196 lilbee manifest. Two other shapes are accepted as a backwards-compat
197 concession for builds already published (whose on-disk layout differs),
198 not as the intended contract: a bare ``<org>/<repo>`` (older builds
199 persisted these into ``config.toml``) resolves to the one quant of that
200 repo that's installed, and a manifest that's missing / unparseable /
201 blob-less falls back to whatever GGUF ``huggingface_hub`` reports the
202 cache holds for that ref. The HF cache layout is stable, so this lets an
203 upgrade keep working without anyone purging their lilbee data dir; it is
204 deliberately the exception here, not a pattern to follow elsewhere.
205 """
206 if is_bare_hf_repo(ref):
207 return self._resolve_repo_only(_validate_hf_repo(ref))
208 hf_repo, gguf_filename = parse_hf_ref(ref)
209 shards = split_shard_filenames(gguf_filename)
210 if len(shards) > 1:
211 return self._resolve_split(ref, hf_repo, shards)
212 manifest = self._read_manifest(hf_repo, gguf_filename)
213 if manifest is not None:
214 backing = self._manifest_backing_file(manifest)
215 if backing is not None:
216 return backing
217 recovered = self._find_cached_gguf(hf_repo, gguf_filename)
218 if recovered is not None:
219 self._reregister_from_cache(hf_repo, gguf_filename, recovered)
220 return recovered
221 if manifest is None:
222 raise KeyError(f"Model {ref} not installed")
223 # Manifest present but neither it nor the cache yields a blob; keep the
224 # specific diagnostic so a corrupted cache stays debuggable.
225 cache_path = self._repo_cache_dir(manifest.hf_repo)
226 if not cache_path.exists():
227 raise KeyError(f"Cache folder missing for {ref}: {cache_path.name}")
228 if manifest.blob is None:
229 raise KeyError(f"Manifest for {ref} has no blob hash; install incomplete")
230 blob_file = cache_path / "blobs" / manifest.blob
231 if blob_file.exists():
232 raise KeyError(
233 f"Blob for {ref} is truncated: {blob_file.stat().st_size} of "
234 f"{manifest.size_bytes} bytes; re-download required"
235 )
236 raise KeyError(f"Blob file missing for {ref}: {manifest.blob}")
238 def _resolve_split(self, ref: str, hf_repo: str, shards: list[str]) -> Path:
239 """Resolve a split GGUF to its first shard's snapshot symlink.
241 llama.cpp loads the whole set from the first shard, locating the siblings
242 by filename next to it. Only the snapshot dir co-locates the shards under
243 their real names (the blobs dir names them by hash), so hand back the
244 symlink, not the blob. Every shard must be present first: the first shard
245 alone used to read as installed, registering an unloadable model that a
246 re-pull then skipped.
247 """
248 if not self._split_shards_present(hf_repo, shards[0]):
249 raise KeyError(f"Split GGUF {ref} is missing shards; re-pull to fetch the full set")
250 first_shard = self._snapshot_gguf_path(hf_repo, shards[0])
251 if first_shard is None:
252 raise KeyError(f"Model {ref} not installed")
253 if self._read_manifest(hf_repo, shards[0]) is None:
254 # Same cache recovery as the single-file path; resolve the symlink so
255 # the manifest records the content-hashed blob, not the link, and pass
256 # the snapshot path so the shard accounting is recovered too.
257 self._reregister_from_cache(
258 hf_repo, shards[0], first_shard.resolve(), snapshot_path=first_shard
259 )
260 return first_shard
262 def _resolve_repo_only(self, hf_repo: str) -> Path:
263 """Resolve a bare ``<org>/<repo>`` ref to the GGUF of that repo on disk.
265 Older builds persisted bare repo refs for the chat / embedding model.
266 Prefers a current-format manifest under the repo; otherwise asks
267 ``huggingface_hub`` what GGUFs the cache holds for the repo and returns
268 the first one (alphabetical for determinism if more than one quant is
269 installed).
270 """
271 manifest_dir = self._manifests_dir / repo_to_dir(hf_repo)
272 if manifest_dir.is_dir():
273 # rglob, like list_installed: a quant-subdir ref writes its manifest one
274 # directory deeper, so a non-recursive scan would miss it and fall through
275 # to the slower huggingface_hub cache recovery.
276 for mf in sorted(manifest_dir.rglob("*.gguf.json")):
277 manifest = self._load_manifest_file(mf)
278 if manifest is None:
279 continue
280 backing = self._manifest_backing_file(manifest)
281 if backing is not None:
282 return backing
283 for filename in sorted(self._cached_gguf_names(hf_repo)):
284 shards = split_shard_filenames(filename)
285 if len(shards) > 1:
286 # A split set in the cache: skip its non-first shards and resolve
287 # the whole set from shard 1 so we hand back the snapshot symlink
288 # (siblings co-located, loadable) with shard accounting, not
289 # shard 1's blob as an unloadable single file.
290 if filename != shards[0]:
291 continue
292 with contextlib.suppress(KeyError):
293 return self._resolve_split(
294 format_native_gguf_ref(hf_repo, filename), hf_repo, shards
295 )
296 continue
297 recovered = self._find_cached_gguf(hf_repo, filename)
298 if recovered is not None:
299 self._reregister_from_cache(hf_repo, filename, recovered)
300 return recovered
301 raise KeyError(f"Model {hf_repo} not installed")
303 def _cached_gguf_names(self, hf_repo: str) -> set[str]:
304 """``.gguf`` filenames the HuggingFace cache holds for *hf_repo*."""
305 if not self._root.is_dir():
306 return set()
307 from huggingface_hub import scan_cache_dir
309 info = scan_cache_dir(self._root)
310 return {
311 f.file_name
312 for repo in info.repos
313 if repo.repo_id == hf_repo
314 for rev in repo.revisions
315 for f in rev.files
316 if f.file_name.endswith(GGUF_SUFFIX)
317 }
319 def _snapshot_gguf_path(self, hf_repo: str, gguf_filename: str) -> Path | None:
320 """Return the snapshot *symlink* path for a cached GGUF, or None.
322 Returns the symlink, not the blob, so a split GGUF loads from a dir where
323 its sibling shards are co-located under their real names.
324 """
325 from huggingface_hub import try_to_load_from_cache
327 hit = try_to_load_from_cache(
328 repo_id=hf_repo, filename=gguf_filename, cache_dir=str(self._root)
329 )
330 candidate: Path | None = None
331 if isinstance(hit, str): # exact repo-relative match
332 candidate = Path(hit)
333 else: # None or the _CACHED_NO_EXIST sentinel: locate the basename instead
334 snapshots = self._repo_cache_dir(hf_repo) / "snapshots"
335 if snapshots.is_dir():
336 basename = Path(gguf_filename).name
337 # Several cached revisions can hold the basename; prefer the most
338 # recently materialized one over an arbitrary lexicographic pick.
339 candidate = max(
340 snapshots.rglob(basename), key=lambda p: p.lstat().st_mtime, default=None
341 )
342 if candidate is None:
343 return None
344 try:
345 validate_path_within(candidate.resolve(), self._root)
346 except ValueError:
347 return None
348 return candidate
350 def _find_cached_gguf(self, hf_repo: str, gguf_filename: str) -> Path | None:
351 """Return the cached blob path for ``hf_repo``/``gguf_filename``, or None.
353 Locates the snapshot symlink (subdir-aware) and resolves it to its blob,
354 bounded to the cache directory.
355 """
356 symlink = self._snapshot_gguf_path(hf_repo, gguf_filename)
357 return symlink.resolve() if symlink is not None else None
359 def _split_shards_present(self, hf_repo: str, gguf_filename: str) -> bool:
360 """True unless *gguf_filename* is a split GGUF missing one of its shards.
362 A single-file GGUF is always present here. For a split set
363 (``<base>-0000N-of-0000M.gguf``) every shard must be cached, since
364 llama.cpp loads the whole set from the first shard but needs them all.
365 """
366 shards = split_shard_filenames(gguf_filename)
367 if len(shards) == 1:
368 return True
369 return all(self._find_cached_gguf(hf_repo, shard) is not None for shard in shards)
371 def shard_paths(self, ref: str) -> list[Path]:
372 """On-disk paths of *ref*'s GGUF shards that exist next to its resolved path.
374 A split GGUF resolves to its first shard's snapshot symlink with the
375 siblings co-located, so every shard is returned; a single-file GGUF
376 resolves to its content-hashed blob, where no sibling exists under the
377 real filename. Raises ``KeyError`` / ``ValueError`` like :meth:`resolve`.
378 """
379 first = self.resolve(ref)
380 _repo, filename = parse_hf_ref(ref)
381 candidates = (
382 first.parent / Path(shard).name for shard in split_shard_filenames(Path(filename).name)
383 )
384 return [path for path in candidates if path.exists()]
386 def _reregister_from_cache(
387 self,
388 hf_repo: str,
389 gguf_filename: str,
390 blob_path: Path,
391 snapshot_path: Path | None = None,
392 ) -> None:
393 """Best-effort manifest write for a cache-recovered model so listings see it.
395 *snapshot_path* is the first shard's snapshot path (siblings co-located);
396 when given, the split-shard accounting is recovered too, so a cache-only
397 split GGUF still frees every shard and reports its full size.
398 """
399 ref = format_native_gguf_ref(hf_repo, gguf_filename)
400 try:
401 task = ModelTask(reclassify_by_name(ref, ModelTask.CHAT))
402 total_size, shard_blobs = (
403 _shard_accounting(snapshot_path) if snapshot_path is not None else (None, [])
404 )
405 self._write_manifest(
406 ModelManifest(
407 hf_repo=hf_repo,
408 gguf_filename=gguf_filename,
409 size_bytes=blob_path.stat().st_size,
410 task=task,
411 downloaded_at=datetime.now(UTC).isoformat(),
412 # Only a file under blobs/ is named by its sha; a symlink-less
413 # snapshot file is not, so no digest is recorded for it.
414 blob=blob_path.name if blob_path.parent.name == "blobs" else None,
415 total_size_bytes=total_size,
416 shard_blobs=shard_blobs,
417 )
418 )
419 log.info("Recovered manifest for %s from the model cache", ref)
420 except Exception: # cache-warming write; the resolve already returned a path
421 log.debug("Could not re-register %s from the model cache", ref, exc_info=True)
423 def is_installed(self, ref: str) -> bool:
424 """Return True if a model is installed and its blob is present."""
425 try:
426 self.resolve(ref)
427 return True
428 except (KeyError, ValueError):
429 return False
431 def install(
432 self,
433 hf_repo: str,
434 gguf_filename: str,
435 source_path: Path,
436 manifest: ModelManifest,
437 ) -> Path:
438 """Write a manifest, copying *source_path* into the HF cache if needed."""
439 digest = _blob_digest(source_path)
440 cache_path = self._repo_cache_dir(hf_repo)
441 blobs_dir = cache_path / "blobs"
442 blob_path = blobs_dir / digest
443 if not blob_path.exists():
444 blobs_dir.mkdir(parents=True, exist_ok=True)
445 _copy_atomic(source_path, blob_path)
447 updated = ModelManifest(
448 hf_repo=hf_repo,
449 gguf_filename=gguf_filename,
450 # Record the size install actually wrote, not the caller's claim,
451 # so the on-disk size check has a trustworthy reference.
452 size_bytes=source_path.stat().st_size,
453 task=manifest.task,
454 downloaded_at=manifest.downloaded_at,
455 blob=digest,
456 # Carry the split-shard accounting through unchanged (computed by the
457 # caller from the full shard set); install only rewrites the primary.
458 total_size_bytes=manifest.total_size_bytes,
459 shard_blobs=manifest.shard_blobs,
460 )
461 self._write_manifest(updated)
462 return blob_path
464 def remove(self, ref: str) -> bool:
465 """Remove a manifest and its backing blob.
467 The blob is shared via SHA-256 digest, so it only goes away
468 when no other installed manifest references the same digest.
469 Empty cache directories (``blobs/``, the per-repo ``models--``
470 folder, and the per-repo manifest folder) are pruned so a
471 deleted model leaves no orphan bytes behind.
472 """
473 try:
474 hf_repo, gguf_filename = parse_hf_ref(ref)
475 except ValueError:
476 return False
477 manifest = self._read_manifest(hf_repo, gguf_filename)
478 if manifest is None:
479 return False
480 # Manifests written before shard accounting existed have no shard_blobs, so
481 # recover them from the cache *before* unlinking (resolve needs the manifest).
482 shard_blobs = manifest.shard_blobs or self._recover_legacy_shard_blobs(ref)
483 manifest_path = self._manifest_path(hf_repo, gguf_filename)
484 manifest_path.unlink()
485 repo_dir = manifest_path.parent
486 if repo_dir.exists() and not any(repo_dir.iterdir()):
487 repo_dir.rmdir()
488 self._unlink_snapshot_entries(manifest)
489 # Free the primary blob and every extra shard blob; a split GGUF has more
490 # than one, and leaving the others orphans them when a sibling quant keeps
491 # the repo cache dir alive. The surviving manifests for this repo are read
492 # once here rather than per digest (list_installed walks the whole tree).
493 # A symlink-less recovery records no digest; GC still runs once so an
494 # unused repo cache dir is pruned.
495 siblings = [m for m in self.list_installed() if m.hf_repo == manifest.hf_repo]
496 digests: list[str | None] = [d for d in [manifest.blob, *shard_blobs] if d is not None]
497 for digest in digests or [None]:
498 self._gc_blob(manifest.hf_repo, digest, siblings=siblings)
499 log.info("Removed model %s", ref)
500 return True
502 def _unlink_snapshot_entries(self, manifest: ModelManifest) -> None:
503 """Drop the snapshot entries for *manifest*'s shards.
505 On a symlink-less cache they hold the bytes and would resurrect the
506 model via cache recovery; on a symlinked cache they would dangle once
507 the blob is gc'd.
508 """
509 for shard in split_shard_filenames(manifest.gguf_filename):
510 snapshot_entry = self._snapshot_gguf_path(manifest.hf_repo, shard)
511 if snapshot_entry is not None:
512 snapshot_entry.unlink(missing_ok=True)
514 def _recover_legacy_shard_blobs(self, ref: str) -> list[str]:
515 """Extra shard blob digests for a pre-accounting split GGUF, best-effort.
517 Older manifests recorded only the first shard, so removing them would
518 orphan the rest. Derive the sibling shards from the cache; empty on any
519 failure or for a single-file model, so removal never breaks.
520 """
521 with contextlib.suppress(Exception):
522 shards = self.shard_paths(ref)
523 return [_blob_digest(path) for path in shards[1:]]
524 return []
526 def _gc_blob(
527 self, hf_repo: str, digest: str | None, *, siblings: list[ModelManifest] | None = None
528 ) -> None:
529 """Drop blob bytes and HuggingFace cache cruft now that *digest*
530 and possibly the whole repo are unused. A None *digest* (a manifest
531 recovered on a symlink-less cache records no digest) only prunes the
532 repo dir when no installed manifest is left.
534 When the per-repo ``models--<repo>/`` directory has no installed
535 manifests left, the whole directory is wiped so HF's ``refs/``,
536 ``snapshots/``, and stale ``blobs/`` all go with it. Otherwise
537 only the specific blob file is removed when no remaining
538 manifest still references its digest.
540 ``siblings`` is the surviving-manifest list for *hf_repo*; callers
541 freeing several blobs at once pass it in so the manifest tree is walked
542 once instead of per digest. Defaults to reading it when omitted.
543 """
544 cache_path = self._repo_cache_dir(hf_repo)
545 try:
546 validate_path_within(cache_path, self._root)
547 except ValueError:
548 log.warning("Refusing to remove cache outside models_dir: %s", cache_path)
549 return
550 if siblings is None:
551 siblings = [m for m in self.list_installed() if m.hf_repo == hf_repo]
552 if not siblings:
553 if cache_path.exists():
554 shutil.rmtree(cache_path)
555 return
556 if digest is None: # no digest recorded (symlink-less recovery): dir-prune only
557 return
558 if any(digest == m.blob or digest in m.shard_blobs for m in siblings):
559 return
560 blob_file = cache_path / "blobs" / digest
561 try:
562 validate_path_within(blob_file, self._root)
563 except ValueError:
564 log.warning("Refusing to remove blob outside models_dir: %s", blob_file)
565 return
566 if blob_file.exists():
567 blob_file.unlink()
569 def list_installed(self) -> list[ModelManifest]:
570 """Return manifests for models whose blob is fully present on disk.
572 A manifest with a null blob field or a missing blob file is the
573 residue of a canceled or partial download. Surfacing it would
574 let the picker offer an unusable selection, so the read filter
575 lives here at the source instead of in every UI caller.
576 """
577 manifests: list[ModelManifest] = []
578 if not self._manifests_dir.exists():
579 return manifests
580 for repo_dir in sorted(self._manifests_dir.iterdir()):
581 if not repo_dir.is_dir():
582 continue
583 # rglob, not glob: a quant-subdir ref (unsloth stores quants under e.g.
584 # Q4_K_S/<model>.gguf) writes its manifest one level deeper, so a
585 # non-recursive scan omitted it from `model list` and /v1/models, and
586 # opencode silently fell back to its own provider.
587 for tag_file in sorted(repo_dir.rglob("*.gguf.json")):
588 manifest = self._load_manifest_file(tag_file)
589 if manifest is not None and self._blob_present(manifest):
590 manifests.append(manifest)
591 return manifests
593 def _blob_present(self, manifest: ModelManifest) -> bool:
594 """True iff *manifest*'s bytes are on disk with a matching size."""
595 return self._manifest_backing_file(manifest) is not None
597 def _manifest_backing_file(self, manifest: ModelManifest) -> Path | None:
598 """The on-disk file holding *manifest*'s bytes, or None if absent/truncated.
600 Normally the content-hashed blob; on a cache without symlinks (the
601 Windows default) the bytes live at the snapshot path, so that is the
602 fallback. ``resolve`` and ``list_installed`` must both gate on this
603 one predicate, or installed-ness disagrees between ``pull`` and ``ask``.
604 """
605 if manifest.blob is not None:
606 blob_file = self._repo_cache_dir(manifest.hf_repo) / "blobs" / manifest.blob
607 if _blob_size_matches(blob_file, manifest.size_bytes):
608 return blob_file
609 recovered = self._find_cached_gguf(manifest.hf_repo, manifest.gguf_filename)
610 if recovered is not None and _blob_size_matches(recovered, manifest.size_bytes):
611 return recovered
612 return None
614 def get_manifest(self, ref: str) -> ModelManifest | None:
615 """Return the manifest for *ref* or None if not installed."""
616 try:
617 hf_repo, gguf_filename = parse_hf_ref(ref)
618 except ValueError:
619 return None
620 return self._read_manifest(hf_repo, gguf_filename)
622 def installed_ref_for_repo(self, hf_repo: str) -> str | None:
623 """Full ``<repo>/<file>.gguf`` ref of an installed quant of *hf_repo*, or None.
625 Alphabetical-first when several quants are installed, matching
626 ``_resolve_repo_only``'s determinism.
627 """
628 refs = sorted(m.ref for m in self.list_installed() if m.hf_repo == hf_repo)
629 return refs[0] if refs else None
631 def _manifest_path(self, hf_repo: str, gguf_filename: str) -> Path:
632 repo = _validate_hf_repo(hf_repo)
633 filename = _validate_gguf_filename(gguf_filename)
634 path = self._manifests_dir / repo_to_dir(repo) / f"{filename}.json"
635 validate_path_within(path, self._manifests_dir)
636 return path
638 def _read_manifest(self, hf_repo: str, gguf_filename: str) -> ModelManifest | None:
639 return self._load_manifest_file(self._manifest_path(hf_repo, gguf_filename))
641 def _write_manifest(self, manifest: ModelManifest) -> None:
642 path = self._manifest_path(manifest.hf_repo, manifest.gguf_filename)
643 path.parent.mkdir(parents=True, exist_ok=True)
644 data = json.dumps(asdict(manifest), indent=2)
645 tmp_path: str | None = None
646 try:
647 with tempfile.NamedTemporaryFile(
648 dir=path.parent, suffix=".tmp", mode="w", encoding="utf-8", delete=False
649 ) as tmp:
650 tmp_path = tmp.name
651 tmp.write(data)
652 os.replace(tmp_path, path)
653 except BaseException:
654 if tmp_path is not None:
655 Path(tmp_path).unlink(missing_ok=True)
656 raise
658 def _load_manifest_file(self, path: Path) -> ModelManifest | None:
659 if not path.exists():
660 return None
661 try:
662 data = json.loads(path.read_text(encoding="utf-8"))
663 return ModelManifest(**data)
664 # UnicodeDecodeError is a ValueError, not a JSONDecodeError.
665 except (json.JSONDecodeError, UnicodeDecodeError, TypeError, KeyError):
666 log.warning("Corrupt manifest: %s", path)
667 return None
670_HF_SNAPSHOTS_DIR = "snapshots"
673def _repo_relative_gguf_name(file_path: Path) -> str:
674 """Recover the repo-relative GGUF filename, keeping any subdir prefix.
676 HF caches a file at ``models--<repo>/snapshots/<rev>/[<subdir>/]<name>``. A
677 subdir-quant giant (unsloth stores quants under e.g. ``Q4_K_M/``) must
678 register under that subdir-relative name so its manifest key round-trips with
679 the ref; ``file_path.name`` alone would drop the subdir. Falls back to the
680 basename when the path is not under a snapshot revision dir.
681 """
682 parts = file_path.parts
683 if _HF_SNAPSHOTS_DIR not in parts:
684 return file_path.name
685 rev_index = parts.index(_HF_SNAPSHOTS_DIR) + 1
686 relative_parts = parts[rev_index + 1 :]
687 return "/".join(relative_parts) if relative_parts else file_path.name
690def _shard_accounting(first_shard_path: Path) -> tuple[int | None, list[str]]:
691 """Total on-disk size and non-primary shard blob digests for a split GGUF.
693 ``(None, [])`` for a single-file model. For a split GGUF, the sibling shards
694 live next to the first shard; ``_blob_digest`` yields each shard's blob digest
695 (the HF cache blob name, or a content hash in copy/non-symlink mode), so the
696 shards are summed and the digests of shards 2..N collected for removal-time
697 garbage collection.
698 """
699 shard_names = split_shard_filenames(first_shard_path.name)
700 if len(shard_names) <= 1:
701 return None, []
702 total = 0
703 shard_blobs: list[str] = []
704 for index, name in enumerate(shard_names):
705 shard_path = first_shard_path.with_name(name)
706 if not shard_path.exists():
707 continue
708 total += shard_path.stat().st_size
709 if index > 0: # the primary blob is tracked separately as manifest.blob
710 shard_blobs.append(_blob_digest(shard_path))
711 return total, shard_blobs
714def register_downloaded_model(entry: CatalogModel, file_path: Path) -> None:
715 """Write a registry manifest for a freshly downloaded GGUF.
717 A failed manifest write is logged, not raised, when the GGUF is still in the
718 HF cache (``resolve`` recovers from it); if it isn't, the download itself is
719 broken and the failure propagates so the caller reports it.
720 """
721 registry = ModelRegistry(cfg.models_dir)
722 gguf_filename = _repo_relative_gguf_name(file_path)
723 total_size, shard_blobs = _shard_accounting(file_path)
724 manifest = ModelManifest(
725 hf_repo=entry.hf_repo,
726 gguf_filename=gguf_filename,
727 size_bytes=file_path.stat().st_size,
728 task=entry.task,
729 downloaded_at=datetime.now(UTC).isoformat(),
730 total_size_bytes=total_size,
731 shard_blobs=shard_blobs,
732 )
733 try:
734 registry.install(entry.hf_repo, gguf_filename, file_path, manifest)
735 log.info("Registered %s/%s in manifest", entry.hf_repo, gguf_filename)
736 except Exception:
737 ref = format_native_gguf_ref(entry.hf_repo, gguf_filename)
738 if not registry.is_installed(ref):
739 raise
740 log.warning(
741 "Manifest write failed for %s; recovered via the model cache", ref, exc_info=True
742 )