# 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_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