fix(setup): weight-aware install-state so truncated model cache isn't read as installed (#622) (#626)
A first-run user whose model download was interrupted after the config/ tokenizer files landed but before the weight shard got stranded on the Models & Engines page: GET /models computed "installed" purely from cache size on disk, so a size-positive-but-weight-less cache reported installed=true, the wizard hid the re-download button, and the model manager (Settings → Models) that could repair it was unreachable behind the wizard gate. Make install-state weight-aware. The boolean weight-floor scan now lives in models.py (the lowest module in the setup import graph) as snapshot_has_weights() + cache_is_complete(); list_models() and recommendations() downgrade a truncated cache to installed=false (+ an explicit incomplete=true on /models), so the existing "install" action re-appears and the user can re-download in-wizard. Fixes the whole class, not just /models: download.py's install-time validator now delegates to the same shared scan (one source of the floors, can't drift), matching the load-time repair in model_manager.py (#581/#606). config_only repos (pyannote/speaker-diarization-3.1 — a pipeline whose real weights live in referenced sub-repos and whose own cache is legitimately tiny) carry a new config_only:true hint in models.yaml and are exempt, so they're not false-flagged as incomplete. Tests: tests/test_mm2_lifecycle.py — snapshot_has_weights truncated-vs-complete, cache_is_complete on a truncated weight repo + config-only exemption, and list_models downgrading a size-positive truncated cache to installed=false / incomplete=true. Full backend suite green (1832 passed). Co-authored-by: mergetest <test@local> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Palash Debnath
co-authored by
mergetest
Claude Opus 4.8
mergetest
Claude Opus 4.8
parent
d21fb765e2
commit
1fe68ba11e
@@ -21,7 +21,17 @@ from pydantic import BaseModel
|
||||
from core import prefs
|
||||
from utils import hf_progress
|
||||
from utils import download_aggregator
|
||||
from .models import KNOWN_MODELS, invalidate_cache
|
||||
# Weight-floor scan (MM2-07 / #352) lives in ``models.py`` — the lowest module in
|
||||
# the setup import graph — so install-time validation here, the first-run
|
||||
# install-state detector (#622), and load-time repair share one set of floors and
|
||||
# can't drift apart. ``_MIN_WEIGHT_BYTES``/``_WEIGHT_FLOORS`` re-exported for tests.
|
||||
from .models import ( # noqa: F401
|
||||
KNOWN_MODELS,
|
||||
invalidate_cache,
|
||||
snapshot_has_weights,
|
||||
_MIN_WEIGHT_BYTES,
|
||||
_WEIGHT_FLOORS,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("omnivoice.setup.download")
|
||||
router = APIRouter()
|
||||
@@ -223,51 +233,26 @@ def _safe_put(queue: asyncio.Queue, event) -> None:
|
||||
# model.safetensors" (#352). 5 MB clears every weight format we ship
|
||||
# (safetensors/bin shards, onnx, pt, gguf) without false-positiving on
|
||||
# config-only aux repos.
|
||||
_MIN_WEIGHT_BYTES = 5 * 1024 * 1024
|
||||
|
||||
# Per-role weight-file floors (MM2-07). A valid model has at least one
|
||||
# recognized weight file at or above its extension's floor. ONNX graphs are
|
||||
# legitimately small (a complete model can be well under 5 MB), so a single
|
||||
# 5 MB rule false-positives on them as "truncated" (#352 over-trigger); give
|
||||
# .onnx a lower floor while still rejecting a 0/KB partial. Tensor formats keep
|
||||
# the original 5 MB floor.
|
||||
_WEIGHT_FLOORS = {
|
||||
".safetensors": _MIN_WEIGHT_BYTES,
|
||||
".bin": _MIN_WEIGHT_BYTES,
|
||||
".ckpt": _MIN_WEIGHT_BYTES,
|
||||
".pt": _MIN_WEIGHT_BYTES,
|
||||
".pth": _MIN_WEIGHT_BYTES,
|
||||
".gguf": _MIN_WEIGHT_BYTES,
|
||||
".onnx": 64 * 1024, # a real ONNX graph is ≥ tens of KB; a truncated one is bytes
|
||||
}
|
||||
|
||||
|
||||
def _validate_snapshot_has_weights(repo_id: str, snapshot_path: str) -> None:
|
||||
"""Raise OSError when a finished snapshot has no plausible weight file —
|
||||
surfaces the truncated-download class (#352) at install time, where the
|
||||
retry loop and the UI's re-download path can deal with it, instead of at
|
||||
first synthesis with an opaque transformers error.
|
||||
|
||||
A snapshot is valid if it contains a recognized weight file meeting its
|
||||
per-extension floor (MM2-07) OR any file ≥ the global 5 MB floor (the
|
||||
original lenient catch — kept so this is never stricter than before)."""
|
||||
Delegates the weight check to ``models.snapshot_has_weights`` (single source of
|
||||
the floors); only the install-time error message lives here."""
|
||||
if snapshot_has_weights(snapshot_path):
|
||||
return
|
||||
biggest = 0
|
||||
try:
|
||||
biggest = 0
|
||||
for root, _dirs, files in os.walk(snapshot_path, followlinks=True):
|
||||
for f in files:
|
||||
try:
|
||||
size = os.path.getsize(os.path.join(root, f))
|
||||
biggest = max(biggest, os.path.getsize(os.path.join(root, f)))
|
||||
except OSError:
|
||||
continue
|
||||
biggest = max(biggest, size)
|
||||
ext = os.path.splitext(f)[1].lower()
|
||||
floor = _WEIGHT_FLOORS.get(ext)
|
||||
if floor is not None and size >= floor:
|
||||
return # a recognized weight file of plausible size
|
||||
if size >= _MIN_WEIGHT_BYTES:
|
||||
return # original lenient catch (non-standard weight names)
|
||||
except OSError:
|
||||
return # can't inspect — don't block the install on the checker itself
|
||||
pass
|
||||
raise OSError(
|
||||
f"{repo_id}: download finished but no model weights were found in the "
|
||||
"snapshot (largest file "
|
||||
|
||||
@@ -146,6 +146,94 @@ def _hub_cache_roots() -> list[str]:
|
||||
return roots
|
||||
|
||||
|
||||
# ── Weight-presence (truncated-cache) detection ─────────────────────────────
|
||||
# A cache that downloaded config/tokenizer files but not the weight shard still
|
||||
# occupies bytes on disk, so a size-only "installed" check (#352/#581/#606) reads
|
||||
# it as installed and the first-run wizard hides the re-download button, stranding
|
||||
# the user (#622). These helpers tell a *complete* snapshot from a truncated one by
|
||||
# checking for a plausible weight file — the same class `download.py` guards at
|
||||
# install time and `model_manager.py` repairs at load time. Shared here (the lowest
|
||||
# module in the setup import graph; `download.py` imports from this module) so the
|
||||
# floors live in exactly one place and can't drift between the three call sites.
|
||||
|
||||
_MIN_WEIGHT_BYTES = 5 * 1024 * 1024 # tensor formats: a real shard is ≥ a few MB
|
||||
|
||||
# Per-extension floors. ONNX graphs are legitimately small (a complete model can be
|
||||
# well under 5 MB), so they get a lower floor that still rejects a bytes-only partial.
|
||||
_WEIGHT_FLOORS = {
|
||||
".safetensors": _MIN_WEIGHT_BYTES,
|
||||
".bin": _MIN_WEIGHT_BYTES,
|
||||
".ckpt": _MIN_WEIGHT_BYTES,
|
||||
".pt": _MIN_WEIGHT_BYTES,
|
||||
".pth": _MIN_WEIGHT_BYTES,
|
||||
".gguf": _MIN_WEIGHT_BYTES,
|
||||
".onnx": 64 * 1024,
|
||||
}
|
||||
|
||||
|
||||
def snapshot_has_weights(snapshot_path: str) -> bool:
|
||||
"""True when a finished snapshot dir holds a plausible weight file.
|
||||
|
||||
A snapshot is complete if it contains a recognized weight file meeting its
|
||||
per-extension floor OR any file ≥ the global 5 MB floor (the lenient catch for
|
||||
non-standard weight names). Returns True when the path can't be inspected — an
|
||||
un-walkable dir must never be reported as truncated, only a confirmed weight-less
|
||||
one. `getsize` follows symlinks, so HF's snapshot→blob links resolve correctly;
|
||||
a broken link (missing blob) raises OSError and is skipped, i.e. counts as absent.
|
||||
"""
|
||||
try:
|
||||
for root, _dirs, files in os.walk(snapshot_path, followlinks=True):
|
||||
for f in files:
|
||||
try:
|
||||
size = os.path.getsize(os.path.join(root, f))
|
||||
except OSError:
|
||||
continue
|
||||
ext = os.path.splitext(f)[1].lower()
|
||||
floor = _WEIGHT_FLOORS.get(ext)
|
||||
if floor is not None and size >= floor:
|
||||
return True
|
||||
if size >= _MIN_WEIGHT_BYTES:
|
||||
return True
|
||||
except OSError:
|
||||
return True # can't inspect — don't mislabel as truncated
|
||||
return False
|
||||
|
||||
|
||||
def _snapshot_dirs(repo_id: str) -> list[str]:
|
||||
"""Existing snapshot revision dirs for a repo across the candidate cache roots."""
|
||||
name = _repo_dir_name(repo_id)
|
||||
dirs: list[str] = []
|
||||
for root in _hub_cache_roots():
|
||||
snaps = os.path.join(root, name, "snapshots")
|
||||
try:
|
||||
for rev in os.listdir(snaps):
|
||||
rev_dir = os.path.join(snaps, rev)
|
||||
if os.path.isdir(rev_dir):
|
||||
dirs.append(rev_dir)
|
||||
except OSError:
|
||||
continue
|
||||
return dirs
|
||||
|
||||
|
||||
def cache_is_complete(model: dict) -> bool:
|
||||
"""True when this model's on-disk cache is usable (not a truncated download).
|
||||
|
||||
Config-only repos (``config_only: true`` in models.yaml — e.g. pyannote's
|
||||
diarisation pipeline, whose real weights live in referenced sub-repos) carry no
|
||||
weight file of their own, so the weight check would false-positive them as
|
||||
incomplete (#622 caveat). They're exempt: cache presence alone means complete.
|
||||
A weight-bearing repo is complete only if at least one of its snapshots has
|
||||
weights; if no snapshot dir is found on disk we can't prove truncation, so we
|
||||
don't downgrade (the size-based caller already decided it's cached).
|
||||
"""
|
||||
if model.get("config_only"):
|
||||
return True
|
||||
dirs = _snapshot_dirs(model["repo_id"])
|
||||
if not dirs:
|
||||
return True
|
||||
return any(snapshot_has_weights(d) for d in dirs)
|
||||
|
||||
|
||||
def _is_cached_on_disk(repo_id: str) -> bool:
|
||||
"""Direct-filesystem fallback for is_cached when scan_cache_dir is unavailable.
|
||||
|
||||
@@ -286,9 +374,15 @@ def list_models():
|
||||
out = []
|
||||
for m in KNOWN_MODELS:
|
||||
cached = cached_by_repo.get(m["repo_id"])
|
||||
on_disk = cached is not None and cached["size_on_disk"] > 0
|
||||
# A size-positive cache can still be a truncated download (config landed,
|
||||
# weight shard didn't). Treat that as not-installed + incomplete so the
|
||||
# wizard re-offers the download instead of stranding the user (#622).
|
||||
incomplete = on_disk and not cache_is_complete(m)
|
||||
out.append({
|
||||
**m,
|
||||
"installed": cached is not None and cached["size_on_disk"] > 0,
|
||||
"installed": on_disk and not incomplete,
|
||||
"incomplete": incomplete,
|
||||
"size_on_disk_bytes": cached["size_on_disk"] if cached else 0,
|
||||
"nb_files": cached["nb_files"] if cached else 0,
|
||||
"supported": _model_supported(m),
|
||||
@@ -383,6 +477,9 @@ def recommendations():
|
||||
entries = []
|
||||
for rid in recommended_ids:
|
||||
meta = known_by_id.get(rid, {})
|
||||
# Mirror /models: a truncated cache (weights missing) is not installed, so
|
||||
# the wizard counts it toward the remaining download instead of "all set".
|
||||
installed = rid in cached_ids and cache_is_complete(meta or {"repo_id": rid})
|
||||
entries.append({
|
||||
"repo_id": rid,
|
||||
"label": meta.get("label", rid),
|
||||
@@ -390,7 +487,7 @@ def recommendations():
|
||||
"size_gb": meta.get("size_gb", 0),
|
||||
"required": bool(meta.get("required", False)),
|
||||
"note": meta.get("note"),
|
||||
"installed": rid in cached_ids,
|
||||
"installed": installed,
|
||||
})
|
||||
|
||||
to_download_gb = sum(e["size_gb"] for e in entries if not e["installed"])
|
||||
|
||||
@@ -14,6 +14,10 @@
|
||||
# required (optional) — true if the app needs this model to function
|
||||
# platforms (optional) — restrict to specific OS+arch tags (e.g. darwin-arm64, cuda)
|
||||
# note (optional) — shown in the UI as a tooltip/footnote
|
||||
# config_only (optional) — true for pipeline repos that ship no weight file of
|
||||
# their own (weights live in referenced sub-repos). Such
|
||||
# a cache is legitimately tiny, so the truncated-download
|
||||
# (weights-missing) detector must NOT flag it incomplete.
|
||||
# ─────────────────────────────────────────────────────────────────────────
|
||||
|
||||
models:
|
||||
@@ -122,6 +126,7 @@ models:
|
||||
label: "pyannote speaker diarisation (multi-speaker videos)"
|
||||
role: Diarisation
|
||||
size_gb: 0.8
|
||||
config_only: true # pipeline repo; real weights live in referenced sub-repos
|
||||
note: "Needs an HF_TOKEN with license accepted."
|
||||
|
||||
# ── Optional TTS ──────────────────────────────────────────────────────
|
||||
|
||||
@@ -161,3 +161,78 @@ def test_truncated_snapshot_still_rejected(tmp_path):
|
||||
def test_large_tensor_weight_passes(tmp_path):
|
||||
(tmp_path / "model.safetensors").write_bytes(b"\0" * (6 * 1024 * 1024))
|
||||
dl._validate_snapshot_has_weights("x/big", str(tmp_path)) # must not raise
|
||||
|
||||
|
||||
# ── #622: install-state detector is weight-aware (truncated cache ≠ installed) ─
|
||||
|
||||
import api.routers.setup.models as models # noqa: E402
|
||||
|
||||
|
||||
def _make_snapshot(cache_root, repo_id, files):
|
||||
"""Build a minimal HF-style snapshots/<rev>/ dir and return its cache root."""
|
||||
name = "models--" + repo_id.replace("/", "--")
|
||||
rev = cache_root / name / "snapshots" / "abc123"
|
||||
rev.mkdir(parents=True)
|
||||
for fname, data in files.items():
|
||||
(rev / fname).write_bytes(data)
|
||||
return rev
|
||||
|
||||
|
||||
def test_snapshot_has_weights_distinguishes_truncated(tmp_path):
|
||||
full = tmp_path / "full"; full.mkdir()
|
||||
(full / "config.json").write_bytes(b"{}")
|
||||
(full / "model.safetensors").write_bytes(b"\0" * (6 * 1024 * 1024))
|
||||
assert models.snapshot_has_weights(str(full)) is True
|
||||
|
||||
trunc = tmp_path / "trunc"; trunc.mkdir()
|
||||
(trunc / "config.json").write_bytes(b"{}")
|
||||
(trunc / "tokenizer.json").write_bytes(b"x" * 4096)
|
||||
assert models.snapshot_has_weights(str(trunc)) is False
|
||||
|
||||
|
||||
def test_cache_is_complete_flags_truncated_weight_repo(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
|
||||
# Weight-bearing repo with config only (interrupted download) → incomplete.
|
||||
_make_snapshot(tmp_path, "k2-fsa/OmniVoice", {"config.json": b"{}"})
|
||||
assert models.cache_is_complete({"repo_id": "k2-fsa/OmniVoice"}) is False
|
||||
# Same repo once the shard lands → complete.
|
||||
_make_snapshot(
|
||||
tmp_path / "ok", "k2-fsa/OmniVoice",
|
||||
{"config.json": b"{}", "model.safetensors": b"\0" * (6 * 1024 * 1024)},
|
||||
)
|
||||
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "ok"))
|
||||
assert models.cache_is_complete({"repo_id": "k2-fsa/OmniVoice"}) is True
|
||||
|
||||
|
||||
def test_cache_is_complete_exempts_config_only_repo(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
|
||||
# pyannote pipeline ships no weight of its own — a tiny cache is legit, not
|
||||
# truncated; the config_only hint must keep it from being flagged incomplete.
|
||||
_make_snapshot(tmp_path, "pyannote/speaker-diarization-3.1", {"config.yaml": b"x"})
|
||||
assert models.cache_is_complete(
|
||||
{"repo_id": "pyannote/speaker-diarization-3.1", "config_only": True}
|
||||
) is True
|
||||
|
||||
|
||||
def test_list_models_downgrades_truncated_cache(tmp_path, monkeypatch):
|
||||
"""A size-positive but weight-less cache must report installed=False so the
|
||||
first-run wizard re-offers the download instead of stranding the user (#622)."""
|
||||
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
|
||||
_make_snapshot(tmp_path, "k2-fsa/OmniVoice", {"config.json": b"{}"})
|
||||
|
||||
class _Repo:
|
||||
def __init__(self, rid, size):
|
||||
self.repo_id, self.size_on_disk = rid, size
|
||||
self.last_accessed, self.nb_files = 0, 1
|
||||
|
||||
class _Info:
|
||||
repos = [_Repo("k2-fsa/OmniVoice", 4096)] # size > 0 (config landed)
|
||||
|
||||
import huggingface_hub
|
||||
monkeypatch.setattr(huggingface_hub, "scan_cache_dir", lambda: _Info())
|
||||
models.invalidate_cache()
|
||||
out = models.list_models()
|
||||
row = next(m for m in out["models"] if m["repo_id"] == "k2-fsa/OmniVoice")
|
||||
assert row["installed"] is False
|
||||
assert row["incomplete"] is True
|
||||
models.invalidate_cache()
|
||||
|
||||
Reference in New Issue
Block a user