diff --git a/tests/test_onnx_runtime.py b/tests/test_onnx_runtime.py index 076afd2eb..63d29abe6 100644 --- a/tests/test_onnx_runtime.py +++ b/tests/test_onnx_runtime.py @@ -1,9 +1,11 @@ +import os import sys from headroom.onnx_runtime import ( ONNX_CPU_ARENA_ENV, cpu_arena_enabled, create_cpu_session_options, + hf_entry_known_absent, ) @@ -104,3 +106,55 @@ def test_create_cpu_session_options_handles_older_session_options(monkeypatch): assert options.intra_op_num_threads is None assert options.inter_op_num_threads is None + + +def _write_fake_hf_cache( + root: str, repo_id: str, revision: str, *, no_exist_files: list[str] +) -> None: + """Build a minimal on-disk HF hub cache layout for a single repo/revision. + + Mirrors the real cache structure closely enough for + ``huggingface_hub.try_to_load_from_cache`` to read it: a ``refs/`` + pointer file, a ``snapshots/`` directory, and a + ``.no_exist//`` marker per file whose absence is cached. + """ + from huggingface_hub.file_download import repo_folder_name + + repo_folder = os.path.join(root, repo_folder_name(repo_id=repo_id, repo_type="model")) + os.makedirs(os.path.join(repo_folder, "refs"), exist_ok=True) + with open(os.path.join(repo_folder, "refs", revision), "w") as f: + f.write("abc123") + os.makedirs(os.path.join(repo_folder, "snapshots", "abc123"), exist_ok=True) + no_exist_dir = os.path.join(repo_folder, ".no_exist", "abc123") + os.makedirs(no_exist_dir, exist_ok=True) + for filename in no_exist_files: + open(os.path.join(no_exist_dir, filename), "w").close() + + +def test_hf_entry_known_absent_true_when_404_was_cached(tmp_path, monkeypatch): + from huggingface_hub import constants + + _write_fake_hf_cache(str(tmp_path), "acme/widget", "main", no_exist_files=["merged.pt"]) + monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HEADROOM_HF_PIN", raising=False) + + assert hf_entry_known_absent("acme/widget", "merged.pt") is True + + +def test_hf_entry_known_absent_false_when_never_checked(tmp_path, monkeypatch): + from huggingface_hub import constants + + _write_fake_hf_cache(str(tmp_path), "acme/widget", "main", no_exist_files=[]) + monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HEADROOM_HF_PIN", raising=False) + + assert hf_entry_known_absent("acme/widget", "merged.pt") is False + + +def test_hf_entry_known_absent_false_when_repo_not_cached_at_all(tmp_path, monkeypatch): + from huggingface_hub import constants + + monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HEADROOM_HF_PIN", raising=False) + + assert hf_entry_known_absent("nobody/nothing", "merged.pt") is False diff --git a/tests/test_transforms/test_kompress_compressor.py b/tests/test_transforms/test_kompress_compressor.py index 37276791d..669923dfb 100644 --- a/tests/test_transforms/test_kompress_compressor.py +++ b/tests/test_transforms/test_kompress_compressor.py @@ -753,3 +753,62 @@ class TestPytorchWeightLoading: for name, param in model.state_dict().items(): assert torch.equal(param, fresh_model.state_dict()[name]) + + def test_cache_only_raises_when_confirmed_absent_but_plain_also_missing( + self, monkeypatch + ) -> None: + """merged.pt confirmed absent, but the plain fallback isn't cached either: + still nothing to load from, so this must defer rather than raise a + confusing lower-level error. + """ + import pytest + + pytest.importorskip("torch") + import headroom.transforms.kompress_compressor as kmod + + def fake_download(model_id, filename, *, allow_network=True, **kwargs): # noqa: ANN001 + raise kmod.LocalEntryNotFoundError(f"{filename} not cached") + + monkeypatch.setattr(kmod, "hf_hub_download_local_first", fake_download) + monkeypatch.setattr(kmod, "hf_entry_known_absent", lambda *a, **k: True) + + with pytest.raises(kmod.KompressModelNotCached): + kmod._load_pytorch_weights( + SimpleNamespace(), "chopratejas/kompress-base", allow_download=False + ) + + def test_genuine_download_failure_propagates_instead_of_falling_back(self, monkeypatch) -> None: + """A real network/download failure (not a 404, not a cache miss under + allow_download=False) must propagate as-is, not be swallowed into a + silent fallback to the plain format. + """ + import pytest + + pytest.importorskip("torch") + import headroom.transforms.kompress_compressor as kmod + + def fake_download(model_id, filename, *, allow_network=True, **kwargs): # noqa: ANN001 + raise OSError("connection reset") + + monkeypatch.setattr(kmod, "hf_hub_download_local_first", fake_download) + + with pytest.raises(OSError, match="connection reset"): + kmod._load_pytorch_weights(SimpleNamespace(), "some/repo", allow_download=True) + + +class TestLoadKompressPytorchCaching: + def test_returns_cached_entry_without_reloading(self, monkeypatch) -> None: + import pytest + + pytest.importorskip("torch") + import headroom.transforms.kompress_compressor as kmod + + sentinel = ("cached-model", "cached-tokenizer", "pytorch") + monkeypatch.setattr(kmod, "_kompress_cache", {"some/repo": sentinel}) + + def boom(*a, **k): # noqa: ANN001, ANN202 + raise AssertionError("should not attempt to reload an already-cached model") + + monkeypatch.setattr(kmod, "_load_pytorch_weights", boom) + + assert kmod._load_kompress_pytorch("some/repo") == sentinel