diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index 6c85dffe..ea42687e 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -51,6 +51,16 @@ def _hf_transport_errors() -> tuple[type[Exception], ...]: _HF_DOWNLOAD_ERRORS = (OSError, RepositoryNotFoundError, ValueError) + _hf_transport_errors() +class CorruptedCacheError(ValueError): + """A cached model file is present but does not match the size recorded for it. + + Distinct from the download failures above so that callers can tell "this snapshot is + unusable, look elsewhere" apart from "the hub could not be reached": only the former + is worth re-fetching a file the hub's existence-only cache check would otherwise keep. + Subclasses ValueError to stay inside _HF_DOWNLOAD_ERRORS for callers that don't care. + """ + + class ModelManagement(Generic[T]): METADATA_FILE = "files_metadata.json" @@ -215,6 +225,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: @@ -226,7 +245,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"] @@ -249,9 +269,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, } @@ -290,18 +311,46 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) if legacy_source is not None: sources.append(legacy_source) + # a corrupt candidate is remembered rather than raised on the spot: the next + # candidate may still be loadable, and if none is, this is the error worth + # reporting, since it is the only one download_model can act on. + corruption: CorruptedCacheError | None = None + for index, source in enumerate(sources): + last_source = index == len(sources) - 1 snapshot_dir = Path(cache_dir) / repo_folder_name( repo_id=source, repo_type="model" ) metadata_file = snapshot_dir / cls.METADATA_FILE 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 one of this model's own files (its weights, or + # whatever else its description lists) means loading would later fail + # with a cryptic error, so the snapshot is not offered at all. A + # mismatch confined to a file nothing asked for 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: + corruption = corruption or CorruptedCacheError( + f"Cached model files for {source} in {cache_dir} do not match " + "the stored metadata; the cache appears corrupted." + ) + if last_source: + raise corruption + logger.info( + f"{source} is corrupted in {cache_dir}, loading " + f"{sources[index + 1]}, cached from the same repo by an older " + "fastembed version." + ) + continue + logger.warning("Local file sizes do not match the metadata.") try: # a legacy source is only ever resolved against the cache: sending it # to the hub would ask for the very redirect this casing avoids @@ -313,7 +362,11 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) **kwargs, ) except _HF_DOWNLOAD_ERRORS: - if index == len(sources) - 1: + if last_source: + # an earlier candidate was corrupt and this one is merely absent: + # report the corruption, the only failure a re-download can fix + if corruption is not None: + raise corruption raise logger.info( f"{source} is not usable in {cache_dir}, loading " @@ -326,9 +379,13 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) # retry or fall back to another source. list_repo_tree and snapshot_download's own # metadata requests can't be given one, but a silent endpoint now fails here first. repo_revision = model_info(hf_source_repo, timeout=constants.HF_HUB_ETAG_TIMEOUT).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 = {Path(pattern).suffix for pattern in allow_patterns} repo_files = ( [ f @@ -567,6 +624,12 @@ def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: An extra_patterns = [model.model_file] extra_patterns.extend(model.additional_files) + # Only a cached file that is present but does not match its recorded size needs + # force_download: the hub's cache check is existence-only, so a plain snapshot_download + # would keep the truncated blob. Every other probe failure (a missing file, an absent + # cache) is served by an ordinary download, which re-uses what is already on disk. + force_download = False + if hf_source: try: cache_kwargs = deepcopy(kwargs) @@ -583,6 +646,8 @@ def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: An (resolved_path / file).exists() for file in extra_patterns ): return resolved_path + except CorruptedCacheError: + force_download = True except Exception: pass finally: @@ -594,21 +659,27 @@ def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: An if hf_source and not local_files_only: # we have already tried loading with `local_files_only=True` via hf and we failed + + # Merged into kwargs so a caller-supplied force_download can't raise a + # duplicate keyword error. Cleared straight away: force_download deletes the + # partially fetched `.incomplete` file, so leaving it on would make every + # retry restart a large download from zero instead of resuming it. + attempt_kwargs = {**kwargs, "force_download": True} if force_download else kwargs + force_download = False try: return Path( cls.download_files_from_huggingface( hf_source, cache_dir=cache_dir, extra_patterns=extra_patterns, - **kwargs, + **attempt_kwargs, ) ) except _HF_DOWNLOAD_ERRORS as e: - if not local_files_only: - logger.error( - f"Could not download model from HuggingFace: {e} " - "Falling back to other sources." - ) + logger.error( + f"Could not download model from HuggingFace: {e} " + "Falling back to other sources." + ) finally: enable_progress_bars() if url_source or local_files_only: diff --git a/tests/test_model_management.py b/tests/test_model_management.py new file mode 100644 index 00000000..4a5e7c44 --- /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)