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
ramisworld 01f7e14988
 Fix Studio custom folders on Linux external drives (#6799) 
* Fix external drive custom folder selection

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

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

* Update studio/backend/tests/test_linux_external_media_paths.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

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

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

* Keep legacy media scan validation strict

* Apply sensitive-dir denylist to legacy folder browser for PR #6799

The legacy /api/models browse endpoint gained the new /run/media mount
roots in its allowlist but not the credential/config guard that scan-folder
registration and the Hub browser already enforce. Filter sensitive names
during enumeration and reject them in _resolve_browse_target so .ssh, .aws,
.config, etc. under allowlisted roots stay unbrowseable, matching the Hub
browser. Add a public contains_sensitive_path_component helper and cover the
legacy resolver with a regression test.

* Trim redundant comments in PR #6799 changes

* Skip sensitive Linux media roots

* Reject sensitive dirs at exact browse roots for PR #6799

Both _resolve_browse_target functions only checked contains_sensitive_path_component
while walking descendant parts, so requesting an allowlisted root itself (empty
relative path) returned it unchecked. A pre-existing scan-folder row under ~/.ssh,
~/.aws, ~/.config, etc. (registerable before the denylist was added) is re-added to
the allowlist on upgrade and could then be browsed. Check the resolved target once
before returning in both the legacy and Hub browsers, and cover the root case in
both test suites.

* fix: avoid unused path helper reexports

* fix: import sensitive path helpers directly

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-07-03 19:10:04 +01:00

3105 lines
102 KiB
Python
Raw Permalink 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)
@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)
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_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