diff --git a/docs/semantic-models.md b/docs/semantic-models.md index 00987d14..9f6ca9ea 100644 --- a/docs/semantic-models.md +++ b/docs/semantic-models.md @@ -51,8 +51,13 @@ Models live in `~/.freshdata/models//`; override the directory with > `fd.models.pull(...)` currently raises `ModelNotPublishedError` with > instructions. You can point `FRESHDATA_MODEL_URL_BASE` at a mirror that > hosts the files, or use the air-gapped path below. The registry pins sha256 -> checksums as artifacts are published; a pinned checksum that does not match -> is refused at download *and* at load time. +> checksums as artifacts are published. Once a model has any pin, every one of +> its files (for example both `model.onnx` and `tokenizer.json`) must have a +> checksum, and a file that does not match is refused at download *and* at +> load time. `fd.models.pull` also verifies files that are already installed +> before it reports success; on a mismatch it raises `ModelChecksumError` and +> leaves the files in place, so re-run it with `force=True` to download them +> again. Downloads use a network timeout of 60 seconds per connection attempt and socket read; override it with `FRESHDATA_MODEL_TIMEOUT` (a positive number of @@ -69,8 +74,9 @@ $FRESHDATA_MODEL_DIR/fd-col-encoder-v1/tokenizer.json ``` `fd.models.status()` detects them. Manually placed files with no pinned hash -load as `unverified` (visible in `status()`); a pinned mismatch refuses to -load. +load as `unverified` (visible in `status()`). For a pinned model every file is +checked: `status()` reports `verified: True` only when all of them match, and +a mismatch in any file refuses to load. ## Using the embedding backend @@ -184,7 +190,7 @@ Requesting `"embedding"` without the extra or the model **never crashes**: | `Semantic backend 'embedding' skipped: optional dependency missing` | `[semantic]` extra not installed | `pip install "freshdata-cleaner[semantic]"` | | `... model 'fd-col-encoder-v1' is not installed` | weights never pulled (they never download automatically) | `fd.models.pull("fd-col-encoder-v1")` or place files in `FRESHDATA_MODEL_DIR` | | `ModelNotPublishedError` from `pull` | no official artifact hosting yet | set `FRESHDATA_MODEL_URL_BASE` to a mirror, or use the air-gapped path | -| `ModelChecksumError` | file does not match the pinned sha256 | re-pull with `--force` / replace the file; FreshData refuses to load mismatches | +| `ModelChecksumError` | a model file does not match its pinned sha256 (or a pinned model has a file without one) | re-pull with `--force` / replace the file; FreshData refuses to load mismatches | | `budget exhausted (...)` fallback event | `semantic_budget` ceiling hit | raise the ceiling or accept the (clean, recorded) early stop | | `calibration_version="uncalibrated"` in metadata | no calibration table found | reinstall (restores the packaged default) or `fd.models.pull("calib-v1")` | diff --git a/src/freshdata/models/download.py b/src/freshdata/models/download.py index 2739d80b..90524509 100644 --- a/src/freshdata/models/download.py +++ b/src/freshdata/models/download.py @@ -14,7 +14,14 @@ import urllib.request from pathlib import Path -from .registry import _sha256_file, check_lib_version, get_config, model_dir +from .registry import ( + _required_checksums, + _sha256_file, + check_lib_version, + get_config, + model_dir, + verify, +) from .types import ModelChecksumError, ModelNotPublishedError _URL_BASE_ENV = "FRESHDATA_MODEL_URL_BASE" @@ -74,8 +81,14 @@ def _fetch(url: str, dest: Path) -> None: def pull(model_id: str, *, force: bool = False) -> Path: """Explicitly download ``model_id`` into the local model directory. - Verifies the registry checksum when one is pinned and refuses to keep a - mismatching artifact. Returns the primary artifact path. Raises + When the model has pinned checksums, every downloaded file is checked + against its own pin before it replaces the installed file, and a + mismatching download is discarded. When every file is already present + and ``force`` is False, the installed files are verified with + :func:`~freshdata.models.registry.verify` before returning; a mismatch + raises :class:`ModelChecksumError` and leaves the files in place (re-run + with ``force=True`` to download them again). Returns the primary + artifact path. Raises :class:`ModelNotPublishedError` while no download location exists — the air-gapped path (drop files into ``FRESHDATA_MODEL_DIR``) always works. """ @@ -90,12 +103,14 @@ def pull(model_id: str, *, force: bool = False) -> Path: "(see docs/semantic-models.md)." ) + pins = _required_checksums(cfg) target_dir = model_dir() / model_id target_dir.mkdir(parents=True, exist_ok=True) base_url = base.rstrip("/") + "/" + cfg.url.rsplit("/", 1)[0] primary = target_dir / cfg.files[0] if primary.is_file() and not force and all((target_dir / f).is_file() for f in cfg.files): + verify(model_id) return primary for name in cfg.files: @@ -105,13 +120,14 @@ def pull(model_id: str, *, force: bool = False) -> Path: tmp = Path(tmp_name) try: _fetch(f"{base_url}/{name}", tmp) - if name == cfg.files[0] and cfg.sha256 is not None: + expected = pins.get(name) + if expected is not None: actual = _sha256_file(tmp) - if actual != cfg.sha256: + if actual != expected: raise ModelChecksumError( - f"Downloaded {model_id!r} failed checksum verification: " - f"expected {cfg.sha256}, got {actual}. The partial download " - "was discarded." + f"Downloaded {model_id!r} file {name!r} failed checksum " + f"verification: expected {expected}, got {actual}. The partial " + "download was discarded." ) tmp.replace(dest) finally: diff --git a/src/freshdata/models/registry.py b/src/freshdata/models/registry.py index e3f4ae87..1eca91ca 100644 --- a/src/freshdata/models/registry.py +++ b/src/freshdata/models/registry.py @@ -108,13 +108,63 @@ def _sha256_file(path: Path) -> str: return digest.hexdigest() +def pinned_checksums(cfg: ModelConfig) -> dict[str, str]: + """Return the pinned sha256 for each file of ``cfg``, keyed by file name. + + Combines ``cfg.file_sha256`` with ``cfg.sha256`` (the primary-file pin). + An empty mapping means the model is unpinned. Raises + :class:`ModelChecksumError` when a pin names a file outside ``cfg.files`` + or when two pins for the same file disagree. + """ + pins: dict[str, str] = {} + for name, digest in cfg.file_sha256: + if name not in cfg.files: + raise ModelChecksumError( + f"Registry entry {cfg.model_id!r} pins a checksum for {name!r}, " + f"which is not one of its files ({', '.join(cfg.files)})." + ) + if pins.get(name, digest) != digest: + raise ModelChecksumError( + f"Registry entry {cfg.model_id!r} pins conflicting checksums for {name!r}." + ) + pins[name] = digest + if cfg.sha256 is not None and cfg.files: + primary = cfg.files[0] + if pins.get(primary, cfg.sha256) != cfg.sha256: + raise ModelChecksumError( + f"Registry entry {cfg.model_id!r} pins conflicting checksums for " + f"its primary file {primary!r} (sha256 and file_sha256 differ)." + ) + pins[primary] = cfg.sha256 + return pins + + +def _required_checksums(cfg: ModelConfig) -> dict[str, str]: + """Return :func:`pinned_checksums`, requiring every file to be pinned once any is. + + Raises :class:`ModelChecksumError` naming the unpinned files of a + partially pinned model, since they could not be checked. + """ + pins = pinned_checksums(cfg) + if pins: + unpinned = [name for name in cfg.files if name not in pins] + if unpinned: + raise ModelChecksumError( + f"Model {cfg.model_id!r} has pinned checksums, but none for " + f"{', '.join(repr(n) for n in unpinned)}; every file of a pinned " + "model must have a checksum." + ) + return pins + + def verify(model_id: str) -> bool: - """Verify the primary artifact's checksum against the registry pin. + """Verify every installed file's checksum against the registry pins. - Returns True when pinned and matching, False when no hash is pinned - (manually placed / unpublished artifacts load as "unverified"). Raises - :class:`ModelChecksumError` on a pinned mismatch and - :class:`ModelNotInstalledError` when files are missing. + Returns True when pinned and every file matches, False when no hash is + pinned (manually placed / unpublished artifacts load as "unverified"). + Raises :class:`ModelChecksumError` on a mismatch or when a pinned model + has a file without a checksum, and :class:`ModelNotInstalledError` when + files are missing. """ cfg = get_config(model_id) if not is_installed(model_id): @@ -122,14 +172,17 @@ def verify(model_id: str) -> bool: f"Model {model_id!r} is not installed. Run fd.models.pull({model_id!r}) " f"or place its files under {model_dir() / model_id}." ) - if cfg.sha256 is None: + pins = _required_checksums(cfg) + if not pins: return False - actual = _sha256_file(_model_files(cfg)[0]) - if actual != cfg.sha256: - raise ModelChecksumError( - f"Checksum mismatch for {model_id!r}: expected {cfg.sha256}, got {actual}. " - "Refusing to load; re-run fd.models.pull(..., force=True) or replace the file." - ) + for name, file in zip(cfg.files, _model_files(cfg)): + actual = _sha256_file(file) + if actual != pins[name]: + raise ModelChecksumError( + f"Checksum mismatch for {model_id!r} file {name!r}: expected " + f"{pins[name]}, got {actual}. Refusing to load; re-run " + f"fd.models.pull({model_id!r}, force=True) or replace the file." + ) return True @@ -205,7 +258,8 @@ def status() -> dict[str, dict[str, Any]]: verified = True note = "using packaged default" else: - note = "not installed" + ("" if cfg.sha256 else " (not yet published)") + pinned = cfg.sha256 or cfg.file_sha256 + note = "not installed" + ("" if pinned else " (not yet published)") out[cfg.model_id] = { "installed": installed, "verified": verified, diff --git a/src/freshdata/models/types.py b/src/freshdata/models/types.py index 582af47c..1abb85c2 100644 --- a/src/freshdata/models/types.py +++ b/src/freshdata/models/types.py @@ -39,9 +39,13 @@ class UnknownModelError(ModelError): class ModelConfig: """Registry metadata for one downloadable (or packaged) model artifact. - ``sha256`` is ``None`` while the artifact is unpublished and no hash has - been pinned yet; installed-but-unpinned models load with a "unverified" - note instead of a checksum guarantee. ``url`` is relative to the download + ``sha256`` pins the primary file (``files[0]``) and is ``None`` while the + artifact is unpublished and no hash has been pinned yet; installed-but- + unpinned models load with a "unverified" note instead of a checksum + guarantee. ``file_sha256`` pins individual files as ``(name, sha256)`` + pairs. Once any pin exists, every file in ``files`` must be pinned (by + either field) and match; see :func:`freshdata.models.registry.pinned_checksums`. + ``url`` is relative to the download base (``FRESHDATA_MODEL_URL_BASE``); an empty base means the model is not yet published and :func:`freshdata.models.pull` explains the manual placement path instead of downloading. @@ -56,3 +60,4 @@ class ModelConfig: quantization: str # "int8" | "fp32" | "json" files: tuple[str, ...] = field(default_factory=tuple) packaged_default: bool = False + file_sha256: tuple[tuple[str, str], ...] = () diff --git a/tests/test_cli_models.py b/tests/test_cli_models.py index 45f2d68e..0a782ca9 100644 --- a/tests/test_cli_models.py +++ b/tests/test_cli_models.py @@ -2,11 +2,15 @@ from __future__ import annotations +import dataclasses +import hashlib + import pandas as pd import pytest from freshdata.enterprise.cli import main from freshdata.models import download as dl +from freshdata.models import registry as reg @pytest.fixture @@ -46,6 +50,27 @@ def test_models_pull_downloads_with_mocked_fetch(model_home, monkeypatch, capsys assert "installed" in capsys.readouterr().out +def test_models_pull_existing_mismatched_files_exits_2(model_home, monkeypatch, capsys): + """#346: pull verifies already-installed files instead of reporting success.""" + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/m") + pinned = dataclasses.replace( + reg.get_config("fd-intent-v1"), sha256=hashlib.sha256(b"GOOD").hexdigest() + ) + monkeypatch.setitem(reg.REGISTRY, "fd-intent-v1", pinned) + target = model_home / "models" / "fd-intent-v1" + target.mkdir(parents=True) + (target / "model.onnx").write_bytes(b"CORRUPT") + + def no_fetch(url, dest): # pragma: no cover - must not run + raise AssertionError("unexpected download") + + monkeypatch.setattr(dl, "_fetch", no_fetch) + assert main(["models", "pull", "fd-intent-v1"]) == 2 + out = capsys.readouterr().out + assert "Checksum mismatch" in out + assert "pulled" not in out + + def test_clean_with_embedding_missing_model_prints_skip(model_home, tmp_path, capsys): df = pd.DataFrame( { diff --git a/tests/test_models_download.py b/tests/test_models_download.py index 981f3905..b2c7091e 100644 --- a/tests/test_models_download.py +++ b/tests/test_models_download.py @@ -84,6 +84,117 @@ def test_pull_checksum_match_keeps_download(model_home, monkeypatch): assert reg.verify("fd-intent-v1") is True +def _sha(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def _no_fetch(url: str, dest: Path) -> None: # pragma: no cover - must not run + raise AssertionError(f"unexpected download of {url}") + + +_ENCODER_FILES = {"model.onnx": b"GENUINE-MODEL", "tokenizer.json": b'{"vocab": []}'} +_ENCODER_PINS = tuple((name, _sha(data)) for name, data in _ENCODER_FILES.items()) + + +def _replace_encoder(monkeypatch, **changes: object) -> None: + cfg = dataclasses.replace(reg.get_config("fd-col-encoder-v1"), **changes) + monkeypatch.setitem(reg.REGISTRY, "fd-col-encoder-v1", cfg) + + +def test_pull_refuses_pinned_secondary_file_mismatch(model_home, monkeypatch): + """#346: every downloaded file is checked, not only the primary.""" + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models") + _replace_encoder(monkeypatch, file_sha256=_ENCODER_PINS) + served = {"model.onnx": _ENCODER_FILES["model.onnx"], "tokenizer.json": b'{"modified": true}'} + monkeypatch.setattr( + dl, "_fetch", lambda url, dest: dest.write_bytes(served[url.rsplit("/", 1)[-1]]) + ) + with pytest.raises(ModelChecksumError, match="tokenizer.json"): + fd.models.pull("fd-col-encoder-v1") + target = model_home / "fd-col-encoder-v1" + assert not (target / "tokenizer.json").exists() + assert reg.is_installed("fd-col-encoder-v1") is False + assert not list(target.glob("*.part")) + + +def test_pull_all_pinned_files_matching_installs(model_home, monkeypatch): + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models") + _replace_encoder( + monkeypatch, + sha256=_sha(_ENCODER_FILES["model.onnx"]), + file_sha256=(("tokenizer.json", _sha(_ENCODER_FILES["tokenizer.json"])),), + ) + monkeypatch.setattr( + dl, "_fetch", lambda url, dest: dest.write_bytes(_ENCODER_FILES[url.rsplit("/", 1)[-1]]) + ) + primary = fd.models.pull("fd-col-encoder-v1") + assert primary.read_bytes() == _ENCODER_FILES["model.onnx"] + assert reg.verify("fd-col-encoder-v1") is True + assert fd.models.status()["fd-col-encoder-v1"]["verified"] is True + + +def test_pull_partially_pinned_model_raises(model_home, monkeypatch): + """#346 repro: only the primary is pinned, so tokenizer.json cannot be checked.""" + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models") + _replace_encoder(monkeypatch, sha256=_sha(_ENCODER_FILES["model.onnx"])) + monkeypatch.setattr(dl, "_fetch", _no_fetch) + with pytest.raises(ModelChecksumError, match="'tokenizer.json'"): + fd.models.pull("fd-col-encoder-v1") + assert list(model_home.iterdir()) == [] + + +def test_pull_early_return_with_corrupt_primary_raises(model_home, monkeypatch): + """#346 repro: already-present files are verified before pull reports success.""" + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models") + monkeypatch.setitem( + reg.REGISTRY, + "fd-intent-v1", + dataclasses.replace(reg.get_config("fd-intent-v1"), sha256=_sha(b"GOOD")), + ) + target = model_home / "fd-intent-v1" + target.mkdir() + (target / "model.onnx").write_bytes(b"CORRUPT") + monkeypatch.setattr(dl, "_fetch", _no_fetch) + with pytest.raises(ModelChecksumError, match="force=True"): + fd.models.pull("fd-intent-v1") + assert (target / "model.onnx").read_bytes() == b"CORRUPT" # left in place + + +def test_pull_early_return_with_corrupt_secondary_raises(model_home, monkeypatch): + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models") + _replace_encoder(monkeypatch, file_sha256=_ENCODER_PINS) + target = model_home / "fd-col-encoder-v1" + target.mkdir() + (target / "model.onnx").write_bytes(_ENCODER_FILES["model.onnx"]) + (target / "tokenizer.json").write_bytes(b'{"modified": true}') + monkeypatch.setattr(dl, "_fetch", _no_fetch) + with pytest.raises(ModelChecksumError, match="tokenizer.json"): + fd.models.pull("fd-col-encoder-v1") + assert (target / "tokenizer.json").read_bytes() == b'{"modified": true}' + + +def test_pull_early_return_with_matching_pins_skips_fetch(model_home, monkeypatch): + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models") + _replace_encoder(monkeypatch, file_sha256=_ENCODER_PINS) + target = model_home / "fd-col-encoder-v1" + target.mkdir() + for name, data in _ENCODER_FILES.items(): + (target / name).write_bytes(data) + monkeypatch.setattr(dl, "_fetch", _no_fetch) + assert fd.models.pull("fd-col-encoder-v1") == target / "model.onnx" + + +def test_pull_early_return_unpinned_is_unchanged(model_home, monkeypatch): + monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models") + target = model_home / "fd-col-encoder-v1" + target.mkdir() + (target / "model.onnx").write_bytes(b"anything") + (target / "tokenizer.json").write_bytes(b"anything") + monkeypatch.setattr(dl, "_fetch", _no_fetch) + assert fd.models.pull("fd-col-encoder-v1") == target / "model.onnx" + assert (target / "model.onnx").read_bytes() == b"anything" + + def test_clean_never_downloads(model_home, monkeypatch): """fd.clean must not touch the network even when embedding is requested.""" diff --git a/tests/test_models_registry.py b/tests/test_models_registry.py index 0da486ef..c51d58e6 100644 --- a/tests/test_models_registry.py +++ b/tests/test_models_registry.py @@ -113,6 +113,81 @@ def test_pinned_checksum_match_and_mismatch(model_home, monkeypatch): assert "mismatch" in status["note"].lower() +def _sha(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def test_status_false_for_mismatched_secondary_file(model_home, monkeypatch): + """#346: status() reports verified=True only when every file matches.""" + base = model_home / "fd-col-encoder-v1" + base.mkdir() + (base / "model.onnx").write_bytes(b"GENUINE-MODEL") + (base / "tokenizer.json").write_bytes(b'{"modified": true}') + pinned = dataclasses.replace( + reg.get_config("fd-col-encoder-v1"), + file_sha256=( + ("model.onnx", _sha(b"GENUINE-MODEL")), + ("tokenizer.json", _sha(b'{"vocab": []}')), + ), + ) + monkeypatch.setitem(reg.REGISTRY, "fd-col-encoder-v1", pinned) + with pytest.raises(ModelChecksumError, match="tokenizer.json"): + reg.verify("fd-col-encoder-v1") + status = fd.models.status()["fd-col-encoder-v1"] + assert status["installed"] is True + assert status["verified"] is False + assert "tokenizer.json" in status["note"] + + (base / "tokenizer.json").write_bytes(b'{"vocab": []}') + assert reg.verify("fd-col-encoder-v1") is True + assert fd.models.status()["fd-col-encoder-v1"]["verified"] is True + + +def test_verify_partially_pinned_model_raises(model_home, monkeypatch): + _place(model_home, "fd-col-encoder-v1", b"weights") + pinned = dataclasses.replace(reg.get_config("fd-col-encoder-v1"), sha256=_sha(b"weights")) + monkeypatch.setitem(reg.REGISTRY, "fd-col-encoder-v1", pinned) + with pytest.raises(ModelChecksumError, match="'tokenizer.json'"): + reg.verify("fd-col-encoder-v1") + assert fd.models.status()["fd-col-encoder-v1"]["verified"] is False + + +def test_pinned_checksums_combines_primary_and_file_pins(): + cfg = dataclasses.replace( + reg.get_config("fd-col-encoder-v1"), + sha256="a" * 64, + file_sha256=(("model.onnx", "a" * 64), ("tokenizer.json", "b" * 64)), + ) + assert reg.pinned_checksums(cfg) == {"model.onnx": "a" * 64, "tokenizer.json": "b" * 64} + unpinned = reg.get_config("fd-col-encoder-v1") + assert reg.pinned_checksums(dataclasses.replace(unpinned, sha256=None)) == {} + + +def test_pinned_checksums_rejects_unknown_file(): + cfg = dataclasses.replace( + reg.get_config("fd-intent-v1"), file_sha256=(("weights.bin", "a" * 64),) + ) + with pytest.raises(ModelChecksumError, match="weights.bin"): + reg.pinned_checksums(cfg) + + +def test_pinned_checksums_rejects_conflicting_primary_pin(): + cfg = dataclasses.replace( + reg.get_config("fd-intent-v1"), + sha256="a" * 64, + file_sha256=(("model.onnx", "b" * 64),), + ) + with pytest.raises(ModelChecksumError, match="conflicting"): + reg.pinned_checksums(cfg) + + +@pytest.mark.parametrize("model_id", sorted(reg.REGISTRY)) +def test_registry_pins_are_empty_or_cover_every_file(model_id): + cfg = reg.REGISTRY[model_id] + pins = reg.pinned_checksums(cfg) + assert pins == {} or set(pins) == set(cfg.files) + + def test_min_lib_version_gate(model_home, monkeypatch): demanding = dataclasses.replace(reg.get_config("fd-intent-v1"), min_lib_version="99.0.0") monkeypatch.setitem(reg.REGISTRY, "fd-intent-v1", demanding)