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