diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index 35e682366..17781e7bc 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -151,6 +151,15 @@ def download_files_from_huggingface( Path: The path to the model directory. """ + def _repo_relative_path(relative_path: Path) -> str: + # Cached files are stored under `snapshots//`. Strip that + # prefix so the remainder matches a repo file's path, which may itself be nested + # (e.g. `onnx/model.onnx`). + parts = relative_path.parts + if len(parts) > 2 and parts[0] == "snapshots": + return "/".join(parts[2:]) + return relative_path.as_posix() + def _verify_files_from_metadata( model_dir: Path, stored_metadata: dict[str, Any], repo_files: list[RepoFile] ) -> bool: @@ -162,7 +171,8 @@ def _verify_files_from_metadata( return False if repo_files: # online verification - file_info = next((f for f in repo_files if f.path == file_path.name), None) + repo_path = _repo_relative_path(Path(rel_path)) + file_info = next((f for f in repo_files if f.path == repo_path), None) if ( not file_info or file_info.size != meta["size"] @@ -185,9 +195,10 @@ def _collect_file_metadata( file_info_map = {f.path: f for f in repo_files} for file_path in model_dir.rglob("*"): if file_path.is_file() and file_path.name != cls.METADATA_FILE: - repo_file = file_info_map.get(file_path.name) + relative_path = file_path.relative_to(model_dir) + repo_file = file_info_map.get(_repo_relative_path(relative_path)) if repo_file: - meta[str(file_path.relative_to(model_dir))] = { + meta[str(relative_path)] = { "size": repo_file.size, "blob_id": repo_file.blob_id, } @@ -218,11 +229,26 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) disable_progress_bars() if metadata_file.exists(): metadata = json.loads(metadata_file.read_text()) - verified = _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[]) - if not verified: - logger.warning( - "Local file sizes do not match the metadata." - ) # do not raise, still make an attempt to load the model + if not _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[]): + # A size mismatch on a model file means ONNX Runtime would later fail + # with a cryptic protobuf error. Raise so download_model's offline probe + # (which wraps this call in `except Exception`) falls through to a forced + # online re-download instead of returning a corrupt cached path. A + # mismatch limited to an auxiliary file (e.g. a config) is left to + # best-effort loading, preserving the previous behavior. + model_files = set(extra_patterns) + model_file_corrupt = any( + _repo_relative_path(Path(rel_path)) in model_files + and (snapshot_dir / rel_path).is_file() + and (snapshot_dir / rel_path).stat().st_size != meta["size"] + for rel_path, meta in metadata.items() + ) + if model_file_corrupt: + raise ValueError( + "Cached model files do not match the stored metadata; " + "the cache appears corrupted." + ) + logger.warning("Local file sizes do not match the metadata.") result = snapshot_download( repo_id=hf_source_repo, allow_patterns=allow_patterns, @@ -233,9 +259,13 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) return result repo_revision = model_info(hf_source_repo).sha - repo_tree = list(list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model")) + repo_tree = list( + list_repo_tree( + hf_source_repo, revision=repo_revision, repo_type="model", recursive=True + ) + ) - allowed_extensions = {".json", ".onnx", ".txt"} + allowed_extensions = {".json", ".onnx", ".txt", ".onnx_data", ".npy", ".vocab"} repo_files = ( [ f @@ -442,7 +472,12 @@ def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: An hf_source, cache_dir=cache_dir, extra_patterns=extra_patterns, - **kwargs, + # The local probe failed (cache missing or corrupt). force_download so + # huggingface_hub re-fetches a present-but-truncated blob instead of + # trusting it (its cache short-circuit is existence-only). Merge into + # kwargs so a caller-supplied force_download can't raise a duplicate + # keyword error. + **{**kwargs, "force_download": True}, ) ) except (EnvironmentError, RepositoryNotFoundError, ValueError) as e: diff --git a/tests/test_model_management.py b/tests/test_model_management.py new file mode 100644 index 000000000..4a5e7c44f --- /dev/null +++ b/tests/test_model_management.py @@ -0,0 +1,92 @@ +"""Offline cache-verification tests for ModelManagement. + +The offline probe must distinguish a corrupt *model* file (which would make ONNX Runtime +fail with a cryptic protobuf error) from a benign size drift on an auxiliary file. A corrupt +model file must raise so ``download_model`` falls through to a forced re-download; an auxiliary +mismatch must be tolerated so an offline caller is not bricked by best-effort-loadable drift. +Model files may live in a subdirectory (e.g. ``onnx/model.onnx``), so matching is done on the +repo-relative path, not the bare filename. +""" + +import json +from pathlib import Path +from typing import Any + +import pytest + +from fastembed.common import model_management +from fastembed.common.model_management import ModelManagement + +REVISION = "0123456789abcdef0123456789abcdef01234567" + + +def _seed_cache(cache: Path, files: dict[str, bytes], metadata: dict[str, Any]) -> Path: + snapshot = cache / "models--qdrant--fake-onnx" + for name, blob in files.items(): + path = snapshot / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(blob) + snapshot.mkdir(parents=True, exist_ok=True) + (snapshot / ModelManagement.METADATA_FILE).write_text(json.dumps(metadata)) + return snapshot + + +def test_offline_probe_raises_when_model_file_is_corrupt(tmp_path: Path) -> None: + cache = tmp_path / "cache" + _seed_cache( + cache, + files={"model.onnx": b"x" * 10}, + metadata={"model.onnx": {"size": 999_999, "blob_id": "deadbeef"}}, + ) + + with pytest.raises(ValueError, match="corrupt"): + ModelManagement.download_files_from_huggingface( + "qdrant/fake-onnx", + cache_dir=str(cache), + extra_patterns=["model.onnx"], + local_files_only=True, + ) + + +def test_offline_probe_raises_when_subdir_model_file_is_corrupt(tmp_path: Path) -> None: + cache = tmp_path / "cache" + key = f"snapshots/{REVISION}/onnx/model.onnx" + _seed_cache( + cache, + files={key: b"x" * 10}, + metadata={key: {"size": 999_999, "blob_id": "deadbeef"}}, + ) + + with pytest.raises(ValueError, match="corrupt"): + ModelManagement.download_files_from_huggingface( + "qdrant/fake-onnx", + cache_dir=str(cache), + extra_patterns=["onnx/model.onnx"], + local_files_only=True, + ) + + +def test_offline_probe_tolerates_auxiliary_file_drift( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cache = tmp_path / "cache" + snapshot = _seed_cache( + cache, + files={"model.onnx": b"x" * 10, "config.json": b"y" * 20}, + metadata={ + "model.onnx": {"size": 10, "blob_id": "modelblob"}, + "config.json": {"size": 999, "blob_id": "configblob"}, + }, + ) + + # The model file matches metadata; only config.json drifted. The probe must NOT raise: + # it falls through to snapshot_download (mocked here to return the cached path). + monkeypatch.setattr(model_management, "snapshot_download", lambda **kwargs: str(snapshot)) + + result = ModelManagement.download_files_from_huggingface( + "qdrant/fake-onnx", + cache_dir=str(cache), + extra_patterns=["model.onnx"], + local_files_only=True, + ) + assert result == str(snapshot)