ml/scripts/promote_bundle.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
promote_bundle.py — promote a candidate model bundle to deploy/.

Steps:
  1. Validate completeness (model.onnx, names.json, postprocess_config.json, model_card.json).
  2. Require a passing gate_report.json (from quality_gate.py) unless --force.
  3. Copy to deploy/<version>/ and refresh deploy/latest/ (real directory, no symlink —
     git and HF Spaces both handle plain dirs reliably).
  4. Regenerate CHECKSUMS.sha256 over the shipped files.
  5. Append an entry to model_registry/registry.json (traceability across versions).

Usage:
  python ml/scripts/promote_bundle.py --candidate artifacts/v20260703 [--version v20260703]
"""
from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import re
import shutil
import sys
from pathlib import Path
from typing import List

VERSION_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")

REQUIRED_FILES = ["model.onnx", "names.json", "postprocess_config.json", "model_card.json"]
OPTIONAL_FILES = ["weight_priors.json", "gate_report.json", "calibration_temp.json", "label_map.json"]


def validate_bundle(candidate: Path) -> List[str]:
    """Return a list of problems; empty list means the bundle is shippable."""
    problems = []
    if not candidate.is_dir():
        return [f"candidate dir not found: {candidate}"]
    for name in REQUIRED_FILES:
        if not (candidate / name).exists():
            problems.append(f"missing required file: {name}")
    # names.json must parse to a non-empty list
    names_p = candidate / "names.json"
    if names_p.exists():
        try:
            names = json.loads(names_p.read_text(encoding="utf-8"))
            if not isinstance(names, (list, dict)) or not names:
                problems.append("names.json is empty or not a list/dict")
        except Exception as e:
            problems.append(f"names.json unparseable: {e}")
    return problems


def gate_passed(candidate: Path) -> bool:
    """A gate report is only valid if it passed AND belongs to exactly this
    model.onnx — a stale report from an earlier export must never promote a
    newer, never-gated model."""
    p = candidate / "gate_report.json"
    if not p.exists():
        return False
    try:
        report = json.loads(p.read_text(encoding="utf-8"))
    except Exception:
        return False
    if not report.get("passed"):
        return False
    report_hash = report.get("model_sha256")
    if not report_hash:
        print("PROMOTE: gate_report.json has no model_sha256 — re-run quality_gate.py "
              "so the report is bound to the current model.onnx.")
        return False
    model_p = candidate / "model.onnx"
    if not model_p.exists():
        return False
    actual = hashlib.sha256(model_p.read_bytes()).hexdigest()
    if actual != report_hash:
        print("PROMOTE: gate_report.json belongs to a DIFFERENT model.onnx "
              "(stale report after re-export). Re-run quality_gate.py.")
        return False
    return True


def write_checksums(bundle_dir: Path) -> None:
    lines = []
    for f in sorted(bundle_dir.iterdir()):
        if f.name == "CHECKSUMS.sha256" or not f.is_file():
            continue
        h = hashlib.sha256(f.read_bytes()).hexdigest()
        lines.append(f"{h}  {f.name}")
    (bundle_dir / "CHECKSUMS.sha256").write_text("\n".join(lines) + "\n", encoding="utf-8")


def copy_bundle(candidate: Path, dest: Path) -> None:
    if dest.exists():
        shutil.rmtree(dest)
    dest.mkdir(parents=True)
    for name in REQUIRED_FILES + OPTIONAL_FILES:
        src = candidate / name
        if src.exists():
            shutil.copy2(src, dest / name)
    write_checksums(dest)


def update_registry(registry_path: Path, version: str, deployed_dir: Path) -> None:
    registry = {"models": []}
    if registry_path.exists():
        try:
            registry = json.loads(registry_path.read_text(encoding="utf-8"))
        except Exception:
            pass
    registry.setdefault("models", [])
    registry["models"].append({
        "model_version": version,
        "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
        "artifacts_dir": str(deployed_dir),
        "files": sorted(p.name for p in deployed_dir.iterdir() if p.is_file()),
        "promoted_by": "promote_bundle.py",
    })
    registry_path.parent.mkdir(parents=True, exist_ok=True)
    registry_path.write_text(json.dumps(registry, indent=2, ensure_ascii=False), encoding="utf-8")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="Promote a candidate bundle to deploy/.")
    ap.add_argument("--candidate", required=True, type=Path)
    ap.add_argument("--version", default=None, help="version name; defaults to the candidate dir name")
    ap.add_argument("--deploy-root", default="deploy", type=Path)
    ap.add_argument("--registry", default="model_registry/registry.json", type=Path)
    ap.add_argument("--force", action="store_true", help="promote even without a passing gate report")
    args = ap.parse_args(argv)

    candidate = args.candidate
    version = args.version or candidate.resolve().name

    # Guard rails: a bad version name must never turn the rmtree in
    # copy_bundle() against the deploy root or an unrelated directory.
    if not version or not VERSION_RE.match(version):
        print(f"PROMOTE FAIL — invalid version name {version!r} "
              "(use e.g. v20260703-1200; pass --version explicitly).")
        return 1
    if version == "latest":
        print("PROMOTE FAIL — version must not be 'latest' (reserved for the active bundle).")
        return 1
    versioned_check = (args.deploy_root / version).resolve()
    if versioned_check == args.deploy_root.resolve() or versioned_check == candidate.resolve():
        print("PROMOTE FAIL — refusing: target directory equals deploy root or the candidate itself.")
        return 1

    problems = validate_bundle(candidate)
    if problems:
        print("PROMOTE FAIL — bundle incomplete:")
        for p in problems:
            print(f"  - {p}")
        return 1

    if not gate_passed(candidate) and not args.force:
        print("PROMOTE FAIL — no passing gate_report.json in candidate. "
              "Run quality_gate.py first, or use --force (not recommended).")
        return 1

    versioned = args.deploy_root / version
    latest = args.deploy_root / "latest"
    copy_bundle(candidate, versioned)
    copy_bundle(candidate, latest)
    update_registry(args.registry, version, versioned)

    print(f"PROMOTED {candidate} -> {versioned} and {latest}")
    print("Next: commit the new bundle + registry, push, and let CI deploy the Space.")
    return 0


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