vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29
unsloth/studio/backend/hub/tests/test_model_services.py
pre-commit-ci[bot] 4ac010b73a [pre-commit.ci] auto fixes from pre-commit.com hooks 
for more information, see https://pre-commit.ci
2026-07-09 05:50:47 +00:00

3244 lines
106 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
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()
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda: [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"]
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_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_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