d6d8b8b1a6
Seven fixes from the latest review round on the Images page and the hub cache inventory. Images page: - The lost-POST settle path built its "already seen" gallery id set inside the catch, after the request failed. By then the earlier runs of the same batch had already prepended their records, so run 2 could accept run 1's image as proof that its own request reached the backend. The set is now captured once before the first POST and grows with every record the batch produces. - settleLostGeneration fell out of its SETTLE_MAX_MS loop and returned normally, so a wedged generation was counted as done and the next run started against a busy backend. It now throws on timeout. - Restoring a recipe cleared the ControlNet selection but left the workflow tab and the init / mask / reference images pointing at whatever was loaded, so the next Generate conditioned on an unrelated image. It now clears all of them and returns to Create. - The download plan omitted the adapter selection the load itself bakes in. A baked LoRA forces the dense build path, so the plan described a different file set than the load that followed and the rest was pulled inline, outside the download manager. Both now derive the list from one helper. Hub cache inventory: - A download for a repo an Images or Video load is staging was allowed to start: only the llama.cpp loader was consulted. Both diffusion backends already expose loading_repo_ids for the delete guard, and the download guard now reads them too. - A companion-only prefetch (pipeline manifest plus VAE and text encoder, no transformer) passed the snapshot-partial check, since every file its manifest expected did arrive, and was advertised as on-device although from_pretrained cannot load it. - The single-file flag never reached the picker through the hub inventory path, so a checkpoint-only diffusion repo read as a full pipeline and failed after the handoff. The two pipeline-shape helpers now live in hub/utils/inventory_scan.py so /api/models/cached and the hub inventory classify the same repos the same way.
4077 lines
135 KiB
Python
4077 lines
135 KiB
Python
# 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 json
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from hub.dependencies import get_hf_token
|
|
from hub.storage import scan_folders
|
|
from hub.services import download_lifecycle
|
|
from hub.services import snapshot_progress
|
|
from hub.services.datasets import downloads as dataset_downloads
|
|
from hub.services.models import (
|
|
cache_inventory,
|
|
common as model_common,
|
|
deletion,
|
|
downloads,
|
|
folder_browser,
|
|
gguf_variants,
|
|
local_inventory,
|
|
ollama,
|
|
)
|
|
from hub.utils import (
|
|
download_manifest,
|
|
download_registry,
|
|
gguf,
|
|
hf_cache_state,
|
|
inventory_scan,
|
|
paths,
|
|
state_dir,
|
|
)
|
|
from hub.workers import hf_download
|
|
|
|
|
|
@pytest.fixture(autouse = True)
|
|
def _denylist_inert(monkeypatch):
|
|
# The browse tests here exercise allowlist containment, symlink safety and
|
|
# the sensitive-name filter, not the system-directory denylist (which has
|
|
# its own suite in tests/test_browse_denylist.py). On macOS tmp_path
|
|
# resolves under /private/var, a denied prefix, so _resolve_browse_target
|
|
# would 403 the fixture dirs before that logic runs. Keep the denylist inert
|
|
# so these assertions hold on every platform. folder_browser binds
|
|
# is_denied_system_path at import, so patch it on that module, not on
|
|
# scan_folders. The "rejects" cases still 403 via the allowlist/sensitive
|
|
# checks, and the non-browse tests never call it.
|
|
monkeypatch.setattr(folder_browser, "is_denied_system_path", lambda _p: False)
|
|
|
|
|
|
def _download_body(**over) -> SimpleNamespace:
|
|
"""A download request with every field the route reads.
|
|
|
|
Built by hand rather than through the schema so these tests stay cheap, so it lists
|
|
the newer scoping fields too: leaving one off reads as AttributeError inside the
|
|
route, not as a missing default.
|
|
"""
|
|
body = {
|
|
"repo_id": "Org/Model",
|
|
"gguf_variant": None,
|
|
"use_xet": False,
|
|
"scope_id": None,
|
|
"files": None,
|
|
}
|
|
body.update(over)
|
|
return SimpleNamespace(**body)
|
|
|
|
|
|
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 _cached_model_row(tmp_path: Path, *, partial: bool, active_cache: bool | None, size_bytes: int):
|
|
path = tmp_path / f"cache-{active_cache}-{partial}-{size_bytes}"
|
|
return model_common._local_model_info(
|
|
scan_path = path,
|
|
load_path = path,
|
|
source = "hf_cache",
|
|
model_format = "safetensors",
|
|
model_id = "Org/Model",
|
|
partial = partial,
|
|
active_cache = active_cache,
|
|
size_bytes = size_bytes,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("reverse", [False, True])
|
|
def test_local_inventory_prefers_complete_previous_cache_copy(tmp_path, reverse):
|
|
active_partial = _cached_model_row(
|
|
tmp_path,
|
|
partial = True,
|
|
active_cache = True,
|
|
size_bytes = 20,
|
|
)
|
|
previous_complete = _cached_model_row(
|
|
tmp_path,
|
|
partial = False,
|
|
active_cache = False,
|
|
size_bytes = 10,
|
|
)
|
|
rows = [active_partial, previous_complete]
|
|
if reverse:
|
|
rows.reverse()
|
|
|
|
result = local_inventory._dedupe_local_models(rows)
|
|
|
|
assert result == [previous_complete]
|
|
|
|
|
|
def test_local_inventory_compares_all_non_active_cache_copies(tmp_path):
|
|
inactive_partial = _cached_model_row(
|
|
tmp_path,
|
|
partial = True,
|
|
active_cache = False,
|
|
size_bytes = 20,
|
|
)
|
|
custom_complete = _cached_model_row(
|
|
tmp_path,
|
|
partial = False,
|
|
active_cache = None,
|
|
size_bytes = 10,
|
|
)
|
|
|
|
assert local_inventory._dedupe_local_models([inactive_partial, custom_complete]) == [
|
|
custom_complete
|
|
]
|
|
|
|
|
|
def test_local_inventory_prefers_active_cache_when_copies_are_equally_complete(tmp_path):
|
|
previous = _cached_model_row(
|
|
tmp_path,
|
|
partial = False,
|
|
active_cache = False,
|
|
size_bytes = 20,
|
|
)
|
|
active = _cached_model_row(
|
|
tmp_path,
|
|
partial = False,
|
|
active_cache = True,
|
|
size_bytes = 10,
|
|
)
|
|
|
|
assert local_inventory._dedupe_local_models([previous, active]) == [active]
|
|
|
|
|
|
def test_loaded_repo_match_accepts_previous_cache_snapshot_path(monkeypatch, tmp_path):
|
|
repo_dir = tmp_path / "old-hub" / "models--Org--Model"
|
|
snapshot = repo_dir / "snapshots" / "revision"
|
|
snapshot.mkdir(parents = True)
|
|
monkeypatch.setattr(deletion, "iter_repo_cache_dirs", lambda *_args: iter([repo_dir]))
|
|
|
|
assert deletion._loaded_id_matches_repo(str(snapshot), "Org/Model") is True
|
|
assert deletion._loaded_id_matches_repo(str(snapshot / "model.gguf"), "Org/Model") is True
|
|
assert deletion._loaded_id_matches_repo(str(tmp_path / "other"), "Org/Model") is False
|
|
|
|
|
|
def test_cached_inventory_loads_previous_cache_copy_by_snapshot(monkeypatch, tmp_path):
|
|
active_hub = tmp_path / "active-hub"
|
|
previous_repo = tmp_path / "previous-hub" / "models--Org--Model"
|
|
snapshot = previous_repo / "snapshots" / "revision"
|
|
snapshot.mkdir(parents = True)
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = active_hub),
|
|
)
|
|
|
|
fields = cache_inventory._cache_inventory_fields(
|
|
"Org/Model",
|
|
"safetensors",
|
|
repo_path = previous_repo,
|
|
snapshot_path = snapshot,
|
|
)
|
|
|
|
assert fields["load_id"] == str(snapshot)
|
|
|
|
|
|
def test_cached_inventory_keeps_repo_id_for_active_cache(monkeypatch, tmp_path):
|
|
active_hub = tmp_path / "active-hub"
|
|
active_repo = active_hub / "models--Org--Model"
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = active_hub),
|
|
)
|
|
|
|
fields = cache_inventory._cache_inventory_fields(
|
|
"Org/Model",
|
|
"safetensors",
|
|
repo_path = active_repo,
|
|
)
|
|
|
|
assert fields["load_id"] == "Org/Model"
|
|
|
|
|
|
def test_cached_inventory_prefers_active_copy_when_completeness_matches():
|
|
previous = {"partial": False, "active_cache": False, "size_bytes": 200}
|
|
active = {"partial": False, "active_cache": True, "size_bytes": 100}
|
|
|
|
assert cache_inventory._prefer_cache_row(active, previous) is True
|
|
assert cache_inventory._prefer_cache_row(previous, active) is False
|
|
|
|
|
|
def test_cached_inventory_prefers_complete_copy_before_active_cache():
|
|
previous = {"partial": False, "active_cache": False, "size_bytes": 100}
|
|
active_partial = {"partial": True, "active_cache": True, "size_bytes": 200}
|
|
|
|
assert cache_inventory._prefer_cache_row(previous, active_partial) is True
|
|
assert cache_inventory._prefer_cache_row(active_partial, previous) is False
|
|
|
|
|
|
def test_inventory_scans_every_dynamic_cache_root(monkeypatch, tmp_path):
|
|
first = tmp_path / "first-hub"
|
|
second = tmp_path / "second-hub"
|
|
unreadable = tmp_path / "unreadable-hub"
|
|
first.mkdir()
|
|
second.mkdir()
|
|
unreadable.mkdir()
|
|
scanned = []
|
|
|
|
monkeypatch.setattr(
|
|
inventory_scan,
|
|
"hf_cache_roots",
|
|
lambda: [first, unreadable, second],
|
|
)
|
|
|
|
def scan_cache(cache_dir):
|
|
path = Path(cache_dir)
|
|
scanned.append(path)
|
|
if path == unreadable:
|
|
raise PermissionError("unreadable")
|
|
return SimpleNamespace(cache_dir = cache_dir)
|
|
|
|
monkeypatch.setattr("huggingface_hub.scan_cache_dir", scan_cache)
|
|
|
|
result = inventory_scan._compute_all_hf_cache_scans()
|
|
|
|
assert scanned == [first, unreadable, second]
|
|
assert [Path(scan.cache_dir) for scan in result] == [first, second]
|
|
|
|
|
|
def test_inventory_applies_download_state_to_its_owning_cache(monkeypatch, tmp_path):
|
|
state_root = tmp_path / "state"
|
|
cache_a = tmp_path / "cache-a"
|
|
cache_b = tmp_path / "cache-b"
|
|
repo_id = "Org/Model"
|
|
repo_name = "models--Org--Model"
|
|
repo_a = cache_a / repo_name
|
|
repo_b = cache_b / repo_name
|
|
snapshot_a = repo_a / "snapshots" / "revision"
|
|
snapshot_b = repo_b / "snapshots" / "revision"
|
|
snapshot_a.mkdir(parents = True)
|
|
snapshot_b.mkdir(parents = True)
|
|
(snapshot_a / "config.json").write_bytes(b"x")
|
|
(snapshot_b / "config.json").write_bytes(b"xx")
|
|
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: state_root)
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = cache_b),
|
|
)
|
|
assert download_manifest.write_manifest(
|
|
"model",
|
|
repo_id,
|
|
None,
|
|
[download_manifest.ExpectedFile(path = "config.json", size = 2)],
|
|
"http",
|
|
hub_cache = cache_a,
|
|
)
|
|
|
|
assert inventory_scan.is_snapshot_partial("model", repo_id, repo_a) is True
|
|
assert inventory_scan.is_snapshot_partial("model", repo_id, repo_b) is False
|
|
assert inventory_scan.partial_transport_for("model", repo_id, None, repo_a) == "http"
|
|
assert inventory_scan.partial_transport_for("model", repo_id, None, repo_b) is None
|
|
|
|
|
|
def test_inventory_scopes_cancel_markers_to_their_owning_cache(monkeypatch, tmp_path):
|
|
state_root = tmp_path / "state"
|
|
cache_a = tmp_path / "cache-a"
|
|
cache_b = tmp_path / "cache-b"
|
|
repo_id = "Org/Model"
|
|
repo_name = "models--Org--Model"
|
|
repo_a = cache_a / repo_name
|
|
repo_b = cache_b / repo_name
|
|
repo_a.mkdir(parents = True)
|
|
repo_b.mkdir(parents = True)
|
|
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: state_root)
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = cache_b),
|
|
)
|
|
assert download_manifest.write_cancel_marker(
|
|
"model",
|
|
repo_id,
|
|
"Q4_K_M",
|
|
"xet",
|
|
hub_cache = cache_a,
|
|
)
|
|
|
|
assert inventory_scan.is_variant_partial(repo_id, "Q4_K_M", repo_cache_dir = repo_a) is True
|
|
assert inventory_scan.is_variant_partial(repo_id, "Q4_K_M", repo_cache_dir = repo_b) is False
|
|
|
|
|
|
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",
|
|
)
|
|
|
|
hub_cache = download_manifest._canonical_hub_cache()
|
|
marker_path = state_dir.marker_path(
|
|
"model",
|
|
repo_id,
|
|
variant,
|
|
hub_cache = hub_cache,
|
|
)
|
|
manifest_path = state_dir.manifest_path(
|
|
"model",
|
|
repo_id,
|
|
variant,
|
|
hub_cache = hub_cache,
|
|
)
|
|
|
|
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)
|
|
]
|
|
|
|
|
|
def test_download_state_isolated_across_hub_cache_switches(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
|
|
cache_a = tmp_path / "cache-a"
|
|
cache_b = tmp_path / "cache-b"
|
|
selected = SimpleNamespace(hub_cache = cache_a)
|
|
|
|
from utils import hf_cache_settings
|
|
|
|
monkeypatch.setattr(hf_cache_settings, "get_hf_cache_paths", lambda: selected)
|
|
expected_a = [download_manifest.ExpectedFile(path = "a.gguf", size = 1)]
|
|
expected_b = [download_manifest.ExpectedFile(path = "b.gguf", size = 2)]
|
|
|
|
assert download_manifest.write_manifest("model", "Owner/Repo", "Q4_K_M", expected_a)
|
|
assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")
|
|
|
|
selected.hub_cache = cache_b
|
|
assert download_manifest.write_manifest("model", "Owner/Repo", "Q4_K_M", expected_b)
|
|
|
|
manifest_b = download_manifest.read_manifest("model", "Owner/Repo", "Q4_K_M")
|
|
manifest_a = download_manifest.read_manifest(
|
|
"model",
|
|
"Owner/Repo",
|
|
"Q4_K_M",
|
|
hub_cache = cache_a,
|
|
)
|
|
|
|
assert manifest_b is not None and manifest_b.expected_files == tuple(expected_b)
|
|
assert manifest_a is not None and manifest_a.expected_files == tuple(expected_a)
|
|
assert not download_manifest.has_cancel_marker("model", "Owner/Repo", "Q4_K_M")
|
|
assert download_manifest.has_cancel_marker(
|
|
"model",
|
|
"Owner/Repo",
|
|
"Q4_K_M",
|
|
hub_cache = cache_a,
|
|
)
|
|
assert len(list((tmp_path / "hub-state" / "manifests").rglob("*.json"))) == 2
|
|
|
|
|
|
def test_legacy_unscoped_download_state_falls_back_only_for_selected_cache(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
|
|
cache_a = tmp_path / "cache-a"
|
|
cache_b = tmp_path / "cache-b"
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = cache_a),
|
|
)
|
|
manifest = state_dir.manifest_path("model", "Owner/Repo", "Q4_K_M")
|
|
marker = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M")
|
|
assert manifest is not None and marker is not None
|
|
manifest.write_text(
|
|
json.dumps(
|
|
{
|
|
"version": 1,
|
|
"repo_id": "Owner/Repo",
|
|
"variant": "Q4_K_M",
|
|
"expected_files": [{"path": "model.gguf", "size": 10}],
|
|
"transport": "http",
|
|
}
|
|
),
|
|
encoding = "utf-8",
|
|
)
|
|
marker.write_text(
|
|
json.dumps({"version": 1, "repo_id": "Owner/Repo", "variant": "Q4_K_M"}),
|
|
encoding = "utf-8",
|
|
)
|
|
|
|
assert download_manifest.read_manifest("model", "Owner/Repo", "Q4_K_M") is not None
|
|
assert download_manifest.has_cancel_marker("model", "Owner/Repo", "Q4_K_M")
|
|
assert list(download_manifest.iter_variant_manifests("model", "Owner/Repo")) == [
|
|
("Q4_K_M", manifest)
|
|
]
|
|
assert list(download_manifest.iter_variant_markers("model", "Owner/Repo")) == [
|
|
("Q4_K_M", marker)
|
|
]
|
|
assert (
|
|
download_manifest.read_manifest(
|
|
"model",
|
|
"Owner/Repo",
|
|
"Q4_K_M",
|
|
hub_cache = cache_b,
|
|
)
|
|
is None
|
|
)
|
|
assert not download_manifest.has_cancel_marker(
|
|
"model",
|
|
"Owner/Repo",
|
|
"Q4_K_M",
|
|
hub_cache = cache_b,
|
|
)
|
|
|
|
|
|
class _RecordingLogger:
|
|
def __init__(self):
|
|
self.warnings = []
|
|
|
|
def warning(self, *args, **kwargs):
|
|
self.warnings.append((args, kwargs))
|
|
|
|
|
|
def test_resolve_browse_target_preserves_allowlist_and_symlink_safety(tmp_path):
|
|
home = tmp_path / "home"
|
|
scan = tmp_path / "scan"
|
|
target = scan / "nested"
|
|
home.mkdir()
|
|
target.mkdir(parents = True)
|
|
(home / "scan-link").symlink_to(scan, target_is_directory = True)
|
|
|
|
resolved = folder_browser._resolve_browse_target(
|
|
str(home / "scan-link" / "nested"),
|
|
[home, scan],
|
|
)
|
|
|
|
assert resolved == target.resolve()
|
|
|
|
|
|
def test_resolve_browse_target_rejects_outside_allowlist(tmp_path):
|
|
allowed = tmp_path / "allowed"
|
|
outside = tmp_path / "outside"
|
|
allowed.mkdir()
|
|
outside.mkdir()
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
folder_browser._resolve_browse_target(str(outside), [allowed])
|
|
|
|
assert exc_info.value.status_code == 403
|
|
|
|
|
|
def test_resolve_browse_target_rejects_sensitive_dir(tmp_path):
|
|
home = tmp_path / "home"
|
|
ssh = home / ".ssh"
|
|
ssh.mkdir(parents = True)
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
folder_browser._resolve_browse_target(str(ssh), [home])
|
|
|
|
assert exc_info.value.status_code == 403
|
|
|
|
|
|
def test_resolve_browse_target_rejects_sensitive_root(tmp_path):
|
|
ssh = tmp_path / "home" / ".ssh"
|
|
ssh.mkdir(parents = True)
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
folder_browser._resolve_browse_target(str(ssh), [ssh])
|
|
|
|
assert exc_info.value.status_code == 403
|
|
|
|
|
|
def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
|
|
home = tmp_path / "home"
|
|
(home / ".ssh").mkdir(parents = True)
|
|
(home / "models").mkdir()
|
|
# Accept and ignore the optional (media_roots, drive_roots) args the caller now passes.
|
|
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home])
|
|
|
|
response = folder_browser.browse_folders_response(str(home), show_hidden = True)
|
|
|
|
names = {entry.name for entry in response.entries}
|
|
assert "models" in names
|
|
assert ".ssh" not in names
|
|
|
|
|
|
def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path):
|
|
home = tmp_path / "home"
|
|
media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB"
|
|
model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6"
|
|
home.mkdir()
|
|
model_dir.mkdir(parents = True)
|
|
monkeypatch.setattr(folder_browser.Path, "home", lambda: home)
|
|
monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root])
|
|
monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf")
|
|
monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: [])
|
|
monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: [])
|
|
|
|
allowlist = folder_browser._build_browse_allowlist()
|
|
|
|
assert media_root.resolve() in allowlist
|
|
assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve()
|
|
|
|
|
|
def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path):
|
|
# The endpoint creates the cache dir on demand so the desktop "Open folder"
|
|
# action works even before the first download.
|
|
target = tmp_path / "hub"
|
|
monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)
|
|
|
|
response = local_inventory.get_models_folder_response()
|
|
|
|
assert response == {"path": str(target)}
|
|
assert target.is_dir()
|
|
|
|
|
|
def test_get_models_folder_response_reports_create_failure(monkeypatch, tmp_path):
|
|
target = tmp_path / "hub"
|
|
target.write_text("not a directory")
|
|
monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
local_inventory.get_models_folder_response()
|
|
|
|
assert exc_info.value.status_code == 500
|
|
assert "Failed to create models folder" in exc_info.value.detail
|
|
|
|
|
|
def test_get_models_folder_response_requires_directory(monkeypatch, tmp_path):
|
|
class MissingPath:
|
|
def __init__(self, value: Path):
|
|
self.value = value
|
|
|
|
def mkdir(self, *, parents: bool, exist_ok: bool):
|
|
assert parents is True
|
|
assert exist_ok is True
|
|
|
|
def is_dir(self):
|
|
return False
|
|
|
|
def __str__(self):
|
|
return str(self.value)
|
|
|
|
target = MissingPath(tmp_path / "hub")
|
|
monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
local_inventory.get_models_folder_response()
|
|
|
|
assert exc_info.value.status_code == 500
|
|
assert "not a directory" in exc_info.value.detail
|
|
|
|
|
|
def test_contained_link_path_confines_to_link_dir(tmp_path):
|
|
link_dir = tmp_path / "ollama" / ".studio_links" / "abc123"
|
|
|
|
legit = ollama._contained_link_path(link_dir, "llama3-latest-Q4_K_M.gguf")
|
|
assert legit == link_dir / "llama3-latest-Q4_K_M.gguf"
|
|
|
|
for unsafe in (
|
|
"",
|
|
".",
|
|
"..",
|
|
"a/b.gguf",
|
|
"../evil.gguf",
|
|
"/etc/passwd",
|
|
"model-tag-../../../pwned.gguf",
|
|
):
|
|
assert ollama._contained_link_path(link_dir, unsafe) is None
|
|
|
|
|
|
def test_make_ollama_blob_link_refuses_escaping_name(tmp_path):
|
|
root = tmp_path / "ollama"
|
|
link_dir = root / ".studio_links" / "abc123"
|
|
blob = root / "blobs" / "sha256-deadbeef"
|
|
blob.parent.mkdir(parents = True)
|
|
blob.write_bytes(b"weights")
|
|
|
|
escaped = ollama._make_ollama_blob_link(link_dir, "model-tag-../../../pwned.gguf", blob)
|
|
assert escaped is None
|
|
assert not list(tmp_path.rglob("pwned.gguf"))
|
|
|
|
safe = ollama._make_ollama_blob_link(link_dir, "model-tag.gguf", blob)
|
|
assert safe == str(link_dir / "model-tag.gguf")
|
|
assert (link_dir / "model-tag.gguf").exists()
|
|
|
|
|
|
def test_cached_gguf_scan_dedupes_and_excludes_mmproj_only(monkeypatch, tmp_path):
|
|
smaller = _repo("Org/Dupe", [_file("Q4_K_M.gguf", 100)], tmp_path / "small")
|
|
larger = _repo(
|
|
"org/dupe",
|
|
[_file("Q4_K_M.gguf", 300), _file("Q8_0.gguf", 200)],
|
|
tmp_path / "large",
|
|
)
|
|
mmproj_only = _repo("Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj")
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [smaller, larger, mmproj_only])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: False,
|
|
)
|
|
|
|
result = {"cached": cache_inventory._scan_cached_gguf()}
|
|
|
|
assert [row["repo_id"] for row in result["cached"]] == ["org/dupe"]
|
|
assert result["cached"][0]["size_bytes"] == 500
|
|
assert result["cached"][0]["model_format"] == "gguf"
|
|
assert result["cached"][0]["capabilities"]["requires_variant"] is True
|
|
|
|
|
|
def test_cached_gguf_scan_preserves_partial_flag(monkeypatch, tmp_path):
|
|
partial = _repo("Org/Partial", [_file("Q4_K_M.gguf", 100)], tmp_path / "partial")
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [partial])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: True,
|
|
)
|
|
|
|
result = {"cached": cache_inventory._scan_cached_gguf()}
|
|
row = result["cached"][0]
|
|
|
|
assert row["partial"] is True
|
|
assert row["partial_transport"] is None
|
|
assert row["capabilities"]["can_chat"] is False
|
|
|
|
|
|
def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
repo_path = tmp_path / "hub" / "models--Org--PartialGguf"
|
|
repo_path.mkdir(parents = True)
|
|
partial = _repo(
|
|
"Org/PartialGguf",
|
|
[_file("config.json", 12)],
|
|
repo_path,
|
|
)
|
|
assert download_manifest.write_manifest(
|
|
"model",
|
|
"Org/PartialGguf",
|
|
"Q4_K_M",
|
|
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)],
|
|
"http",
|
|
hub_cache = repo_path.parent,
|
|
)
|
|
assert download_manifest.write_cancel_marker(
|
|
"model",
|
|
"Org/PartialGguf",
|
|
"Q4_K_M",
|
|
"http",
|
|
hub_cache = repo_path.parent,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [partial])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: True,
|
|
)
|
|
|
|
result = {"cached": cache_inventory._scan_cached_gguf()}
|
|
row = result["cached"][0]
|
|
|
|
assert row["repo_id"] == "Org/PartialGguf"
|
|
assert row["model_format"] == "gguf"
|
|
assert row["size_bytes"] == 4096
|
|
assert row["partial"] is True
|
|
assert row["capabilities"]["requires_variant"] is True
|
|
|
|
|
|
def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, tmp_path):
|
|
probe = _repo(
|
|
"ggml-org/models",
|
|
[_file("tinyllamas/stories260K.gguf", 1_200_000)],
|
|
tmp_path / "probe",
|
|
)
|
|
embedder = _repo(
|
|
"unsloth/bge-small-en-v1.5-GGUF",
|
|
[_file("bge-small-en-v1.5-f16.gguf", 60_000_000)],
|
|
tmp_path / "embedder",
|
|
)
|
|
chat = _repo("Org/Chat-GGUF", [_file("Q4_K_M.gguf", 100)], tmp_path / "chat")
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [probe, embedder, chat])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: False,
|
|
)
|
|
|
|
result = {"cached": cache_inventory._scan_cached_gguf()}
|
|
|
|
assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat-GGUF"]
|
|
|
|
|
|
def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
embedder = _repo(
|
|
"unsloth/bge-small-en-v1.5-GGUF",
|
|
[
|
|
_file("bge-small-en-v1.5-f16.gguf", 60_000_000),
|
|
_file("bge-small-en-v1.5-Q8_0.gguf", 35_000_000),
|
|
],
|
|
tmp_path / "embedder",
|
|
)
|
|
# Variant manifests only exist for user Hub downloads, not auto-downloads.
|
|
assert download_manifest.write_manifest(
|
|
"model",
|
|
"unsloth/bge-small-en-v1.5-GGUF",
|
|
"Q8_0",
|
|
[download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)],
|
|
"http",
|
|
hub_cache = Path(embedder.repo_path).parent,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [embedder])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: False,
|
|
)
|
|
|
|
result = {"cached": cache_inventory._scan_cached_gguf()}
|
|
|
|
assert [row["repo_id"] for row in result["cached"]] == ["unsloth/bge-small-en-v1.5-GGUF"]
|
|
assert result["cached"][0]["capabilities"]["can_chat"] is False
|
|
|
|
|
|
def test_cached_models_scan_hides_non_gguf_embedder(monkeypatch, tmp_path):
|
|
embedder_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5"
|
|
embedder_path.mkdir(parents = True)
|
|
embedder = _repo(
|
|
"unsloth/bge-small-en-v1.5",
|
|
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
|
|
embedder_path,
|
|
)
|
|
chat_path = tmp_path / "hub" / "models--Org--Chat"
|
|
chat_path.mkdir(parents = True)
|
|
chat = _repo(
|
|
"Org/Chat",
|
|
[_file("config.json", 12), _file("model.safetensors", 100)],
|
|
chat_path,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [embedder, chat])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda _kind, _repo_id, _path: False,
|
|
)
|
|
|
|
result = {"cached": cache_inventory._scan_cached_models()}
|
|
|
|
assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat"]
|
|
|
|
|
|
def _diffusion_scan(monkeypatch, tmp_path, repo_id: str, files: list, *, task: str):
|
|
"""One cached diffusion repo through _scan_cached_models, with the download-partial signal
|
|
forced off so only the pipeline-shape checks can flag the row."""
|
|
repo_path = tmp_path / f"hub/models--{repo_id.replace('/', '--')}"
|
|
repo_path.mkdir(parents = True)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [_repo(repo_id, files, repo_path)])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda _kind, _repo_id, _path: False,
|
|
)
|
|
monkeypatch.setattr(cache_inventory, "_cached_row_task", lambda _repo, gguf: task)
|
|
rows = cache_inventory._scan_cached_models()
|
|
assert len(rows) == 1
|
|
return rows[0]
|
|
|
|
|
|
def test_cached_models_scan_marks_a_companion_only_pipeline_partial(monkeypatch, tmp_path):
|
|
"""A GGUF image load prefetches the base repo's manifest + VAE + text encoder and skips the
|
|
multi-GB transformer. Every file its manifest expected arrived, so the download-partial check
|
|
passes it, but from_pretrained cannot load it -- the picker must not advertise it as on-device
|
|
(same rule /api/models/cached applies)."""
|
|
row = _diffusion_scan(
|
|
monkeypatch,
|
|
tmp_path,
|
|
"Org/Pipeline-Companions-Only",
|
|
[
|
|
_file("model_index.json", 900),
|
|
_file("vae/diffusion_pytorch_model.safetensors", 300_000_000),
|
|
_file("text_encoder/model.safetensors", 900_000_000),
|
|
],
|
|
task = "text-to-image",
|
|
)
|
|
|
|
assert row["partial"] is True
|
|
# A companion-only snapshot arrived intact, so it has no Resume / Redownload story.
|
|
assert row["partial_transport"] is None
|
|
|
|
|
|
def test_cached_models_scan_keeps_a_complete_pipeline_loadable(monkeypatch, tmp_path):
|
|
row = _diffusion_scan(
|
|
monkeypatch,
|
|
tmp_path,
|
|
"Org/Pipeline-Complete",
|
|
[
|
|
_file("model_index.json", 900),
|
|
_file("vae/diffusion_pytorch_model.safetensors", 300_000_000),
|
|
_file("text_encoder/model.safetensors", 900_000_000),
|
|
_file("transformer/diffusion_pytorch_model-00001-of-00002.safetensors", 4_000_000_000),
|
|
],
|
|
task = "text-to-image",
|
|
)
|
|
|
|
assert row["partial"] is False
|
|
assert row["single_file"] is False
|
|
|
|
|
|
def test_cached_models_scan_flags_a_single_file_diffusion_checkpoint(monkeypatch, tmp_path):
|
|
"""No root model_index.json: loadable only through from_single_file + a filename. The picker
|
|
gates on this flag, and before it was carried here every hub-sourced row read as a full
|
|
pipeline -- so a checkpoint-only repo was offered as a pipeline load and failed after the
|
|
handoff."""
|
|
row = _diffusion_scan(
|
|
monkeypatch,
|
|
tmp_path,
|
|
"Org/Single-File-Checkpoint",
|
|
[_file("config.json", 12), _file("z-image-turbo-fp8.safetensors", 6_000_000_000)],
|
|
task = "text-to-image",
|
|
)
|
|
|
|
assert row["single_file"] is True
|
|
assert row["partial"] is False
|
|
|
|
|
|
def test_cached_models_scan_leaves_chat_repos_unflagged(monkeypatch, tmp_path):
|
|
"""The flag is a diffusion-picker concern: a chat repo (no task) never carries it, so a plain
|
|
safetensors model is not mistaken for a single-file checkpoint."""
|
|
row = _diffusion_scan(
|
|
monkeypatch,
|
|
tmp_path,
|
|
"Org/Chat-Model",
|
|
[_file("config.json", 12), _file("model.safetensors", 100)],
|
|
task = None,
|
|
)
|
|
|
|
assert row["single_file"] is False
|
|
|
|
|
|
def test_cached_scans_hide_embedders_configured_by_cache_path(monkeypatch, tmp_path):
|
|
from core.rag import config as rag_config
|
|
|
|
gguf_path = tmp_path / "hub" / "models--Org--PathEmbedder-GGUF"
|
|
gguf_path.mkdir(parents = True)
|
|
gguf = _repo(
|
|
"Org/PathEmbedder-GGUF",
|
|
[_file("model-F16.gguf", 60_000_000)],
|
|
gguf_path,
|
|
)
|
|
model_path = tmp_path / "hub" / "models--Org--PathEmbedder"
|
|
model_path.mkdir(parents = True)
|
|
model = _repo(
|
|
"Org/PathEmbedder",
|
|
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
|
|
model_path,
|
|
)
|
|
monkeypatch.setattr(
|
|
rag_config,
|
|
"effective_embedding_model",
|
|
lambda: str(model_path),
|
|
)
|
|
monkeypatch.setattr(
|
|
rag_config,
|
|
"effective_gguf_repo",
|
|
lambda: str(gguf_path),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [gguf, model])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda _kind, _repo_id, _path: False,
|
|
)
|
|
|
|
assert cache_inventory._scan_cached_gguf() == []
|
|
assert cache_inventory._scan_cached_models() == []
|
|
|
|
|
|
def test_cached_scans_hide_embedders_configured_by_snapshot_path(monkeypatch, tmp_path):
|
|
from core.rag import config as rag_config
|
|
|
|
gguf_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder-GGUF"
|
|
gguf_snapshot = gguf_path / "snapshots" / "gguf-revision"
|
|
gguf_snapshot.mkdir(parents = True)
|
|
gguf = _repo(
|
|
"Org/SnapshotEmbedder-GGUF",
|
|
[_file("model-F16.gguf", 60_000_000)],
|
|
gguf_path,
|
|
)
|
|
model_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder"
|
|
model_snapshot = model_path / "snapshots" / "model-revision"
|
|
model_snapshot.mkdir(parents = True)
|
|
model = _repo(
|
|
"Org/SnapshotEmbedder",
|
|
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
|
|
model_path,
|
|
)
|
|
monkeypatch.setattr(
|
|
rag_config,
|
|
"effective_embedding_model",
|
|
lambda: str(model_snapshot),
|
|
)
|
|
monkeypatch.setattr(
|
|
rag_config,
|
|
"effective_gguf_repo",
|
|
lambda: str(gguf_snapshot),
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [gguf, model])],
|
|
)
|
|
|
|
def _resolve_snapshot(repo_path):
|
|
return str(
|
|
{
|
|
gguf_path: gguf_snapshot,
|
|
model_path: model_snapshot,
|
|
}.get(Path(repo_path), Path(repo_path))
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"resolve_hf_cache_realpath",
|
|
_resolve_snapshot,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda _kind, _repo_id, _path: False,
|
|
)
|
|
|
|
assert cache_inventory._scan_cached_gguf() == []
|
|
assert cache_inventory._scan_cached_models() == []
|
|
|
|
|
|
def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder(
|
|
monkeypatch, tmp_path
|
|
):
|
|
# A custom embedder with a generic basename ("org/model") must be hidden by
|
|
# EXACT repo-id match only. An unrelated cached chat model whose id merely
|
|
# contains "model" (e.g. "user/model-chat") must stay on device: substring
|
|
# basename matching used to drop real chat models from the inventory.
|
|
from core.rag import config as rag_config
|
|
|
|
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
|
|
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
|
|
|
|
def _model_repo(repo_id: str):
|
|
path = tmp_path / "hub" / f"models--{repo_id.replace('/', '--')}"
|
|
path.mkdir(parents = True)
|
|
return _repo(
|
|
repo_id,
|
|
[_file("config.json", 12), _file("model.safetensors", 100)],
|
|
path,
|
|
)
|
|
|
|
embedder = _model_repo("org/model")
|
|
chat = _model_repo("user/model-chat")
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [embedder, chat])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda _kind, _repo_id, _path: False,
|
|
)
|
|
|
|
result = {"cached": cache_inventory._scan_cached_models()}
|
|
|
|
assert [row["repo_id"] for row in result["cached"]] == ["user/model-chat"]
|
|
|
|
|
|
def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypatch, tmp_path):
|
|
from core.rag import config as rag_config
|
|
|
|
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom")
|
|
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF")
|
|
|
|
gguf = _repo(
|
|
"unsloth/bge-small-en-v1.5-GGUF",
|
|
[_file("bge-small-en-v1.5-f16.gguf", 60_000_000)],
|
|
tmp_path / "default-gguf",
|
|
)
|
|
weights_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5"
|
|
weights_path.mkdir(parents = True)
|
|
weights = _repo(
|
|
"unsloth/bge-small-en-v1.5",
|
|
[_file("config.json", 12), _file("model.safetensors", 130_000_000)],
|
|
weights_path,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory,
|
|
"all_hf_cache_scans",
|
|
lambda: [SimpleNamespace(repos = [gguf, weights])],
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda _repo_id, _path: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
cache_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda _kind, _repo_id, _path: False,
|
|
)
|
|
|
|
assert cache_inventory._scan_cached_gguf() == []
|
|
assert cache_inventory._scan_cached_models() == []
|
|
|
|
|
|
def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj():
|
|
requirements = gguf_variants._build_gguf_variant_requirements(
|
|
[
|
|
_sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),
|
|
_sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"),
|
|
_sibling("mmproj-BF16.gguf", 7, "mm-bf16"),
|
|
_sibling("mmproj-F16.gguf", 5, "mm-f16"),
|
|
]
|
|
)
|
|
|
|
req = requirements["q4_k_m"]
|
|
|
|
assert req.main_size_bytes == 30
|
|
assert req.download_size_bytes == 35
|
|
assert req.main_hashes == frozenset({"main-a", "main-b"})
|
|
assert req.required_hashes == frozenset({"main-a", "main-b", "mm-f16"})
|
|
assert req.companion_hashes == frozenset({"mm-f16"})
|
|
assert req.mmproj_hashes == frozenset({"mm-bf16", "mm-f16"})
|
|
assert req.target_filenames == (
|
|
"model-Q4_K_M-00001-of-00002.gguf",
|
|
"model-Q4_K_M-00002-of-00002.gguf",
|
|
"mmproj-F16.gguf",
|
|
)
|
|
|
|
|
|
def test_gguf_variant_requirements_skip_big_endian_sibling():
|
|
requirements = gguf_variants._build_gguf_variant_requirements(
|
|
[
|
|
_sibling("model-Q4_K_M-be.gguf", 100, "main-be"),
|
|
_sibling("model-Q4_K_M.gguf", 10, "main-le"),
|
|
]
|
|
)
|
|
|
|
req = requirements["q4_k_m"]
|
|
|
|
assert req.main_size_bytes == 10
|
|
assert req.main_hashes == frozenset({"main-le"})
|
|
assert req.main_filenames == frozenset({"model-Q4_K_M.gguf"})
|
|
assert req.target_filenames == ("model-Q4_K_M.gguf",)
|
|
|
|
|
|
def test_worker_gguf_variant_plan_matches_service_requirement(monkeypatch):
|
|
siblings = [
|
|
_sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),
|
|
_sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"),
|
|
_sibling("mmproj-BF16.gguf", 7, "mm-bf16"),
|
|
_sibling("mmproj-F16.gguf", 5, "mm-f16"),
|
|
]
|
|
monkeypatch.setattr(
|
|
hf_download,
|
|
"_model_info_with_retry",
|
|
lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings),
|
|
)
|
|
|
|
service_req = gguf_variants._build_gguf_variant_requirements(siblings)["q4_k_m"]
|
|
worker_plan = hf_download._gguf_variant_target_plan("Org/Vision", "Q4_K_M", None)
|
|
|
|
assert worker_plan == service_req
|
|
|
|
|
|
def test_gguf_variant_blob_hashes_accept_dict_lfs_fallback(monkeypatch):
|
|
with gguf_variants._VARIANT_HASH_LOCK:
|
|
gguf_variants._VARIANT_HASH_CACHE.clear()
|
|
gguf_variants._VARIANT_REQUIREMENT_CACHE.clear()
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"huggingface_hub",
|
|
SimpleNamespace(
|
|
HfApi = lambda *_args, **_kwargs: SimpleNamespace(
|
|
model_info = lambda *_a, **_k: SimpleNamespace(
|
|
siblings = [
|
|
_sibling("model-Q4_K_M.gguf", 10, "main-dict"),
|
|
_sibling("model-Q8_0.gguf", 20, "other"),
|
|
_sibling("mmproj-F16.gguf", 5, "mmproj"),
|
|
]
|
|
)
|
|
)
|
|
),
|
|
)
|
|
|
|
result = gguf_variants.gguf_variant_blob_hashes("Org/DictLfs", "Q4_K_M", None)
|
|
main_only = gguf_variants.gguf_variant_blob_hashes(
|
|
"Org/DictLfs",
|
|
"Q4_K_M",
|
|
None,
|
|
include_companions = False,
|
|
)
|
|
|
|
assert result == frozenset({"main-dict", "mmproj"})
|
|
assert main_only == frozenset({"main-dict"})
|
|
|
|
|
|
def test_gguf_variant_blob_hashes_skip_missing_rfilename(monkeypatch):
|
|
with gguf_variants._VARIANT_HASH_LOCK:
|
|
gguf_variants._VARIANT_HASH_CACHE.clear()
|
|
gguf_variants._VARIANT_REQUIREMENT_CACHE.clear()
|
|
siblings = [
|
|
SimpleNamespace(rfilename = None, size = 1, lfs = {"sha256": "bad"}),
|
|
_sibling("model-Q4_K_M.gguf", 10, "main"),
|
|
]
|
|
monkeypatch.setattr(
|
|
gguf_variants,
|
|
"_fetch_gguf_variant_requirements",
|
|
lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements(siblings),
|
|
)
|
|
|
|
result = gguf_variants.gguf_variant_blob_hashes("Org/Malformed", "Q4_K_M", None)
|
|
|
|
assert result == frozenset({"main"})
|
|
|
|
|
|
def test_worker_gguf_variant_targets_skip_missing_rfilename(monkeypatch):
|
|
monkeypatch.setattr(
|
|
hf_download,
|
|
"_model_info_with_retry",
|
|
lambda *_args, **_kwargs: SimpleNamespace(
|
|
siblings = [
|
|
SimpleNamespace(rfilename = None, size = 1),
|
|
_sibling("model-Q4_K_M.gguf", 10, "main"),
|
|
_sibling("mmproj-F16.gguf", 5, "mm"),
|
|
]
|
|
),
|
|
)
|
|
|
|
result = hf_download._gguf_variant_target_plan("Org/Malformed", "Q4_K_M", None)
|
|
|
|
assert list(result.target_filenames) == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
|
|
|
|
|
|
def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_path):
|
|
prepare_calls = []
|
|
snapshot_calls = []
|
|
written = []
|
|
verified = []
|
|
|
|
monkeypatch.setattr(
|
|
hf_download,
|
|
"_model_info_with_retry",
|
|
lambda *_args, **_kwargs: SimpleNamespace(
|
|
siblings = [
|
|
_sibling("model-Q4_K_M.gguf", 10, "q4-main"),
|
|
_sibling("model-Q8_0.gguf", 20, "q8-main"),
|
|
_sibling("mmproj-F16.gguf", 5, "shared-mmproj"),
|
|
]
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
|
|
)
|
|
monkeypatch.setattr(
|
|
download_registry,
|
|
"prepare_cache_for_transport",
|
|
lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0,
|
|
)
|
|
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
|
|
monkeypatch.setattr(
|
|
download_manifest, "write_manifest", lambda *args: written.append(args) or True
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"huggingface_hub",
|
|
SimpleNamespace(
|
|
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
|
|
),
|
|
)
|
|
|
|
hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http")
|
|
|
|
assert prepare_calls == [
|
|
(
|
|
("model", "Org/Vision", "http", "Q4_K_M"),
|
|
{
|
|
"only_blob_hashes": frozenset({"q4-main"}),
|
|
"companion_blob_hashes": frozenset({"shared-mmproj"}),
|
|
"protected_blob_hashes": frozenset(),
|
|
},
|
|
)
|
|
]
|
|
assert [file.path for file in written[0][3]] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
|
|
assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
|
|
assert verified == [("model", "Org/Vision", "Q4_K_M", str(tmp_path))]
|
|
|
|
|
|
def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(monkeypatch, tmp_path):
|
|
prepare_calls = []
|
|
snapshot_calls = []
|
|
|
|
def _metadata_unavailable(*_args, **_kwargs):
|
|
raise RuntimeError("metadata down")
|
|
|
|
manifest = download_manifest.Manifest(
|
|
repo_type = "model",
|
|
repo_id = "Org/Vision",
|
|
variant = "Q4_K_M",
|
|
started_at = "",
|
|
expected_files = (
|
|
download_manifest.ExpectedFile(
|
|
path = "model-Q4_K_M.gguf",
|
|
size = 10,
|
|
sha256 = "q4-main",
|
|
),
|
|
download_manifest.ExpectedFile(
|
|
path = "mmproj-F16.gguf",
|
|
size = 5,
|
|
sha256 = "shared-mmproj",
|
|
),
|
|
),
|
|
transport = "http",
|
|
)
|
|
monkeypatch.setattr(
|
|
hf_download,
|
|
"_gguf_variant_target_plan",
|
|
_metadata_unavailable,
|
|
)
|
|
monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args: manifest)
|
|
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
|
|
monkeypatch.setattr(
|
|
download_registry,
|
|
"prepare_cache_for_transport",
|
|
lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0,
|
|
)
|
|
monkeypatch.setattr(hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"huggingface_hub",
|
|
SimpleNamespace(
|
|
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
|
|
),
|
|
)
|
|
|
|
hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http")
|
|
|
|
assert prepare_calls == [
|
|
(
|
|
("model", "Org/Vision", "http", "Q4_K_M"),
|
|
{
|
|
"only_blob_hashes": frozenset({"q4-main"}),
|
|
"companion_blob_hashes": frozenset({"shared-mmproj"}),
|
|
"protected_blob_hashes": frozenset(),
|
|
},
|
|
)
|
|
]
|
|
assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
|
|
|
|
|
|
def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch, tmp_path):
|
|
metadata_calls = []
|
|
written = []
|
|
cleared = []
|
|
verified = []
|
|
|
|
def _metadata(*_args, **_kwargs):
|
|
metadata_calls.append(True)
|
|
if len(metadata_calls) == 1:
|
|
raise RuntimeError("metadata down")
|
|
return SimpleNamespace(siblings = [SimpleNamespace(rfilename = "config.json", size = 12)])
|
|
|
|
monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
|
|
monkeypatch.setattr(
|
|
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
|
|
)
|
|
monkeypatch.setattr(
|
|
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
|
|
)
|
|
monkeypatch.setattr(
|
|
download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args)
|
|
)
|
|
monkeypatch.setattr(
|
|
download_manifest, "write_manifest", lambda *args: written.append(args) or True
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"huggingface_hub",
|
|
SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
|
|
)
|
|
|
|
hf_download._download_snapshot("Org/Model", None, "http")
|
|
|
|
assert len(metadata_calls) == 2
|
|
assert cleared == [("model", "Org/Model", None)]
|
|
assert written[0][0:3] == ("model", "Org/Model", None)
|
|
assert written[0][3][0].path == "config.json"
|
|
assert verified == [("model", "Org/Model", None, str(tmp_path))]
|
|
|
|
|
|
def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_path):
|
|
metadata_calls = []
|
|
snapshot_calls = []
|
|
written = []
|
|
cleared = []
|
|
verified = []
|
|
|
|
def _metadata(*_args, **_kwargs):
|
|
metadata_calls.append(True)
|
|
raise RuntimeError("metadata down")
|
|
|
|
monkeypatch.setattr(hf_download, "_dataset_info_with_retry", _metadata)
|
|
monkeypatch.setattr(
|
|
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
|
|
)
|
|
monkeypatch.setattr(
|
|
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
|
|
)
|
|
monkeypatch.setattr(
|
|
download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args)
|
|
)
|
|
monkeypatch.setattr(
|
|
download_manifest, "write_manifest", lambda *args: written.append(args) or True
|
|
)
|
|
monkeypatch.setattr(
|
|
hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: False
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"huggingface_hub",
|
|
SimpleNamespace(
|
|
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
|
|
),
|
|
)
|
|
|
|
hf_download._download_dataset("Org/Data", None, "http")
|
|
|
|
assert len(metadata_calls) == 2
|
|
assert cleared == [("dataset", "Org/Data", None)]
|
|
assert written == []
|
|
assert snapshot_calls == [
|
|
{
|
|
"repo_id": "Org/Data",
|
|
"token": False,
|
|
"repo_type": "dataset",
|
|
"max_workers": 1,
|
|
}
|
|
]
|
|
assert verified == [("dataset", "Org/Data", None, str(tmp_path))]
|
|
|
|
|
|
def test_download_snapshot_fails_when_metadata_unavailable_and_partial_remains(
|
|
monkeypatch, tmp_path
|
|
):
|
|
"""No prior manifest + metadata unavailable + leftover .incomplete blobs means
|
|
a cached partial was returned without downloading: the worker must exit 1, not
|
|
derive a self-certifying manifest from the finalized subset."""
|
|
written = []
|
|
verified = []
|
|
|
|
def _metadata(*_args, **_kwargs):
|
|
raise RuntimeError("metadata down")
|
|
|
|
monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
|
|
monkeypatch.setattr(
|
|
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
|
|
)
|
|
monkeypatch.setattr(
|
|
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
|
|
)
|
|
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
|
|
monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args, **_kwargs: None)
|
|
monkeypatch.setattr(
|
|
download_manifest, "write_manifest", lambda *args: written.append(args) or True
|
|
)
|
|
monkeypatch.setattr(
|
|
hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: True
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"huggingface_hub",
|
|
SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
|
|
)
|
|
|
|
with pytest.raises(SystemExit) as excinfo:
|
|
hf_download._download_snapshot("Org/Model", None, "http")
|
|
|
|
assert excinfo.value.code == 1
|
|
assert written == []
|
|
assert verified == []
|
|
|
|
|
|
def test_purge_repo_cache_dirs_skips_top_level_symlink(monkeypatch, tmp_path):
|
|
root = tmp_path / "hub"
|
|
target = tmp_path / "target"
|
|
root.mkdir()
|
|
target.mkdir()
|
|
link = root / "models--Org--Repo"
|
|
link.symlink_to(target, target_is_directory = True)
|
|
monkeypatch.setattr(hf_cache_state, "hf_cache_roots", lambda: [root])
|
|
|
|
removed = hf_cache_state.purge_repo_cache_dirs("model", "Org/Repo")
|
|
|
|
assert removed is False
|
|
assert link.is_symlink()
|
|
assert target.is_dir()
|
|
|
|
|
|
def test_gguf_download_progress_fallback_logs_warning(monkeypatch):
|
|
token = "hf_12345678901234567890"
|
|
logger = _RecordingLogger()
|
|
|
|
async def _run_inline(fn, *args, **kwargs):
|
|
return fn(*args, **kwargs)
|
|
|
|
def _raise_permission_error(*_args, **_kwargs):
|
|
raise PermissionError(f"denied {token}")
|
|
|
|
monkeypatch.setattr(snapshot_progress, "logger", logger)
|
|
monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
|
|
monkeypatch.setattr(
|
|
downloads.gguf_variants,
|
|
"gguf_variant_requirements",
|
|
lambda *_args, **_kwargs: None,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloads.gguf_variants,
|
|
"gguf_variant_blob_hashes",
|
|
lambda *_args, **_kwargs: frozenset(),
|
|
)
|
|
monkeypatch.setattr(
|
|
snapshot_progress,
|
|
"preferred_repo_cache_dirs",
|
|
_raise_permission_error,
|
|
)
|
|
monkeypatch.setattr(
|
|
downloads,
|
|
"_registry",
|
|
SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
|
|
)
|
|
|
|
result = asyncio.run(
|
|
downloads.get_gguf_download_progress_response(
|
|
"Org/Model",
|
|
variant = "Q4_K_M",
|
|
expected_bytes = -1,
|
|
hf_token = token,
|
|
)
|
|
)
|
|
|
|
assert result == {
|
|
"downloaded_bytes": 0,
|
|
"completed_bytes": 0,
|
|
"complete_on_disk": False,
|
|
"expected_bytes": 0,
|
|
"progress": 0,
|
|
"cache_path": None,
|
|
}
|
|
assert logger.warnings
|
|
args, kwargs = logger.warnings[0]
|
|
assert args[:4] == (
|
|
"Error checking %s download progress for %s: %s: %s",
|
|
"model",
|
|
"Org/Model",
|
|
"PermissionError",
|
|
)
|
|
assert token not in args[4]
|
|
assert "***" in args[4]
|
|
assert kwargs == {}
|
|
|
|
|
|
def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch, tmp_path):
|
|
"""A finished mmproj companion keeps counting toward progress once the caller
|
|
supplies expected bytes; resolving the variant requirement credits it."""
|
|
entry = tmp_path / "models--Org--Model-GGUF"
|
|
snap = entry / "snapshots" / "rev0"
|
|
blobs = entry / "blobs"
|
|
snap.mkdir(parents = True)
|
|
blobs.mkdir(parents = True)
|
|
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
|
|
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 30)
|
|
(blobs / "mainhash").write_bytes(b"x" * 100)
|
|
(blobs / "mmprojhash").write_bytes(b"y" * 30)
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
assert download_manifest.write_manifest(
|
|
"model",
|
|
"Org/Model-GGUF",
|
|
"Q4_K_M",
|
|
[
|
|
download_manifest.ExpectedFile(
|
|
path = "model-Q4_K_M.gguf",
|
|
size = 100,
|
|
sha256 = "mainhash",
|
|
),
|
|
download_manifest.ExpectedFile(
|
|
path = "mmproj-F16.gguf",
|
|
size = 30,
|
|
sha256 = "mmprojhash",
|
|
),
|
|
],
|
|
"http",
|
|
hub_cache = entry.parent,
|
|
)
|
|
|
|
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",
|
|
hub_cache = entry.parent,
|
|
)
|
|
|
|
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",
|
|
hub_cache = entry.parent,
|
|
)
|
|
|
|
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",
|
|
hub_cache = cache_dir,
|
|
)
|
|
assert download_manifest.write_cancel_marker(
|
|
"model",
|
|
"Org/PartialGguf",
|
|
"Q4_K_M",
|
|
"http",
|
|
hub_cache = cache_dir,
|
|
)
|
|
monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
|
|
monkeypatch.setattr(
|
|
local_inventory.hf_cache_scan,
|
|
"is_snapshot_partial",
|
|
lambda *_args, **_kwargs: False,
|
|
)
|
|
monkeypatch.setattr(
|
|
local_inventory.hf_cache_scan,
|
|
"is_gguf_repo_partial",
|
|
lambda *_args, **_kwargs: True,
|
|
)
|
|
monkeypatch.setattr(
|
|
local_inventory.hf_cache_scan,
|
|
"resolve_hf_cache_realpath",
|
|
lambda *_args, **_kwargs: None,
|
|
)
|
|
|
|
rows = local_inventory._scan_hf_cache(cache_dir)
|
|
|
|
assert len(rows) == 1
|
|
assert rows[0].model_id == "Org/PartialGguf"
|
|
assert rows[0].source == "hf_cache"
|
|
assert rows[0].model_format == "gguf"
|
|
assert rows[0].partial is True
|
|
assert rows[0].size_bytes == 8192
|
|
assert rows[0].capabilities.requires_variant is True
|
|
|
|
|
|
def test_local_inventory_filters_custom_embedder_hf_cache_row(monkeypatch, tmp_path):
|
|
from core.rag import config as rag_config
|
|
|
|
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/embedder")
|
|
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF")
|
|
|
|
def _row(repo_id: str):
|
|
repo_path = tmp_path / f"models--{repo_id.replace('/', '--')}"
|
|
return model_common._local_model_info(
|
|
scan_path = repo_path,
|
|
load_path = repo_path,
|
|
source = "hf_cache",
|
|
model_format = "safetensors",
|
|
model_id = repo_id,
|
|
)
|
|
|
|
rows = local_inventory._filter_hidden_models([_row("org/embedder"), _row("org/chat-model")])
|
|
|
|
assert [row.model_id for row in rows] == ["org/chat-model"]
|
|
|
|
|
|
def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatch, tmp_path):
|
|
from core.rag import config as rag_config
|
|
|
|
embedder_path = tmp_path / "hub" / "models--org--embedder"
|
|
embedder_snapshot = embedder_path / "snapshots" / "revision"
|
|
embedder_snapshot.mkdir(parents = True)
|
|
chat_path = tmp_path / "hub" / "models--org--chat-model"
|
|
chat_path.mkdir(parents = True)
|
|
monkeypatch.setattr(
|
|
rag_config,
|
|
"effective_embedding_model",
|
|
lambda: str(embedder_snapshot),
|
|
)
|
|
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF")
|
|
monkeypatch.setattr(
|
|
local_inventory.hf_cache_scan,
|
|
"resolve_hf_cache_realpath",
|
|
lambda path: str(embedder_snapshot) if Path(path) == embedder_path else str(path),
|
|
)
|
|
|
|
def _row(repo_id: str, repo_path: Path):
|
|
return model_common._local_model_info(
|
|
scan_path = repo_path,
|
|
load_path = repo_path,
|
|
source = "hf_cache",
|
|
model_format = "safetensors",
|
|
model_id = repo_id,
|
|
)
|
|
|
|
rows = local_inventory._filter_hidden_models(
|
|
[_row("org/embedder", embedder_path), _row("org/chat-model", chat_path)]
|
|
)
|
|
|
|
assert [row.model_id for row in rows] == ["org/chat-model"]
|
|
|
|
|
|
def test_model_download_job_helpers_preserve_idle_shape():
|
|
key = downloads._download_job_key("Org/Model", None)
|
|
status = downloads._job_status(key)
|
|
|
|
assert key == "org/model::"
|
|
assert status.state == "idle"
|
|
assert status.error is None
|
|
|
|
|
|
def test_gguf_repo_partial_treats_completed_disk_variant_as_clean(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
snapshot = tmp_path / "cache" / "models--Org--Repo" / "snapshots" / "abc"
|
|
snapshot.mkdir(parents = True)
|
|
(snapshot / "model-Q8_0.gguf").write_bytes(b"complete")
|
|
assert download_manifest.write_cancel_marker("model", "Org/Repo", "Q4_K_M", "xet")
|
|
monkeypatch.setattr(
|
|
inventory_scan,
|
|
"resolve_snapshot_dir_for_scan",
|
|
lambda *_args: snapshot,
|
|
)
|
|
|
|
assert inventory_scan.is_gguf_repo_partial("Org/Repo", snapshot.parents[1]) is False
|
|
|
|
|
|
def test_gguf_repo_partial_flags_vision_variant_missing_mmproj(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
snapshot = tmp_path / "cache" / "models--Org--Vision" / "snapshots" / "abc"
|
|
snapshot.mkdir(parents = True)
|
|
(snapshot / "model-Q4_K_M.gguf").write_bytes(b"complete-weight")
|
|
assert download_manifest.write_manifest(
|
|
"model",
|
|
"Org/Vision",
|
|
"Q4_K_M",
|
|
[
|
|
download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 15),
|
|
download_manifest.ExpectedFile(path = "mmproj-F16.gguf", size = 8),
|
|
],
|
|
"http",
|
|
)
|
|
monkeypatch.setattr(
|
|
inventory_scan,
|
|
"resolve_snapshot_dir_for_scan",
|
|
lambda *_args: snapshot,
|
|
)
|
|
|
|
assert inventory_scan.is_gguf_repo_partial("Org/Vision") is True
|
|
|
|
|
|
def test_cancel_worker_leaves_exited_process_to_watcher():
|
|
calls: list = []
|
|
|
|
class _Registry:
|
|
def get_process(self, _key):
|
|
return SimpleNamespace(poll = lambda: 1)
|
|
|
|
def get_job(self, _key):
|
|
return SimpleNamespace(state = "running")
|
|
|
|
def mark_pending_cancel(self, key, generation):
|
|
calls.append(("pending", key, generation))
|
|
return True
|
|
|
|
def request_cancel(self, key, proc, generation):
|
|
calls.append(("request", key, generation))
|
|
return True
|
|
|
|
def cancel_requested(self, _key):
|
|
return False
|
|
|
|
state = download_lifecycle.cancel_worker(
|
|
_Registry(),
|
|
"org/model::",
|
|
generation = 3,
|
|
label = "Org/Model",
|
|
logger = SimpleNamespace(warning = lambda *_a, **_k: None),
|
|
)
|
|
|
|
assert state == "running"
|
|
assert calls == []
|
|
|
|
|
|
def test_completed_gguf_split_variant_requires_all_shards(tmp_path):
|
|
snapshot = tmp_path / "snapshot"
|
|
snapshot.mkdir()
|
|
first = snapshot / "model-Q8_0-00001-of-00002.gguf"
|
|
second = snapshot / "model-Q8_0-00002-of-00002.gguf"
|
|
first.write_bytes(b"first")
|
|
|
|
assert "Q8_0" not in inventory_scan._completed_gguf_variants(snapshot)
|
|
|
|
second.write_bytes(b"second")
|
|
assert "Q8_0" in inventory_scan._completed_gguf_variants(snapshot)
|
|
|
|
|
|
def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
|
|
|
|
assert inventory_scan.is_variant_partial(
|
|
"Org/Repo",
|
|
"Q4_K_M",
|
|
incomplete_blob_hashes = {"main-q4", "main-q8"},
|
|
variant_blob_hashes = frozenset({"main-q4"}),
|
|
)
|
|
assert not inventory_scan.is_variant_partial(
|
|
"Org/Repo",
|
|
"Q5_K_M",
|
|
incomplete_blob_hashes = {"main-q4"},
|
|
variant_blob_hashes = frozenset({"main-q5"}),
|
|
)
|
|
|
|
|
|
def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path):
|
|
"""A verified GGUF update can prune an older snapshot and make that old
|
|
directory the newest by mtime. The variant is still complete when another
|
|
snapshot satisfies its manifest."""
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
repo_dir = tmp_path / "cache" / "models--Org--Repo"
|
|
old_snapshot = repo_dir / "snapshots" / "old"
|
|
new_snapshot = repo_dir / "snapshots" / "new"
|
|
old_snapshot.mkdir(parents = True)
|
|
new_snapshot.mkdir(parents = True)
|
|
(old_snapshot / "model-Q8_0.gguf").write_bytes(b"sibling")
|
|
(new_snapshot / "model-Q4_K_M.gguf").write_bytes(b"new")
|
|
assert download_manifest.write_manifest(
|
|
"model",
|
|
"Org/Repo",
|
|
"Q4_K_M",
|
|
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 3)],
|
|
"http",
|
|
)
|
|
|
|
assert not inventory_scan.is_variant_partial(
|
|
"Org/Repo",
|
|
"Q4_K_M",
|
|
snapshot_dir = old_snapshot,
|
|
repo_cache_dir = repo_dir,
|
|
)
|
|
|
|
|
|
def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path):
|
|
async def _run_inline(fn, *args, **kwargs):
|
|
return fn(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline)
|
|
assert download_manifest.write_cancel_marker("model", "Org/PartialRepo", "Q4_K_M", "http")
|
|
snapshot = tmp_path / "cache" / "models--Org--PartialRepo" / "snapshots" / "rev0"
|
|
snapshot.mkdir(parents = True)
|
|
(snapshot / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
|
|
|
|
monkeypatch.setattr(
|
|
gguf_variants,
|
|
"list_gguf_variants",
|
|
lambda *_args, **_kwargs: (
|
|
[
|
|
SimpleNamespace(
|
|
filename = "model-Q4_K_M.gguf",
|
|
quant = "Q4_K_M",
|
|
display_label = None,
|
|
size_bytes = 100,
|
|
)
|
|
],
|
|
False,
|
|
None,
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
gguf_variants,
|
|
"iter_hf_cache_snapshots",
|
|
lambda _repo_id, root = None: [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_gguf_variants_scopes_partial_state_to_requested_cache(monkeypatch, tmp_path):
|
|
async def _run_inline(fn, *args, **kwargs):
|
|
return fn(*args, **kwargs)
|
|
|
|
repo_id = "Org/SharedRepo"
|
|
repo_name = "models--Org--SharedRepo"
|
|
cache_a = tmp_path / "cache-a"
|
|
cache_b = tmp_path / "cache-b"
|
|
repo_a = cache_a / repo_name
|
|
snapshot_a = repo_a / "snapshots" / "revision"
|
|
snapshot_a.mkdir(parents = True)
|
|
(snapshot_a / "model-Q8_0.gguf").write_bytes(b"complete")
|
|
blobs_b = cache_b / repo_name / "blobs"
|
|
blobs_b.mkdir(parents = True)
|
|
(blobs_b / "q8-hash.incomplete").write_bytes(b"partial")
|
|
|
|
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
|
|
monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline)
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = cache_b),
|
|
)
|
|
assert download_manifest.write_cancel_marker(
|
|
"model",
|
|
repo_id,
|
|
"Q8_0",
|
|
"http",
|
|
hub_cache = cache_b,
|
|
)
|
|
monkeypatch.setattr(
|
|
gguf_variants,
|
|
"list_gguf_variants",
|
|
lambda *_args, **_kwargs: (
|
|
[
|
|
SimpleNamespace(
|
|
filename = "model-Q8_0.gguf",
|
|
quant = "Q8_0",
|
|
display_label = None,
|
|
size_bytes = 8,
|
|
)
|
|
],
|
|
False,
|
|
[
|
|
SimpleNamespace(
|
|
rfilename = "model-Q8_0.gguf",
|
|
size = 8,
|
|
lfs = SimpleNamespace(sha256 = "q8-hash"),
|
|
)
|
|
],
|
|
),
|
|
)
|
|
monkeypatch.setattr(cache_inventory, "all_hf_cache_scans", lambda: [])
|
|
|
|
result = asyncio.run(
|
|
gguf_variants.get_gguf_variants_response(
|
|
repo_id,
|
|
local_path = str(repo_a),
|
|
)
|
|
)
|
|
|
|
assert result.variants[0].downloaded is True
|
|
assert result.variants[0].partial is False
|
|
|
|
|
|
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 test_prepare_cache_for_transport_uses_captured_root(monkeypatch, tmp_path):
|
|
cache_a = tmp_path / "cache-a"
|
|
cache_b = tmp_path / "cache-b"
|
|
repo_name = "models--Org--Repo"
|
|
partial_a = cache_a / repo_name / "blobs" / "blob.incomplete"
|
|
partial_b = cache_b / repo_name / "blobs" / "blob.incomplete"
|
|
partial_a.parent.mkdir(parents = True)
|
|
partial_b.parent.mkdir(parents = True)
|
|
partial_a.write_bytes(b"a")
|
|
partial_b.write_bytes(b"b")
|
|
monkeypatch.setattr(
|
|
download_registry,
|
|
"hf_cache_root",
|
|
lambda create = False, root = None: root or cache_b,
|
|
)
|
|
|
|
purged = download_registry.prepare_cache_for_transport(
|
|
"model",
|
|
"Org/Repo",
|
|
download_registry.TRANSPORT_HTTP,
|
|
root = cache_a,
|
|
)
|
|
|
|
assert purged == 1
|
|
assert not partial_a.exists()
|
|
assert partial_b.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(_download_body(gguf_variant = "Q4_K_M")))
|
|
|
|
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(_download_body(gguf_variant = "Q4_K_M")))
|
|
|
|
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_orphan_reaper_uses_worker_cache_root_after_setting_changes(monkeypatch, tmp_path):
|
|
workers = tmp_path / "workers"
|
|
workers.mkdir()
|
|
cache_a = tmp_path / "cache-a" / "hub"
|
|
cache_b = tmp_path / "cache-b" / "hub"
|
|
partial = cache_a / "models--Org--Model" / "blobs" / "abc.incomplete"
|
|
partial.parent.mkdir(parents = True)
|
|
partial.write_bytes(b"partial")
|
|
cache_b.mkdir(parents = True)
|
|
monkeypatch.setattr(state_dir, "workers_dir", lambda: workers)
|
|
monkeypatch.setattr(download_registry, "_process_alive", lambda _pid: False)
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = cache_b),
|
|
)
|
|
markers = []
|
|
monkeypatch.setattr(
|
|
download_registry,
|
|
"persist_cancel_marker",
|
|
lambda *args, **kwargs: markers.append(args),
|
|
)
|
|
metadata = download_registry.DownloadMetadata(
|
|
repo_type = "model",
|
|
repo_id = "Org/Model",
|
|
variant = None,
|
|
transport = download_registry.TRANSPORT_HTTP,
|
|
hub_cache = str(cache_a),
|
|
xet_cache = str(tmp_path / "cache-a" / "xet"),
|
|
)
|
|
download_registry.write_worker_breadcrumb("org/model", 1234, metadata)
|
|
[breadcrumb] = list(workers.iterdir())
|
|
payload = json.loads(breadcrumb.read_text(encoding = "utf-8"))
|
|
assert payload["hub_cache"] == str(cache_a)
|
|
assert payload["xet_cache"] == str(tmp_path / "cache-a" / "xet")
|
|
|
|
download_registry.reap_orphan_workers()
|
|
|
|
assert markers == [("model", "Org/Model", None, "http")]
|
|
assert list(workers.iterdir()) == []
|
|
|
|
|
|
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(_download_body()))
|
|
|
|
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(_download_body()))
|
|
|
|
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(_download_body(gguf_variant = "Q4_K_M")),
|
|
downloads.download_model_response(_download_body(gguf_variant = "Q8_0")),
|
|
)
|
|
|
|
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, hub_cache = None):
|
|
monkeypatch.setattr(
|
|
deletion.download_manifest,
|
|
"purge_state",
|
|
lambda *_args, **_kwargs: False,
|
|
)
|
|
# The repo under test lives in this cache; make it the active one so the
|
|
# delete scopes to it (default target root is the active hub cache).
|
|
if hub_cache is not None:
|
|
monkeypatch.setattr(
|
|
"utils.hf_cache_settings.get_hf_cache_paths",
|
|
lambda: SimpleNamespace(hub_cache = hub_cache),
|
|
)
|
|
|
|
|
|
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, tmp_path)
|
|
|
|
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, tmp_path)
|
|
|
|
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, tmp_path)
|
|
|
|
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
|