ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/hub/tests/test_model_services.py
Michael Han 74d1a284eb
 Studio: hide the RAG embedder and llama.cpp probe from the hub cached inventory (#7018) 
* Studio: hide infra models from the hub cached inventory

The hub inventory scans behind /api/hub/cached-gguf and /api/hub/cached-models
returned the llama.cpp install validation probe (ggml-org/models) and the RAG
embedder (unsloth/bge-small-en-v1.5[-GGUF]) as on-device models. Share the
hidden-model check from routes/models.py via utils/models/hidden_models.py and
apply it in both scans. A GGUF infra repo stays visible when the user
explicitly downloaded a variant through the Hub, since variant manifests only
exist for user-initiated downloads.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: make On Device trust the hub inventory, match repo ids exactly, lighten the hidden-model import

Follow-up on the hub cached-inventory hidden-model change, addressing the review.

On Device now trusts the Hub inventory API for cached rows. The backend already
hides the RAG embedder and the llama.cpp probe and re-includes a GGUF infra repo
once the user downloads a variant through the Hub, but the frontend was
re-hiding it by repo id, so the user-downloaded variant never appeared in the On
Device list or the count. isVisibleInventoryRow now short-circuits cached rows
(kind === "cache") to visible and keeps client-side needle hiding only for local
filesystem rows and Discover.

is_hidden_model matches Hub repo ids exactly (case-insensitive) against the probe
plus the effective embedder and its GGUF companion, instead of substring
matching the configured-embedder basename. A custom embedder with a generic
basename like org/model no longer hides unrelated cached repos such as
user/model-chat or org/model-instruct. The probe filename and local-path
embedders keep exact matching.

The helper moves to utils/hidden_models.py and is imported at module scope in the
hub cache scanner, so it no longer pulls in utils/models/__init__ (the eager
model-config/checkpoint stack) and a broken import fails at startup instead of
being swallowed per-repo and silently emptying the inventory. routes.models
keeps the _is_hidden_model and _safe_resolve aliases and drops the unused
_HF_REPO_ID_RE re-export that was failing source lint.

Tests: exact repo-id matching with a custom embedder, the cached-models scan
keeping an unrelated repo, and a clean-interpreter check that the helper imports
without the model-config stack.

* Studio: match the llama.cpp probe filename on both path separators

The hidden-model check compared the probe's on-disk filename with
Path(value).name, which on a POSIX interpreter does not split a Windows-style
path ("...\stories260K.gguf") and would let the probe through. Split on both
separators so the probe is matched regardless of which OS produced the path,
matching the tolerance of the previous substring check. Adds a Windows-path
assertion to the probe test.

* Studio: harden hidden infra model handling

* Fix hidden cache row confirmation

* Fix hidden local rows and confirmed hint merges

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Handle snapshot-configured hidden models

* Hide basename-only default embedders

* Fix dynamic embedder inventory filtering

* Studio: hide the configured RAG embedder from Discover and feed rows

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothshared@gmail.com>
Co-authored-by: Daniel Han <23090290+danielhanchen@users.noreply.github.com>
2026-07-19 03:20:56 -07:00

3504 lines
115 KiB
Python
Raw Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
from fastapi import HTTPException
from hub.dependencies import get_hf_token
from hub.storage import scan_folders
from hub.services import download_lifecycle
from hub.services import snapshot_progress
from hub.services.datasets import downloads as dataset_downloads
from hub.services.models import (
cache_inventory,
common as model_common,
deletion,
downloads,
folder_browser,
gguf_variants,
local_inventory,
ollama,
)
from hub.utils import (
download_manifest,
download_registry,
gguf,
hf_cache_state,
inventory_scan,
paths,
state_dir,
)
from hub.workers import hf_download
@pytest.fixture(autouse = True)
def _denylist_inert(monkeypatch):
# The browse tests here exercise allowlist containment, symlink safety and
# the sensitive-name filter, not the system-directory denylist (which has
# its own suite in tests/test_browse_denylist.py). On macOS tmp_path
# resolves under /private/var, a denied prefix, so _resolve_browse_target
# would 403 the fixture dirs before that logic runs. Keep the denylist inert
# so these assertions hold on every platform. folder_browser binds
# is_denied_system_path at import, so patch it on that module, not on
# scan_folders. The "rejects" cases still 403 via the allowlist/sensitive
# checks, and the non-browse tests never call it.
monkeypatch.setattr(folder_browser, "is_denied_system_path", lambda _p: False)
def _repo(repo_id: str, files: list[SimpleNamespace], repo_path: Path):
return SimpleNamespace(
repo_id = repo_id,
repo_type = "model",
repo_path = repo_path,
revisions = [SimpleNamespace(files = files)],
)
def _file(
name: str,
size: int,
blob_path: str | None = None,
):
return SimpleNamespace(file_name = name, size_on_disk = size, blob_path = blob_path)
def _sibling(name: str, size: int, sha: str):
return SimpleNamespace(rfilename = name, size = size, lfs = {"sha256": sha})
class TestExtractQuantToken:
def test_trailing_precision_is_kept(self):
assert gguf.extract_quant_token("model-it-F16.gguf") == "F16"
assert gguf.extract_quant_token("model-BF16.gguf") == "BF16"
def test_real_quant_wins_over_infix_precision(self):
assert gguf.extract_quant_token("Foo-BF16-Q4_K_M.gguf") == "Q4_K_M"
assert gguf.extract_quant_token("Foo-F16-Q8_0.gguf") == "Q8_0"
assert gguf.extract_quant_token("Foo-F32-IQ4_XS.gguf") == "IQ4_XS"
def test_ud_prefix_preserved(self):
assert gguf.extract_quant_token("Foo-BF16-UD-Q4_K_XL.gguf") == "UD-Q4_K_XL"
def test_precision_infix_variants_do_not_collapse(self):
labels = {
gguf.extract_quant_label("Foo-BF16-Q4_K_M.gguf"),
gguf.extract_quant_label("Foo-BF16-Q8_0.gguf"),
}
assert labels == {"Q4_K_M", "Q8_0"}
def test_big_endian_detection_ignores_model_name_be_token():
assert gguf.is_big_endian_gguf_path("model-Q4_K_M-be.gguf", "Q4_K_M")
assert gguf.is_big_endian_gguf_path("model-Q4_K_M_be_infill.gguf", "Q4_K_M")
assert not gguf.is_big_endian_gguf_path("foo-be-Q4_K_M.gguf", "Q4_K_M")
assert not gguf.is_big_endian_gguf_path("Q4_K_M/foo-be.gguf", "Q4_K_M")
assert gguf.pick_best_gguf(["model-Q4_K_M-be.gguf", "model-Q4_K_M.gguf"]) == (
"model-Q4_K_M.gguf"
)
def test_list_local_gguf_variants_skips_big_endian_sibling(tmp_path):
(tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 100)
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"y" * 10)
variants, has_vision = gguf.list_local_gguf_variants(str(tmp_path))
assert has_vision is False
assert [(v.quant, v.filename, v.size_bytes) for v in variants] == [
("Q4_K_M", "model-Q4_K_M.gguf", 10)
]
@pytest.mark.parametrize("repo_id", ["bert-base-uncased", "owner/repo"])
def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id):
assert paths.is_valid_repo_id(repo_id)
def test_repo_id_validation_accepts_max_length_namespaced_repo():
assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}")
@pytest.mark.parametrize(
"repo_id",
[
"datasets/foo/bar",
".repo",
"repo.git",
"foo..bar",
"foo--bar",
"../repo",
"owner/../repo",
],
)
def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id):
assert not paths.is_valid_repo_id(repo_id)
def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M")
assert path is not None
assert path.name == "models--owner--repo--variant--q4_k_m.json"
@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64])
def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
repo_id = f"{'a' * 96}/{'b' * 96}"
assert paths.is_valid_repo_id(repo_id)
assert download_manifest.write_cancel_marker("model", repo_id, variant, "http")
assert download_manifest.write_manifest(
"model",
repo_id,
variant,
[download_manifest.ExpectedFile(path = "model.gguf", size = 1)],
"http",
)
marker_path = state_dir.marker_path("model", repo_id, variant)
manifest_path = state_dir.manifest_path("model", repo_id, variant)
assert marker_path is not None
assert manifest_path is not None
assert "--sha256-" in marker_path.name
assert len(marker_path.name.encode("utf-8")) <= 255
assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255
assert download_manifest.has_cancel_marker("model", repo_id, variant)
assert download_manifest.read_manifest("model", repo_id, variant) is not None
assert list(download_manifest.iter_variant_markers("model", repo_id)) == [
(variant, marker_path)
]
assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [
(variant, manifest_path)
]
class _RecordingLogger:
def __init__(self):
self.warnings = []
def warning(self, *args, **kwargs):
self.warnings.append((args, kwargs))
def test_resolve_browse_target_preserves_allowlist_and_symlink_safety(tmp_path):
home = tmp_path / "home"
scan = tmp_path / "scan"
target = scan / "nested"
home.mkdir()
target.mkdir(parents = True)
(home / "scan-link").symlink_to(scan, target_is_directory = True)
resolved = folder_browser._resolve_browse_target(
str(home / "scan-link" / "nested"),
[home, scan],
)
assert resolved == target.resolve()
def test_resolve_browse_target_rejects_outside_allowlist(tmp_path):
allowed = tmp_path / "allowed"
outside = tmp_path / "outside"
allowed.mkdir()
outside.mkdir()
with pytest.raises(HTTPException) as exc_info:
folder_browser._resolve_browse_target(str(outside), [allowed])
assert exc_info.value.status_code == 403
def test_resolve_browse_target_rejects_sensitive_dir(tmp_path):
home = tmp_path / "home"
ssh = home / ".ssh"
ssh.mkdir(parents = True)
with pytest.raises(HTTPException) as exc_info:
folder_browser._resolve_browse_target(str(ssh), [home])
assert exc_info.value.status_code == 403
def test_resolve_browse_target_rejects_sensitive_root(tmp_path):
ssh = tmp_path / "home" / ".ssh"
ssh.mkdir(parents = True)
with pytest.raises(HTTPException) as exc_info:
folder_browser._resolve_browse_target(str(ssh), [ssh])
assert exc_info.value.status_code == 403
def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
home = tmp_path / "home"
(home / ".ssh").mkdir(parents = True)
(home / "models").mkdir()
# Accept and ignore the optional (media_roots, drive_roots) args the caller now passes.
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home])
response = folder_browser.browse_folders_response(str(home), show_hidden = True)
names = {entry.name for entry in response.entries}
assert "models" in names
assert ".ssh" not in names
def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path):
home = tmp_path / "home"
media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB"
model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6"
home.mkdir()
model_dir.mkdir(parents = True)
monkeypatch.setattr(folder_browser.Path, "home", lambda: home)
monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root])
monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf")
monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: [])
monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: [])
allowlist = folder_browser._build_browse_allowlist()
assert media_root.resolve() in allowlist
assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve()
def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path):
# The endpoint creates the cache dir on demand so the desktop "Open folder"
# action works even before the first download.
target = tmp_path / "hub"
monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)
response = local_inventory.get_models_folder_response()
assert response == {"path": str(target)}
assert target.is_dir()
def test_get_models_folder_response_reports_create_failure(monkeypatch, tmp_path):
target = tmp_path / "hub"
target.write_text("not a directory")
monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)
with pytest.raises(HTTPException) as exc_info:
local_inventory.get_models_folder_response()
assert exc_info.value.status_code == 500
assert "Failed to create models folder" in exc_info.value.detail
def test_get_models_folder_response_requires_directory(monkeypatch, tmp_path):
class MissingPath:
def __init__(self, value: Path):
self.value = value
def mkdir(self, *, parents: bool, exist_ok: bool):
assert parents is True
assert exist_ok is True
def is_dir(self):
return False
def __str__(self):
return str(self.value)
target = MissingPath(tmp_path / "hub")
monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)
with pytest.raises(HTTPException) as exc_info:
local_inventory.get_models_folder_response()
assert exc_info.value.status_code == 500
assert "not a directory" in exc_info.value.detail
def test_contained_link_path_confines_to_link_dir(tmp_path):
link_dir = tmp_path / "ollama" / ".studio_links" / "abc123"
legit = ollama._contained_link_path(link_dir, "llama3-latest-Q4_K_M.gguf")
assert legit == link_dir / "llama3-latest-Q4_K_M.gguf"
for unsafe in (
"",
".",
"..",
"a/b.gguf",
"../evil.gguf",
"/etc/passwd",
"model-tag-../../../pwned.gguf",
):
assert ollama._contained_link_path(link_dir, unsafe) is None
def test_make_ollama_blob_link_refuses_escaping_name(tmp_path):
root = tmp_path / "ollama"
link_dir = root / ".studio_links" / "abc123"
blob = root / "blobs" / "sha256-deadbeef"
blob.parent.mkdir(parents = True)
blob.write_bytes(b"weights")
escaped = ollama._make_ollama_blob_link(link_dir, "model-tag-../../../pwned.gguf", blob)
assert escaped is None
assert not list(tmp_path.rglob("pwned.gguf"))
safe = ollama._make_ollama_blob_link(link_dir, "model-tag.gguf", blob)
assert safe == str(link_dir / "model-tag.gguf")
assert (link_dir / "model-tag.gguf").exists()
def test_cached_gguf_scan_dedupes_and_excludes_mmproj_only(monkeypatch, tmp_path):
smaller = _repo("Org/Dupe", [_file("Q4_K_M.gguf", 100)], tmp_path / "small")
larger = _repo(
"org/dupe",
[_file("Q4_K_M.gguf", 300), _file("Q8_0.gguf", 200)],
tmp_path / "large",
)
mmproj_only = _repo("Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj")
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [smaller, larger, mmproj_only])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: False,
)
result = {"cached": cache_inventory._scan_cached_gguf()}
assert [row["repo_id"] for row in result["cached"]] == ["org/dupe"]
assert result["cached"][0]["size_bytes"] == 500
assert result["cached"][0]["model_format"] == "gguf"
assert result["cached"][0]["capabilities"]["requires_variant"] is True
def test_cached_gguf_scan_preserves_partial_flag(monkeypatch, tmp_path):
partial = _repo("Org/Partial", [_file("Q4_K_M.gguf", 100)], tmp_path / "partial")
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [partial])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: True,
)
result = {"cached": cache_inventory._scan_cached_gguf()}
row = result["cached"][0]
assert row["partial"] is True
assert row["partial_transport"] is None
assert row["capabilities"]["can_chat"] is False
def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
repo_path = tmp_path / "hub" / "models--Org--PartialGguf"
repo_path.mkdir(parents = True)
partial = _repo(
"Org/PartialGguf",
[_file("config.json", 12)],
repo_path,
)
assert download_manifest.write_manifest(
"model",
"Org/PartialGguf",
"Q4_K_M",
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)],
"http",
)
assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [partial])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: True,
)
result = {"cached": cache_inventory._scan_cached_gguf()}
row = result["cached"][0]
assert row["repo_id"] == "Org/PartialGguf"
assert row["model_format"] == "gguf"
assert row["size_bytes"] == 4096
assert row["partial"] is True
assert row["capabilities"]["requires_variant"] is True
def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, tmp_path):
probe = _repo(
"ggml-org/models",
[_file("tinyllamas/stories260K.gguf", 1_200_000)],
tmp_path / "probe",
)
embedder = _repo(
"unsloth/bge-small-en-v1.5-GGUF",
[_file("bge-small-en-v1.5-f16.gguf", 60_000_000)],
tmp_path / "embedder",
)
chat = _repo("Org/Chat-GGUF", [_file("Q4_K_M.gguf", 100)], tmp_path / "chat")
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [probe, embedder, chat])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: False,
)
result = {"cached": cache_inventory._scan_cached_gguf()}
assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat-GGUF"]
def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
embedder = _repo(
"unsloth/bge-small-en-v1.5-GGUF",
[
_file("bge-small-en-v1.5-f16.gguf", 60_000_000),
_file("bge-small-en-v1.5-Q8_0.gguf", 35_000_000),
],
tmp_path / "embedder",
)
# Variant manifests only exist for user Hub downloads, not auto-downloads.
assert download_manifest.write_manifest(
"model",
"unsloth/bge-small-en-v1.5-GGUF",
"Q8_0",
[download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)],
"http",
)
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [embedder])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: False,
)
result = {"cached": cache_inventory._scan_cached_gguf()}
assert [row["repo_id"] for row in result["cached"]] == ["unsloth/bge-small-en-v1.5-GGUF"]
assert result["cached"][0]["capabilities"]["can_chat"] is False
def test_cached_models_scan_hides_non_gguf_embedder(monkeypatch, tmp_path):
embedder_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5"
embedder_path.mkdir(parents = True)
embedder = _repo(
"unsloth/bge-small-en-v1.5",
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
embedder_path,
)
chat_path = tmp_path / "hub" / "models--Org--Chat"
chat_path.mkdir(parents = True)
chat = _repo(
"Org/Chat",
[_file("config.json", 12), _file("model.safetensors", 100)],
chat_path,
)
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [embedder, chat])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda _kind, _repo_id, _path: False,
)
result = {"cached": cache_inventory._scan_cached_models()}
assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat"]
def test_cached_scans_hide_embedders_configured_by_cache_path(monkeypatch, tmp_path):
from core.rag import config as rag_config
gguf_path = tmp_path / "hub" / "models--Org--PathEmbedder-GGUF"
gguf_path.mkdir(parents = True)
gguf = _repo(
"Org/PathEmbedder-GGUF",
[_file("model-F16.gguf", 60_000_000)],
gguf_path,
)
model_path = tmp_path / "hub" / "models--Org--PathEmbedder"
model_path.mkdir(parents = True)
model = _repo(
"Org/PathEmbedder",
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
model_path,
)
monkeypatch.setattr(
rag_config,
"effective_embedding_model",
lambda: str(model_path),
)
monkeypatch.setattr(
rag_config,
"effective_gguf_repo",
lambda: str(gguf_path),
)
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [gguf, model])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: False,
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda _kind, _repo_id, _path: False,
)
assert cache_inventory._scan_cached_gguf() == []
assert cache_inventory._scan_cached_models() == []
def test_cached_scans_hide_embedders_configured_by_snapshot_path(monkeypatch, tmp_path):
from core.rag import config as rag_config
gguf_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder-GGUF"
gguf_snapshot = gguf_path / "snapshots" / "gguf-revision"
gguf_snapshot.mkdir(parents = True)
gguf = _repo(
"Org/SnapshotEmbedder-GGUF",
[_file("model-F16.gguf", 60_000_000)],
gguf_path,
)
model_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder"
model_snapshot = model_path / "snapshots" / "model-revision"
model_snapshot.mkdir(parents = True)
model = _repo(
"Org/SnapshotEmbedder",
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
model_path,
)
monkeypatch.setattr(
rag_config,
"effective_embedding_model",
lambda: str(model_snapshot),
)
monkeypatch.setattr(
rag_config,
"effective_gguf_repo",
lambda: str(gguf_snapshot),
)
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [gguf, model])],
)
def _resolve_snapshot(repo_path):
return str(
{
gguf_path: gguf_snapshot,
model_path: model_snapshot,
}.get(Path(repo_path), Path(repo_path))
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"resolve_hf_cache_realpath",
_resolve_snapshot,
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: False,
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda _kind, _repo_id, _path: False,
)
assert cache_inventory._scan_cached_gguf() == []
assert cache_inventory._scan_cached_models() == []
def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder(
monkeypatch, tmp_path
):
# A custom embedder with a generic basename ("org/model") must be hidden by
# EXACT repo-id match only. An unrelated cached chat model whose id merely
# contains "model" (e.g. "user/model-chat") must stay on device: substring
# basename matching used to drop real chat models from the inventory.
from core.rag import config as rag_config
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
def _model_repo(repo_id: str):
path = tmp_path / "hub" / f"models--{repo_id.replace('/', '--')}"
path.mkdir(parents = True)
return _repo(
repo_id,
[_file("config.json", 12), _file("model.safetensors", 100)],
path,
)
embedder = _model_repo("org/model")
chat = _model_repo("user/model-chat")
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [embedder, chat])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda _kind, _repo_id, _path: False,
)
result = {"cached": cache_inventory._scan_cached_models()}
assert [row["repo_id"] for row in result["cached"]] == ["user/model-chat"]
def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypatch, tmp_path):
from core.rag import config as rag_config
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom")
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF")
gguf = _repo(
"unsloth/bge-small-en-v1.5-GGUF",
[_file("bge-small-en-v1.5-f16.gguf", 60_000_000)],
tmp_path / "default-gguf",
)
weights_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5"
weights_path.mkdir(parents = True)
weights = _repo(
"unsloth/bge-small-en-v1.5",
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
weights_path,
)
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [gguf, weights])],
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda _repo_id, _path: False,
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda _kind, _repo_id, _path: False,
)
assert cache_inventory._scan_cached_gguf() == []
assert cache_inventory._scan_cached_models() == []
def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj():
requirements = gguf_variants._build_gguf_variant_requirements(
[
_sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),
_sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"),
_sibling("mmproj-BF16.gguf", 7, "mm-bf16"),
_sibling("mmproj-F16.gguf", 5, "mm-f16"),
]
)
req = requirements["q4_k_m"]
assert req.main_size_bytes == 30
assert req.download_size_bytes == 35
assert req.main_hashes == frozenset({"main-a", "main-b"})
assert req.required_hashes == frozenset({"main-a", "main-b", "mm-f16"})
assert req.companion_hashes == frozenset({"mm-f16"})
assert req.mmproj_hashes == frozenset({"mm-bf16", "mm-f16"})
assert req.target_filenames == (
"model-Q4_K_M-00001-of-00002.gguf",
"model-Q4_K_M-00002-of-00002.gguf",
"mmproj-F16.gguf",
)
def test_gguf_variant_requirements_skip_big_endian_sibling():
requirements = gguf_variants._build_gguf_variant_requirements(
[
_sibling("model-Q4_K_M-be.gguf", 100, "main-be"),
_sibling("model-Q4_K_M.gguf", 10, "main-le"),
]
)
req = requirements["q4_k_m"]
assert req.main_size_bytes == 10
assert req.main_hashes == frozenset({"main-le"})
assert req.main_filenames == frozenset({"model-Q4_K_M.gguf"})
assert req.target_filenames == ("model-Q4_K_M.gguf",)
def test_worker_gguf_variant_plan_matches_service_requirement(monkeypatch):
siblings = [
_sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),
_sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"),
_sibling("mmproj-BF16.gguf", 7, "mm-bf16"),
_sibling("mmproj-F16.gguf", 5, "mm-f16"),
]
monkeypatch.setattr(
hf_download,
"_model_info_with_retry",
lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings),
)
service_req = gguf_variants._build_gguf_variant_requirements(siblings)["q4_k_m"]
worker_plan = hf_download._gguf_variant_target_plan("Org/Vision", "Q4_K_M", None)
assert worker_plan == service_req
def test_gguf_variant_blob_hashes_accept_dict_lfs_fallback(monkeypatch):
with gguf_variants._VARIANT_HASH_LOCK:
gguf_variants._VARIANT_HASH_CACHE.clear()
gguf_variants._VARIANT_REQUIREMENT_CACHE.clear()
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(
HfApi = lambda *_args, **_kwargs: SimpleNamespace(
model_info = lambda *_a, **_k: SimpleNamespace(
siblings = [
_sibling("model-Q4_K_M.gguf", 10, "main-dict"),
_sibling("model-Q8_0.gguf", 20, "other"),
_sibling("mmproj-F16.gguf", 5, "mmproj"),
]
)
)
),
)
result = gguf_variants.gguf_variant_blob_hashes("Org/DictLfs", "Q4_K_M", None)
main_only = gguf_variants.gguf_variant_blob_hashes(
"Org/DictLfs",
"Q4_K_M",
None,
include_companions = False,
)
assert result == frozenset({"main-dict", "mmproj"})
assert main_only == frozenset({"main-dict"})
def test_gguf_variant_blob_hashes_skip_missing_rfilename(monkeypatch):
with gguf_variants._VARIANT_HASH_LOCK:
gguf_variants._VARIANT_HASH_CACHE.clear()
gguf_variants._VARIANT_REQUIREMENT_CACHE.clear()
siblings = [
SimpleNamespace(rfilename = None, size = 1, lfs = {"sha256": "bad"}),
_sibling("model-Q4_K_M.gguf", 10, "main"),
]
monkeypatch.setattr(
gguf_variants,
"_fetch_gguf_variant_requirements",
lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements(siblings),
)
result = gguf_variants.gguf_variant_blob_hashes("Org/Malformed", "Q4_K_M", None)
assert result == frozenset({"main"})
def test_worker_gguf_variant_targets_skip_missing_rfilename(monkeypatch):
monkeypatch.setattr(
hf_download,
"_model_info_with_retry",
lambda *_args, **_kwargs: SimpleNamespace(
siblings = [
SimpleNamespace(rfilename = None, size = 1),
_sibling("model-Q4_K_M.gguf", 10, "main"),
_sibling("mmproj-F16.gguf", 5, "mm"),
]
),
)
result = hf_download._gguf_variant_target_plan("Org/Malformed", "Q4_K_M", None)
assert list(result.target_filenames) == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_path):
prepare_calls = []
snapshot_calls = []
written = []
verified = []
monkeypatch.setattr(
hf_download,
"_model_info_with_retry",
lambda *_args, **_kwargs: SimpleNamespace(
siblings = [
_sibling("model-Q4_K_M.gguf", 10, "q4-main"),
_sibling("model-Q8_0.gguf", 20, "q8-main"),
_sibling("mmproj-F16.gguf", 5, "shared-mmproj"),
]
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
)
monkeypatch.setattr(
download_registry,
"prepare_cache_for_transport",
lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0,
)
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
),
)
hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http")
assert prepare_calls == [
(
("model", "Org/Vision", "http", "Q4_K_M"),
{
"only_blob_hashes": frozenset({"q4-main"}),
"companion_blob_hashes": frozenset({"shared-mmproj"}),
"protected_blob_hashes": frozenset(),
},
)
]
assert [file.path for file in written[0][3]] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
assert verified == [("model", "Org/Vision", "Q4_K_M", str(tmp_path))]
def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(monkeypatch, tmp_path):
prepare_calls = []
snapshot_calls = []
def _metadata_unavailable(*_args, **_kwargs):
raise RuntimeError("metadata down")
manifest = download_manifest.Manifest(
repo_type = "model",
repo_id = "Org/Vision",
variant = "Q4_K_M",
started_at = "",
expected_files = (
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 10,
sha256 = "q4-main",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 5,
sha256 = "shared-mmproj",
),
),
transport = "http",
)
monkeypatch.setattr(
hf_download,
"_gguf_variant_target_plan",
_metadata_unavailable,
)
monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args: manifest)
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
monkeypatch.setattr(
download_registry,
"prepare_cache_for_transport",
lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0,
)
monkeypatch.setattr(hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
),
)
hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http")
assert prepare_calls == [
(
("model", "Org/Vision", "http", "Q4_K_M"),
{
"only_blob_hashes": frozenset({"q4-main"}),
"companion_blob_hashes": frozenset({"shared-mmproj"}),
"protected_blob_hashes": frozenset(),
},
)
]
assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch, tmp_path):
metadata_calls = []
written = []
cleared = []
verified = []
def _metadata(*_args, **_kwargs):
metadata_calls.append(True)
if len(metadata_calls) == 1:
raise RuntimeError("metadata down")
return SimpleNamespace(siblings = [SimpleNamespace(rfilename = "config.json", size = 12)])
monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
)
monkeypatch.setattr(
download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args)
)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
)
hf_download._download_snapshot("Org/Model", None, "http")
assert len(metadata_calls) == 2
assert cleared == [("model", "Org/Model", None)]
assert written[0][0:3] == ("model", "Org/Model", None)
assert written[0][3][0].path == "config.json"
assert verified == [("model", "Org/Model", None, str(tmp_path))]
def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_path):
metadata_calls = []
snapshot_calls = []
written = []
cleared = []
verified = []
def _metadata(*_args, **_kwargs):
metadata_calls.append(True)
raise RuntimeError("metadata down")
monkeypatch.setattr(hf_download, "_dataset_info_with_retry", _metadata)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
)
monkeypatch.setattr(
download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args)
)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
monkeypatch.setattr(
hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: False
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
),
)
hf_download._download_dataset("Org/Data", None, "http")
assert len(metadata_calls) == 2
assert cleared == [("dataset", "Org/Data", None)]
assert written == []
assert snapshot_calls == [
{
"repo_id": "Org/Data",
"token": False,
"repo_type": "dataset",
"max_workers": 1,
}
]
assert verified == [("dataset", "Org/Data", None, str(tmp_path))]
def test_download_snapshot_fails_when_metadata_unavailable_and_partial_remains(
monkeypatch, tmp_path
):
"""No prior manifest + metadata unavailable + leftover .incomplete blobs means
a cached partial was returned without downloading: the worker must exit 1, not
derive a self-certifying manifest from the finalized subset."""
written = []
verified = []
def _metadata(*_args, **_kwargs):
raise RuntimeError("metadata down")
monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
)
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args, **_kwargs: None)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
monkeypatch.setattr(
hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: True
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
)
with pytest.raises(SystemExit) as excinfo:
hf_download._download_snapshot("Org/Model", None, "http")
assert excinfo.value.code == 1
assert written == []
assert verified == []
def test_purge_repo_cache_dirs_skips_top_level_symlink(monkeypatch, tmp_path):
root = tmp_path / "hub"
target = tmp_path / "target"
root.mkdir()
target.mkdir()
link = root / "models--Org--Repo"
link.symlink_to(target, target_is_directory = True)
monkeypatch.setattr(hf_cache_state, "hf_cache_roots", lambda: [root])
removed = hf_cache_state.purge_repo_cache_dirs("model", "Org/Repo")
assert removed is False
assert link.is_symlink()
assert target.is_dir()
def test_gguf_download_progress_fallback_logs_warning(monkeypatch):
token = "hf_12345678901234567890"
logger = _RecordingLogger()
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
def _raise_permission_error(*_args, **_kwargs):
raise PermissionError(f"denied {token}")
monkeypatch.setattr(snapshot_progress, "logger", logger)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: None,
)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda *_args, **_kwargs: frozenset(),
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
_raise_permission_error,
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model",
variant = "Q4_K_M",
expected_bytes = -1,
hf_token = token,
)
)
assert result == {
"downloaded_bytes": 0,
"completed_bytes": 0,
"complete_on_disk": False,
"expected_bytes": 0,
"progress": 0,
"cache_path": None,
}
assert logger.warnings
args, kwargs = logger.warnings[0]
assert args[:4] == (
"Error checking %s download progress for %s: %s: %s",
"model",
"Org/Model",
"PermissionError",
)
assert token not in args[4]
assert "***" in args[4]
assert kwargs == {}
def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch, tmp_path):
"""A finished mmproj companion keeps counting toward progress once the caller
supplies expected bytes; resolving the variant requirement credits it."""
entry = tmp_path / "models--Org--Model-GGUF"
snap = entry / "snapshots" / "rev0"
blobs = entry / "blobs"
snap.mkdir(parents = True)
blobs.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 30)
(blobs / "mainhash").write_bytes(b"x" * 100)
(blobs / "mmprojhash").write_bytes(b"y" * 30)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
assert download_manifest.write_manifest(
"model",
"Org/Model-GGUF",
"Q4_K_M",
[
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 30,
sha256 = "mmprojhash",
),
],
"http",
)
requirement = gguf_variants._GgufVariantRequirement(
main_filenames = frozenset({"model-Q4_K_M.gguf"}),
target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
main_hashes = frozenset({"mainhash"}),
required_hashes = frozenset({"mainhash", "mmprojhash"}),
companion_hashes = frozenset({"mmprojhash"}),
mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
mmproj_hashes = frozenset({"mmprojhash"}),
expected_files = (
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 30,
sha256 = "mmprojhash",
),
),
main_size_bytes = 100,
download_size_bytes = 130,
)
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: requirement,
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "idle")),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model-GGUF",
variant = "Q4_K_M",
expected_bytes = 130,
)
)
assert result["completed_bytes"] == 130
assert result["downloaded_bytes"] == 130
assert result["complete_on_disk"] is True
assert result["progress"] == 1.0
def test_gguf_progress_subtracts_new_job_completed_baseline(monkeypatch, tmp_path):
entry = tmp_path / "models--Org--Model-GGUF"
snap = entry / "snapshots" / "rev0"
blobs = entry / "blobs"
snap.mkdir(parents = True)
blobs.mkdir(parents = True)
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 30)
(blobs / "mmprojhash").write_bytes(b"y" * 30)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
assert download_manifest.write_manifest(
"model",
"Org/Model-GGUF",
"Q4_K_M",
[
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 30,
sha256 = "mmprojhash",
),
],
"http",
)
requirement = gguf_variants._GgufVariantRequirement(
main_filenames = frozenset({"model-Q4_K_M.gguf"}),
target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
main_hashes = frozenset({"mainhash"}),
required_hashes = frozenset({"mainhash", "mmprojhash"}),
companion_hashes = frozenset({"mmprojhash"}),
mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
mmproj_hashes = frozenset({"mmprojhash"}),
expected_files = (
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 30,
sha256 = "mmprojhash",
),
),
main_size_bytes = 100,
download_size_bytes = 130,
)
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: requirement,
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(
get_job = lambda _key: SimpleNamespace(state = "running"),
get_job_metadata = lambda _key: SimpleNamespace(
completed_baseline_bytes = 30,
),
),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model-GGUF",
variant = "Q4_K_M",
expected_bytes = 130,
)
)
assert result["completed_bytes"] == 0
assert result["downloaded_bytes"] == 0
assert result["expected_bytes"] == 100
assert result["complete_on_disk"] is False
assert result["progress"] == 0
def test_gguf_progress_shows_main_when_companion_left_the_count(monkeypatch, tmp_path):
# The mmproj companion that seeded the baseline is gone, so completed_bytes
# is main-only and below the baseline; it must not be subtracted to 0.
entry = tmp_path / "models--Org--Model-GGUF"
blobs = entry / "blobs"
blobs.mkdir(parents = True)
(blobs / "mainhash").write_bytes(b"x" * 20)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
requirement = gguf_variants._GgufVariantRequirement(
main_filenames = frozenset({"model-Q4_K_M.gguf"}),
target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
main_hashes = frozenset({"mainhash"}),
required_hashes = frozenset({"mainhash", "mmprojhash"}),
companion_hashes = frozenset({"mmprojhash"}),
mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
mmproj_hashes = frozenset({"mmprojhash"}),
expected_files = (
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 30,
sha256 = "mmprojhash",
),
),
main_size_bytes = 100,
download_size_bytes = 130,
)
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: requirement,
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(
get_job = lambda _key: SimpleNamespace(state = "running"),
get_job_metadata = lambda _key: SimpleNamespace(
completed_baseline_bytes = 30,
),
),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model-GGUF",
variant = "Q4_K_M",
expected_bytes = 130,
)
)
assert result["completed_bytes"] == 20
assert result["downloaded_bytes"] == 20
assert result["expected_bytes"] == 130
assert result["complete_on_disk"] is False
def test_gguf_progress_complete_on_disk_ignores_full_baseline(monkeypatch, tmp_path):
# A variant already complete on disk carries a baseline equal to its full
# size; subtracting it would report 0/0 for a finished variant (frontend
# evicts it as gone). Once complete_on_disk is verified, the full figures
# must survive.
entry = tmp_path / "models--Org--Model-GGUF"
snap = entry / "snapshots" / "rev0"
blobs = entry / "blobs"
snap.mkdir(parents = True)
blobs.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 30)
(blobs / "mainhash").write_bytes(b"x" * 100)
(blobs / "mmprojhash").write_bytes(b"y" * 30)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
assert download_manifest.write_manifest(
"model",
"Org/Model-GGUF",
"Q4_K_M",
[
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 30,
sha256 = "mmprojhash",
),
],
"http",
)
requirement = gguf_variants._GgufVariantRequirement(
main_filenames = frozenset({"model-Q4_K_M.gguf"}),
target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
main_hashes = frozenset({"mainhash"}),
required_hashes = frozenset({"mainhash", "mmprojhash"}),
companion_hashes = frozenset({"mmprojhash"}),
mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
mmproj_hashes = frozenset({"mmprojhash"}),
expected_files = (
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
download_manifest.ExpectedFile(
path = "mmproj-F16.gguf",
size = 30,
sha256 = "mmprojhash",
),
),
main_size_bytes = 100,
download_size_bytes = 130,
)
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: requirement,
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(
get_job = lambda _key: SimpleNamespace(state = "running"),
get_job_metadata = lambda _key: SimpleNamespace(
completed_baseline_bytes = 130,
),
),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model-GGUF",
variant = "Q4_K_M",
expected_bytes = 130,
)
)
assert result["complete_on_disk"] is True
assert result["completed_bytes"] == 130
assert result["downloaded_bytes"] == 130
assert result["expected_bytes"] == 130
assert result["progress"] == 1.0
def test_gguf_progress_scoped_hashes_exclude_sibling_quant(monkeypatch, tmp_path):
# The "instant ~900 MB" bug: a sibling quant is fully cached when a different
# variant starts. With this variant's hashes resolved, progress counts ONLY
# its in-progress blob, never the sibling's finalized bytes in the shared
# blobs/ dir.
entry = tmp_path / "models--Org--Model-GGUF"
blobs = entry / "blobs"
blobs.mkdir(parents = True)
(blobs / "siblinghash").write_bytes(b"z" * 900) # other quant, complete
(blobs / "mainhash.incomplete").write_bytes(b"x" * 5) # this variant, started
requirement = gguf_variants._GgufVariantRequirement(
main_filenames = frozenset({"model-Q4_K_M.gguf"}),
target_filenames = ("model-Q4_K_M.gguf",),
main_hashes = frozenset({"mainhash"}),
required_hashes = frozenset({"mainhash"}),
companion_hashes = frozenset(),
mmproj_filenames = frozenset(),
mmproj_hashes = frozenset(),
expected_files = (
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
),
),
main_size_bytes = 100,
download_size_bytes = 100,
)
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: requirement,
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(
get_job = lambda _key: SimpleNamespace(state = "running"),
get_job_metadata = lambda _key: SimpleNamespace(
completed_baseline_bytes = 0,
),
),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model-GGUF",
variant = "Q4_K_M",
expected_bytes = 100,
)
)
assert result["completed_bytes"] == 0
assert result["downloaded_bytes"] == 5
def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(monkeypatch, tmp_path):
# With a variant's hashes unresolved (metadata flaked, no manifest), the
# shared blobs/ dir's FINALIZED blobs must NOT be counted wholesale: a cached
# sibling quant (``siblinghash``) alongside is the "instant ~900 MB" bug.
# With no .incomplete present, downloaded must be 0.
entry = tmp_path / "models--Org--Model-GGUF"
snap = entry / "snapshots" / "rev0"
blobs = entry / "blobs"
snap.mkdir(parents = True)
blobs.mkdir(parents = True)
(blobs / "mainhash").write_bytes(b"x" * 100)
(blobs / "mmprojhash").write_bytes(b"y" * 30)
(blobs / "siblinghash").write_bytes(b"z" * 900)
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: None,
)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda *_args, **_kwargs: frozenset(),
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model-GGUF",
variant = "Q4_K_M",
expected_bytes = 130,
)
)
assert result["completed_bytes"] == 0
assert result["downloaded_bytes"] == 0
assert result["complete_on_disk"] is False
def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(monkeypatch, tmp_path):
# With hashes unresolved, an .incomplete in the shared blobs/ dir can't be
# attributed to this variant (it may be a concurrent sibling's active write),
# so it is dropped, mirroring the finalized-blob guard. In production the
# worker writes the manifest before any .incomplete exists, so hashes resolve
# via the manifest backstop and this window never suppresses real progress.
entry = tmp_path / "models--Org--Model-GGUF"
blobs = entry / "blobs"
blobs.mkdir(parents = True)
(blobs / "activehash.incomplete").write_bytes(b"x" * 50) # unattributable
(blobs / "siblinghash").write_bytes(b"z" * 900) # finalized sibling
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: None,
)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda *_args, **_kwargs: frozenset(),
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"Org/Model-GGUF",
variant = "Q4_K_M",
expected_bytes = 1000,
)
)
assert result["downloaded_bytes"] == 0 # unscoped .incomplete not leaked
assert result["completed_bytes"] == 0 # finalized sibling still ignored
def test_gguf_progress_unknown_hashes_no_backward_dip_when_variant_finalizes(monkeypatch, tmp_path):
# Regression for the two-variant dip: with hashes unresolved, the first quant
# finalizes while the sibling still writes its .incomplete. The sibling's
# bytes used to leak into this numerator, dipping the bar ~99% -> ~78% for
# one poll. The unscoped .incomplete must be dropped so the reading stays 0.
entry = tmp_path / "models--unsloth--SmolLM2-360M-Instruct-GGUF"
blobs = entry / "blobs"
snap = entry / "snapshots" / "rev0"
blobs.mkdir(parents = True)
snap.mkdir(parents = True)
own_total = 218_673_760 # Q2_K finished blob size (denominator)
sibling_total = 234_686_560 # Q3_K_M total
def _sparse_file(path: Path, size: int) -> None:
with path.open("wb") as handle:
handle.truncate(size)
own_finalized = blobs / "q2hash"
_sparse_file(own_finalized, own_total)
# ~72.7% of the sibling => sibling_partial / own_total == 0.78 pre-fix.
sibling_incomplete = blobs / "q3hash.incomplete"
_sparse_file(sibling_incomplete, int(sibling_total * 0.727))
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_requirements",
lambda *_args, **_kwargs: None,
)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda *_args, **_kwargs: frozenset(),
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda *_args, **_kwargs: [entry],
)
monkeypatch.setattr(
downloads,
"_registry",
SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
)
result = asyncio.run(
downloads.get_gguf_download_progress_response(
"unsloth/SmolLM2-360M-Instruct-GGUF",
variant = "Q2_K",
expected_bytes = own_total,
)
)
assert result["downloaded_bytes"] == 0 # sibling .incomplete did not leak
assert result["progress"] == 0 # no ~0.78 backward dip
def test_hf_cache_model_file_probe_is_bounded(monkeypatch, tmp_path):
snapshot = tmp_path / "snapshot"
snapshot.mkdir()
first = tmp_path / "README.md"
second = tmp_path / "notes.txt"
model = tmp_path / "model.safetensors"
first.write_text("readme", encoding = "utf-8")
second.write_text("notes", encoding = "utf-8")
model.write_bytes(b"weights")
entries = [first, second, model]
monkeypatch.setattr(model_common.Path, "rglob", lambda _self, _pattern: iter(entries))
monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 2)
bounded = model_common._iter_hf_cache_model_files(snapshot)
assert bounded == [first, second]
monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 3)
unbounded = model_common._iter_hf_cache_model_files(snapshot)
assert unbounded == [first, second, model]
def test_download_state_lookup_is_repo_case_insensitive(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
assert download_manifest.write_manifest(
"model",
"Owner/Repo",
None,
[download_manifest.ExpectedFile(path = "config.json", size = 12)],
)
assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")
manifest = download_manifest.read_manifest("model", "owner/repo", None)
assert manifest is not None
assert manifest.repo_id == "Owner/Repo"
assert manifest.expected_files[0].path == "config.json"
assert download_manifest.has_cancel_marker("model", "owner/repo", "Q4_K_M")
assert (
download_manifest.read_cancel_marker_transport(
"model",
"owner/repo",
"Q4_K_M",
)
== "http"
)
assert [
variant
for variant, _path in download_manifest.iter_variant_markers(
"model",
"owner/repo",
)
] == ["Q4_K_M"]
assert download_manifest.purge_all_state_for_repo("model", "owner/repo") == 2
assert download_manifest.read_manifest("model", "owner/repo", None) is None
def test_hf_cache_scan_fallback_row_uses_local_model_info_alias(monkeypatch, tmp_path):
cache_dir = tmp_path / "hub"
repo_dir = cache_dir / "models--Org--Broken"
blobs_dir = repo_dir / "blobs"
blobs_dir.mkdir(parents = True)
(blobs_dir / "blob").write_bytes(b"content")
monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"resolve_hf_cache_realpath",
lambda *_args, **_kwargs: None,
)
rows = local_inventory._scan_hf_cache(cache_dir)
assert len(rows) == 1
assert rows[0].model_id == "Org/Broken"
assert rows[0].source == "hf_cache"
assert rows[0].model_format == "unknown"
def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
cache_dir = tmp_path / "hub"
repo_dir = cache_dir / "models--Org--PartialGguf"
blobs_dir = repo_dir / "blobs"
blobs_dir.mkdir(parents = True)
(blobs_dir / "partial").write_bytes(b"content")
assert download_manifest.write_manifest(
"model",
"Org/PartialGguf",
"Q4_K_M",
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 8192)],
"http",
)
assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"is_gguf_repo_partial",
lambda *_args, **_kwargs: True,
)
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"resolve_hf_cache_realpath",
lambda *_args, **_kwargs: None,
)
rows = local_inventory._scan_hf_cache(cache_dir)
assert len(rows) == 1
assert rows[0].model_id == "Org/PartialGguf"
assert rows[0].source == "hf_cache"
assert rows[0].model_format == "gguf"
assert rows[0].partial is True
assert rows[0].size_bytes == 8192
assert rows[0].capabilities.requires_variant is True
def test_local_inventory_filters_custom_embedder_hf_cache_row(monkeypatch, tmp_path):
from core.rag import config as rag_config
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/embedder")
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF")
def _row(repo_id: str):
repo_path = tmp_path / f"models--{repo_id.replace('/', '--')}"
return model_common._local_model_info(
scan_path = repo_path,
load_path = repo_path,
source = "hf_cache",
model_format = "safetensors",
model_id = repo_id,
)
rows = local_inventory._filter_hidden_models([_row("org/embedder"), _row("org/chat-model")])
assert [row.model_id for row in rows] == ["org/chat-model"]
def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatch, tmp_path):
from core.rag import config as rag_config
embedder_path = tmp_path / "hub" / "models--org--embedder"
embedder_snapshot = embedder_path / "snapshots" / "revision"
embedder_snapshot.mkdir(parents = True)
chat_path = tmp_path / "hub" / "models--org--chat-model"
chat_path.mkdir(parents = True)
monkeypatch.setattr(
rag_config,
"effective_embedding_model",
lambda: str(embedder_snapshot),
)
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF")
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"resolve_hf_cache_realpath",
lambda path: str(embedder_snapshot) if Path(path) == embedder_path else str(path),
)
def _row(repo_id: str, repo_path: Path):
return model_common._local_model_info(
scan_path = repo_path,
load_path = repo_path,
source = "hf_cache",
model_format = "safetensors",
model_id = repo_id,
)
rows = local_inventory._filter_hidden_models(
[_row("org/embedder", embedder_path), _row("org/chat-model", chat_path)]
)
assert [row.model_id for row in rows] == ["org/chat-model"]
def test_model_download_job_helpers_preserve_idle_shape():
key = downloads._download_job_key("Org/Model", None)
status = downloads._job_status(key)
assert key == "org/model::"
assert status.state == "idle"
assert status.error is None
def test_gguf_repo_partial_treats_completed_disk_variant_as_clean(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
snapshot = tmp_path / "cache" / "models--Org--Repo" / "snapshots" / "abc"
snapshot.mkdir(parents = True)
(snapshot / "model-Q8_0.gguf").write_bytes(b"complete")
assert download_manifest.write_cancel_marker("model", "Org/Repo", "Q4_K_M", "xet")
monkeypatch.setattr(
inventory_scan,
"resolve_snapshot_dir_for_scan",
lambda *_args: snapshot,
)
assert inventory_scan.is_gguf_repo_partial("Org/Repo", snapshot.parents[1]) is False
def test_gguf_repo_partial_flags_vision_variant_missing_mmproj(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
snapshot = tmp_path / "cache" / "models--Org--Vision" / "snapshots" / "abc"
snapshot.mkdir(parents = True)
(snapshot / "model-Q4_K_M.gguf").write_bytes(b"complete-weight")
assert download_manifest.write_manifest(
"model",
"Org/Vision",
"Q4_K_M",
[
download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 15),
download_manifest.ExpectedFile(path = "mmproj-F16.gguf", size = 8),
],
"http",
)
monkeypatch.setattr(
inventory_scan,
"resolve_snapshot_dir_for_scan",
lambda *_args: snapshot,
)
assert inventory_scan.is_gguf_repo_partial("Org/Vision") is True
def test_cancel_worker_leaves_exited_process_to_watcher():
calls: list = []
class _Registry:
def get_process(self, _key):
return SimpleNamespace(poll = lambda: 1)
def get_job(self, _key):
return SimpleNamespace(state = "running")
def mark_pending_cancel(self, key, generation):
calls.append(("pending", key, generation))
return True
def request_cancel(self, key, proc, generation):
calls.append(("request", key, generation))
return True
def cancel_requested(self, _key):
return False
state = download_lifecycle.cancel_worker(
_Registry(),
"org/model::",
generation = 3,
label = "Org/Model",
logger = SimpleNamespace(warning = lambda *_a, **_k: None),
)
assert state == "running"
assert calls == []
def test_completed_gguf_split_variant_requires_all_shards(tmp_path):
snapshot = tmp_path / "snapshot"
snapshot.mkdir()
first = snapshot / "model-Q8_0-00001-of-00002.gguf"
second = snapshot / "model-Q8_0-00002-of-00002.gguf"
first.write_bytes(b"first")
assert "Q8_0" not in inventory_scan._completed_gguf_variants(snapshot)
second.write_bytes(b"second")
assert "Q8_0" in inventory_scan._completed_gguf_variants(snapshot)
def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
assert inventory_scan.is_variant_partial(
"Org/Repo",
"Q4_K_M",
incomplete_blob_hashes = {"main-q4", "main-q8"},
variant_blob_hashes = frozenset({"main-q4"}),
)
assert not inventory_scan.is_variant_partial(
"Org/Repo",
"Q5_K_M",
incomplete_blob_hashes = {"main-q4"},
variant_blob_hashes = frozenset({"main-q5"}),
)
def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path):
"""A verified GGUF update can prune an older snapshot and make that old
directory the newest by mtime. The variant is still complete when another
snapshot satisfies its manifest."""
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
repo_dir = tmp_path / "cache" / "models--Org--Repo"
old_snapshot = repo_dir / "snapshots" / "old"
new_snapshot = repo_dir / "snapshots" / "new"
old_snapshot.mkdir(parents = True)
new_snapshot.mkdir(parents = True)
(old_snapshot / "model-Q8_0.gguf").write_bytes(b"sibling")
(new_snapshot / "model-Q4_K_M.gguf").write_bytes(b"new")
assert download_manifest.write_manifest(
"model",
"Org/Repo",
"Q4_K_M",
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 3)],
"http",
)
assert not inventory_scan.is_variant_partial(
"Org/Repo",
"Q4_K_M",
snapshot_dir = old_snapshot,
repo_cache_dir = repo_dir,
)
def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline)
assert download_manifest.write_cancel_marker("model", "Org/PartialRepo", "Q4_K_M", "http")
snapshot = tmp_path / "cache" / "models--Org--PartialRepo" / "snapshots" / "rev0"
snapshot.mkdir(parents = True)
(snapshot / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
monkeypatch.setattr(
gguf_variants,
"list_gguf_variants",
lambda *_args, **_kwargs: (
[
SimpleNamespace(
filename = "model-Q4_K_M.gguf",
quant = "Q4_K_M",
display_label = None,
size_bytes = 100,
)
],
False,
None,
),
)
monkeypatch.setattr(
gguf_variants,
"iter_hf_cache_snapshots",
lambda _repo_id: [snapshot],
)
monkeypatch.setattr(
gguf_variants,
"_gguf_all_variant_requirements",
lambda *_args, **_kwargs: {},
)
monkeypatch.setattr(
gguf_variants.download_registry,
"incomplete_blob_hashes",
lambda *_args, **_kwargs: set(),
)
result = asyncio.run(gguf_variants.get_gguf_variants_response("Org/PartialRepo"))
assert result.variants[0].downloaded is False
assert result.variants[0].partial is True
def test_download_registry_repo_keys_are_case_insensitive():
registry = download_registry.DownloadRegistry()
claimed, state = registry.claim(
"Org/Repo::Q8_0",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q8_0",
)
# The same variant under a different-cased repo id resolves to the same
# job, so the second claim attaches to the running one instead of starting
# a duplicate.
duplicate_claimed, duplicate_state = registry.claim(
"org/repo::Q8_0",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "org/repo",
variant = "Q8_0",
)
assert claimed is True
assert state == "running"
assert duplicate_claimed is False
assert duplicate_state == "running"
assert registry.active_jobs("ORG/REPO") == {"org/repo::Q8_0": "running"}
def test_download_registry_allows_disjoint_gguf_variant_downloads():
registry = download_registry.DownloadRegistry()
claimed, state = registry.claim(
"Org/Repo::Q8_0",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q8_0",
blob_hashes = frozenset({"q8-main"}),
progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
)
second_claimed, second_state = registry.claim(
"Org/Repo::Q4_K_M",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q4_K_M",
blob_hashes = frozenset({"q4-main"}),
progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
)
assert claimed is True
assert state == "running"
assert second_claimed is True
assert second_state == "running"
assert registry.active_jobs("org/repo") == {
"org/repo::Q8_0": "running",
"org/repo::Q4_K_M": "running",
}
def test_download_registry_allows_overlapping_same_transport_variant_downloads():
# Two variants sharing one mmproj blob still download together on one
# transport: huggingface_hub's per-blob lock serializes the shared write and
# prepare_cache_for_transport never purges a blob a peer is writing.
registry = download_registry.DownloadRegistry()
claimed, state = registry.claim(
"Org/Repo::Q8_0",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q8_0",
blob_hashes = frozenset({"q8-main"}),
progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
)
second_claimed, second_state = registry.claim(
"Org/Repo::Q4_K_M",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q4_K_M",
blob_hashes = frozenset({"q4-main"}),
progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
)
assert claimed is True
assert state == "running"
assert second_claimed is True
assert second_state == "running"
def test_download_registry_variant_delete_does_not_block_sibling_download():
# Deleting one quant's partial must be allowed while a different quant of the
# same repo is downloading, and must protect every blob the live sibling is
# writing (including a shared mmproj companion).
registry = download_registry.DownloadRegistry()
registry.claim(
"Org/Repo::Q8_0",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q8_0",
blob_hashes = frozenset({"q8-main"}),
progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
)
# A sibling variant delete is allowed; deleting the in-flight variant is not.
assert registry.begin_delete("Org/Repo", "Q4_K_M") is True
assert registry.begin_delete("Org/Repo", "Q8_0") is False
# A whole-repo delete still waits for every active download.
assert registry.begin_delete("Org/Repo") is False
# The live sibling is detected so the delete keeps the shared companion.
assert registry.has_active_peer_variant("Org/Repo", "Q4_K_M") is True
assert registry.has_active_peer_variant("Org/Repo", "Q8_0") is False
# While Q4_K_M is being deleted, re-downloading it is blocked but an
# untouched third variant may still start.
blocked, blocked_state = registry.claim(
"Org/Repo::Q4_K_M",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q4_K_M",
)
assert blocked is False
assert blocked_state == "deleting"
started, started_state = registry.claim(
"Org/Repo::Q5_K_M",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q5_K_M",
)
assert started is True
assert started_state == "running"
registry.end_delete("Org/Repo", "Q4_K_M")
assert registry.begin_delete("Org/Repo", "Q4_K_M") is True
def test_partial_gguf_reconstruction_dedupes_variant_casing(monkeypatch):
# The manifest keeps original casing while the marker is lowercased; offline
# reconstruction must collapse them to ONE entry (manifest's casing), not two.
monkeypatch.setattr(
download_manifest,
"iter_variant_manifests",
lambda _repo_type, _repo_id: iter([("Q4_K_M", Path("manifest.json"))]),
)
monkeypatch.setattr(
download_manifest,
"iter_variant_markers",
lambda _repo_type, _repo_id: iter([("q4_k_m", Path("marker.json"))]),
)
monkeypatch.setattr(download_manifest, "read_manifest", lambda *_a, **_k: None)
result = gguf.list_partial_gguf_variants_from_state("Org/Repo")
assert result is not None
variants, _has_vision = result
assert [variant.quant for variant in variants] == ["Q4_K_M"]
def test_download_registry_serializes_cross_transport_variant_downloads():
# An HTTP append-resume and an XET rewrite of the same shared blob would
# corrupt each other, so different-transport variants are serialized.
registry = download_registry.DownloadRegistry()
claimed, state = registry.claim(
"Org/Repo::Q8_0",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q8_0",
blob_hashes = frozenset({"q8-main"}),
progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
)
second_claimed, second_state = registry.claim(
"Org/Repo::Q4_K_M",
download_registry.TRANSPORT_XET,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q4_K_M",
blob_hashes = frozenset({"q4-main"}),
progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
)
assert claimed is True
assert state == "running"
assert second_claimed is False
assert second_state == "running"
def test_download_registry_allows_unknown_hash_gguf_variant_downloads():
# Resolved blob hashes are NOT required to run two same-transport variants
# concurrently: on-disk safety comes from each worker purging only its own
# main-quant blobs plus huggingface_hub's per-etag lock. Requiring them here
# used to reject the second variant whenever a metadata fetch flaked.
registry = download_registry.DownloadRegistry()
claimed, state = registry.claim(
"Org/Repo::Q8_0",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q8_0",
)
second_claimed, second_state = registry.claim(
"Org/Repo::Q4_K_M",
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q4_K_M",
blob_hashes = frozenset({"q4-main"}),
progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
)
assert claimed is True
assert state == "running"
assert second_claimed is True
assert second_state == "running"
assert registry.active_jobs("org/repo") == {
"org/repo::Q8_0": "running",
"org/repo::Q4_K_M": "running",
}
def test_finalize_worker_exit_never_kills_a_healthy_worker(monkeypatch, tmp_path):
# finalize_worker_exit relies solely on the worker's exit code and never kills
# a live process: huggingface_hub already bounds reads with timeouts, so a
# liveness kill could only false-cancel a healthy download.
import inspect
import io
import logging
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
class _Proc:
pid = 4242
def __init__(self):
self.killed = False
self.stderr = io.BytesIO(b"")
def poll(self):
return 0
def wait(self, timeout = None):
return 0
def kill(self):
self.killed = True
registry = download_registry.DownloadRegistry()
proc = _Proc()
key = "Org/Repo::Q4_K_M"
registry.claim(
key,
download_registry.TRANSPORT_HTTP,
repo_type = "model",
repo_id = "Org/Repo",
variant = "Q4_K_M",
)
registry.register_process(key, proc)
download_lifecycle.finalize_worker_exit(
registry,
key,
proc,
hf_token = None,
label = "Org/Repo [Q4_K_M]",
log_prefix = "Download",
logger = logging.getLogger("test"),
repo_type = "model",
repo_id = "Org/Repo",
transport = "http",
)
assert proc.killed is False
assert registry.get_job(key).state == "complete"
# The stall-watchdog knob is gone entirely; no caller may re-enable it.
assert (
"enable_stall_watchdog"
not in inspect.signature(download_lifecycle.finalize_worker_exit).parameters
)
def test_prepare_cache_for_transport_purges_only_requested_hashes(monkeypatch, tmp_path):
root = tmp_path / "hub"
blobs = root / "models--Org--Repo" / "blobs"
blobs.mkdir(parents = True)
(blobs / "variant-main.incomplete").write_bytes(b"x")
(blobs / "shared-mmproj.incomplete").write_bytes(b"y")
monkeypatch.setattr(download_registry, "hf_cache_root", lambda create = False: root)
purged = download_registry.prepare_cache_for_transport(
"model",
"Org/Repo",
download_registry.TRANSPORT_XET,
"Q4_K_M",
frozenset({"variant-main"}),
)
assert purged == 1
assert not (blobs / "variant-main.incomplete").exists()
assert (blobs / "shared-mmproj.incomplete").exists()
def _vision_cache_root(monkeypatch, tmp_path):
root = tmp_path / "hub"
blobs = root / "models--Org--Vision" / "blobs"
blobs.mkdir(parents = True)
monkeypatch.setattr(download_registry, "hf_cache_root", lambda create = False: root)
return blobs
def test_prepare_cache_for_transport_purges_cross_transport_companion(monkeypatch, tmp_path):
blobs = _vision_cache_root(monkeypatch, tmp_path)
companion = frozenset({"shared-mmproj"})
# An interrupted XET download stamps the companion marker "xet" and leaves a
# sparse partial. A later HTTP download of a different variant must purge it,
# else the HTTP resumer appends to the sparse bytes and corrupts the blob.
download_registry.prepare_cache_for_transport(
"model",
"Org/Vision",
download_registry.TRANSPORT_XET,
"Q4_K_M",
only_blob_hashes = frozenset({"q4-main"}),
companion_blob_hashes = companion,
)
(blobs / "shared-mmproj.incomplete").write_bytes(b"sparse")
purged = download_registry.prepare_cache_for_transport(
"model",
"Org/Vision",
download_registry.TRANSPORT_HTTP,
"Q8_0",
only_blob_hashes = frozenset({"q8-main"}),
companion_blob_hashes = companion,
)
assert purged == 1
assert not (blobs / "shared-mmproj.incomplete").exists()
def test_prepare_cache_for_transport_preserves_same_transport_companion(monkeypatch, tmp_path):
blobs = _vision_cache_root(monkeypatch, tmp_path)
companion = frozenset({"shared-mmproj"})
download_registry.prepare_cache_for_transport(
"model",
"Org/Vision",
download_registry.TRANSPORT_HTTP,
"Q4_K_M",
only_blob_hashes = frozenset({"q4-main"}),
companion_blob_hashes = companion,
)
(blobs / "shared-mmproj.incomplete").write_bytes(b"resumable")
purged = download_registry.prepare_cache_for_transport(
"model",
"Org/Vision",
download_registry.TRANSPORT_HTTP,
"Q4_K_M",
only_blob_hashes = frozenset({"q4-main"}),
companion_blob_hashes = companion,
)
assert purged == 0
assert (blobs / "shared-mmproj.incomplete").exists()
def test_prepare_cache_for_transport_protects_peer_companion(monkeypatch, tmp_path):
blobs = _vision_cache_root(monkeypatch, tmp_path)
companion = frozenset({"shared-mmproj"})
download_registry.prepare_cache_for_transport(
"model",
"Org/Vision",
download_registry.TRANSPORT_XET,
"Q4_K_M",
only_blob_hashes = frozenset({"q4-main"}),
companion_blob_hashes = companion,
)
(blobs / "shared-mmproj.incomplete").write_bytes(b"sparse")
purged = download_registry.prepare_cache_for_transport(
"model",
"Org/Vision",
download_registry.TRANSPORT_HTTP,
"Q8_0",
only_blob_hashes = frozenset({"q8-main"}),
companion_blob_hashes = companion,
protected_blob_hashes = companion,
)
assert purged == 0
assert (blobs / "shared-mmproj.incomplete").exists()
def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypatch, tmp_path):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, repo_type = "model": repo_id,
)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: (
frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"})
),
)
monkeypatch.setattr(
downloads.download_registry,
"completed_blob_bytes",
lambda *_args, **_kwargs: 30,
)
class _Registry:
claim_kwargs = None
def claim(self, _key, _transport, **kwargs):
self.claim_kwargs = kwargs
return True, "running"
def current_generation(self, _key):
return 1
def get_job(self, _key):
return SimpleNamespace(state = "running")
def register_process(self, _key, _proc):
return False
def peer_blob_hashes(self, _key):
return frozenset()
class _Proc:
pid = 123
stderr = None
def poll(self):
return None
def kill(self):
return None
def wait(self, timeout = None):
return 0
registry = _Registry()
monkeypatch.setattr(downloads, "_registry", registry)
monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc())
asyncio.run(
downloads.download_model_response(
SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False)
)
)
assert registry.claim_kwargs["blob_hashes"] == frozenset({"mainhash"})
assert registry.claim_kwargs["progress_blob_hashes"] == frozenset({"mainhash", "mmprojhash"})
assert registry.claim_kwargs["completed_baseline_bytes"] == 30
def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state(
monkeypatch, tmp_path
):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
assert download_manifest.write_manifest(
"model",
"Org/Model",
"Q4_K_M",
[
download_manifest.ExpectedFile(
path = "model-Q4_K_M.gguf",
size = 100,
sha256 = "mainhash",
)
],
"http",
)
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, repo_type = "model": repo_id,
)
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: (
frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"})
),
)
monkeypatch.setattr(
downloads.download_registry,
"completed_blob_bytes",
lambda *_args, **_kwargs: 30,
)
class _Registry:
claim_kwargs = None
def claim(self, _key, _transport, **kwargs):
self.claim_kwargs = kwargs
return True, "running"
def current_generation(self, _key):
return 1
def get_job(self, _key):
return SimpleNamespace(state = "running")
def register_process(self, _key, _proc):
return False
def peer_blob_hashes(self, _key):
return frozenset()
class _Proc:
pid = 123
stderr = None
def poll(self):
return None
def kill(self):
return None
def wait(self, timeout = None):
return 0
registry = _Registry()
monkeypatch.setattr(downloads, "_registry", registry)
monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc())
asyncio.run(
downloads.download_model_response(
SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False)
)
)
assert registry.claim_kwargs["completed_baseline_bytes"] == 0
def test_model_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry())
assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")
status = asyncio.run(downloads.get_download_status_response("owner/repo", "Q4_K_M"))
assert status.state == "cancelled"
assert status.error is None
def test_shutdown_kills_all_workers_before_shared_deadline_reap(monkeypatch):
events = []
now = [100.0]
class _Proc:
def __init__(self, name):
self.name = name
def poll(self):
return None
def kill(self):
events.append(("kill", self.name))
def wait(self, timeout):
events.append(("wait", self.name, timeout))
now[0] += 7.0
registry = download_registry.DownloadRegistry()
proc_a = _Proc("a")
proc_b = _Proc("b")
registry.claim(
"Org/A",
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "Org/A",
)
registry.claim(
"Org/B",
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "Org/B",
)
assert registry.register_process("org/a", proc_a)
assert registry.register_process("org/b", proc_b)
monkeypatch.setattr(
download_registry,
"persist_cancel_marker",
lambda *args, **kwargs: events.append(("marker", args[1])),
)
monkeypatch.setattr(download_registry.time, "monotonic", lambda: now[0])
registry.terminate_all("dataset download")
assert events == [
("kill", "a"),
("kill", "b"),
("wait", "a", 10.0),
("marker", "Org/A"),
("wait", "b", 3.0),
("marker", "Org/B"),
]
def test_shutdown_skips_marker_for_worker_that_exits_cleanly(monkeypatch):
markers = []
class _Proc:
def __init__(self, final_rc):
self._final_rc = final_rc
self._exited = False
def poll(self):
return self._final_rc if self._exited else None
def kill(self):
pass
def wait(self, timeout):
self._exited = True
registry = download_registry.DownloadRegistry()
clean = _Proc(0)
interrupted = _Proc(-9)
registry.claim(
"Org/Clean",
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "Org/Clean",
)
registry.claim(
"Org/Cut",
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "Org/Cut",
)
assert registry.register_process("org/clean", clean)
assert registry.register_process("org/cut", interrupted)
monkeypatch.setattr(
download_registry,
"persist_cancel_marker",
lambda *args, **kwargs: markers.append(args[1]),
)
registry.terminate_all("dataset download")
assert markers == ["Org/Cut"]
def test_model_claim_register_cancel_uses_registry_marker_owner(monkeypatch):
killed = []
class _Registry:
def claim(self, *_args, **_kwargs):
return True, "running"
def current_generation(self, _key):
return 1
def register_process(self, _key, _proc):
return False
def persist_cancel_for_key(self, *_args, **_kwargs):
raise AssertionError("register_process owns pending-cancel markers")
def get_job(self, _key):
return SimpleNamespace(state = "cancelled", error = None)
monkeypatch.setattr(downloads, "_registry", _Registry())
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, **_kwargs: repo_id,
)
monkeypatch.setattr(
downloads.download_registry,
"download_transport_unavailable_reason",
lambda _transport: None,
)
monkeypatch.setattr(
downloads,
"_spawn_download_worker",
lambda *_args, **_kwargs: object(),
)
monkeypatch.setattr(
downloads.download_lifecycle,
"kill_and_reap_process",
lambda proc, **_kwargs: killed.append(proc),
)
result = asyncio.run(
downloads.download_model_response(
SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False)
)
)
assert result["state"] == "cancelled"
assert killed
def test_model_cancel_registered_worker_requests_and_kills(monkeypatch):
events = []
class _Proc:
def poll(self):
return None
def kill(self):
events.append(("kill",))
class _Registry:
def get_process(self, _key):
return _Proc()
def request_cancel(self, key, _proc, generation):
events.append(("request", key, generation))
return True
def persist_cancel_for_key(self, *_args, **_kwargs):
raise AssertionError(
"cancel_worker must leave marker persistence to the exit watcher; "
"an eager persist races a clean completion and strands a stale marker"
)
def get_job(self, _key):
return SimpleNamespace(state = "running")
monkeypatch.setattr(downloads, "_registry", _Registry())
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, **_kwargs: repo_id,
)
result = asyncio.run(
downloads.cancel_download_model_response(
SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", generation = 7)
)
)
assert result == {
"job_key": downloads._download_job_key("Org/Model", "Q4_K_M"),
"state": "cancelling",
}
assert events == [
("request", downloads._download_job_key("Org/Model", "Q4_K_M"), 7),
("kill",),
]
def test_model_download_watcher_invalidates_hf_cache_scan(monkeypatch):
invalidated = []
class _Registry:
def claim(self, *_args, **_kwargs):
return True, "running"
def current_generation(self, _key):
return 1
def register_process(self, _key, _proc):
return True
def get_job(self, _key):
return SimpleNamespace(state = "complete", error = None)
class _ImmediateThread:
def __init__(self, *, target, **_kwargs):
self._target = target
def start(self):
self._target()
monkeypatch.setattr(downloads, "_registry", _Registry())
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, **_kwargs: repo_id,
)
monkeypatch.setattr(
downloads.download_registry,
"download_transport_unavailable_reason",
lambda _transport: None,
)
monkeypatch.setattr(
downloads.download_lifecycle,
"finalize_worker_exit",
lambda *_args, **_kwargs: None,
)
monkeypatch.setattr(
downloads,
"_spawn_download_worker",
lambda *_args, **_kwargs: object(),
)
monkeypatch.setattr(downloads.download_lifecycle.threading, "Thread", _ImmediateThread)
monkeypatch.setattr(
downloads.hf_cache_scan,
"invalidate_hf_cache_scans",
lambda: invalidated.append(True),
)
async def _inline_to_thread(func, *args, **kwargs):
return func(*args, **kwargs)
monkeypatch.setattr(downloads.asyncio, "to_thread", _inline_to_thread)
result = asyncio.run(
downloads.download_model_response(
SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False)
)
)
assert result["accepted"] is True
assert invalidated == [True]
def test_two_concurrent_same_repo_variants_both_complete(monkeypatch, tmp_path):
# End-to-end proof that two GGUF variants of ONE repo download concurrently
# without cancelling each other, with real registry/finalize/subprocess/watch
# threads exercising the claim gate, register, finalize funnel, and
# classify_exit under true concurrency.
import subprocess
import time
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(
downloads,
"_registry",
download_registry.DownloadRegistry(),
)
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, **_k: repo_id,
)
monkeypatch.setattr(
downloads.download_registry,
"download_transport_unavailable_reason",
lambda _transport: None,
)
# Per-variant blob hashes (distinct main shard, shared mmproj companion).
monkeypatch.setattr(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda _repo, variant, _token = None, include_companions = True, **_k: (
frozenset({f"{variant.lower()}-main", "shared-mmproj"})
if include_companions
else frozenset({f"{variant.lower()}-main"})
),
)
monkeypatch.setattr(
downloads.download_registry,
"completed_blob_bytes",
lambda *_a, **_k: 0,
)
monkeypatch.setattr(
downloads.hf_cache_scan,
"invalidate_hf_cache_scans",
lambda: None,
)
# Real subprocess that exits 0 immediately, with a stderr pipe to drain.
spawned: list[subprocess.Popen] = []
def _fake_spawn(*_args, **_kwargs):
proc = subprocess.Popen(
[sys.executable, "-c", "import sys; sys.exit(0)"],
stderr = subprocess.PIPE,
)
spawned.append(proc)
return proc
monkeypatch.setattr(downloads, "_spawn_download_worker", _fake_spawn)
async def _run_both():
return await asyncio.gather(
downloads.download_model_response(
SimpleNamespace(
repo_id = "Org/Model",
gguf_variant = "Q4_K_M",
use_xet = False,
)
),
downloads.download_model_response(
SimpleNamespace(
repo_id = "Org/Model",
gguf_variant = "Q8_0",
use_xet = False,
)
),
)
results = asyncio.run(_run_both())
assert all(r["accepted"] is True for r in results), results
registry = downloads._registry
key_q4 = downloads._download_job_key("Org/Model", "Q4_K_M")
key_q8 = downloads._download_job_key("Org/Model", "Q8_0")
deadline = time.monotonic() + 10.0
while time.monotonic() < deadline:
s4 = registry.get_job(key_q4).state
s8 = registry.get_job(key_q8).state
if s4 in download_registry.TERMINAL_STATES and s8 in download_registry.TERMINAL_STATES:
break
time.sleep(0.02)
for p in spawned:
try:
p.wait(timeout = 5)
except Exception:
pass
assert registry.get_job(key_q4).state == "complete"
assert registry.get_job(key_q8).state == "complete"
def test_download_registry_factories_reuse_service_singletons():
registry_module = downloads.download_registry
before_count = len(registry_module._REGISTRIES)
assert registry_module.get_models_registry() is downloads.registry
assert registry_module.get_models_registry() is downloads.registry
assert registry_module.get_datasets_registry() is dataset_downloads.registry
assert registry_module.get_datasets_registry() is dataset_downloads.registry
assert len(registry_module._REGISTRIES) == before_count
def test_hub_hf_token_header_uses_namespaced_header_only():
assert get_hf_token("new-token") == "new-token"
assert get_hf_token(None) is None
def test_scan_folder_rejects_credential_directories(tmp_path):
sensitive_dir = tmp_path / ".ssh" / "models"
sensitive_dir.mkdir(parents = True)
with pytest.raises(ValueError, match = "Credential or configuration"):
scan_folders.add_scan_folder(str(sensitive_dir))
def _build_variant_cache_repo(repo_dir, blob_specs, snapshot_links):
"""Build a HF cache repo dir with blobs + snapshot symlinks for the
per-variant deletion path. blob_specs: {blob_name: bytes_payload};
snapshot_links: list of (revision, filename, blob_name)."""
blobs_dir = repo_dir / "blobs"
blobs_dir.mkdir(parents = True)
for blob_name, payload in blob_specs.items():
(blobs_dir / blob_name).write_bytes(payload)
files = []
for revision, filename, blob_name in snapshot_links:
snap_dir = repo_dir / "snapshots" / revision
snap_dir.mkdir(parents = True, exist_ok = True)
blob = blobs_dir / blob_name
link = snap_dir / filename
link.symlink_to(blob)
files.append(
SimpleNamespace(
file_name = filename,
file_path = str(link),
blob_path = str(blob),
size_on_disk = blob.stat().st_size,
)
)
repo = SimpleNamespace(
repo_id = "Org/Repo-GGUF",
repo_type = "model",
repo_path = repo_dir,
revisions = [SimpleNamespace(commit_hash = "rev1", files = files)],
)
return repo
def _patch_variant_delete_side_effects(monkeypatch):
monkeypatch.setattr(
deletion.download_manifest,
"purge_state",
lambda *_args, **_kwargs: False,
)
def test_snapshot_progress_filters_stale_blobs(monkeypatch, tmp_path):
"""Exclude superseded-revision blobs; count an in-progress blob only when its
hash belongs to the target."""
entry = tmp_path / "datasets--Org--Data"
blobs = entry / "blobs"
blobs.mkdir(parents = True)
(blobs / "keep1").write_bytes(b"a" * 100)
(blobs / "stale").write_bytes(b"b" * 500)
(blobs / "keep2.incomplete").write_bytes(b"c" * 40)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda _repo_type, _repo_id, force_active = False: [entry],
)
result = snapshot_progress.compute_snapshot_progress(
repo_type = "dataset",
repo_id = "Org/Data",
job_key = "org/data",
expected_bytes = 0,
hf_token = None,
registry = SimpleNamespace(
get_job = lambda _key: SimpleNamespace(state = "running"),
),
metadata_resolver = lambda _repo_id, _hf_token: (
140,
frozenset({"keep1", "keep2"}),
),
)
assert result["completed_bytes"] == 100
assert result["downloaded_bytes"] == 140
assert result["complete_on_disk"] is False
assert result["expected_bytes"] == 140
def test_snapshot_progress_confirms_complete_only_with_verified_snapshot(monkeypatch, tmp_path):
entry = tmp_path / "models--Org--Model"
blobs = entry / "blobs"
snap = entry / "snapshots" / "rev0"
blobs.mkdir(parents = True)
snap.mkdir(parents = True)
(blobs / "keep1").write_bytes(b"a" * 100)
(snap / "model.safetensors").write_bytes(b"a" * 100)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda _repo_type, _repo_id, force_active = False: [entry],
)
monkeypatch.setattr(
snapshot_progress.download_manifest,
"has_cancel_marker",
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
snapshot_progress.download_manifest,
"read_manifest",
lambda *_args, **_kwargs: SimpleNamespace(),
)
monkeypatch.setattr(
snapshot_progress.download_manifest,
"verify_against_disk",
lambda *_args, **_kwargs: SimpleNamespace(ok = True),
)
result = snapshot_progress.compute_snapshot_progress(
repo_type = "model",
repo_id = "Org/Model",
job_key = "org/model::",
expected_bytes = 100,
hf_token = None,
registry = SimpleNamespace(
get_job = lambda _key: SimpleNamespace(state = "idle"),
),
metadata_resolver = lambda _repo_id, _hf_token: (
100,
frozenset({"keep1"}),
),
)
assert result["completed_bytes"] == 100
assert result["complete_on_disk"] is True
def test_expected_files_from_snapshot_dir_records_relative_paths_and_sizes(tmp_path):
snap = tmp_path / "snapshots" / "rev0"
(snap / "nested").mkdir(parents = True)
(snap / "model.safetensors").write_bytes(b"a" * 12)
(snap / "nested" / "config.json").write_bytes(b"b" * 3)
files = download_manifest.expected_files_from_snapshot_dir(snap)
by_path = {f.path: f for f in files}
assert by_path["model.safetensors"].size == 12
assert by_path["nested/config.json"].size == 3
assert all(f.sha256 is None for f in files)
def test_snapshot_progress_complete_with_manifest_synthesized_from_disk(monkeypatch, tmp_path):
"""A finished snapshot whose only manifest was synthesized from on-disk files
still verifies as complete, so a refresh finalizes it instead of capping at
99% and evicting it as gone."""
entry = tmp_path / "models--Org--Model"
blobs = entry / "blobs"
snap = entry / "snapshots" / "rev0"
blobs.mkdir(parents = True)
snap.mkdir(parents = True)
(blobs / "keep1").write_bytes(b"a" * 100)
(snap / "model.safetensors").write_bytes(b"a" * 100)
synthesized = download_manifest.expected_files_from_snapshot_dir(snap)
manifest = download_manifest.Manifest(
repo_type = "model",
repo_id = "Org/Model",
variant = None,
started_at = "",
expected_files = tuple(synthesized),
)
monkeypatch.setattr(
snapshot_progress,
"preferred_repo_cache_dirs",
lambda _repo_type, _repo_id, force_active = False: [entry],
)
monkeypatch.setattr(
snapshot_progress.download_manifest,
"has_cancel_marker",
lambda *_args, **_kwargs: False,
)
monkeypatch.setattr(
snapshot_progress.download_manifest,
"read_manifest",
lambda *_args, **_kwargs: manifest,
)
result = snapshot_progress.compute_snapshot_progress(
repo_type = "model",
repo_id = "Org/Model",
job_key = "org/model::",
expected_bytes = 100,
hf_token = None,
registry = SimpleNamespace(
get_job = lambda _key: SimpleNamespace(state = "idle"),
),
metadata_resolver = lambda _repo_id, _hf_token: (
100,
frozenset({"keep1"}),
),
)
assert result["complete_on_disk"] is True
assert result["progress"] == 1.0
def test_delete_variant_keeps_blob_shared_with_other_snapshot(monkeypatch, tmp_path):
"""A blob still referenced by a non-target snapshot symlink survives so that
symlink doesn't dangle (which the scanner reports as partial)."""
repo_dir = tmp_path / "models--Org--Repo-GGUF"
repo = _build_variant_cache_repo(
repo_dir,
blob_specs = {"sharedblob": b"x" * 200, "q8blob": b"y" * 300},
snapshot_links = [
("rev1", "model-Q4_K_M.gguf", "sharedblob"),
("rev1", "model-Q8_0.gguf", "q8blob"),
# An unrelated file that happens to share Q4's blob content.
("rev1", "extra-copy.gguf", "sharedblob"),
],
)
monkeypatch.setattr(
deletion.cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
_patch_variant_delete_side_effects(monkeypatch)
result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)
assert result["status"] == "deleted"
# Q4 snapshot link gone, but its blob survives (extra-copy still links it).
assert not (repo_dir / "snapshots" / "rev1" / "model-Q4_K_M.gguf").exists()
assert (repo_dir / "blobs" / "sharedblob").exists()
extra = repo_dir / "snapshots" / "rev1" / "extra-copy.gguf"
assert extra.is_symlink() and extra.exists() # not dangling
def test_delete_variant_unlinks_unshared_blob(monkeypatch, tmp_path):
repo_dir = tmp_path / "models--Org--Repo-GGUF"
repo = _build_variant_cache_repo(
repo_dir,
blob_specs = {"q4blob": b"x" * 200, "q8blob": b"y" * 300},
snapshot_links = [
("rev1", "model-Q4_K_M.gguf", "q4blob"),
("rev1", "model-Q8_0.gguf", "q8blob"),
],
)
monkeypatch.setattr(
deletion.cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
_patch_variant_delete_side_effects(monkeypatch)
result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)
assert result["status"] == "deleted"
assert not (repo_dir / "blobs" / "q4blob").exists()
# Untouched sibling variant remains fully intact.
assert (repo_dir / "blobs" / "q8blob").exists()
q8 = repo_dir / "snapshots" / "rev1" / "model-Q8_0.gguf"
assert q8.is_symlink() and q8.exists()
def test_delete_variant_surfaces_locked_file_as_conflict(monkeypatch, tmp_path):
"""A blob unlink that fails (e.g. a Windows file lock on a loaded model)
must raise a clear 409, not report a misleading success."""
repo_dir = tmp_path / "models--Org--Repo-GGUF"
repo = _build_variant_cache_repo(
repo_dir,
blob_specs = {"lockedblob": b"x" * 200},
snapshot_links = [("rev1", "model-Q4_K_M.gguf", "lockedblob")],
)
monkeypatch.setattr(
deletion.cache_inventory,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
_patch_variant_delete_side_effects(monkeypatch)
real_unlink = Path.unlink
def fake_unlink(self, *args, **kwargs):
if self.name == "lockedblob":
raise PermissionError("file in use")
return real_unlink(self, *args, **kwargs)
monkeypatch.setattr(Path, "unlink", fake_unlink)
with pytest.raises(HTTPException) as exc_info:
deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)
assert exc_info.value.status_code == 409
def test_download_snapshot_writes_manifest_for_xet(monkeypatch, tmp_path):
written = []
verified = []
monkeypatch.setattr(
hf_download,
"_model_info_with_retry",
lambda *_args, **_kwargs: SimpleNamespace(
siblings = [SimpleNamespace(rfilename = "config.json", size = 12)]
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
)
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
)
hf_download._download_snapshot("Org/Model", None, "xet")
assert written, "XET snapshot download must still record a manifest"
assert written[0][0:3] == ("model", "Org/Model", None)
assert written[0][3][0].path == "config.json"
assert verified == [("model", "Org/Model", None, str(tmp_path))]
def test_download_gguf_variant_writes_manifest_for_xet(monkeypatch, tmp_path):
written = []
verified = []
monkeypatch.setattr(
hf_download,
"_model_info_with_retry",
lambda *_args, **_kwargs: SimpleNamespace(
siblings = [_sibling("model-Q4_K_M.gguf", 10, "main")]
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
)
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
)
hf_download._download_gguf_variant("Org/Model", "Q4_K_M", None, "xet")
assert written, "XET GGUF variant download must still record a manifest"
assert written[0][0:3] == ("model", "Org/Model", "Q4_K_M")
assert written[0][3][0].path == "model-Q4_K_M.gguf"
assert verified == [("model", "Org/Model", "Q4_K_M", str(tmp_path))]
def test_download_dataset_writes_manifest_for_xet(monkeypatch, tmp_path):
written = []
verified = []
monkeypatch.setattr(
hf_download,
"_dataset_info_with_retry",
lambda *_args, **_kwargs: SimpleNamespace(
siblings = [SimpleNamespace(rfilename = "data.parquet", size = 30)]
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
)
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
)
hf_download._download_dataset("Org/Data", None, "xet")
assert written, "XET dataset download must still record a manifest"
assert written[0][0:3] == ("dataset", "Org/Data", None)
assert written[0][3][0].path == "data.parquet"
assert verified == [("dataset", "Org/Data", None, str(tmp_path))]
def test_dataset_status_includes_generation(monkeypatch):
class _Registry:
def get_job(self, _key):
return SimpleNamespace(state = "running", error = None)
def current_generation(self, _key):
return 4
monkeypatch.setattr(dataset_downloads, "_registry", _Registry())
monkeypatch.setattr(
dataset_downloads,
"resolve_cached_repo_id_case",
lambda repo_id, **_kwargs: repo_id,
)
result = asyncio.run(dataset_downloads.get_dataset_download_status_response("Org/Data"))
assert result.state == "running"
assert result.generation == 4
Reference in a new issue View git blame Copy permalink