ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge 9abe3d7823 into c70c1d2d89

Browse source
This commit is contained in:
Souravrajvi0 2026-07-29 09:04:15 +01:00 • committed by GitHub
commit 740e7a1565
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 2389 additions and 115 deletions
Download patch file Download diff file Expand all files Collapse all files

198
studio/backend/hub/services/models/cache_inventory.py
View file

@ -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):

26
studio/backend/hub/tests/test_model_services.py
View file

@ -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() == []

33
studio/backend/hub/utils/gguf.py
View file

@ -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(

521
studio/backend/hub/utils/inventory_scan.py
View file

@ -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/<branch>`` file names a
commit with no ``snapshots/<commit>/`` 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/<sha>/`` 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/<revision> 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/<revision>`` 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:

23
studio/backend/routes/models.py
View file

@ -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]:

1080
studio/backend/tests/test_hf_cache_dangling_refs.py Normal file
View file

File diff suppressed because it is too large Load diff

122
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -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<string>();
// 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.");
}

11
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -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(

20
tests/studio/load_freeze/test_load_orchestrator.py
View file

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

470
tests/studio/test_chat_autoload_failure_gate.py Normal file
View file

@ -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<string, unknown>;
export type Scenario = {
ggufRepos: any[];
modelRepos: any[];
variants: Record<string, any>;
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