scripts/install_models.py

"""Fetch the artefacts the committed model cards describe.

Run by a person, once, before the service starts. **Not** a runtime download:
the service never calls this, and a model that is not on disk leaves its
capability unavailable rather than triggering a fetch (brief §30, ADR 0005).

The cards are the input, not this file. Each one already names its source URL
and the sha256 of the bytes that were reviewed, so this script has no list of
its own to drift out of date — it reads what review approved and refuses
anything else.

    .venv/bin/python scripts/install_models.py
    .venv/bin/python scripts/install_models.py --check   # verify, never fetch
"""

from __future__ import annotations

import argparse
import hashlib
import json
import ssl
import sys
import urllib.request
from pathlib import Path

MODELS_DIR = Path(__file__).resolve().parent.parent / "models"
CHUNK = 1 << 20
TIMEOUT_SECONDS = 300


def _digest(path: Path) -> str:
    sha = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(CHUNK):
            sha.update(chunk)
    return sha.hexdigest()


def _ssl_context() -> ssl.SSLContext:
    """Verify certificates, using certifi's bundle where the platform has none.

    A python.org build on macOS ships no CA store, so `urlopen` fails on every
    HTTPS URL. Disabling verification would be the quick fix and the wrong one:
    this script's whole job is fetching a binary that will later run against
    farm photographs.
    """
    try:
        import certifi
    except ImportError:
        return ssl.create_default_context()
    return ssl.create_default_context(cafile=certifi.where())


def _download(url: str, destination: Path) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    # Downloaded beside the target and renamed, so an interrupted fetch cannot
    # leave a truncated file that looks installed.
    partial = destination.with_suffix(destination.suffix + ".partial")
    request = urllib.request.Request(url, headers={"User-Agent": "animap-inference"})
    with urllib.request.urlopen(
        request, timeout=TIMEOUT_SECONDS, context=_ssl_context()
    ) as response, partial.open("wb") as handle:
        while chunk := response.read(CHUNK):
            handle.write(chunk)
    partial.replace(destination)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--check", action="store_true",
                        help="Report what is missing or mismatched; download nothing.")
    args = parser.parse_args(argv)

    # Deeper than `discover()` looks on purpose. Cards under `alternates/` are
    # installed and verified but never registered, so the backend that keeps the
    # licence decision reversible stays runnable without ever being able to
    # answer a request by accident.
    cards = sorted(MODELS_DIR.glob("**/model_card.json"))
    if not cards:
        print(f"No model cards under {MODELS_DIR}.")
        return 1

    failures = 0
    seen: dict[Path, str] = {}

    for card_path in cards:
        card = json.loads(card_path.read_text())
        target = (card_path.parent / card["artefact"]).resolve()
        expected = card["sha256"]
        label = f"{card_path.parent.name} → {target.name}"

        # Several capabilities share one detector. Two cards naming the same
        # file with different checksums is a review error, and the download
        # would silently make one of them right.
        if target in seen and seen[target] != expected:
            print(f"  ✗ {label}: two cards claim different checksums for this file.")
            failures += 1
            continue
        seen[target] = expected

        if target.is_file():
            actual = _digest(target)
            if actual == expected:
                print(f"  ✓ {label}")
                continue
            print(f"  ✗ {label}: on disk is {actual[:12]}…, card says {expected[:12]}…")
            failures += 1
            continue

        if args.check:
            print(f"  ✗ {label}: not installed.")
            failures += 1
            continue

        # Some artefacts have no URL to fetch because they are produced here —
        # an ONNX export of a published checkpoint is not a file the publisher
        # hosts. `source` still records where the weights came from, because
        # that is the provenance question; `produced_by` records how the bytes
        # on disk were made from them. Downloading `source` would write a model
        # card page into a .onnx file, so this refuses instead.
        producer = card.get("produced_by")
        if producer:
            print(f"  ✗ {label}: not installed, and it is a local export rather "
                  f"than a download. Run: {producer}")
            failures += 1
            continue

        print(f"  … {label}: fetching {card['source']}")
        _download(card["source"], target)
        actual = _digest(target)
        if actual != expected:
            target.unlink()
            print(f"  ✗ {label}: downloaded {actual[:12]}…, card says {expected[:12]}…. "
                  f"Removed.")
            failures += 1
            continue
        print(f"  ✓ {label} ({target.stat().st_size / 1e6:.1f} MB, {card['license']})")

    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())