diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 807ec70991..14f03d9da4 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -269,7 +269,14 @@ def _cache_inventory_fields( active_hub_cache: Optional[Path] = None, partial: bool = False, requires_variant: bool = False, + payload_snapshots: Optional[frozenset[str]] = None, ) -> dict: + # *snapshot_path* becomes the load identity whenever the repo id will not + # resolve, so callers must pass a snapshot that holds the payload this row + # advertises, not merely the newest one. *payload_snapshots* are every + # snapshot that does hold it, used to judge where the repo id would land; + # None means the caller does not track payloads and *snapshot_path* is + # taken on trust. load_id = repo_id active_cache = True if repo_path is not None: @@ -285,6 +292,27 @@ def _cache_inventory_fields( except (OSError, RuntimeError, ValueError): active_cache = False load_id = str(snapshot_path or repo_path) + # Only pin a snapshot that is known to hold the payload. When none does the + # caller falls back to the newest snapshot, and handing that out names a + # directory the load cannot use; the repo id at least still completes the + # missing files from the hub. + if ( + load_id == repo_id + and snapshot_path is not None + and repo_path is not None + and (payload_snapshots is None or str(snapshot_path) in payload_snapshots) + ): + default_snapshot = hf_cache_scan.default_ref_snapshot(repo_path) + # No usable refs/main: from_pretrained(repo_id) would fail offline and + # would fetch the current upstream HEAD online, ignoring the snapshot + # already on disk. A refs/main that resolves is no better when it lands + # on a revision without the payload this row advertises: a metadata + # probe against a moved commit leaves exactly such a snapshot, and it + # is the newest, so it wins the ref. Point the load at the payload. + if default_snapshot is None or ( + payload_snapshots and str(default_snapshot) not in payload_snapshots + ): + load_id = str(snapshot_path) return { "inventory_id": _local_inventory_id("cache", model_format, repo_id), "load_id": load_id, @@ -344,9 +372,15 @@ def _scan_cached_gguf() -> list[dict]: continue if total_size == 0 and not has_variant_state: continue + # Walks each snapshot's quants, so it runs only for repos that + # made it past the skips above. Computed before the partial + # walk, which has to be judged against the snapshot this row + # hands out rather than the newest one. + gguf_snapshot, gguf_payload_snapshots = _repo_gguf_payload_snapshots(repo_info) partial = hf_cache_scan.is_gguf_repo_partial( repo_id, repo_path, + snapshot_dir = gguf_snapshot, ) if total_size == 0 and not partial: continue @@ -370,10 +404,11 @@ def _scan_cached_gguf() -> list[dict]: repo_id, "gguf", repo_path = repo_path, - snapshot_path = snapshot_path, + snapshot_path = gguf_snapshot or snapshot_path, active_hub_cache = active_hub_cache, partial = bool(row["partial"]), requires_variant = True, + payload_snapshots = gguf_payload_snapshots, ) ) if _repo_has_mmproj(repo_info): @@ -418,6 +453,92 @@ class _CachedNonGgufPayload(NamedTuple): has_runnable_weights: bool model_format: ModelFormat last_modified: float + payload_snapshot: Optional[Path] + payload_snapshots: frozenset[str] + + +# Keys mirror _classify_non_gguf_model_format's keyword arguments so a revision's +# flags can be classified on their own, exactly as the whole repo's are. +_PAYLOAD_FLAGS = ( + "has_config", + "has_adapter_config", + "has_adapter_weights", + "has_safetensors", + "has_transformers_safetensors", + "has_checkpoint_weights", +) + + +def _newest_snapshot_dir(candidates) -> Optional[Path]: + """Newest of *candidates* by directory mtime, or None when there are none. + + Ordering and resolution match ``_cached_model_snapshot_path`` and + ``hub.utils.gguf.iter_hf_cache_snapshots`` so every consumer names the same + directory by the same string. + """ + best: Optional[tuple[float, Path]] = None + for candidate in candidates: + path = Path(candidate) + try: + mtime = path.stat().st_mtime + except OSError: + mtime = 0.0 + if best is None or mtime > best[0]: + best = (mtime, path) + if best is None: + return None + try: + return best[1].resolve() + except OSError: + return best[1] + + +def _resolved_snapshot_ids(candidates) -> frozenset[str]: + """The same strings ``_newest_snapshot_dir`` would return, for membership.""" + resolved: set[str] = set() + for candidate in candidates: + path = Path(candidate) + try: + resolved.add(str(path.resolve())) + except OSError: + resolved.add(str(path)) + return frozenset(resolved) + + +def _repo_gguf_payload_snapshots(repo_info) -> tuple[Optional[Path], frozenset[str]]: + """Snapshot dirs a GGUF load can actually use, plus the newest of them. + + The row's size sums quants over every revision, while local variant + resolution reads only the directory handed out as ``load_id``, so the two + must agree or an advertised quant resolves to nothing. A snapshot holding + only part of a split quant is not usable either: the picker still offers + that quant and the generated command asks for shards that are absent, so + prefer one holding a whole quant exactly as ``_repo_gguf_load_id`` does. A + snapshot that mixes a whole quant with an interrupted split one still counts, + because the lister trims its offer to the completed subset; demanding the + whole directory be complete would hide that finished quant behind an older + revision's larger one. Fall back to any primary GGUF when nothing is + complete, which is what shipped before. + """ + # Matched on the snapshot-relative path, not ``file_name``: huggingface_hub + # sets that to the bare name for a nested file (the recovered mirror copies + # it), and the ``MTP/`` drafters unsloth ships are only recognisable as + # companions from their directory. Matching the bare name lets a snapshot + # holding nothing but a drafter win the load id, where the variant lister + # -- which does relativise -- then offers no quant at all. + with_gguf = [ + snapshot + for revision in repo_info.revisions + if (snapshot := getattr(revision, "snapshot_path", None)) is not None + and any(_is_main_gguf_filename(_cached_repo_file_name(f)) for f in revision.files) + ] + complete = [ + snapshot + for snapshot in with_gguf + if hf_cache_scan.snapshot_has_complete_variants(str(snapshot)) + ] + usable = complete or with_gguf + return _newest_snapshot_dir(usable), _resolved_snapshot_ids(usable) def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: @@ -425,12 +546,8 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: adapter_blobs: dict[str, tuple[int, float]] = {} safetensors_blobs: dict[str, tuple[int, float]] = {} checkpoint_blobs: dict[str, tuple[int, float]] = {} - has_config = False - has_adapter_config = False - has_adapter_weights = False - has_safetensors = False - has_transformers_safetensors = False - has_checkpoint = False + repo_flags = dict.fromkeys(_PAYLOAD_FLAGS, False) + revision_flags: list[tuple[Path, dict[str, bool]]] = [] def _record_blob( target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str @@ -444,6 +561,7 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: for revision in repo_info.revisions: rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) + flags = dict.fromkeys(_PAYLOAD_FLAGS, False) for f in revision.files: file_name = str(f.file_name) lower = file_name.lower() @@ -451,37 +569,34 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: if _is_gguf_filename(lower): continue if name == "config.json": - has_config = True + flags["has_config"] = True continue if name == "adapter_config.json": - has_adapter_config = True + flags["has_adapter_config"] = True continue is_adapter = _is_adapter_weight_name(name) is_safetensors = name.endswith(".safetensors") and not is_adapter is_checkpoint = _is_checkpoint_weight_name(name) if is_adapter: - has_adapter_weights = True + flags["has_adapter_weights"] = True _record_blob(adapter_blobs, f, rev_id, file_name) if is_safetensors: - has_safetensors = True + flags["has_safetensors"] = True if _is_transformers_safetensors_weight_name(name): - has_transformers_safetensors = True + flags["has_transformers_safetensors"] = True _record_blob(safetensors_blobs, f, rev_id, file_name) if is_checkpoint: - has_checkpoint = True + flags["has_checkpoint_weights"] = True _record_blob(checkpoint_blobs, f, rev_id, file_name) + snapshot = getattr(revision, "snapshot_path", None) + if snapshot is not None: + revision_flags.append((Path(snapshot), flags)) + for key, seen in flags.items(): + if seen: + repo_flags[key] = True model_format = ( - _classify_non_gguf_model_format( - has_config = has_config, - has_adapter_config = has_adapter_config, - has_adapter_weights = has_adapter_weights, - has_safetensors = has_safetensors, - has_transformers_safetensors = has_transformers_safetensors, - has_checkpoint_weights = has_checkpoint, - trusted_hf_cache_repo = True, - ) - or "unknown" + _classify_non_gguf_model_format(**repo_flags, trusted_hf_cache_repo = True) or "unknown" ) if model_format == "adapter": selected_blobs = adapter_blobs @@ -492,11 +607,27 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: else: selected_blobs = all_weight_blobs + # Weights are pooled across revisions, so the newest snapshot need not hold + # any: the load id has to name one that classifies the same way on its own, + # otherwise the load fails on a fully cached model. Untrusted here on + # purpose: the repo-level classification may rest on transformer-named + # weights alone, but a pinned load id names ONE directory and + # from_pretrained needs config.json inside it. A snapshot that only + # qualifies through that trust is not self-contained, so it must not become + # the load id; the row then keeps the repo id, which can still fill the + # config in from the hub. + payload_snapshots = [ + snapshot + for snapshot, flags in revision_flags + if _classify_non_gguf_model_format(**flags, trusted_hf_cache_repo = False) == model_format + ] return _CachedNonGgufPayload( size_bytes = sum(size for size, _mtime in selected_blobs.values()), has_runnable_weights = model_format != "unknown", model_format = model_format, last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0), + payload_snapshot = _newest_snapshot_dir(payload_snapshots), + payload_snapshots = _resolved_snapshot_ids(payload_snapshots), ) @@ -553,8 +684,13 @@ def _read_model_card_frontmatter(path: Path) -> dict: return {} -def _cached_model_local_metadata(repo_path: Path) -> dict: - snapshot = _cached_model_snapshot_path(repo_path) +def _cached_model_local_metadata(repo_path: Path, snapshot: Optional[Path] = None) -> dict: + # Describe the directory the row hands out, not merely the newest one: the + # metadata probe that strands the payload in an older snapshot carries + # neither the quantization config nor the model card. Fall back to the + # newest snapshot when no revision holds the payload on its own. + if snapshot is None: + snapshot = _cached_model_snapshot_path(repo_path) if snapshot is None: return {} @@ -624,14 +760,21 @@ def _scan_cached_models() -> list[dict]: continue key = repo_id.lower() existing = seen_lower.get(key) - local_metadata = _cached_model_local_metadata(repo_path) + local_metadata = _cached_model_local_metadata( + repo_path, + payload.payload_snapshot, + ) if local_metadata.pop("_hidden_stt", False): skipped_stt += 1 continue + # Scoped to the snapshot the row will advertise: an incomplete + # newer revision must not flip can_chat off for the complete + # one this row actually hands out as its load id. snapshot_partial = hf_cache_scan.is_snapshot_partial( "model", repo_id, repo_path, + snapshot_dir = payload.payload_snapshot, ) row = { "repo_id": repo_id, @@ -660,9 +803,10 @@ def _scan_cached_models() -> list[dict]: repo_id, payload.model_format, repo_path = repo_path, - snapshot_path = snapshot_path, + snapshot_path = payload.payload_snapshot or snapshot_path, active_hub_cache = active_hub_cache, partial = bool(row["partial"]), + payload_snapshots = payload.payload_snapshots, ) ) if _prefer_cache_row(row, existing): diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index fa5862a13c..e0e955be23 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -702,7 +702,7 @@ def test_cached_gguf_scan_dedupes_and_excludes_mmproj_only(monkeypatch, tmp_path monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: False, + lambda _repo_id, _path, **_kw: False, ) result = {"cached": cache_inventory._scan_cached_gguf()} @@ -723,7 +723,7 @@ def test_cached_gguf_scan_preserves_partial_flag(monkeypatch, tmp_path): monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: True, + lambda _repo_id, _path, **_kw: True, ) result = {"cached": cache_inventory._scan_cached_gguf()} @@ -766,7 +766,7 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: True, + lambda _repo_id, _path, **_kw: True, ) result = {"cached": cache_inventory._scan_cached_gguf()} @@ -799,7 +799,7 @@ def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: False, + lambda _repo_id, _path, **_kw: False, ) result = {"cached": cache_inventory._scan_cached_gguf()} @@ -834,7 +834,7 @@ def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypa monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: False, + lambda _repo_id, _path, **_kw: False, ) result = {"cached": cache_inventory._scan_cached_gguf()} @@ -866,7 +866,7 @@ def test_cached_models_scan_hides_non_gguf_embedder(monkeypatch, tmp_path): monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_snapshot_partial", - lambda _kind, _repo_id, _path: False, + lambda _kind, _repo_id, _path, **_kw: False, ) result = {"cached": cache_inventory._scan_cached_models()} @@ -909,12 +909,12 @@ def test_cached_scans_hide_embedders_configured_by_cache_path(monkeypatch, tmp_p monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: False, + lambda _repo_id, _path, **_kw: False, ) monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_snapshot_partial", - lambda _kind, _repo_id, _path: False, + lambda _kind, _repo_id, _path, **_kw: False, ) assert cache_inventory._scan_cached_gguf() == [] @@ -972,12 +972,12 @@ def test_cached_scans_hide_embedders_configured_by_snapshot_path(monkeypatch, tm monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: False, + lambda _repo_id, _path, **_kw: False, ) monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_snapshot_partial", - lambda _kind, _repo_id, _path: False, + lambda _kind, _repo_id, _path, **_kw: False, ) assert cache_inventory._scan_cached_gguf() == [] @@ -1015,7 +1015,7 @@ def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder( monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_snapshot_partial", - lambda _kind, _repo_id, _path: False, + lambda _kind, _repo_id, _path, **_kw: False, ) result = {"cached": cache_inventory._scan_cached_models()} @@ -1049,12 +1049,12 @@ def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypat monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_gguf_repo_partial", - lambda _repo_id, _path: False, + lambda _repo_id, _path, **_kw: False, ) monkeypatch.setattr( cache_inventory.hf_cache_scan, "is_snapshot_partial", - lambda _kind, _repo_id, _path: False, + lambda _kind, _repo_id, _path, **_kw: False, ) assert cache_inventory._scan_cached_gguf() == [] diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py index eb768db5d6..121e7816a7 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -316,16 +316,43 @@ def list_empty_gguf_variant_dirs(repo_id: str, root: Optional[Path] = None) -> s def list_gguf_variants_from_hf_cache( repo_id: str, root: Optional[Path] = None ) -> Optional[tuple[list[GgufVariantInfo], bool]]: + # Imported here, not at module scope: inventory_scan imports this module. + from hub.utils.inventory_scan import complete_snapshot_variants + snapshots = ( iter_hf_cache_snapshots(repo_id, root = root) if root is not None else iter_hf_cache_snapshots(repo_id) ) + # The inventory row hands out one snapshot as its load id and a local load + # reads only that one directory, so this walk has to land on the same one. + # A plain newest-first walk reports a half-downloaded split quant from a + # newer snapshot as downloaded while /load points elsewhere, so the newest + # snapshot holding at least one whole quant wins, exactly as + # _repo_gguf_payload_snapshots does, and only that snapshot's completed + # subset is offered. Requiring the whole directory to be complete instead + # would skip a newer snapshot that mixes a finished quant with an + # interrupted one and hide the finished quant behind an older revision's + # larger one, which auto-load may not have the memory for. + # + # The vision flag is OR-ed in because a skipped newer snapshot can be a + # projector fetched on its own. When no snapshot has a whole quant the first + # one with anything wins, as it did before. + any_vision = False + fallback: Optional[tuple[list[GgufVariantInfo], bool]] = None for snapshot in snapshots: variants, has_vision = list_local_gguf_variants(str(snapshot)) - if variants or has_vision: - return variants, has_vision - return None + any_vision = any_vision or has_vision + if variants: + complete = complete_snapshot_variants(str(snapshot)) + # A quant with no label cannot be judged, so it is kept rather than + # dropped; with nothing complete the list stays as it was. + usable = [v for v in variants if not v.quant or v.quant in complete] + if usable: + return usable, any_vision + if fallback is None and (variants or has_vision): + fallback = (variants, has_vision) + return fallback def list_partial_gguf_variants_from_state( diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py index 058fdf9b65..250d95d1f8 100644 --- a/studio/backend/hub/utils/inventory_scan.py +++ b/studio/backend/hub/utils/inventory_scan.py @@ -17,7 +17,7 @@ import hashlib import re import threading import time -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Callable, Optional @@ -125,18 +125,301 @@ def all_hf_cache_scans() -> list: flight.event.set() +# huggingface_hub skips these when walking refs/ and snapshots/; mirrored so a +# stray OS helper file is not mistaken for cache corruption. +_CACHE_ENTRIES_TO_IGNORE = frozenset({".DS_Store"}) +_HF_REPO_TYPES = frozenset({"model", "dataset", "space"}) + + +# Recovered entries deliberately mirror huggingface_hub's CachedFileInfo / +# CachedRevisionInfo / CachedRepoInfo field-for-field, but are constructed here +# rather than imported so a field added or removed upstream cannot break the +# call. test_hf_cache_dangling_refs asserts the surfaces stay in step. Frozen +# (so hashable) because HFCacheInfo.delete_revisions() keys a dict by repo and +# takes a set difference over ``revisions``. +@dataclass(frozen = True) +class _RecoveredFileInfo: + file_name: str + file_path: Path + size_on_disk: int + blob_path: Path + blob_last_accessed: float + blob_last_modified: float + + +@dataclass(frozen = True) +class _RecoveredRevisionInfo: + commit_hash: str + snapshot_path: Path + size_on_disk: int + files: frozenset + refs: frozenset + last_modified: float + + +@dataclass(frozen = True) +class _RecoveredRepoInfo: + repo_id: str + repo_type: str + repo_path: Path + size_on_disk: int + nb_files: int + revisions: frozenset + last_accessed: float + last_modified: float + + @property + def refs(self) -> dict: + return {ref: rev for rev in self.revisions for ref in rev.refs} + + +def _hf_repo_identity(repo_dir_name: str) -> Optional[tuple[str, str]]: + """``models--Org--Model`` -> ``("model", "Org/Model")``, as huggingface_hub parses it.""" + if "--" not in repo_dir_name: + return None + repo_type, _, repo_id = repo_dir_name.partition("--") + repo_type = repo_type[:-1] + if repo_type not in _HF_REPO_TYPES or not repo_id: + return None + return repo_type, repo_id.replace("--", "/") + + +def _read_refs_by_commit(refs_dir: Path) -> Optional[dict[str, set[str]]]: + """Map commit hash -> ref names under ``refs/``. None if unreadable.""" + refs_by_commit: dict[str, set[str]] = {} + if not refs_dir.exists(): + return refs_by_commit + if refs_dir.is_file(): + return None + try: + entries = sorted(refs_dir.rglob("*")) + except OSError: + return None + for ref_path in entries: + try: + if ref_path.is_dir() or ref_path.name in _CACHE_ENTRIES_TO_IGNORE: + continue + commit = ref_path.read_text(encoding = "utf-8") + except (OSError, UnicodeDecodeError): + return None + # PurePath keeps the separator platform-native; huggingface_hub stores + # ref names the same way, so no manual posix normalisation here. + refs_by_commit.setdefault(commit, set()).add(str(ref_path.relative_to(refs_dir))) + return refs_by_commit + + +def _recover_repo_hidden_by_dangling_refs(repo_dir: Path) -> Optional[_RecoveredRepoInfo]: + """Rebuild the scan entry for a repo dropped *solely* over leftover refs. + + ``_scan_cached_repo`` assembles every revision successfully and only then + raises ``CorruptedCacheException`` because a ``refs/`` file names a + commit with no ``snapshots//`` dir, so ``scan_cache_dir`` omits an + entirely intact repo from ``.repos``. Studio creates that state itself: + ``snapshot_download`` writes ``refs/main`` at the live upstream sha *before* + fetching the first file and never creates ``snapshots//`` itself, and no + Studio caller pins ``revision``. So any repo re-uploaded since it was + downloaded goes invisible to every inventory endpoint the moment a refresh + starts, while the model picker's plain directory walk still lists it. No race + is needed: when the allow/ignore patterns match nothing the download returns + normally having written only the ref. + + This reads the same directories huggingface_hub reads and writes nothing: + the ref file that upstream's assertion trips over is left exactly as it is. + Returns None whenever anything *other* than leftover refs would have failed + the upstream scan, so a genuinely corrupt repo stays omitted as before. + """ + identity = _hf_repo_identity(repo_dir.name) + if identity is None: + return None + repo_type, repo_id = identity + snapshots_dir = repo_dir / "snapshots" + refs_by_commit = _read_refs_by_commit(repo_dir / "refs") + if refs_by_commit is None: + return None + try: + if not snapshots_dir.is_dir(): + return None + snapshot_entries = sorted(snapshots_dir.iterdir()) + except OSError: + return None + + blob_stats: dict[Path, object] = {} + revisions: set[_RecoveredRevisionInfo] = set() + dangling = dict(refs_by_commit) + for snapshot in snapshot_entries: + if snapshot.name in _CACHE_ENTRIES_TO_IGNORE: + continue + try: + if not snapshot.is_dir(): + # Upstream treats a file here as corruption; defer to it. + return None + entries = sorted(snapshot.rglob("*")) + except OSError: + return None + files: set[_RecoveredFileInfo] = set() + for entry in entries: + try: + if entry.is_dir(): + continue + blob_path = entry.resolve() + stat = blob_stats.get(blob_path) or blob_path.stat() + except OSError: + # Broken symlink / unreadable blob: upstream raises here too. + return None + blob_stats[blob_path] = stat + files.add( + _RecoveredFileInfo( + file_name = entry.name, + file_path = entry, + size_on_disk = stat.st_size, + blob_path = blob_path, + blob_last_accessed = stat.st_atime, + blob_last_modified = stat.st_mtime, + ) + ) + try: + last_modified = ( + max(f.blob_last_modified for f in files) if files else snapshot.stat().st_mtime + ) + except OSError: + return None + revisions.add( + _RecoveredRevisionInfo( + commit_hash = snapshot.name, + snapshot_path = snapshot, + size_on_disk = sum(blob_stats[blob].st_size for blob in {f.blob_path for f in files}), + files = frozenset(files), + refs = frozenset(dangling.pop(snapshot.name, set())), + last_modified = last_modified, + ) + ) + # A download writes refs/ before fetching its first file, so a + # repo with no snapshot yet is not downloaded and must not be reported as + # such -- that is the "already have it" lie this whole fix is about. + if not revisions: + return None + # Every ref resolved, so upstream did not drop this repo over leftover refs + # and either already returned it or failed for a reason we must not paper over. + if not dangling: + return None + try: + repo_stats = repo_dir.stat() + except OSError: + return None + return _RecoveredRepoInfo( + repo_id = repo_id, + repo_type = repo_type, + repo_path = repo_dir, + size_on_disk = sum(stat.st_size for stat in blob_stats.values()), + nb_files = len(blob_stats), + revisions = frozenset(revisions), + last_accessed = ( + max((stat.st_atime for stat in blob_stats.values()), default = repo_stats.st_atime) + ), + last_modified = ( + max((stat.st_mtime for stat in blob_stats.values()), default = repo_stats.st_mtime) + ), + ) + + +def _with_repos_hidden_by_dangling_refs(scan, cache_root: Path): + """Add back the repos ``scan_cache_dir`` dropped over a dangling ref.""" + try: + repo_dirs = sorted(entry for entry in cache_root.iterdir() if "--" in entry.name) + except OSError: + return scan + known = getattr(scan, "repos", ()) + scanned: set[str] = set() + for repo in known: + try: + scanned.add(str(Path(repo.repo_path).resolve(strict = False))) + except (AttributeError, OSError, RuntimeError, TypeError, ValueError): + continue + recovered: list[_RecoveredRepoInfo] = [] + for repo_dir in repo_dirs: + try: + if str(repo_dir.resolve(strict = False)) in scanned: + continue + entry = _recover_repo_hidden_by_dangling_refs(repo_dir) + except (OSError, RuntimeError, ValueError): + continue + if entry is None: + continue + logger.info( + "Recovered HF cache repo %s hidden by a dangling ref (%d revision(s) on disk)", + entry.repo_id, + len(entry.revisions), + ) + recovered.append(entry) + if not recovered: + return scan + try: + return replace( + scan, + repos = frozenset(known) | frozenset(recovered), + size_on_disk = getattr(scan, "size_on_disk", 0) + + sum(entry.size_on_disk for entry in recovered), + ) + except (AttributeError, TypeError, ValueError) as exc: + # A scan shape we cannot rebuild is left untouched rather than dropped. + logger.debug("Could not attach recovered HF cache repos: %s", exc) + return scan + + def _compute_all_hf_cache_scans() -> list: from huggingface_hub import scan_cache_dir scans: list = [] for cache_root in hf_cache_roots(): try: - scans.append(scan_cache_dir(cache_dir = str(cache_root))) + scan = scan_cache_dir(cache_dir = str(cache_root)) + # Only a warned-about scan can be hiding a repo, so a healthy cache + # is never walked twice. getattr: a scan object without .warnings + # must not take the whole cache root down with it. + if getattr(scan, "warnings", None): + scan = _with_repos_hidden_by_dangling_refs(scan, cache_root) + scans.append(scan) except Exception as exc: logger.warning("Could not scan HF cache %s: %s", cache_root, exc) return scans +def default_ref_snapshot(repo_dir: Path) -> Optional[Path]: + """Snapshot dir that ``refs/main`` names in *repo_dir*, or ``None``. + + This is the directory ``from_pretrained(repo_id)`` ends up in, so callers + compare it against the snapshot the inventory row advertises: naming the + repo id is only safe when the two hold the same payload. + """ + ref_path = repo_dir / "refs" / "main" + try: + # No strip: huggingface_hub matches the raw ref contents against the + # snapshot dir name, so a ref with stray whitespace resolves nowhere. + commit = ref_path.read_text(encoding = "utf-8") + except (OSError, UnicodeDecodeError): + return None + if not commit: + return None + snapshot = repo_dir / "snapshots" / commit + try: + if not snapshot.is_dir(): + return None + return snapshot.resolve() + except (OSError, ValueError): + return None + + +def default_ref_resolves_on_disk(repo_dir: Path) -> bool: + """Whether ``refs/main`` names a snapshot that exists in *repo_dir*. + + When it does not, ``from_pretrained(repo_id)`` has nothing to resolve: an + offline load fails outright and an online one silently fetches the current + upstream HEAD instead of the snapshot already on disk. Callers use this to + hand out the snapshot path as the load identity instead of the repo id. + """ + return default_ref_snapshot(repo_dir) is not None + + def token_fingerprint(hf_token: Optional[str]) -> str: """16-char SHA256 prefix used as a cache-key qualifier for gated repos. @@ -236,19 +519,21 @@ def _repo_cache_dir_incomplete_hashes(repo_cache_dir: Path) -> set[str]: return hashes -def _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir: Path) -> bool: - latest = latest_snapshot_dir(repo_cache_dir) - if latest is None: +def _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks( + repo_cache_dir: Path, snapshot_dir: Optional[Path] = None +) -> bool: + target = snapshot_dir if snapshot_dir is not None else latest_snapshot_dir(repo_cache_dir) + if target is None: return False try: - entries = list(latest.rglob("*")) + entries = list(target.rglob("*")) except OSError: return False for entry in entries: try: if not entry.is_symlink() or entry.exists(): continue - rel = entry.relative_to(latest).as_posix() + rel = entry.relative_to(target).as_posix() if is_gguf_filename(rel): continue return True @@ -257,6 +542,65 @@ def _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir: Path) return False +def _is_latest_snapshot(repo_cache_dir: Path, snapshot_dir: Path) -> bool: + latest = latest_snapshot_dir(repo_cache_dir) + if latest is None: + return False + try: + return latest.resolve() == snapshot_dir.resolve() + except OSError: + return latest == snapshot_dir + + +def _default_ref_names_an_absent_snapshot(repo_cache_dir: Path) -> bool: + """Whether ``refs/main`` is present and names a commit with no snapshot dir. + + ``snapshot_download`` rewrites ``refs/`` with the resolved commit + *before* it fetches a single file, and the snapshot directory is only + created once the first file lands, so this state is the window between the + two. A missing ``refs/main`` is not the same thing: a repo fetched by + commit hash never gets one, so it carries no evidence either way. + """ + ref_path = repo_cache_dir / "refs" / "main" + try: + commit = ref_path.read_text(encoding = "utf-8") + except (OSError, UnicodeDecodeError): + return False + if not commit: + return False + try: + return not (repo_cache_dir / "snapshots" / commit).is_dir() + except (OSError, ValueError): + return False + + +def _repo_signal_applies_to_snapshot( + repo_cache_dir: Optional[Path], snapshot_dir: Optional[Path] +) -> bool: + """Whether a repo-wide partial signal describes *snapshot_dir*. + + A cancel marker is cleared at every download start and again on success, so + one that is present records the most recent attempt; an ``.incomplete`` blob + likewise belongs to the revision a download is writing. Both attach to the + newest snapshot, so a row advertising an older, already complete one must + not inherit them and lose ``can_chat``. With nothing to attribute against, + the signal is kept rather than dropped. + + A ``refs/main`` naming a commit with no directory pins that attempt to a + revision that is not on disk at all, so no snapshot here may inherit it. + The downloader never passes a revision, so every attempt it starts rewrites + that ref first; leaving the signal on the newest snapshot instead charged an + interrupted update to the previous, complete payload and hid a model that + still loads. This is the very state the dangling-ref recovery restores rows + from, so the recovered row would arrive unusable. + """ + if repo_cache_dir is None or snapshot_dir is None: + return True + if _default_ref_names_an_absent_snapshot(repo_cache_dir): + return False + return _is_latest_snapshot(repo_cache_dir, snapshot_dir) + + def _gguf_variant_manifest_blob_hashes( repo_id: str, repo_cache_dir: Optional[Path] = None ) -> frozenset[str]: @@ -284,18 +628,30 @@ def _gguf_variant_manifest_blob_hashes( def _repo_cache_dir_has_snapshot_legacy_partial( - repo_cache_dir: Path, *, ignored_blob_hashes: frozenset[str] + repo_cache_dir: Path, + *, + ignored_blob_hashes: frozenset[str], + snapshot_dir: Optional[Path] = None, ) -> bool: - incomplete_hashes = _repo_cache_dir_incomplete_hashes(repo_cache_dir) - if any(blob_hash not in ignored_blob_hashes for blob_hash in incomplete_hashes): + if _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir, snapshot_dir): return True - return _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir) + # ``.incomplete`` blobs sit in ``blobs/`` with no revision of their own, so + # they can only be charged to the revision a download is currently writing, + # which is the newest one. A row that advertises an older, already complete + # snapshot must not go partial (and lose ``can_chat``) over a separate fetch. + if snapshot_dir is not None and not _repo_signal_applies_to_snapshot( + repo_cache_dir, snapshot_dir + ): + return False + incomplete_hashes = _repo_cache_dir_incomplete_hashes(repo_cache_dir) + return any(blob_hash not in ignored_blob_hashes for blob_hash in incomplete_hashes) def _snapshot_legacy_partial( repo_type: str, repo_id: str, repo_cache_dir: Optional[Path] = None, + snapshot_dir: Optional[Path] = None, ) -> bool: if repo_type != "model": return _legacy_partial(repo_type, repo_id, repo_cache_dir) @@ -304,7 +660,10 @@ def _snapshot_legacy_partial( return _repo_cache_dir_has_snapshot_legacy_partial( repo_cache_dir, ignored_blob_hashes = ignored_hashes, + snapshot_dir = snapshot_dir, ) + # Without a repo dir the snapshot cannot be attributed to one of the roots + # below, so the repo-wide signal is kept rather than applied to the wrong dir. return any( _repo_cache_dir_has_snapshot_legacy_partial( entry, @@ -350,6 +709,62 @@ def _completed_gguf_variants(snapshot_dir: Optional[Path]) -> set[str]: return complete +def snapshot_variants_all_complete(snapshot: str) -> bool: + """True when every quant the variant lister would advertise from *snapshot* is + fully on disk. + + One complete quant is not enough: the picker enumerates the whole directory, so a + half-downloaded split quant sitting beside a good one still gets offered and the + generated command asks llama-server for shards that are absent. Both sides derive + their labels from ``extract_quant_label`` over paths relative to the snapshot, so + the sets are directly comparable. + """ + from hub.utils.gguf import list_local_gguf_variants + try: + variants, _ = list_local_gguf_variants(snapshot) + offered = {v.quant for v in variants if getattr(v, "quant", None)} + if not offered: + return False + return offered <= _completed_gguf_variants(Path(snapshot)) + except Exception: + return False + + +def snapshot_has_complete_variants(snapshot: str) -> bool: + """True when at least one quant the variant lister would advertise from + *snapshot* is fully on disk. + + Deliberately weaker than ``snapshot_variants_all_complete``: a snapshot that + mixes a whole quant with an interrupted split one is still loadable for the + whole quant, and the lister trims the offer to that completed subset. Skipping + such a snapshot outright hides a fully downloaded quant behind an older + revision's larger one, which auto-load may not have the memory for. + + Snapshot selection and the offered variants have to agree on one directory, so + every caller that pins a load id uses this predicate and the lister uses the + matching subset. + """ + from hub.utils.gguf import list_local_gguf_variants + try: + variants, _ = list_local_gguf_variants(snapshot) + offered = {v.quant for v in variants if getattr(v, "quant", None)} + if not offered: + return False + return bool(offered & _completed_gguf_variants(Path(snapshot))) + except Exception: + return False + + +def complete_snapshot_variants(snapshot: str) -> set[str]: + """Quant labels in *snapshot* whose files are all on disk. Same labels as + ``snapshot_variants_all_complete`` compares, for callers that need the subset + rather than the all-or-nothing answer.""" + try: + return _completed_gguf_variants(Path(snapshot)) + except (OSError, RuntimeError, ValueError): + return set() + + def _manifest_partial( repo_type: RepoType, repo_id: str, @@ -427,31 +842,49 @@ def is_snapshot_partial( repo_type: RepoType, repo_id: str, repo_cache_dir: Optional[Path] = None, + snapshot_dir: Optional[Path] = None, ) -> bool: """Repo-row partial flag for snapshot-style downloads (full-snapshot models — safetensors/adapter/checkpoint — and all datasets). Composes three signals, cheapest first: - 1. Cancel marker (single stat). + 1. Cancel marker (single stat), charged to the newest snapshot. 2. Snapshot-attributed legacy .incomplete blob / broken-symlink check. 3. Manifest walk (stat per expected file under the latest snapshot). A manifest without a resolvable snapshot is partial: the worker got - far enough to record expectations but did not leave a usable snapshot.""" + far enough to record expectations but did not leave a usable snapshot. + + *snapshot_dir* pins both the legacy and the manifest walk to the snapshot the + row will hand out as its load identity. Without it they use the newest + snapshot, so a weightless metadata-only revision beside a complete download + flags the row partial and ``can_chat`` goes false for a model that loads + fine. + + The repo-wide manifest carries no revision either, so it gets the same + attribution as the marker and the ``.incomplete`` blobs: it describes the + revision the last attempt was writing, which is the newest. Verifying it + against an older pinned snapshot compares one revision's file list with + another's payload, and any rename or size change between the two flagged a + complete, loadable row partial.""" from hub.utils import download_manifest + + repo_signal_applies = _repo_signal_applies_to_snapshot(repo_cache_dir, snapshot_dir) return _compose_partial( - lambda: download_manifest.has_cancel_marker( + lambda: repo_signal_applies + and download_manifest.has_cancel_marker( repo_type, repo_id, None, hub_cache = _hub_cache_for_repo_dir(repo_cache_dir), ), - lambda: _snapshot_legacy_partial(repo_type, repo_id, repo_cache_dir), - lambda: _manifest_partial( + lambda: _snapshot_legacy_partial(repo_type, repo_id, repo_cache_dir, snapshot_dir), + lambda: repo_signal_applies + and _manifest_partial( repo_type, repo_id, None, - None, + snapshot_dir, repo_cache_dir, ), ) @@ -465,6 +898,7 @@ def is_variant_partial( incomplete_blob_hashes: Optional[set[str]] = None, variant_blob_hashes: Optional[frozenset[str]] = None, repo_cache_dir: Optional[Path] = None, + repo_signal_applies: bool = True, ) -> bool: """Per-variant partial detection. Owns its manifest, owns its marker. Used by the GGUF variants endpoint to flag a specific quant as broken @@ -472,10 +906,19 @@ def is_variant_partial( snapshot_dir is an optional hint to avoid re-walking the cache when a caller is checking many variants of the same repo (see - is_gguf_repo_partial for that usage).""" + is_gguf_repo_partial for that usage). + + ``repo_signal_applies`` is the same attribution the repo-wide signals get. + The marker and the manifest are both keyed by (repo, variant) with no + revision and are both overwritten by the next attempt, so a caller that + pinned *snapshot_dir* to an older revision than that attempt was writing + passes False rather than judge the quant it verifies there by another + revision's file list. Defaults True so the per-variant endpoint keeps + reporting a cancelled or unfinished quant as broken.""" from hub.utils import download_manifest return _compose_partial( - lambda: download_manifest.has_cancel_marker( + lambda: repo_signal_applies + and download_manifest.has_cancel_marker( "model", repo_id, variant, @@ -486,7 +929,8 @@ def is_variant_partial( and variant_blob_hashes and incomplete_blob_hashes.intersection(variant_blob_hashes) ), - lambda: _manifest_partial( + lambda: repo_signal_applies + and _manifest_partial( "model", repo_id, variant, @@ -496,7 +940,12 @@ def is_variant_partial( ) -def is_gguf_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> bool: +def is_gguf_repo_partial( + repo_id: str, + repo_cache_dir: Optional[Path] = None, + *, + snapshot_dir: Optional[Path] = None, +) -> bool: """Repo-row partial flag for a GGUF repo. The inventory shows ONE row per GGUF repo (requires_variant=True); per-variant detail lives in GET /api/models/gguf-variants and uses is_variant_partial. @@ -515,16 +964,28 @@ def is_gguf_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> Composes signals: 1. Cheap legacy fast-path (.incomplete blobs / broken symlinks). 2. Per-variant manifest + marker enumeration, gated on "all broken". + + *snapshot_dir* pins all three to the snapshot the row hands out as its load + id. Without it the newest snapshot supplies the completed quants while the + legacy walk and the per-variant cancel markers stay repo-wide, so an + interrupted re-download flips can_chat off for the older complete quant the + row actually advertises. """ from hub.utils import download_manifest - has_legacy_partial = _legacy_partial("model", repo_id, repo_cache_dir) - snapshot_dir = resolve_snapshot_dir_for_scan( - "model", - repo_id, - repo_cache_dir, - ) - variants: set[str] = set(_completed_gguf_variants(snapshot_dir)) + if snapshot_dir is None: + snapshot_dir = resolve_snapshot_dir_for_scan( + "model", + repo_id, + repo_cache_dir, + ) + # Same attribution as is_snapshot_partial: an .incomplete blob or a broken + # symlink belongs to the revision a download is writing, which is the newest. + # Variant cancel markers carry no revision either, so they get it too. + repo_signal_applies = _repo_signal_applies_to_snapshot(repo_cache_dir, snapshot_dir) + has_legacy_partial = repo_signal_applies and _legacy_partial("model", repo_id, repo_cache_dir) + complete_here = _completed_gguf_variants(snapshot_dir) + variants: set[str] = set(complete_here) hub_cache = _hub_cache_for_repo_dir(repo_cache_dir) for variant, _path in download_manifest.iter_variant_manifests( "model", @@ -563,6 +1024,10 @@ def is_gguf_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> variant, snapshot_dir, repo_cache_dir = repo_cache_dir, + # A quant whose files are all in the pinned snapshot is loadable + # from it whatever a newer attempt's marker or manifest says, and + # that attempt's manifest lists another revision's files. + repo_signal_applies = repo_signal_applies or variant not in complete_here, ): has_broken = True else: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 6e587c18e8..8363500e0b 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3052,26 +3052,11 @@ def _repo_gguf_last_modified(repo_info) -> float: def snapshot_variants_all_complete(snapshot: str) -> bool: - """True when every quant the variant lister would advertise from *snapshot* is - fully on disk. - - One complete quant is not enough: the picker enumerates the whole directory, so a - half-downloaded split quant sitting beside a good one still gets offered and the - generated command asks llama-server for shards that are absent. Both sides derive - their labels from ``extract_quant_label`` over paths relative to the snapshot, so - the sets are directly comparable. - """ + """Re-exported for callers that already import it from here; the scan-side + cache inventory needs the same predicate, so it lives beside the completed + variant walk it is built on.""" from hub.utils import inventory_scan - from hub.utils.gguf import list_local_gguf_variants - - try: - variants, _ = list_local_gguf_variants(snapshot) - offered = {v.quant for v in variants if getattr(v, "quant", None)} - if not offered: - return False - return offered <= inventory_scan._completed_gguf_variants(Path(snapshot)) - except Exception: - return False + return inventory_scan.snapshot_variants_all_complete(snapshot) def _repo_gguf_load_id(repo_info, active_root: Optional[Path]) -> Optional[str]: diff --git a/studio/backend/tests/test_hf_cache_dangling_refs.py b/studio/backend/tests/test_hf_cache_dangling_refs.py new file mode 100644 index 0000000000..d311b00d4a --- /dev/null +++ b/studio/backend/tests/test_hf_cache_dangling_refs.py @@ -0,0 +1,1080 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""A dangling ``refs/`` must not hide an intact repo from the scan. + +``scan_cache_dir`` raises CorruptedCacheException for a repo whose ref names a +commit with no ``snapshots//`` directory and omits it from ``.repos``, +so the model stays visible in the model picker (a plain directory walk) while +disappearing from every Hub inventory endpoint that feeds chat auto-load. + +The repair is read-only: the hidden repo is rebuilt from the same directories +huggingface_hub reads and the ref file is left exactly as it is, because +``_cache_commit_hash_for_specific_revision`` writes refs with an unlocked +in-place ``write_text``, so no external process can delete one race-free. +""" + +from __future__ import annotations + +import dataclasses +import os +import stat +from pathlib import Path + +import pytest + +from hub.utils import inventory_scan + +SNAPSHOT = "a" * 40 +UPSTREAM_HEAD = "b" * 40 + + +def _build_repo( + cache_root: Path, + *, + ref: str | None = UPSTREAM_HEAD, + extra_refs: dict[str, str] | None = None, + name: str = "models--Org--Model", + payload: bytes = b"\0" * 11, + snapshots: tuple[str, ...] = (SNAPSHOT,), +) -> Path: + """A cache repo shaped like HF's, using regular files rather than symlinks. + + ``_scan_cached_repo`` resolves each snapshot entry to its blob, and a + regular file resolves to itself, so this exercises the real scanner while + staying runnable on Windows without the symlink privilege. + """ + repo_dir = cache_root / name + refs = repo_dir / "refs" + refs.mkdir(parents = True, exist_ok = True) + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + for commit in snapshots: + snapshot = repo_dir / "snapshots" / commit + snapshot.mkdir(parents = True, exist_ok = True) + (snapshot / "model.safetensors").write_bytes(payload) + if ref is not None: + (refs / "main").write_text(ref, encoding = "utf-8") + for ref_name, commit in (extra_refs or {}).items(): + (refs / ref_name).write_text(commit, encoding = "utf-8") + return repo_dir + + +def _ref_names(repo_dir: Path) -> list[str]: + return sorted(entry.name for entry in (repo_dir / "refs").rglob("*") if entry.is_file()) + + +def _scan(cache_root: Path, monkeypatch) -> list: + monkeypatch.setattr(inventory_scan, "hf_cache_roots", lambda: [cache_root]) + return inventory_scan._compute_all_hf_cache_scans() + + +def _scanned_repo_ids(cache_root: Path, monkeypatch) -> list[str]: + return sorted(repo.repo_id for scan in _scan(cache_root, monkeypatch) for repo in scan.repos) + + +def _only_repo(cache_root: Path, monkeypatch): + repos = [repo for scan in _scan(cache_root, monkeypatch) for repo in scan.repos] + assert [repo.repo_id for repo in repos] == ["Org/Model"], "repo hidden from the scan" + return repos[0] + + +# --- the #7374 symptom ------------------------------------------------------- + + +def test_huggingface_hub_really_hides_a_repo_behind_a_dangling_ref(tmp_path): + """Baseline for the bug: the snapshot is intact, yet the repo is dropped.""" + from huggingface_hub import scan_cache_dir + + repo_dir = _build_repo(tmp_path) + + raw = scan_cache_dir(cache_dir = str(tmp_path)) + + assert [repo.repo_id for repo in raw.repos] == [] + assert raw.warnings + assert (repo_dir / "snapshots" / SNAPSHOT / "model.safetensors").is_file() + + +def test_dangling_ref_no_longer_hides_an_intact_repo(tmp_path, monkeypatch): + repo_dir = _build_repo(tmp_path) + + repo = _only_repo(tmp_path, monkeypatch) + + assert repo.repo_id == "Org/Model" + assert repo.repo_type == "model" + # routes/models.py does repo_path.parent, so this must stay a real Path. + assert isinstance(repo.repo_path, Path) and repo.repo_path == repo_dir + assert repo.size_on_disk == 11 + # The recovered revision keeps the identity the snapshot is loadable by. + revision = next(iter(repo.revisions)) + assert revision.commit_hash == SNAPSHOT + assert revision.snapshot_path == repo_dir / "snapshots" / SNAPSHOT + assert {f.file_name for f in revision.files} == {"model.safetensors"} + # The dangling ref resolves to nothing, so it maps to no revision... + assert revision.refs == frozenset() + # ...and, crucially, is still on disk: the repair never writes to the cache. + assert _ref_names(repo_dir) == ["main"] + assert (repo_dir / "refs" / "main").read_text(encoding = "utf-8") == UPSTREAM_HEAD + + +def test_recovery_reads_a_multi_file_multi_revision_repo(tmp_path, monkeypatch): + other = "c" * 40 + repo_dir = _build_repo(tmp_path, snapshots = (SNAPSHOT, other)) + (repo_dir / "snapshots" / SNAPSHOT / "nested").mkdir() + (repo_dir / "snapshots" / SNAPSHOT / "nested" / "extra.json").write_bytes(b"{}") + + repo = _only_repo(tmp_path, monkeypatch) + + assert {rev.commit_hash for rev in repo.revisions} == {SNAPSHOT, other} + files = {f.file_path for rev in repo.revisions for f in rev.files} + assert repo_dir / "snapshots" / SNAPSHOT / "nested" / "extra.json" in files + # _resolve_cached_model_path relativises file_path against snapshot_path. + for rev in repo.revisions: + for f in rev.files: + assert f.file_path.relative_to(rev.snapshot_path) + + +def test_a_still_resolvable_ref_keeps_its_revision_mapping(tmp_path, monkeypatch): + """One good ref plus one stale ref: hub drops the repo, we keep the mapping.""" + repo_dir = _build_repo(tmp_path, ref = SNAPSHOT, extra_refs = {"stale": UPSTREAM_HEAD}) + + repo = _only_repo(tmp_path, monkeypatch) + + assert next(iter(repo.revisions)).refs == frozenset({"main"}) + assert _ref_names(repo_dir) == ["main", "stale"] + + +# --- the repair must not widen past the leftover-refs assertion -------------- + + +def test_a_healthy_cache_is_returned_untouched(tmp_path, monkeypatch): + import huggingface_hub + + repo_dir = _build_repo(tmp_path, ref = SNAPSHOT, extra_refs = {"v1.0": SNAPSHOT}) + calls: list[str] = [] + real_scan = huggingface_hub.scan_cache_dir + + def counting_scan(cache_dir = None): + calls.append(str(cache_dir)) + return real_scan(cache_dir = cache_dir) + + monkeypatch.setattr(huggingface_hub, "scan_cache_dir", counting_scan) + + assert _scanned_repo_ids(tmp_path, monkeypatch) == ["Org/Model"] + assert len(calls) == 1 + assert _ref_names(repo_dir) == ["main", "v1.0"] + + +def test_a_download_that_has_only_written_its_ref_is_not_invented(tmp_path, monkeypatch): + """snapshot_download writes refs/ before fetching the first file. + + There is no snapshot to recover yet, so nothing must be reported as + downloaded -- that would be the very "already have it" lie #7374 is about. + """ + repo_dir = tmp_path / "models--Org--Model" + (repo_dir / "snapshots").mkdir(parents = True) + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text(UPSTREAM_HEAD, encoding = "utf-8") + + assert _scanned_repo_ids(tmp_path, monkeypatch) == [] + assert _ref_names(repo_dir) == ["main"] + + +def test_a_repo_corrupted_beyond_a_dangling_ref_stays_omitted(tmp_path, monkeypatch): + """A broken snapshot symlink is corruption hub rejects for its own reasons.""" + repo_dir = _build_repo(tmp_path) + broken = repo_dir / "snapshots" / SNAPSHOT / "weights.bin" + try: + os.symlink(repo_dir / "blobs" / "missing", broken) + except (NotImplementedError, OSError): + pytest.skip("symlinks unavailable (Windows without developer mode)") + + assert _scanned_repo_ids(tmp_path, monkeypatch) == [] + + +def test_an_unrelated_repo_is_never_disturbed(tmp_path, monkeypatch): + hidden = _build_repo(tmp_path) + healthy = _build_repo(tmp_path, ref = SNAPSHOT, name = "models--Org--Healthy") + + assert _scanned_repo_ids(tmp_path, monkeypatch) == ["Org/Healthy", "Org/Model"] + assert _ref_names(hidden) == ["main"] + assert _ref_names(healthy) == ["main"] + + +def test_an_unreadable_repo_does_not_abort_the_recovery(tmp_path): + """One unreadable repo must not stop the others being recovered. + + Scoped to the recovery pass because scan_cache_dir itself raises on an + unreadable repo dir, which is upstream of this code and unchanged here. + """ + from huggingface_hub import HFCacheInfo + + hidden = _build_repo(tmp_path) + locked = _build_repo(tmp_path, ref = SNAPSHOT, name = "models--Org--Locked") + locked.chmod(0) + if os.access(locked / "refs", os.R_OK): + pytest.skip("filesystem does not enforce directory permissions") + try: + scan = HFCacheInfo(size_on_disk = 0, repos = frozenset(), warnings = []) + merged = inventory_scan._with_repos_hidden_by_dangling_refs(scan, tmp_path) + assert sorted(repo.repo_id for repo in merged.repos) == ["Org/Model"] + assert _ref_names(hidden) == ["main"] + finally: + locked.chmod(stat.S_IRWXU) + + +def test_a_non_repo_directory_is_ignored(tmp_path, monkeypatch): + _build_repo(tmp_path) + (tmp_path / ".locks").mkdir() + (tmp_path / "notarepo").mkdir() + (tmp_path / "spaces--Org--Thing").mkdir() + + assert _scanned_repo_ids(tmp_path, monkeypatch) == ["Org/Model"] + + +def test_a_scan_object_without_warnings_still_reaches_the_caller(tmp_path, monkeypatch): + """The gate must read .warnings defensively: an AttributeError here lands in + the per-root ``except Exception`` and silently blanks the whole cache.""" + from types import SimpleNamespace + + import huggingface_hub + + monkeypatch.setattr( + huggingface_hub, + "scan_cache_dir", + lambda cache_dir = None: SimpleNamespace(cache_dir = cache_dir), + ) + + assert len(_scan(tmp_path, monkeypatch)) == 1 + + +# --- version robustness ------------------------------------------------------ + + +def test_recovered_entries_match_the_huggingface_hub_field_surface(): + """The recovered entries are duck-typed rather than built with hub's own + constructors, so nothing breaks when a field is added or removed upstream. + This is the tripwire that says the surfaces have drifted.""" + from huggingface_hub import CachedFileInfo, CachedRepoInfo, CachedRevisionInfo + + pairs = ( + (inventory_scan._RecoveredFileInfo, CachedFileInfo), + (inventory_scan._RecoveredRevisionInfo, CachedRevisionInfo), + (inventory_scan._RecoveredRepoInfo, CachedRepoInfo), + ) + for ours, theirs in pairs: + missing = {f.name for f in dataclasses.fields(theirs)} - { + f.name for f in dataclasses.fields(ours) + } + assert not missing, f"{ours.__name__} is missing {sorted(missing)}" + + +def test_a_recovered_repo_survives_delete_revisions(tmp_path, monkeypatch): + """Deleting a recovered model routes through HFCacheInfo.delete_revisions, + which keys a dict by repo and takes a set difference over .revisions.""" + repo_dir = _build_repo(tmp_path) + scan = _scan(tmp_path, monkeypatch)[0] + + strategy = scan.delete_revisions(SNAPSHOT) + + assert strategy.repos == frozenset({repo_dir}) + assert strategy.expected_freed_size == 11 + + +# --- load identity for a recovered snapshot ---------------------------------- + + +def _autoload_rows( + cache_root: Path, + monkeypatch, + *, + gguf: bool = False, +) -> list[dict]: + """What chat auto-load sees: GET /api/hub/cached-models, or with *gguf* set + GET /api/hub/cached-gguf.""" + from hub.services.models import cache_inventory + from types import SimpleNamespace + + monkeypatch.setattr(inventory_scan, "hf_cache_roots", lambda: [cache_root]) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = cache_root), + ) + inventory_scan.invalidate_hf_cache_scans() + try: + if gguf: + return cache_inventory._scan_cached_gguf() + return cache_inventory._scan_cached_models() + finally: + inventory_scan.invalidate_hf_cache_scans() + + +def test_auto_load_sees_a_model_hidden_behind_a_dangling_ref(tmp_path, monkeypatch): + """The #7374 symptom end to end: the snapshot is on disk and the picker + lists it, but auto-load's inventory reported nothing downloaded and the app + fell through to a fresh download.""" + repo_dir = _build_repo(tmp_path) + snapshot = repo_dir / "snapshots" / SNAPSHOT + (snapshot / "config.json").write_text("{}", encoding = "utf-8") + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + # Recovered rows must carry the snapshot as their load identity: refs/main + # dangles, so from_pretrained("Org/Model") would fail offline and download + # the current upstream HEAD online instead of using what is already here. + assert rows[0]["load_id"] == str(snapshot) + assert rows[0]["active_cache"] is True + + +def test_a_resolvable_repo_still_loads_by_repo_id(tmp_path, monkeypatch): + repo_dir = _build_repo(tmp_path, ref = SNAPSHOT) + (repo_dir / "snapshots" / SNAPSHOT / "config.json").write_text("{}", encoding = "utf-8") + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + assert rows[0]["load_id"] == "Org/Model" + + +def test_default_ref_resolves_only_when_main_names_a_snapshot(tmp_path): + dangling = _build_repo(tmp_path, name = "models--Org--Dangling") + resolved = _build_repo(tmp_path, ref = SNAPSHOT, name = "models--Org--Resolved") + detached = _build_repo(tmp_path, ref = None, name = "models--Org--Detached") + + assert inventory_scan.default_ref_resolves_on_disk(dangling) is False + assert inventory_scan.default_ref_resolves_on_disk(resolved) is True + assert inventory_scan.default_ref_resolves_on_disk(detached) is False + + +# --- the load id must name a snapshot that holds the advertised payload ------ + +OLDER = "d" * 40 +NEWER = "e" * 40 + + +def _age(path: Path, seconds: float) -> None: + """Backdate a snapshot dir; snapshot selection orders by directory mtime.""" + stamp = os.stat(path).st_mtime - seconds + os.utime(path, (stamp, stamp)) + + +def _autoload_gguf_rows(cache_root: Path, monkeypatch) -> list[dict]: + """What chat auto-load sees for GGUF: GET /api/hub/cached-gguf.""" + return _autoload_rows(cache_root, monkeypatch, gguf = True) + + +def _two_snapshot_repo( + cache_root: Path, + older_files: dict, + newer_files: dict, + *, + ref: str | None = UPSTREAM_HEAD, +) -> Path: + """A repo whose payload sits in the older of two snapshots. + + Realistic because a metadata probe (config.json only) against a commit that + has moved on materialises a newer, weightless snapshot beside the download. + + ``ref = None`` leaves no ``refs/main`` at all, which is what a fetch pinned + to a commit hash leaves behind: ``snapshot_download`` only writes a ref for + a branch or tag. + """ + repo_dir = cache_root / "models--Org--Model" + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + (repo_dir / "refs").mkdir(parents = True, exist_ok = True) + if ref is not None: + (repo_dir / "refs" / "main").write_text(ref, encoding = "utf-8") + for commit, files in ((OLDER, older_files), (NEWER, newer_files)): + snapshot = repo_dir / "snapshots" / commit + snapshot.mkdir(parents = True, exist_ok = True) + for name, payload in files.items(): + # Keys may name a subdir ("MTP/...") the way real GGUF repos ship + # their companions; forward slashes work on Windows too. + (snapshot / name).parent.mkdir(parents = True, exist_ok = True) + (snapshot / name).write_bytes(payload) + _age(repo_dir / "snapshots" / OLDER, 600) + return repo_dir + + +def test_load_id_names_the_snapshot_holding_the_safetensors_payload(tmp_path, monkeypatch): + """The row aggregates weights over every revision, so the payload it + advertises can live in an older snapshot than the newest directory.""" + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"config.json": b"{}", "model.safetensors": b"\0" * 11}, + newer_files = {"config.json": b"{}"}, + ) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + assert rows[0]["model_format"] == "safetensors" + load_dir = Path(rows[0]["load_id"]) + # Behavioural: from_pretrained(load_id) has to find the weights the row + # advertised, otherwise auto-load fails on a model that is fully cached. + assert any( + entry.suffix == ".safetensors" for entry in load_dir.iterdir() + ), f"load_id {load_dir.name} holds no weights; payload is in {OLDER[:8]}" + assert load_dir == repo_dir / "snapshots" / OLDER + + +def test_load_id_names_the_snapshot_holding_the_advertised_gguf_quant(tmp_path, monkeypatch): + """Same for GGUF: the row's size sums quants across revisions, while local + variant resolution only ever reads the one directory in ``load_id``.""" + from hub.utils.gguf import list_local_gguf_variants + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"config.json": b"{}", "Model-Q4_K_M.gguf": b"\0" * 32}, + newer_files = {"config.json": b"{}"}, + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + assert rows[0]["size_bytes"] == 32 + load_dir = Path(rows[0]["load_id"]) + variants, _has_vision = list_local_gguf_variants(str(load_dir)) + assert [v.quant for v in variants] == [ + "Q4_K_M" + ], f"no variant resolves under load_id {load_dir.name}; the quant is in {OLDER[:8]}" + assert load_dir == repo_dir / "snapshots" / OLDER + + +def test_load_id_still_prefers_the_newest_snapshot_that_holds_the_payload(tmp_path, monkeypatch): + """The payload rule must not pin loads to stale revisions: when both + snapshots are runnable, the newest one still wins.""" + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"config.json": b"{}", "model.safetensors": b"\0" * 11}, + newer_files = {"config.json": b"{}", "model.safetensors": b"\0" * 13}, + ) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == repo_dir / "snapshots" / NEWER + + +def test_load_id_leaves_the_payload_snapshot_when_main_resolves_elsewhere(tmp_path, monkeypatch): + """A ``refs/main`` that resolves is not enough: the metadata probe that + strands the weights in an older snapshot also repoints ``refs/main`` at the + weightless one, so loading by repo id lands on a revision with no weights + and the app downloads the model again.""" + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"config.json": b"{}", "model.safetensors": b"\0" * 11}, + newer_files = {"config.json": b"{}"}, + ref = NEWER, + ) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + load_id = rows[0]["load_id"] + # Behavioural: follow the load identity to a directory the way a load would + # (repo id resolves through refs/main) and require the weights to be there. + resolved = ( + repo_dir / "snapshots" / (repo_dir / "refs" / "main").read_text(encoding = "utf-8") + if load_id == "Org/Model" + else Path(load_id) + ) + assert any( + entry.suffix == ".safetensors" for entry in resolved.iterdir() + ), f"load_id {load_id} resolves to {resolved.name}, which holds no weights" + assert Path(load_id) == repo_dir / "snapshots" / OLDER + + +def test_load_id_stays_the_repo_id_when_main_resolves_onto_the_payload(tmp_path, monkeypatch): + """The rule must stay narrow: when ``refs/main`` names a snapshot that does + hold the payload, the pinned revision keeps winning and the row keeps the + repo id, even though a newer snapshot is runnable too.""" + _two_snapshot_repo( + tmp_path, + older_files = {"config.json": b"{}", "model.safetensors": b"\0" * 11}, + newer_files = {"config.json": b"{}", "model.safetensors": b"\0" * 13}, + ref = OLDER, + ) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert rows[0]["load_id"] == "Org/Model" + + +# --- everything the row advertises must resolve under its load id ------------ + + +def _local_gguf_variants_for_autoload(row: dict, cache_root: Path) -> list[str]: + """The quants chat auto-load is offered: GET /api/models/gguf-variants with + ``preferLocalCache`` and the row's ``cache_path``, exactly as chat-adapter + calls it before handing ``load_id`` to /load.""" + import asyncio + + from hub.services.models import gguf_variants + + response = asyncio.run( + gguf_variants.get_gguf_variants_response( + row["repo_id"], + prefer_local_cache = True, + local_path = row["cache_path"], + ) + ) + return [variant.quant for variant in response.variants if variant.downloaded] + + +def test_a_half_split_quant_shadows_neither_the_load_id_nor_the_variants(tmp_path, monkeypatch): + """The newest snapshot can hold shard 1 of an interrupted split download. + The picker still offers that quant, so the generated command asks + llama-server for a shard that is absent while a complete quant sits in an + older snapshot. + + The load id therefore names the newest snapshot whose quants are all on + disk, but the variant lookup still walked snapshots newest-first, so the + half downloaded split quant was offered as downloaded while /load pointed at + an older snapshot that does not hold it. Auto-load then failed on a cache + that has a perfectly usable quant. Both ends are asserted here because they + are only correct together: the load id and the quants offered under it have + to agree on one directory.""" + from hub.utils.gguf import list_local_gguf_variants + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M.gguf": b"\0" * 32}, + newer_files = {"Model-Q8_0-00001-of-00002.gguf": b"\0" * 16}, + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + load_dir = Path(rows[0]["load_id"]) + held = {v.quant for v in list_local_gguf_variants(str(load_dir))[0] if v.quant} + complete = inventory_scan._completed_gguf_variants(load_dir) + assert held and held <= complete, ( + f"load_id {load_dir.name} offers {sorted(held)} with only " + f"{sorted(complete)} complete; the usable quant is in {OLDER[:8]}" + ) + assert load_dir == repo_dir / "snapshots" / OLDER + offered = _local_gguf_variants_for_autoload(rows[0], tmp_path) + resolvable = {v.quant for v in list_local_gguf_variants(str(load_dir))[0]} + # Behavioural: every quant offered as downloaded has to resolve under the + # load id, and the complete one must not be shadowed by the broken one. + assert set(offered) <= resolvable, ( + f"auto-load is offered {sorted(offered)} but load_id {load_dir.name[:8]} " + f"resolves only {sorted(resolvable)}" + ) + assert offered == ["Q4_K_M"] + + +@pytest.mark.parametrize( + "newer_files, offered", + [ + # With nothing complete anywhere the newest snapshot holding quants is + # still reported, which is what shipped before, and it still agrees with + # the load id. + pytest.param( + {"Model-Q8_0-00001-of-00002.gguf": b"\0" * 16}, + ["Q8_0"], + id = "nothing-complete-anywhere", + ), + # When that snapshot holds a whole quant beside the half-downloaded one, + # offering both as downloaded shadowed the usable one: auto-load takes + # only the smallest and a rejected load now suppresses the default + # download, so chat was left with no model at all. + pytest.param( + { + "Model-Q8_0.gguf": b"\0" * 64, + "Model-Q4_K_M-00001-of-00002.gguf": b"\0" * 16, + }, + ["Q8_0"], + id = "one-whole-quant-beside-a-half-one", + ), + ], +) +def test_gguf_variants_still_list_when_no_snapshot_is_complete( + newer_files, offered, tmp_path, monkeypatch +): + """The completeness preference must not empty the list, and must not offer a + quant whose shards are missing while a whole one sits beside it.""" + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M-00001-of-00002.gguf": b"\0" * 32}, + newer_files = newer_files, + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + load_dir = Path(rows[0]["load_id"]) + assert load_dir == repo_dir / "snapshots" / NEWER + assert _local_gguf_variants_for_autoload(rows[0], tmp_path) == offered + + +def test_a_whole_quant_in_a_mixed_newest_snapshot_beats_an_older_larger_one(tmp_path, monkeypatch): + """A whole small quant can sit in the newest snapshot beside an interrupted + split one while an older snapshot holds nothing but a whole larger quant. + + Requiring the whole directory to be complete skipped that newest snapshot + outright, so both the load id and the offered quants fell back to the older + revision and the fully downloaded small quant disappeared. Auto-load takes + only the smallest quant offered, so it would then spend its one attempt on + the larger one and, if that does not fit in memory, leave chat with no model + while a usable quant sat on disk. The snapshot counts as usable and its + completed subset is what gets offered, so both ends still name one directory + and the interrupted quant is still withheld.""" + from hub.utils.gguf import list_local_gguf_variants + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q6_K.gguf": b"\0" * 96}, + newer_files = { + "Model-Q4_K_M.gguf": b"\0" * 16, + "Model-Q8_0-00001-of-00002.gguf": b"\0" * 8, + }, + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + load_dir = Path(rows[0]["load_id"]) + assert load_dir == repo_dir / "snapshots" / NEWER + offered = _local_gguf_variants_for_autoload(rows[0], tmp_path) + assert offered == ["Q4_K_M"] + # The pair still has to agree on one directory, and the interrupted split + # quant must not be advertised from it. + resolvable = {v.quant for v in list_local_gguf_variants(str(load_dir))[0]} + assert set(offered) <= resolvable, ( + f"auto-load is offered {sorted(offered)} but load_id {load_dir.name[:8]} " + f"resolves only {sorted(resolvable)}" + ) + assert "Q8_0" not in offered + # Pinning the snapshot that holds the interrupted download must not flip the + # row partial: a whole quant is loadable from it. + assert rows[0].get("partial") is False + assert rows[0].get("capabilities", {}).get("can_chat") is True + + +def test_load_id_is_not_pinned_to_a_snapshot_that_has_no_config(tmp_path, monkeypatch): + """The repo-level format is allowed to rest on transformer-named weights + alone, but a pinned load id names one directory and from_pretrained needs + ``config.json`` inside it. Pinning a weight-only snapshot advertised the row + as loadable while the load could only fail; keep the repo id, which can + still fill the config in from the hub.""" + _two_snapshot_repo( + tmp_path, + older_files = {"model.safetensors": b"\0" * 11}, + newer_files = {"config.json": b"{}"}, + ref = NEWER, + ) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + assert rows[0]["model_format"] == "safetensors" + load_id = rows[0]["load_id"] + assert load_id == "Org/Model", ( + f"load_id {Path(load_id).name[:8]} holds no config.json, so a local " + "from_pretrained on it cannot resolve the architecture" + ) + + +def test_no_snapshot_holds_the_payload_so_a_dangling_ref_pins_nothing(tmp_path, monkeypatch): + """Same repo, but with the dangling ``refs/main`` this branch exists for. + + The rule above is enforced by leaving the row on the repo id, and the + dangling-ref arm used to overrule it and pin the fallback newest snapshot + anyway, which is the one already known not to hold the payload. Pinning a + directory the load cannot use is worse than the repo id, which can still + complete the config from the hub.""" + _two_snapshot_repo( + tmp_path, + older_files = {"model.safetensors": b"\0" * 11}, + newer_files = {"config.json": b"{}"}, + ref = UPSTREAM_HEAD, + ) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + load_id = rows[0]["load_id"] + if load_id != "Org/Model": + # Behavioural: a pinned directory is only useful if from_pretrained can + # read an architecture and weights out of it. + pinned = Path(load_id) + held = sorted(entry.name for entry in pinned.iterdir()) + assert (pinned / "config.json").is_file() and any( + entry.endswith(".safetensors") for entry in held + ), f"load_id pins {pinned.name[:8]}, which holds only {held}" + assert load_id == "Org/Model" + + +# --- the metadata must describe the snapshot the row hands out --------------- + +QUANTIZED_CONFIG = b'{"quantization_config": {"quant_method": "bitsandbytes"}}' +MODEL_CARD = b"---\npipeline_tag: text-generation\nlibrary_name: transformers\n---\n" + + +@pytest.mark.parametrize( + "older_files, newer_files, pinned", + [ + # ``quant_method`` and the model-card fields were read from the newest + # snapshot while the load id names the payload one, so the row described + # a directory it does not hand out. The metadata probe that strands the + # payload carries neither the quantization config nor the model card, so + # the quant chip and the On Device type filter judged the model on + # absent data. + pytest.param( + { + "config.json": QUANTIZED_CONFIG, + "model.safetensors": b"\0" * 11, + "README.md": MODEL_CARD, + }, + {"config.json": b"{}"}, + True, + id = "payload-snapshot-supplies-the-row", + ), + # The rule stays narrow: with no self-contained payload snapshot there + # is nothing to scope to, so the newest snapshot still supplies the row. + pytest.param( + {"model.safetensors": b"\0" * 11}, + {"config.json": QUANTIZED_CONFIG, "README.md": MODEL_CARD}, + False, + id = "newest-snapshot-fallback", + ), + ], +) +def test_metadata_describes_the_snapshot_the_row_hands_out( + older_files, newer_files, pinned, tmp_path, monkeypatch +): + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = older_files, + newer_files = newer_files, + ref = NEWER, + ) + + rows = _autoload_rows(tmp_path, monkeypatch) + + expected_load_id = str(repo_dir / "snapshots" / OLDER) if pinned else "Org/Model" + assert rows[0]["load_id"] == expected_load_id + assert rows[0].get("quant_method") == "bitsandbytes" + assert rows[0].get("pipeline_tag") == "text-generation" + assert rows[0].get("library_name") == "transformers" + + +# --- the signals paired with the pinned snapshot ------------------------------ + + +def test_a_companion_only_snapshot_is_not_a_gguf_payload(tmp_path, monkeypatch): + """Payload selection matched GGUF companions by bare file name, but the + ``MTP/`` drafters unsloth ships are only recognisable from the path + relative to the snapshot (``huggingface_hub`` sets ``file_name`` to the + bare name for nested files, and the recovered mirror matches it). A + snapshot holding nothing but a drafter therefore counted as holding the + payload and won the load id, while the variant lister -- which does look at + the relative path -- offers nothing from it.""" + from hub.utils.gguf import list_local_gguf_variants + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M-00001-of-00002.gguf": b"\0" * 32}, + newer_files = {"MTP/Model-Q8_0-MTP.gguf": b"\0" * 64}, + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert [row["repo_id"] for row in rows] == ["Org/Model"] + load_dir = Path(rows[0]["load_id"]) + variants, _has_vision = list_local_gguf_variants(str(load_dir)) + assert [v.quant for v in variants], ( + f"load_id {load_dir.name[:8]} offers no quant at all; it holds only a " "companion drafter" + ) + assert load_dir == repo_dir / "snapshots" / OLDER + + +def test_a_repo_root_drafter_still_leaves_a_real_quant_selectable(tmp_path, monkeypatch): + """The rule must stay narrow: a snapshot that holds a drafter *and* a real + quant is still a payload snapshot.""" + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M.gguf": b"\0" * 32}, + newer_files = {"mtp-Model-Q8_0.gguf": b"\0" * 64, "Model-Q8_0.gguf": b"\0" * 128}, + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == repo_dir / "snapshots" / NEWER + + +def _write_repo_wide_signal(kind: str, hub_cache: Path) -> None: + """Either repo-wide partial signal, neither of which records a revision. + + The manifest names a file the pinned older snapshot holds at a different + size, which is what a revision that renamed or resized its weights leaves + behind for the previous one. + """ + from hub.utils import download_manifest + + if kind == "marker": + download_manifest.write_cancel_marker( + "model", "Org/Model", None, "http", hub_cache = hub_cache + ) + return + download_manifest.write_manifest( + "model", + "Org/Model", + None, + [download_manifest.ExpectedFile("model.safetensors", 99)], + "http", + hub_cache = hub_cache, + ) + + +@pytest.mark.parametrize("signal", ["marker", "manifest"]) +@pytest.mark.parametrize( + "newer_files, ref, advertised, partial", + [ + # The signal belongs to the newest snapshot while the row advertises an + # older, complete one, so inheriting it turned ``can_chat`` off for a + # model that loads fine. + pytest.param({"config.json": b"{}"}, NEWER, OLDER, False, id = "pinned-older-snapshot"), + # Negative side: when the row advertises the newest snapshot the signal + # does describe it, and a cancelled download with no ``.incomplete`` + # blob left behind has nothing else to give it away. No ``refs/main`` + # at all, which is what a commit-pinned fetch leaves: the ref carries no + # evidence about the attempt, so the newest snapshot still owns it. + pytest.param( + {"config.json": b"{}", "model.safetensors": b"\0" * 13}, + None, + NEWER, + True, + id = "advertised-snapshot", + ), + # A ``refs/main`` naming a commit with no directory does carry evidence: + # the ref is rewritten before the first file lands, so the attempt that + # left the signal never materialised a snapshot and the complete payload + # already on disk must not inherit it. + pytest.param( + {"config.json": b"{}", "model.safetensors": b"\0" * 13}, + UPSTREAM_HEAD, + NEWER, + False, + id = "unmaterialised-attempt", + ), + ], +) +def test_repo_wide_partial_signals_are_charged_to_the_newest_snapshot( + signal, newer_files, ref, advertised, partial, tmp_path, monkeypatch +): + """A cancel marker and a repo-wide manifest both record only the *last* + attempt (the marker is cleared at every download start and on success; the + manifest is overwritten), so both belong to the newest snapshot and neither + may be verified against an older revision's payload.""" + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"config.json": b"{}", "model.safetensors": b"\0" * 11}, + newer_files = newer_files, + ref = ref, + ) + _write_repo_wide_signal(signal, tmp_path) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == repo_dir / "snapshots" / advertised + assert rows[0].get("partial") is partial + assert rows[0]["capabilities"].get("can_chat") is not partial + + +def test_gguf_partial_is_judged_against_the_snapshot_the_row_advertises(tmp_path, monkeypatch): + """The GGUF row picked its payload snapshot after computing ``partial``, + and that walk used the repo's blobs plus the newest snapshot. An + interrupted re-download therefore flipped ``can_chat`` off for the older, + complete quant the row hands out as its load id.""" + from hub.utils.gguf import list_local_gguf_variants + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M.gguf": b"\0" * 32}, + newer_files = {"Model-Q8_0-00001-of-00002.gguf": b"\0" * 16}, + ) + (repo_dir / "blobs" / ("a" * 40 + ".incomplete")).write_bytes(b"\0" * 3) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + load_dir = Path(rows[0]["load_id"]) + assert load_dir == repo_dir / "snapshots" / OLDER + variants, _has_vision = list_local_gguf_variants(str(load_dir)) + assert [v.quant for v in variants] == ["Q4_K_M"] + assert rows[0].get("partial") is False + assert rows[0]["capabilities"].get("can_chat") is True + + +def test_a_gguf_download_interrupted_in_its_own_snapshot_is_still_partial(tmp_path, monkeypatch): + """Negative side of the same rule: with no complete quant anywhere the row + falls back to the newest snapshot, the ``.incomplete`` blob does belong to + it, and the row must still be partial.""" + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M-00001-of-00002.gguf": b"\0" * 32}, + newer_files = {"Model-Q8_0-00001-of-00002.gguf": b"\0" * 16}, + ref = None, + ) + (repo_dir / "blobs" / ("a" * 40 + ".incomplete")).write_bytes(b"\0" * 3) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == repo_dir / "snapshots" / NEWER + assert rows[0].get("partial") is True + assert rows[0]["capabilities"].get("can_chat") is False + + +@pytest.mark.parametrize("signal", ["marker", "manifest"]) +def test_an_update_that_never_materialised_leaves_the_cached_payload_chattable( + signal, tmp_path, monkeypatch +): + """The recovered row's own scenario, and the reason it must not arrive + partial. ``snapshot_download`` rewrites ``refs/main`` with the new commit + before it fetches a byte, and the downloader writes the manifest earlier + still, so an update interrupted before the first file leaves the previous + complete snapshot as the only payload on disk under a ref that resolves + nowhere. That is exactly the state this branch recovers rows from, and + charging the interrupted attempt's signals to the cached payload handed + chat auto-load a model it had to skip.""" + repo_dir = _build_repo(tmp_path, ref = UPSTREAM_HEAD) + snapshot = repo_dir / "snapshots" / SNAPSHOT + (snapshot / "config.json").write_text("{}", encoding = "utf-8") + _write_repo_wide_signal(signal, tmp_path) + + rows = _autoload_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == snapshot + assert rows[0].get("partial") is False + assert rows[0]["capabilities"].get("can_chat") is True + + +@pytest.mark.parametrize( + "older_files, newer_files, ref, advertised, partial", + [ + # A same-variant re-download that stops before materialising its + # snapshot leaves a manifest listing the new revision's files. Verifying + # it against the older snapshot the row pins fails on any rename or size + # change, so the one complete quant on disk looked broken. + pytest.param( + {"Model-Q4_K_M.gguf": b"\0" * 32}, + {"config.json": b"{}"}, + NEWER, + OLDER, + False, + id = "pinned-older-snapshot", + ), + # Negative side: the quant the manifest names is not complete under the + # pinned snapshot, so the manifest is the only thing that can judge it + # and the row stays partial. + pytest.param( + {"config.json": b"{}"}, + {"Model-Q4_K_M.gguf": b"\0" * 32}, + None, + NEWER, + True, + id = "advertised-snapshot", + ), + ], +) +def test_a_gguf_variant_manifest_is_scoped_to_the_snapshot_the_row_pins( + older_files, newer_files, ref, advertised, partial, tmp_path, monkeypatch +): + from hub.utils import download_manifest + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = older_files, + newer_files = newer_files, + ref = ref, + ) + download_manifest.write_manifest( + "model", + "Org/Model", + "Q4_K_M", + [download_manifest.ExpectedFile("Model-Q4_K_M.gguf", 999)], + "http", + hub_cache = tmp_path, + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == repo_dir / "snapshots" / advertised + assert rows[0].get("partial") is partial + assert rows[0]["capabilities"].get("can_chat") is not partial + + +def test_a_gguf_variant_marker_from_a_newer_attempt_does_not_disable_the_pinned_quant( + tmp_path, monkeypatch +): + """The GGUF twin of the repo-wide marker rule. A cancel marker is keyed by + (repo, variant) with no revision, so re-downloading the quant the row + already holds and cancelling it marked the complete copy in the older, + advertised snapshot broken and On Device hid a loadable GGUF.""" + from hub.utils import download_manifest + from hub.utils.gguf import list_local_gguf_variants + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M.gguf": b"\0" * 32}, + newer_files = {"config.json": b"{}"}, + ref = NEWER, + ) + download_manifest.write_cancel_marker( + "model", "Org/Model", "Q4_K_M", "http", hub_cache = tmp_path + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + load_dir = Path(rows[0]["load_id"]) + assert load_dir == repo_dir / "snapshots" / OLDER + # The quant the marker names does resolve under the load id, so the row has + # something to chat with. + assert [v.quant for v in list_local_gguf_variants(str(load_dir))[0]] == ["Q4_K_M"] + assert rows[0].get("partial") is False + assert rows[0]["capabilities"].get("can_chat") is True + + +def test_a_gguf_variant_marker_against_the_advertised_snapshot_is_still_partial( + tmp_path, monkeypatch +): + """Negative side of the same rule: when the row advertises the newest + snapshot the marker does describe the attempt that wrote it, so the only + quant stays broken and the row keeps its resume affordance.""" + from hub.utils import download_manifest + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"config.json": b"{}"}, + newer_files = {"Model-Q4_K_M.gguf": b"\0" * 32}, + ref = None, + ) + download_manifest.write_cancel_marker( + "model", "Org/Model", "Q4_K_M", "http", hub_cache = tmp_path + ) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == repo_dir / "snapshots" / NEWER + assert rows[0].get("partial") is True + assert rows[0]["capabilities"].get("can_chat") is False + + +def test_a_marker_for_another_quant_still_leaves_the_pinned_one_chattable(tmp_path, monkeypatch): + """The Q8+Q4 mixed-state rule still holds across snapshots: a cancelled + quant that is not the one the load id resolves must not veto the clean one.""" + from hub.utils import download_manifest + + repo_dir = _two_snapshot_repo( + tmp_path, + older_files = {"Model-Q4_K_M.gguf": b"\0" * 32}, + newer_files = {"config.json": b"{}"}, + ref = NEWER, + ) + download_manifest.write_cancel_marker("model", "Org/Model", "Q8_0", "http", hub_cache = tmp_path) + + rows = _autoload_gguf_rows(tmp_path, monkeypatch) + + assert Path(rows[0]["load_id"]) == repo_dir / "snapshots" / OLDER + assert rows[0].get("partial") is False + assert rows[0]["capabilities"].get("can_chat") is True diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 5f6c6cc589..030372e8db 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1445,9 +1445,29 @@ function isAutoLoadableGgufVariant(variant: GgufVariantDetail | null): boolean { return !hasBigEndianGgufMarker(filename, variant.quant); } +/** Whether a cache row is a model the user can actually chat with. + * + * The cache endpoints also return rows that exist only for the resume/delete + * affordances (an interrupted or cancelled download), and those carry + * partial: true with can_chat: false. Auto-load used to attempt them anyway + * and fall through on the rejection; now that a rejection suppresses the + * default download, attempting one would leave chat with no model at all. + * Both fields are optional so an older backend that omits them keeps its + * current behaviour of trying the row. */ +function isChattableCachedRepo(repo: { + partial?: boolean; + capabilities?: { can_chat?: boolean } | null; +}): boolean { + return repo.partial !== true && repo.capabilities?.can_chat !== false; +} + async function autoLoadSmallestModel(): Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean; + /** A specific load failure was already reported, so callers must not replace + * it with their generic "no model loaded" advice. Optional so every other + * return keeps its existing shape. */ + loadFailureReported?: boolean; }> { if (await tryAdoptServerActiveModel()) { return { loaded: true, blockedByTrustRemoteCode: false }; @@ -1494,6 +1514,26 @@ async function autoLoadSmallestModel(): Promise<{ let hadNonTrustFailure = false; let loadAttempts = 0; const skippedAutoLoadCandidates = new Set(); + // Why the last attempted load failed, set only when /api/inference/load + // itself rejected, so enumeration hiccups (variant listing, validate) keep + // their existing fall-through behaviour. Boxed because a plain `let` assigned + // only inside a nested function narrows to `null` under control-flow + // analysis, which would make every read below a `never`. + const loadFailure: { current: { label: string; detail: string } | null } = { + current: null, + }; + + function noteLoadFailure(label: string, error: unknown): void { + const detail = + error instanceof Error && error.message.trim() ? error.message.trim() : ""; + loadFailure.current = { + label, + // Older backends (and non-Error throws) carry no detail; still name the + // model that failed rather than silently fetching a different one. + detail: + detail || "The server did not report a reason. Check the Studio logs.", + }; + } async function canAutoLoad(payload: { model_path: string; @@ -1541,6 +1581,9 @@ async function autoLoadSmallestModel(): Promise<{ } const currentStore = useChatRuntimeStore.getState(); const modelPath = candidate.loadId ?? candidate.id; + const failureLabel = candidate.ggufVariant + ? `${candidate.id} (${candidate.ggufVariant})` + : candidate.id; const { config } = resolveInitialConfig(candidate.id, candidate.ggufVariant); const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ modelId: candidate.id, @@ -1642,6 +1685,12 @@ async function autoLoadSmallestModel(): Promise<{ n_parallel: config.nParallel ?? null, } : {}), + }).catch((error: unknown) => { + // The sweep's parameterless catches discard this error, which is what let + // a genuine load failure fall through to the "no downloaded models" Hub + // download. Rethrowing keeps the awaited type and their control flow. + noteLoadFailure(failureLabel, error); + throw error; }); // Only persist the global preference when the value came from the global // settings. A per-model config's choice must stay load-local, or autoloading @@ -1768,10 +1817,14 @@ async function autoLoadSmallestModel(): Promise<{ return true; } try { - const [ggufRepos, modelRepos] = await Promise.all([ + const [allGgufRepos, allModelRepos] = await Promise.all([ listCachedGguf().catch(() => []), listCachedModels().catch(() => []), ]); + // Filtered once, so the last-used lookup below sees the same set as the + // sweeps and neither can spend a load attempt on a resume-only row. + const ggufRepos = allGgufRepos.filter(isChattableCachedRepo); + const modelRepos = allModelRepos.filter(isChattableCachedRepo); if (lastLoaded) { if (lastLoaded.kind === "gguf") { @@ -1923,12 +1976,24 @@ async function autoLoadSmallestModel(): Promise<{ // Cap also gates the default download, so total /api/inference/load // budget across cached + fallback is MAX_AUTO_LOAD_ATTEMPTS, not +1. - if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) { + // A cached model that was tried and failed stops here too: the user has + // models on disk, so the reason is the useful answer and pulling an + // unrelated default off the Hub is not. A device with nothing cached never + // sets loadFailure and still falls through to the download below. + if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS || loadFailure.current) { toast.dismiss(toastId); + if (loadFailure.current) { + toast.error(`Could not load ${loadFailure.current.label}`, { + description: loadFailure.current.detail, + duration: 10000, + closeButton: true, + }); + } return { loaded: false, blockedByTrustRemoteCode: blockedByTrustRemoteCode && !hadNonTrustFailure, + loadFailureReported: loadFailure.current !== null, }; } @@ -2091,19 +2156,23 @@ export function createOpenAIStreamAdapter( await waitForModelReady(abortSignal); } if (!useChatRuntimeStore.getState().params.checkpoint) { - const { loaded, blockedByTrustRemoteCode } = + const { loaded, blockedByTrustRemoteCode, loadFailureReported } = await autoLoadSmallestModel(); if (!loaded) { - toast.error( - blockedByTrustRemoteCode - ? "This model needs custom code approval" - : "No model loaded", - { - description: blockedByTrustRemoteCode - ? "Select it from the top bar to review and approve its custom code, or pick another model." - : "Pick a model in the top bar, then retry.", - }, - ); + // A reported load failure already names the model and the reason, + // so the generic advice would only bury it. + if (!loadFailureReported) { + toast.error( + blockedByTrustRemoteCode + ? "This model needs custom code approval" + : "No model loaded", + { + description: blockedByTrustRemoteCode + ? "Select it from the top bar to review and approve its custom code, or pick another model." + : "Pick a model in the top bar, then retry.", + }, + ); + } throw new Error("Load a model first."); } } @@ -2383,24 +2452,29 @@ export function createOpenAIStreamAdapter( // Prefer a model already loaded by the CLI/API before auto-loading. let loaded: boolean; let blockedByTrustRemoteCode: boolean; + let loadFailureReported: boolean | undefined; try { - ({ loaded, blockedByTrustRemoteCode } = + ({ loaded, blockedByTrustRemoteCode, loadFailureReported } = await autoLoadSmallestModel()); } catch (error) { clearSelectedImageEditReference(); throw error; } if (!loaded) { - toast.error( - blockedByTrustRemoteCode - ? "This model needs custom code approval" - : "No model loaded", - { - description: blockedByTrustRemoteCode - ? "Select it from the top bar to review and approve its custom code, or pick another model." - : "Pick a model in the top bar, then retry.", - }, - ); + // A reported load failure already names the model and the reason, so + // the generic advice would only bury it. + if (!loadFailureReported) { + toast.error( + blockedByTrustRemoteCode + ? "This model needs custom code approval" + : "No model loaded", + { + description: blockedByTrustRemoteCode + ? "Select it from the top bar to review and approve its custom code, or pick another model." + : "Pick a model in the top bar, then retry.", + }, + ); + } clearSelectedImageEditReference(); throw new Error("Load a model first."); } diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 8ad2691391..7288d8dc6b 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -292,6 +292,15 @@ export interface CachedGgufRepo { /** True when the repo ships an mmproj adapter (image inputs). Optional for * older-backend compatibility. */ has_vision?: boolean; + partial?: boolean; + capabilities?: CachedRepoCapabilities | null; +} + +/** The subset of the row's capabilities auto-load acts on. The backend sends + * the whole block on both cache endpoints; the rest is only read by the Hub + * view models, which have their own wider type. */ +export interface CachedRepoCapabilities { + can_chat?: boolean; } export async function getGgufDownloadProgress( @@ -411,6 +420,8 @@ export interface CachedModelRepo { /** Owning cache dir; sent so a delete targets this copy, not the active * cache. Optional for older-backend compatibility. */ cache_path?: string | null; + partial?: boolean; + capabilities?: CachedRepoCapabilities | null; } export async function listCachedModels( diff --git a/tests/studio/load_freeze/test_load_orchestrator.py b/tests/studio/load_freeze/test_load_orchestrator.py index a1f4caa309..5cec40653b 100644 --- a/tests/studio/load_freeze/test_load_orchestrator.py +++ b/tests/studio/load_freeze/test_load_orchestrator.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import importlib.util import os import re import socket @@ -49,7 +50,24 @@ import logging as _logging # noqa: E402 _loggers_stub = types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: _logging.getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) -sys.modules.setdefault("structlog", types.ModuleType("structlog")) +# structlog is a hard studio.txt requirement, but it is only imported lazily, so a +# bare setdefault here used to park an empty placeholder BEFORE anything imported +# the real package -- and it then shadowed it for the rest of the session. Every +# later file importing a studio module that calls structlog.get_logger at module +# scope (routes.inference -> core.inference.external_provider, utils.mlx_repair) +# blew up with AttributeError, but only when this file was collected first, so the +# same test passed alone and failed under `pytest tests/studio`. Only stub when the +# package is genuinely missing, and give the stub the attribute those callers use. +# Guard on sys.modules FIRST: another test module may have parked its own bare +# stub, and find_spec() raises ValueError on a module whose __spec__ is None. +# Anything already there (real or stub) is left alone; only a genuinely absent +# package gets stubbed. +if "structlog" not in sys.modules and importlib.util.find_spec("structlog") is None: + _structlog_stub = types.ModuleType("structlog") + _structlog_stub.get_logger = lambda *args, **kwargs: _logging.getLogger( + args[0] if args else "structlog" + ) + sys.modules["structlog"] = _structlog_stub import httpx # noqa: E402 diff --git a/tests/studio/test_chat_autoload_failure_gate.py b/tests/studio/test_chat_autoload_failure_gate.py new file mode 100644 index 0000000000..1b9070cca5 --- /dev/null +++ b/tests/studio/test_chat_autoload_failure_gate.py @@ -0,0 +1,470 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""A failed auto-load of a cached model must not become a Hub download. + +Runs the real ``autoLoadSmallestModel`` from chat-adapter.ts under node with the +module boundary stubbed, so these assert behaviour (which /api/inference/load +calls happen, what the user is told) rather than source text. The sweep's +catches are parameterless, so before the fix a cached repo whose load rejected +fell straight through to fetching an unrelated default model and reported +success for it. +""" + +import json +import os +import shutil +import subprocess +import tempfile +import textwrap +from pathlib import Path + +import pytest + +WORKDIR = Path(__file__).resolve().parents[2] + + +def _source_path(relative_path: str) -> Path: + direct = WORKDIR / relative_path + if direct.exists(): + return direct + return WORKDIR / "unsloth_repo" / relative_path + + +ADAPTER = _source_path("studio/frontend/src/features/chat/api/chat-adapter.ts") +TEMP = WORKDIR / "temp" / "chat_autoload_failure_gate" +DEFAULT_MODEL = "unsloth/Qwen3.5-4B-MTP-GGUF" +GEMMA_REPO = "unsloth/gemma-4-26B-A4B-it-qat-GGUF" + +# Stubs for everything autoLoadSmallestModel imports. Each scenario supplies the +# cache inventory and how /validate and /load answer for a given model_path. +PREAMBLE = """ +type LastLocalModelKind = "gguf" | "model"; +type GgufVariantDetail = { + quant?: string | null; + filename?: string | null; + downloaded?: boolean; + size_bytes: number; +}; +type ChatModelSummary = Record; + +export type Scenario = { + ggufRepos: any[]; + modelRepos: any[]; + variants: Record; + lastLoaded: any; + validate: (payload: any) => any; + load: (payload: any) => any; +}; + +export const EVENTS: any[] = []; +let SCENARIO: Scenario; +export function setScenario(scenario: Scenario) { + SCENARIO = scenario; + EVENTS.length = 0; + STORE = makeStore(); +} + +const GPU_LAYERS_AUTO = -1; + +function makeStore(): any { + const state: any = { + hfToken: null, + params: { maxSeqLength: 4096, checkpoint: "" }, + activeGgufVariant: null, + activePresetSource: null, + gpuMemoryMode: "auto", + selectedGpuIds: null, + models: [], + setCheckpoint: () => {}, + setModelRequiresTrustRemoteCode: () => {}, + setParams: (p: any) => { state.params = p; }, + setModels: (m: any[]) => { state.models = m; }, + }; + return state; +} +let STORE: any = makeStore(); +const useChatRuntimeStore = { + getState: () => STORE, + setState: (_p: any) => {}, +}; + +function createLoadingToastIcon() { return null; } +const toast: any = Object.assign( + (_msg: string, _opts?: any) => "toast-id", + { + message: (msg: string, opts?: any) => { + EVENTS.push({ kind: "toast.message", msg, description: opts?.description }); + return "toast-id"; + }, + success: (msg: string) => EVENTS.push({ kind: "toast.success", msg }), + error: (msg: string, opts?: any) => + EVENTS.push({ kind: "toast.error", msg, description: opts?.description }), + dismiss: () => EVENTS.push({ kind: "toast.dismiss" }), + info: (msg: string) => EVENTS.push({ kind: "toast.info", msg }), + }, +); + +async function tryAdoptServerActiveModel() { return false; } +function resolveSpeculativeSettingsForLoad() { + return { speculativeType: null, specDraftNMax: 0 }; +} +function readLastLocalModelLoad() { return SCENARIO.lastLoaded; } +function recordLastLocalModelLoad(_x: any) {} +function resolveInitialConfig(_id: string, _variant: any) { + return { config: { + customContextLength: null, maxSeqLength: null, gpuMemoryMode: null, + gpuLayers: null, nCpuMoe: null, selectedGpuIds: undefined, + speculativeType: null, specDraftNMax: null, chatTemplateOverride: null, + kvCacheDtype: null, tensorParallel: false, + } }; +} +function resolveLoadMaxSeqLength(args: any) { return args.maxSeqLength ?? 0; } +function resolveFitMaxSeqLength(..._a: any[]) { return 0; } +function resolveManualAutoCtxPin(..._a: any[]) { return null; } +async function ensureGpuDeviceCache() {} +function reconcilePersistedGpuIds(ids: any) { return ids; } +function saveSpeculativeType(_x: any) {} +function persistGpuMemoryModeOnLoad(..._a: any[]) {} +function reasoningCapsFromLoad(_x: any) { return {}; } +function resolveToolsEnabledOnLoad(_x: any) { return {}; } +function loadedGpuMemoryFields(_x: any) { return {}; } +function resolveLoadedSpeculativeSettings(_x: any) { return {}; } +function isMultimodalResponse(_x: any) { return false; } + +async function listCachedGguf() { return SCENARIO.ggufRepos as any; } +async function listCachedModels() { return SCENARIO.modelRepos as any; } +async function listGgufVariants(repoId: string, _b?: any, _c?: any) { + const entry = SCENARIO.variants[repoId]; + if (entry === "throw") throw new Error("variant listing failed"); + return entry ?? { variants: [] }; +} +async function validateModel(payload: any) { + const result = SCENARIO.validate(payload); + if (result instanceof Error) throw result; + return result; +} +async function loadModel(payload: any) { + const result = SCENARIO.load(payload); + EVENTS.push({ + kind: "loadModel", + model_path: payload.model_path, + gguf_variant: payload.gguf_variant ?? null, + rejected: result instanceof Error, + }); + if (result instanceof Error) throw result; + return result; +} +""" + +SCENARIO_HELPERS = """ + const GEMMA = { + repo_id: "unsloth/gemma-4-26B-A4B-it-qat-GGUF", + load_id: "unsloth/gemma-4-26B-A4B-it-qat-GGUF", + cache_path: + "/home/john-doe/.cache/huggingface/hub/models--unsloth--gemma-4-26B-A4B-it-qat-GGUF", + size_bytes: 15800000000, + }; + const GEMMA_VARIANTS = { + variants: [{ + quant: "UD-Q4_K_XL", + filename: "UD-Q4_K_XL/gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf", + downloaded: true, + size_bytes: 15800000000, + }], + }; + const OOM = + "Failed to load model: llama-server was stopped by the operating system " + + "(signal 9), most likely out of memory."; + const VALIDATE_OK = () => ({ + requires_trust_remote_code: false, + requires_security_review: false, + requires_transformers_upgrade: false, + }); + const LOADED = (payload) => ({ + model: payload.model_path, + is_gguf: true, + context_length: 32768, + }); + const scenario = (over) => ({ + ggufRepos: [], + modelRepos: [], + variants: {}, + lastLoaded: null, + validate: VALIDATE_OK, + load: LOADED, + ...over, + }); +""" + + +def _require_node(): + if shutil.which("node") is None: + pytest.skip("node not available") + if not ADAPTER.exists(): + pytest.skip("studio chat sources not present") + result = subprocess.run( + ["node", "--experimental-strip-types", "--version"], + capture_output = True, + text = True, + timeout = 5, + ) + if result.returncode != 0: + pytest.skip("node --experimental-strip-types not available") + + +def _build_harness(run_dir: Path): + """Slice autoLoadSmallestModel and its helpers verbatim out of the adapter.""" + lines = ADAPTER.read_text(encoding = "utf-8").splitlines() + start = next( + (i for i, line in enumerate(lines) if line.startswith("const MAX_AUTO_LOAD_ATTEMPTS")), + None, + ) + end = next( + ( + i + for i, line in enumerate(lines) + if line.startswith("export function createOpenAIStreamAdapter") + ), + None, + ) + assert ( + start is not None and end is not None and start < end + ), "could not locate the auto-load region in chat-adapter.ts" + body = "\n".join(lines[start:end]) + assert "async function autoLoadSmallestModel" in body + (run_dir / "harness.ts").write_text( + "// @ts-nocheck\n" + PREAMBLE + "\n" + body + "\nexport { autoLoadSmallestModel };\n", + encoding = "utf-8", + ) + + +def _run(scenario_expr: str) -> dict: + _require_node() + # Its own directory per invocation, harness included: sharing either file + # lets a concurrent runner read one that another is mid-rewrite. + TEMP.mkdir(parents = True, exist_ok = True) + run_dir = Path(tempfile.mkdtemp(prefix = "run", dir = TEMP)) + _build_harness(run_dir) + script = ( + textwrap.dedent( + """ + // @ts-nocheck + import { autoLoadSmallestModel, setScenario, EVENTS } from "./harness.ts"; + """ + ) + + SCENARIO_HELPERS + + textwrap.dedent( + f""" + setScenario({scenario_expr}); + const result = await autoLoadSmallestModel(); + console.log(JSON.stringify({{ result, events: EVENTS }})); + """ + ) + ) + (run_dir / "run.mts").write_text(script, encoding = "utf-8") + completed = subprocess.run( + ["node", "--experimental-strip-types", "--no-warnings", "run.mts"], + cwd = str(run_dir), + capture_output = True, + text = True, + timeout = 60, + env = dict(os.environ, NODE_NO_WARNINGS = "1"), + ) + assert completed.returncode == 0, f"stderr: {completed.stderr}\nstdout: {completed.stdout}" + last = [line for line in completed.stdout.strip().splitlines() if line.strip()][-1] + return json.loads(last) + + +def _loaded_paths(out: dict) -> list[str]: + return [event["model_path"] for event in out["events"] if event["kind"] == "loadModel"] + + +def _toasts(out: dict, kind: str) -> list[dict]: + return [event for event in out["events"] if event["kind"] == kind] + + +def test_failed_cached_load_does_not_download_the_default_model(): + """The reported case: the only cached repo is enumerated fine but its load + OOMs, and the default GGUF would load. Auto-load must stop at the failure + instead of fetching a model the user never asked for.""" + out = _run( + "scenario({ ggufRepos: [GEMMA], variants: { [GEMMA.repo_id]: GEMMA_VARIANTS }," + " load: (p) => p.model_path === GEMMA.repo_id ? new Error(OOM) : LOADED(p) })" + ) + + assert _loaded_paths(out) == ["unsloth/gemma-4-26B-A4B-it-qat-GGUF"] + assert DEFAULT_MODEL not in _loaded_paths(out) + assert out["result"]["loaded"] is False + assert _toasts(out, "toast.success") == [] + assert not any( + "Downloading a small model" in event["msg"] for event in _toasts(out, "toast.message") + ) + + +def test_failed_cached_load_surfaces_the_backend_reason(): + out = _run( + "scenario({ ggufRepos: [GEMMA], variants: { [GEMMA.repo_id]: GEMMA_VARIANTS }," + " load: () => new Error(OOM) })" + ) + + [error] = _toasts(out, "toast.error") + assert error["msg"] == "Could not load unsloth/gemma-4-26B-A4B-it-qat-GGUF (UD-Q4_K_XL)" + assert "out of memory" in error["description"] + + +def test_load_rejection_without_a_message_still_names_the_model(): + """Old backends and non-Error throws carry no detail; the model that failed + must still be named rather than swapped for the default.""" + out = _run( + "scenario({ ggufRepos: [GEMMA], variants: { [GEMMA.repo_id]: GEMMA_VARIANTS }," + " load: () => new Error('') })" + ) + + assert DEFAULT_MODEL not in _loaded_paths(out) + [error] = _toasts(out, "toast.error") + assert "unsloth/gemma-4-26B-A4B-it-qat-GGUF (UD-Q4_K_XL)" in error["msg"] + assert error["description"] + + +def test_empty_device_still_downloads_the_default_model(): + """Nothing cached means nothing failed, so the download path is untouched.""" + out = _run("scenario({})") + + assert _loaded_paths(out) == [DEFAULT_MODEL] + assert out["result"]["loaded"] is True + assert _toasts(out, "toast.error") == [] + assert [event["msg"] for event in _toasts(out, "toast.success")] == [ + "Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)" + ] + + +def test_enumeration_failure_still_downloads_the_default_model(): + """A cached repo whose variants cannot be listed never reached /load, so it + keeps falling through: only a real load rejection changes behaviour.""" + out = _run("scenario({ ggufRepos: [GEMMA], variants: { [GEMMA.repo_id]: 'throw' } })") + + assert _loaded_paths(out) == [DEFAULT_MODEL] + assert out["result"]["loaded"] is True + + +def test_consent_gated_candidate_still_downloads_the_default_model(): + """trust_remote_code / security review block the load before it is attempted, + which is a deferral rather than a failure.""" + out = _run( + "scenario({ ggufRepos: [GEMMA], variants: { [GEMMA.repo_id]: GEMMA_VARIANTS }," + " validate: (p) => p.model_path === GEMMA.repo_id" + " ? { requires_trust_remote_code: true, requires_security_review: false," + " requires_transformers_upgrade: false } : VALIDATE_OK() })" + ) + + assert _loaded_paths(out) == [DEFAULT_MODEL] + assert out["result"]["loaded"] is True + + +def test_attempt_cap_still_gates_the_default_download(): + """Four broken cached repos: the sweep keeps trying smaller candidates, the + cap stops it at three attempts, and no fifth load goes to the Hub.""" + out = _run( + "scenario({ ggufRepos: [1, 2, 3, 4].map((i) => ({ ...GEMMA, repo_id: `r${i}`," + " load_id: `r${i}`, size_bytes: i }))," + " variants: Object.fromEntries([1, 2, 3, 4].map((i) => [`r${i}`, GEMMA_VARIANTS]))," + " load: () => new Error(OOM) })" + ) + + assert _loaded_paths(out) == ["r1", "r2", "r3"] + assert DEFAULT_MODEL not in _loaded_paths(out) + + +def test_a_later_cached_model_can_still_load_after_an_earlier_failure(): + """One broken repo must not veto a working one: the sweep continues and the + failure toast is only for a sweep that ends with nothing loaded.""" + out = _run( + "scenario({ ggufRepos: [1, 2].map((i) => ({ ...GEMMA, repo_id: `r${i}`," + " load_id: `r${i}`, size_bytes: i }))," + " variants: { r1: GEMMA_VARIANTS, r2: GEMMA_VARIANTS }," + " load: (p) => p.model_path === 'r1' ? new Error(OOM) : LOADED(p) })" + ) + + assert _loaded_paths(out) == ["r1", "r2"] + assert out["result"]["loaded"] is True + assert _toasts(out, "toast.error") == [] + + +def test_reported_failure_is_flagged_so_callers_drop_the_generic_advice(): + """Both send paths show a generic "No model loaded" toast whenever the sweep + returns loaded: false. Without a flag that lands after the detailed toast, + making the retry advice the last thing the user sees.""" + out = _run( + "scenario({ ggufRepos: [GEMMA], variants: { [GEMMA.repo_id]: GEMMA_VARIANTS }," + " load: () => new Error(OOM) })" + ) + + assert out["result"]["loaded"] is False + assert out["result"]["loadFailureReported"] is True + + +def test_empty_device_does_not_flag_a_reported_failure(): + """Nothing was attempted, so the callers must keep their generic advice.""" + out = _run("scenario({ ggufRepos: [], models: [] })") + + assert out["result"].get("loadFailureReported") is not True + + +def test_a_resume_only_cached_row_is_skipped_so_the_default_still_downloads(): + """An interrupted download leaves a row the backend already marks + partial/can_chat=false, kept for the resume and delete affordances. The + sweep attempted it anyway, and with the failure gate above that rejection + suppressed the default download, so a half-finished cache left chat with no + model at all.""" + out = _run( + "scenario({ modelRepos: [{ repo_id: 'org/half', load_id: 'org/half'," + " size_bytes: 1, partial: true, capabilities: { can_chat: false } }]," + " load: (p) => p.model_path === 'org/half'" + " ? new Error('config.json not found') : LOADED(p) })" + ) + + assert _loaded_paths(out) == [DEFAULT_MODEL] + assert out["result"]["loaded"] is True + assert _toasts(out, "toast.error") == [] + + +def test_a_can_chat_false_cached_row_is_skipped_on_its_own(): + """The two fields are set independently on the row, so can_chat alone (a + weightless config-only repo) has to be enough to skip it.""" + out = _run( + "scenario({ ggufRepos: [{ ...GEMMA, capabilities: { can_chat: false } }]," + " variants: { [GEMMA.repo_id]: GEMMA_VARIANTS }," + " load: (p) => p.model_path === GEMMA.repo_id ? new Error(OOM) : LOADED(p) })" + ) + + assert _loaded_paths(out) == [DEFAULT_MODEL] + assert out["result"]["loaded"] is True + + +def test_the_last_used_model_is_skipped_when_its_row_went_partial(): + """The last-used shortcut reads the same rows, so an update that was + cancelled over the model the user last chatted with must not spend the + attempt either.""" + out = _run( + "scenario({ ggufRepos: [{ ...GEMMA, partial: true }]," + " variants: { [GEMMA.repo_id]: GEMMA_VARIANTS }," + " lastLoaded: { id: GEMMA.repo_id, kind: 'gguf', ggufVariant: 'UD-Q4_K_XL' }," + " load: (p) => p.model_path === GEMMA.repo_id ? new Error(OOM) : LOADED(p) })" + ) + + assert _loaded_paths(out) == [DEFAULT_MODEL] + assert out["result"]["loaded"] is True + + +def test_a_complete_cached_row_is_still_attempted(): + """Guard on the filter itself: a row with the fields present and healthy + must still be swept, and a backend that omits them entirely (older Studio) + keeps its current behaviour.""" + out = _run( + "scenario({ ggufRepos: [{ ...GEMMA, partial: false, capabilities: { can_chat: true } }]," + " variants: { [GEMMA.repo_id]: GEMMA_VARIANTS } })" + ) + + assert _loaded_paths(out) == [GEMMA_REPO] + assert out["result"]["loaded"] is True