Coverage for src/lilbee/modelhub/registry.py: 100%
397 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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
35from lilbee.providers import gguf_meta
36from lilbee.providers.base import ProviderError
38if TYPE_CHECKING:
39 from lilbee.catalog.models import CatalogModel
41log = logging.getLogger(__name__)
43_HASH_CHUNK_SIZE = 8192 # bytes read per iteration when hashing
44_REPO_SEGMENT_RE = re.compile(r"^[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+$")
45# A GGUF filename, optionally under repo subdirectories (unsloth stores quants
46# in e.g. ``Q4_K_M/Model-...-00001-of-00003.gguf``). Path separators are allowed;
47# ``..`` and absolute paths are rejected in the validator to stay inside the repo.
48_FILENAME_RE = re.compile(r"^[a-zA-Z0-9._/-]+\.gguf$")
50REPO_DIR_SEPARATOR = "--"
51_MANIFEST_SUFFIX = f"{GGUF_SUFFIX}.json"
54def _validate_hf_repo(hf_repo: str) -> str:
55 """Validate that a HuggingFace repo id has the form ``org/name``."""
56 if not hf_repo or not _REPO_SEGMENT_RE.match(hf_repo) or ".." in hf_repo:
57 raise ValueError(f"Invalid hf_repo: {hf_repo!r}")
58 return hf_repo
61def _validate_gguf_filename(filename: str) -> str:
62 """Validate a ``.gguf`` filename, allowing repo subdirectories but no traversal."""
63 if (
64 not filename
65 or not _FILENAME_RE.match(filename)
66 or ".." in filename
67 or filename.startswith("/")
68 ):
69 raise ValueError(f"Invalid gguf_filename: {filename!r}")
70 return filename
73_REF_SHAPE_HINT = "Use '<org>/<repo>/<filename>.gguf'."
76def parse_hf_ref(ref: str) -> tuple[str, str]:
77 """Split ``<org>/<repo>/<file>.gguf`` into ``(hf_repo, gguf_filename)``.
79 The repo is always the first two segments (``<org>/<repo>``); everything
80 after is the filename, which may include repo subdirectories (unsloth stores
81 quants under e.g. ``Q4_K_M/Model-...-00001-of-00003.gguf``).
82 """
83 if not ref.endswith(GGUF_SUFFIX) or ref.count("/") < NATIVE_GGUF_REF_MIN_SLASHES:
84 raise ValueError(f"Model ref {ref!r} is not a HuggingFace ref. {_REF_SHAPE_HINT}")
85 parts = ref.split("/")
86 hf_repo = "/".join(parts[:NATIVE_GGUF_REF_MIN_SLASHES])
87 gguf_filename = "/".join(parts[NATIVE_GGUF_REF_MIN_SLASHES:])
88 return _validate_hf_repo(hf_repo), _validate_gguf_filename(gguf_filename)
91def repo_to_dir(hf_repo: str) -> str:
92 """Encode an HF repo for use as a directory name (HF cache convention)."""
93 return hf_repo.replace("/", REPO_DIR_SEPARATOR)
96@dataclass
97class ModelManifest:
98 """One installed model's metadata. Identity: ``(hf_repo, gguf_filename)``."""
100 hf_repo: str
101 gguf_filename: str
102 size_bytes: int # primary (first-shard) blob size; validated against the blob on disk
103 task: ModelTask
104 downloaded_at: str # ISO 8601
105 blob: str | None = None # SHA-256 hex of the blob in the HF cache; None pre-install
106 # A split GGUF has further shard blobs beyond ``blob``. ``total_size_bytes``
107 # is the sum across every shard (None = single file, use ``size_bytes``);
108 # ``shard_blobs`` are the non-primary shard digests, so removal frees them all.
109 total_size_bytes: int | None = None
110 shard_blobs: list[str] = field(default_factory=list)
112 @property
113 def ref(self) -> str:
114 return format_native_gguf_ref(self.hf_repo, self.gguf_filename)
116 @property
117 def disk_size_bytes(self) -> int:
118 """Total bytes this model occupies on disk, across all shards."""
119 return self.total_size_bytes if self.total_size_bytes is not None else self.size_bytes
122def _copy_atomic(source_path: Path, blob_path: Path) -> None:
123 """Copy *source_path* to *blob_path* via a temp file + atomic rename.
125 A crash mid-copy leaves only the temp file, never a partial blob at
126 the final digest path that callers would treat as complete.
127 """
128 fd, tmp_name = tempfile.mkstemp(dir=str(blob_path.parent), suffix=".part")
129 tmp_path = Path(tmp_name)
130 try:
131 with os.fdopen(fd, "wb") as dst, source_path.open("rb") as src:
132 shutil.copyfileobj(src, dst)
133 os.replace(tmp_path, blob_path)
134 except BaseException:
135 tmp_path.unlink(missing_ok=True)
136 raise
139def _blob_size_matches(blob_file: Path, expected_size: int) -> bool:
140 """True iff *blob_file* exists and its byte size equals *expected_size*.
142 A blob shorter than the manifest's recorded size is a truncated /
143 interrupted download and must not count as installed.
144 """
145 try:
146 return blob_file.stat().st_size == expected_size
147 except OSError:
148 return False
151def _sha256_file(path: Path) -> str:
152 """Compute SHA-256 hex digest of a file."""
153 h = hashlib.sha256()
154 with path.open("rb") as f:
155 while True:
156 chunk = f.read(_HASH_CHUNK_SIZE)
157 if not chunk:
158 break
159 h.update(chunk)
160 return h.hexdigest()
163_SHA256_HEX = re.compile(r"[0-9a-f]{64}")
164_BLOB_SHARD_LEN = 2
167def _blob_digest(source_path: Path) -> str:
168 """Digest for *source_path*, reusing the HF-cache blob name when possible.
170 huggingface_hub names cache blobs by their sha256, so a snapshot path that
171 resolves into a ``blobs/`` dir already carries its digest. Re-hashing a
172 multi-GB GGUF only to recompute that name is slow and, on network volumes,
173 I/O-fragile enough to fail registration outright. Plain files still hash.
174 """
175 real = source_path.resolve()
176 name = _cache_blob_name(real)
177 if name is not None and _SHA256_HEX.fullmatch(real.name):
178 return name
179 return _sha256_file(source_path)
182def _cache_blob_name(real: Path) -> str | None:
183 """Name of *real* relative to its cache ``blobs/`` dir, or None outside one.
185 huggingface_hub writes flat blobs (``blobs/<sha>``), and its shared blob
186 store shards them under a two-character prefix (``blobs/a0/<sha>``). The name
187 keeps that prefix, so ``cache_path / "blobs" / name`` finds the file in both.
188 """
189 if real.parent.name == "blobs":
190 return real.name
191 if real.parent.parent.name == "blobs" and len(real.parent.name) == _BLOB_SHARD_LEN:
192 return f"{real.parent.name}/{real.name}"
193 return None
196def _reraise_walk_error(error: OSError) -> None:
197 """Re-raise a directory-read fault; a directory that vanished is not one."""
198 if not isinstance(error, FileNotFoundError):
199 raise error
202def _manifest_files(repo_dir: Path) -> list[Path]:
203 """Sorted manifest paths under *repo_dir*, quant subdirectories included.
205 Raises ``OSError`` when a directory under *repo_dir* cannot be read, so an
206 unreadable tree never reads as an empty one.
207 """
208 found: list[Path] = []
209 for dirpath, _dirnames, filenames in os.walk(repo_dir, onerror=_reraise_walk_error):
210 found.extend(Path(dirpath) / name for name in filenames if name.endswith(_MANIFEST_SUFFIX))
211 return sorted(found)
214class ModelRegistry:
215 """Read/write manifests and resolve refs to blobs in the HF cache."""
217 def __init__(self, models_dir: Path) -> None:
218 self._root = models_dir
219 self._manifests_dir = models_dir / "manifests"
221 def _repo_cache_dir(self, hf_repo: str) -> Path:
222 """The HuggingFace cache directory for *hf_repo* under this registry root."""
223 return self._root / f"models--{repo_to_dir(hf_repo)}"
225 def resolve(self, ref: str) -> Path:
226 """Return the loadable GGUF path for *ref*; ``KeyError`` if not installed.
228 A single-file GGUF resolves to its content-hashed blob (the snapshot
229 file on a cache without symlinks); a split GGUF to its first shard's
230 snapshot symlink (so llama.cpp finds the sibling shards).
232 The canonical *ref* is ``<org>/<repo>/<file>.gguf`` resolved via the
233 lilbee manifest. Two other shapes are accepted as a backwards-compat
234 concession for builds already published (whose on-disk layout differs),
235 not as the intended contract: a bare ``<org>/<repo>`` (older builds
236 persisted these into ``config.toml``) resolves to the one quant of that
237 repo that's installed, and a manifest that's missing / unparseable /
238 blob-less falls back to whatever GGUF ``huggingface_hub`` reports the
239 cache holds for that ref. The HF cache layout is stable, so this lets an
240 upgrade keep working without anyone purging their lilbee data dir; it is
241 deliberately the exception here, not a pattern to follow elsewhere.
242 """
243 if is_bare_hf_repo(ref):
244 return self._resolve_repo_only(_validate_hf_repo(ref))
245 hf_repo, gguf_filename = parse_hf_ref(ref)
246 shards = split_shard_filenames(gguf_filename)
247 if len(shards) > 1:
248 return self._resolve_split(ref, hf_repo, shards)
249 manifest = self._read_manifest(hf_repo, gguf_filename)
250 if manifest is not None:
251 backing = self._manifest_backing_file(manifest)
252 if backing is not None:
253 return backing
254 recovered = self._find_cached_gguf(hf_repo, gguf_filename)
255 if recovered is not None:
256 self._reregister_from_cache(hf_repo, gguf_filename, recovered)
257 return recovered
258 if manifest is None:
259 raise KeyError(f"Model {ref} not installed")
260 # Manifest present but neither it nor the cache yields a blob; keep the
261 # specific diagnostic so a corrupted cache stays debuggable.
262 cache_path = self._repo_cache_dir(manifest.hf_repo)
263 if not cache_path.exists():
264 raise KeyError(f"Cache folder missing for {ref}: {cache_path.name}")
265 if manifest.blob is None:
266 raise KeyError(f"Manifest for {ref} has no blob hash; install incomplete")
267 blob_file = cache_path / "blobs" / manifest.blob
268 if blob_file.exists():
269 raise KeyError(
270 f"Blob for {ref} is truncated: {blob_file.stat().st_size} of "
271 f"{manifest.size_bytes} bytes; re-download required"
272 )
273 raise KeyError(f"Blob file missing for {ref}: {manifest.blob}")
275 def _resolve_split(self, ref: str, hf_repo: str, shards: list[str]) -> Path:
276 """Resolve a split GGUF to its first shard's snapshot symlink.
278 llama.cpp loads the whole set from the first shard, locating the siblings
279 by filename next to it. Only the snapshot dir co-locates the shards under
280 their real names (the blobs dir names them by hash), so hand back the
281 symlink, not the blob. Every shard must be present first: the first shard
282 alone used to read as installed, registering an unloadable model that a
283 re-pull then skipped.
284 """
285 if not self._split_shards_present(hf_repo, shards[0]):
286 raise KeyError(f"Split GGUF {ref} is missing shards; re-pull to fetch the full set")
287 first_shard = self._snapshot_gguf_path(hf_repo, shards[0])
288 if first_shard is None:
289 raise KeyError(f"Model {ref} not installed")
290 if self._read_manifest(hf_repo, shards[0]) is None:
291 # Same cache recovery as the single-file path; resolve the symlink so
292 # the manifest records the content-hashed blob, not the link, and pass
293 # the snapshot path so the shard accounting is recovered too.
294 self._reregister_from_cache(
295 hf_repo, shards[0], first_shard.resolve(), snapshot_path=first_shard
296 )
297 return first_shard
299 def _resolve_repo_only(self, hf_repo: str) -> Path:
300 """Resolve a bare ``<org>/<repo>`` ref to the GGUF of that repo on disk.
302 Older builds persisted bare repo refs for the chat / embedding model.
303 Prefers a current-format manifest under the repo; otherwise asks
304 ``huggingface_hub`` what GGUFs the cache holds for the repo and returns
305 the first one (alphabetical for determinism if more than one quant is
306 installed).
307 """
308 manifest_dir = self._manifests_dir / repo_to_dir(hf_repo)
309 if manifest_dir.is_dir():
310 for mf in _manifest_files(manifest_dir):
311 manifest = self._load_manifest_file(mf)
312 if manifest is None:
313 continue
314 backing = self._manifest_backing_file(manifest)
315 if backing is not None:
316 return backing
317 for filename in sorted(self._cached_gguf_names(hf_repo)):
318 shards = split_shard_filenames(filename)
319 if len(shards) > 1:
320 # A split set in the cache: skip its non-first shards and resolve
321 # the whole set from shard 1 so we hand back the snapshot symlink
322 # (siblings co-located, loadable) with shard accounting, not
323 # shard 1's blob as an unloadable single file.
324 if filename != shards[0]:
325 continue
326 with contextlib.suppress(KeyError):
327 return self._resolve_split(
328 format_native_gguf_ref(hf_repo, filename), hf_repo, shards
329 )
330 continue
331 recovered = self._find_cached_gguf(hf_repo, filename)
332 if recovered is not None:
333 self._reregister_from_cache(hf_repo, filename, recovered)
334 return recovered
335 raise KeyError(f"Model {hf_repo} not installed")
337 def _cached_gguf_names(self, hf_repo: str) -> set[str]:
338 """``.gguf`` filenames the HuggingFace cache holds for *hf_repo*."""
339 if not self._root.is_dir():
340 return set()
341 from huggingface_hub import scan_cache_dir
343 info = scan_cache_dir(self._root)
344 return {
345 f.file_name
346 for repo in info.repos
347 if repo.repo_id == hf_repo
348 for rev in repo.revisions
349 for f in rev.files
350 if f.file_name.endswith(GGUF_SUFFIX)
351 }
353 def _snapshot_gguf_path(self, hf_repo: str, gguf_filename: str) -> Path | None:
354 """Return the snapshot *symlink* path for a cached GGUF, or None.
356 Returns the symlink, not the blob, so a split GGUF loads from a dir where
357 its sibling shards are co-located under their real names.
358 """
359 from huggingface_hub import try_to_load_from_cache
361 hit = try_to_load_from_cache(
362 repo_id=hf_repo, filename=gguf_filename, cache_dir=str(self._root)
363 )
364 candidate: Path | None = None
365 if isinstance(hit, str): # exact repo-relative match
366 candidate = Path(hit)
367 else: # None or the _CACHED_NO_EXIST sentinel: locate the basename instead
368 snapshots = self._repo_cache_dir(hf_repo) / "snapshots"
369 if snapshots.is_dir():
370 basename = Path(gguf_filename).name
371 # Several cached revisions can hold the basename; prefer the most
372 # recently materialized one over an arbitrary lexicographic pick.
373 candidate = max(
374 snapshots.rglob(basename), key=lambda p: p.lstat().st_mtime, default=None
375 )
376 if candidate is None:
377 return None
378 try:
379 validate_path_within(candidate.resolve(), self._root)
380 except ValueError:
381 return None
382 return candidate
384 def _find_cached_gguf(self, hf_repo: str, gguf_filename: str) -> Path | None:
385 """Return the cached blob path for ``hf_repo``/``gguf_filename``, or None.
387 Locates the snapshot symlink (subdir-aware) and resolves it to its blob,
388 bounded to the cache directory.
389 """
390 symlink = self._snapshot_gguf_path(hf_repo, gguf_filename)
391 return symlink.resolve() if symlink is not None else None
393 def _split_shards_present(self, hf_repo: str, gguf_filename: str) -> bool:
394 """True unless *gguf_filename* is a split GGUF missing one of its shards.
396 A single-file GGUF is always present here. For a split set
397 (``<base>-0000N-of-0000M.gguf``) every shard must be cached, since
398 llama.cpp loads the whole set from the first shard but needs them all.
399 """
400 shards = split_shard_filenames(gguf_filename)
401 if len(shards) == 1:
402 return True
403 return all(self._find_cached_gguf(hf_repo, shard) is not None for shard in shards)
405 def shard_paths(self, ref: str) -> list[Path]:
406 """On-disk paths of *ref*'s GGUF shards that exist next to its resolved path.
408 A split GGUF resolves to its first shard's snapshot symlink with the
409 siblings co-located, so every shard is returned; a single-file GGUF
410 resolves to its content-hashed blob, where no sibling exists under the
411 real filename. Raises ``KeyError`` / ``ValueError`` like :meth:`resolve`.
412 """
413 first = self.resolve(ref)
414 _repo, filename = parse_hf_ref(ref)
415 candidates = (
416 first.parent / Path(shard).name for shard in split_shard_filenames(Path(filename).name)
417 )
418 return [path for path in candidates if path.exists()]
420 def _reregister_from_cache(
421 self,
422 hf_repo: str,
423 gguf_filename: str,
424 blob_path: Path,
425 snapshot_path: Path | None = None,
426 ) -> None:
427 """Best-effort manifest write for a cache-recovered model so listings see it.
429 *snapshot_path* is the first shard's snapshot path (siblings co-located);
430 when given, the split-shard accounting is recovered too, so a cache-only
431 split GGUF still frees every shard and reports its full size.
432 """
433 ref = format_native_gguf_ref(hf_repo, gguf_filename)
434 try:
435 task = ModelTask(reclassify_by_name(ref, ModelTask.CHAT))
436 total_size, shard_blobs = (
437 _shard_accounting(snapshot_path) if snapshot_path is not None else (None, [])
438 )
439 self._write_manifest(
440 ModelManifest(
441 hf_repo=hf_repo,
442 gguf_filename=gguf_filename,
443 size_bytes=blob_path.stat().st_size,
444 task=task,
445 downloaded_at=datetime.now(UTC).isoformat(),
446 # Only a file under blobs/ is named by its sha; a symlink-less
447 # snapshot file is not, so no digest is recorded for it.
448 blob=_cache_blob_name(blob_path),
449 total_size_bytes=total_size,
450 shard_blobs=shard_blobs,
451 )
452 )
453 log.info("Recovered manifest for %s from the model cache", ref)
454 except Exception: # cache-warming write; the resolve already returned a path
455 log.debug("Could not re-register %s from the model cache", ref, exc_info=True)
457 def has_vision_projector(self, ref: str) -> bool:
458 """True when *ref* is installed with an mmproj projector beside its weights.
460 A vision model is a chat model with a projector, so the manifest task
461 cannot tell the two apart. The vision picker and the vision-role check
462 both ask this question. The lookup walks the HF cache layout, so it
463 finds the projector from a content-hashed blob path too.
464 """
465 try:
466 gguf_meta.find_mmproj_for_model(self.resolve(ref))
467 except (KeyError, ValueError, ProviderError):
468 return False
469 return True
471 def is_installed(self, ref: str) -> bool:
472 """Return True if a model is installed and its blob is present."""
473 try:
474 self.resolve(ref)
475 return True
476 except (KeyError, ValueError):
477 return False
479 def has_manifest(self, ref: str) -> bool:
480 """True when a manifest, not only a cached blob, records *ref*.
482 ``is_installed`` also answers True for a GGUF found in the HF cache,
483 which the model list does not show until a manifest exists.
484 """
485 if is_bare_hf_repo(ref):
486 return self.installed_ref_for_repo(ref) is not None
487 return self.get_manifest(ref) is not None
489 def install(
490 self,
491 hf_repo: str,
492 gguf_filename: str,
493 source_path: Path,
494 manifest: ModelManifest,
495 ) -> Path:
496 """Write a manifest, copying *source_path* into the HF cache if needed."""
497 digest = _blob_digest(source_path)
498 cache_path = self._repo_cache_dir(hf_repo)
499 blobs_dir = cache_path / "blobs"
500 blob_path = blobs_dir / digest
501 if not blob_path.exists():
502 blob_path.parent.mkdir(parents=True, exist_ok=True)
503 _copy_atomic(source_path, blob_path)
505 updated = ModelManifest(
506 hf_repo=hf_repo,
507 gguf_filename=gguf_filename,
508 # Record the size install actually wrote, not the caller's claim,
509 # so the on-disk size check has a trustworthy reference.
510 size_bytes=source_path.stat().st_size,
511 task=manifest.task,
512 downloaded_at=manifest.downloaded_at,
513 blob=digest,
514 # Carry the split-shard accounting through unchanged (computed by the
515 # caller from the full shard set); install only rewrites the primary.
516 total_size_bytes=manifest.total_size_bytes,
517 shard_blobs=manifest.shard_blobs,
518 )
519 self._write_manifest(updated)
520 return blob_path
522 def remove(self, ref: str) -> bool:
523 """Remove a manifest and its backing blob.
525 The blob is shared via SHA-256 digest, so it only goes away
526 when no other installed manifest references the same digest.
527 Empty cache directories (``blobs/``, the per-repo ``models--``
528 folder, and the per-repo manifest folder) are pruned so a
529 deleted model leaves no orphan bytes behind.
530 """
531 try:
532 hf_repo, gguf_filename = parse_hf_ref(ref)
533 except ValueError:
534 return False
535 manifest = self._read_manifest(hf_repo, gguf_filename)
536 if manifest is None:
537 return False
538 # Manifests written before shard accounting existed have no shard_blobs, so
539 # recover them from the cache *before* unlinking (resolve needs the manifest).
540 shard_blobs = manifest.shard_blobs or self._recover_legacy_shard_blobs(ref)
541 manifest_path = self._manifest_path(hf_repo, gguf_filename)
542 manifest_path.unlink()
543 repo_dir = manifest_path.parent
544 if repo_dir.exists() and not any(repo_dir.iterdir()):
545 repo_dir.rmdir()
546 self._unlink_snapshot_entries(manifest)
547 # Free the primary blob and every extra shard blob; a split GGUF has more
548 # than one, and leaving the others orphans them when a sibling quant keeps
549 # the repo cache dir alive. The surviving manifests for this repo are read
550 # once here rather than per digest (list_installed walks the whole tree).
551 # A symlink-less recovery records no digest; GC still runs once so an
552 # unused repo cache dir is pruned.
553 siblings = [m for m in self.list_installed() if m.hf_repo == manifest.hf_repo]
554 digests: list[str | None] = [d for d in [manifest.blob, *shard_blobs] if d is not None]
555 for digest in digests or [None]:
556 self._gc_blob(manifest.hf_repo, digest, siblings=siblings)
557 log.info("Removed model %s", ref)
558 return True
560 def _unlink_snapshot_entries(self, manifest: ModelManifest) -> None:
561 """Drop the snapshot entries for *manifest*'s shards.
563 On a symlink-less cache they hold the bytes and would resurrect the
564 model via cache recovery; on a symlinked cache they would dangle once
565 the blob is gc'd.
566 """
567 for shard in split_shard_filenames(manifest.gguf_filename):
568 snapshot_entry = self._snapshot_gguf_path(manifest.hf_repo, shard)
569 if snapshot_entry is not None:
570 snapshot_entry.unlink(missing_ok=True)
572 def _recover_legacy_shard_blobs(self, ref: str) -> list[str]:
573 """Extra shard blob digests for a pre-accounting split GGUF, best-effort.
575 Older manifests recorded only the first shard, so removing them would
576 orphan the rest. Derive the sibling shards from the cache; empty on any
577 failure or for a single-file model, so removal never breaks.
578 """
579 with contextlib.suppress(Exception):
580 shards = self.shard_paths(ref)
581 return [_blob_digest(path) for path in shards[1:]]
582 return []
584 def _gc_blob(
585 self, hf_repo: str, digest: str | None, *, siblings: list[ModelManifest] | None = None
586 ) -> None:
587 """Drop blob bytes and HuggingFace cache cruft now that *digest*
588 and possibly the whole repo are unused. A None *digest* (a manifest
589 recovered on a symlink-less cache records no digest) only prunes the
590 repo dir when no installed manifest is left.
592 When the per-repo ``models--<repo>/`` directory has no installed
593 manifests left, the whole directory is wiped so HF's ``refs/``,
594 ``snapshots/``, and stale ``blobs/`` all go with it. Otherwise
595 only the specific blob file is removed when no remaining
596 manifest still references its digest.
598 ``siblings`` is the surviving-manifest list for *hf_repo*; callers
599 freeing several blobs at once pass it in so the manifest tree is walked
600 once instead of per digest. Defaults to reading it when omitted.
601 """
602 cache_path = self._repo_cache_dir(hf_repo)
603 try:
604 validate_path_within(cache_path, self._root)
605 except ValueError:
606 log.warning("Refusing to remove cache outside models_dir: %s", cache_path)
607 return
608 if siblings is None:
609 siblings = [m for m in self.list_installed() if m.hf_repo == hf_repo]
610 if not siblings:
611 if cache_path.exists():
612 shutil.rmtree(cache_path)
613 return
614 if digest is None: # no digest recorded (symlink-less recovery): dir-prune only
615 return
616 if any(digest == m.blob or digest in m.shard_blobs for m in siblings):
617 return
618 blob_file = cache_path / "blobs" / digest
619 try:
620 validate_path_within(blob_file, self._root)
621 except ValueError:
622 log.warning("Refusing to remove blob outside models_dir: %s", blob_file)
623 return
624 if blob_file.exists():
625 blob_file.unlink()
627 def list_installed(self) -> list[ModelManifest]:
628 """Return manifests for models whose blob is fully present on disk.
630 A manifest with a null blob field or a missing blob file is the
631 residue of a canceled or partial download. Surfacing it would
632 let the picker offer an unusable selection, so the read filter
633 lives here at the source instead of in every UI caller.
635 An absent manifest tree is an empty list. A tree that is present but
636 cannot be read raises ``OSError``, because "unknown" is not "none".
637 """
638 manifests: list[ModelManifest] = []
639 try:
640 repo_dirs = sorted(self._manifests_dir.iterdir())
641 except FileNotFoundError:
642 return manifests
643 for repo_dir in repo_dirs:
644 if not repo_dir.is_dir():
645 continue
646 for tag_file in _manifest_files(repo_dir):
647 manifest = self._load_manifest_file(tag_file)
648 if manifest is not None and self._blob_present(manifest):
649 manifests.append(manifest)
650 return manifests
652 def _blob_present(self, manifest: ModelManifest) -> bool:
653 """True iff *manifest*'s bytes are on disk with a matching size."""
654 return self._manifest_backing_file(manifest) is not None
656 def _manifest_backing_file(self, manifest: ModelManifest) -> Path | None:
657 """The on-disk file holding *manifest*'s bytes, or None if absent/truncated.
659 Normally the content-hashed blob; on a cache without symlinks (the
660 Windows default) the bytes live at the snapshot path, so that is the
661 fallback. ``resolve`` and ``list_installed`` must both gate on this
662 one predicate, or installed-ness disagrees between ``pull`` and ``ask``.
663 """
664 if manifest.blob is not None:
665 blob_file = self._repo_cache_dir(manifest.hf_repo) / "blobs" / manifest.blob
666 if _blob_size_matches(blob_file, manifest.size_bytes):
667 return blob_file
668 recovered = self._find_cached_gguf(manifest.hf_repo, manifest.gguf_filename)
669 if recovered is not None and _blob_size_matches(recovered, manifest.size_bytes):
670 return recovered
671 return None
673 def get_manifest(self, ref: str) -> ModelManifest | None:
674 """Return the manifest for *ref* or None if not installed."""
675 try:
676 hf_repo, gguf_filename = parse_hf_ref(ref)
677 except ValueError:
678 return None
679 return self._read_manifest(hf_repo, gguf_filename)
681 def installed_ref_for_repo(self, hf_repo: str) -> str | None:
682 """Full ``<repo>/<file>.gguf`` ref of an installed quant of *hf_repo*, or None.
684 Alphabetical-first when several quants are installed, matching
685 ``_resolve_repo_only``'s determinism.
686 """
687 refs = sorted(m.ref for m in self.list_installed() if m.hf_repo == hf_repo)
688 return refs[0] if refs else None
690 def _manifest_path(self, hf_repo: str, gguf_filename: str) -> Path:
691 repo = _validate_hf_repo(hf_repo)
692 filename = _validate_gguf_filename(gguf_filename)
693 path = self._manifests_dir / repo_to_dir(repo) / f"{filename}.json"
694 validate_path_within(path, self._manifests_dir)
695 return path
697 def _read_manifest(self, hf_repo: str, gguf_filename: str) -> ModelManifest | None:
698 return self._load_manifest_file(self._manifest_path(hf_repo, gguf_filename))
700 def _write_manifest(self, manifest: ModelManifest) -> None:
701 path = self._manifest_path(manifest.hf_repo, manifest.gguf_filename)
702 path.parent.mkdir(parents=True, exist_ok=True)
703 data = json.dumps(asdict(manifest), indent=2)
704 tmp_path: str | None = None
705 try:
706 with tempfile.NamedTemporaryFile(
707 dir=path.parent, suffix=".tmp", mode="w", encoding="utf-8", delete=False
708 ) as tmp:
709 tmp_path = tmp.name
710 tmp.write(data)
711 os.replace(tmp_path, path)
712 except BaseException:
713 if tmp_path is not None:
714 Path(tmp_path).unlink(missing_ok=True)
715 raise
717 def _load_manifest_file(self, path: Path) -> ModelManifest | None:
718 if not path.exists():
719 return None
720 try:
721 data = json.loads(path.read_text(encoding="utf-8"))
722 return ModelManifest(**data)
723 # UnicodeDecodeError is a ValueError, not a JSONDecodeError.
724 except (json.JSONDecodeError, UnicodeDecodeError, TypeError, KeyError):
725 log.warning("Corrupt manifest: %s", path)
726 return None
729_HF_SNAPSHOTS_DIR = "snapshots"
732def _repo_relative_gguf_name(file_path: Path) -> str:
733 """Recover the repo-relative GGUF filename, keeping any subdir prefix.
735 HF caches a file at ``models--<repo>/snapshots/<rev>/[<subdir>/]<name>``. A
736 subdir-quant giant (unsloth stores quants under e.g. ``Q4_K_M/``) must
737 register under that subdir-relative name so its manifest key round-trips with
738 the ref; ``file_path.name`` alone would drop the subdir. Falls back to the
739 basename when the path is not under a snapshot revision dir.
740 """
741 parts = file_path.parts
742 if _HF_SNAPSHOTS_DIR not in parts:
743 return file_path.name
744 rev_index = parts.index(_HF_SNAPSHOTS_DIR) + 1
745 relative_parts = parts[rev_index + 1 :]
746 return "/".join(relative_parts) if relative_parts else file_path.name
749def _shard_accounting(first_shard_path: Path) -> tuple[int | None, list[str]]:
750 """Total on-disk size and non-primary shard blob digests for a split GGUF.
752 ``(None, [])`` for a single-file model. For a split GGUF, the sibling shards
753 live next to the first shard; ``_blob_digest`` yields each shard's blob digest
754 (the HF cache blob name, or a content hash in copy/non-symlink mode), so the
755 shards are summed and the digests of shards 2..N collected for removal-time
756 garbage collection.
757 """
758 shard_names = split_shard_filenames(first_shard_path.name)
759 if len(shard_names) <= 1:
760 return None, []
761 total = 0
762 shard_blobs: list[str] = []
763 for index, name in enumerate(shard_names):
764 shard_path = first_shard_path.with_name(name)
765 if not shard_path.exists():
766 continue
767 total += shard_path.stat().st_size
768 if index > 0: # the primary blob is tracked separately as manifest.blob
769 shard_blobs.append(_blob_digest(shard_path))
770 return total, shard_blobs
773def register_downloaded_model(entry: CatalogModel, file_path: Path) -> None:
774 """Write a registry manifest for a freshly downloaded GGUF.
776 A failed manifest write raises ``RuntimeError`` with a plain message. The
777 model list reads manifests only, so a pull that reported success without one
778 would leave a model that is installed but not listed.
779 """
780 registry = ModelRegistry(cfg.models_dir)
781 gguf_filename = _repo_relative_gguf_name(file_path)
782 total_size, shard_blobs = _shard_accounting(file_path)
783 manifest = ModelManifest(
784 hf_repo=entry.hf_repo,
785 gguf_filename=gguf_filename,
786 size_bytes=file_path.stat().st_size,
787 task=entry.task,
788 downloaded_at=datetime.now(UTC).isoformat(),
789 total_size_bytes=total_size,
790 shard_blobs=shard_blobs,
791 )
792 try:
793 registry.install(entry.hf_repo, gguf_filename, file_path, manifest)
794 log.info("Registered %s/%s in manifest", entry.hf_repo, gguf_filename)
795 except OSError as exc:
796 ref = format_native_gguf_ref(entry.hf_repo, gguf_filename)
797 raise RuntimeError(
798 f"{ref} downloaded, but lilbee cannot save its record in "
799 f"{cfg.models_dir}: {exc}. Fix the folder permissions or free disk space, "
800 "then pull again."
801 ) from exc