scripts/import_data.py

#!/usr/bin/env python3
"""Build and seal the private train/eval/import input archives on an HF CPU Job."""

from __future__ import annotations

import argparse
import json
import shutil
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from repro_control.archives import create_deterministic_tar_gz
from repro_control.artifacts import finalize_attempt, verify_read_back
from repro_control.checkpoints import load_neutral_checkpoint
from repro_control.data import (
    C5_SIZES,
    build_mnist_mask_bank,
    canonical_maze_identity,
    generate_c5_size,
    validate_smoke_fixture_registry,
    write_mnist_mask_bank,
    write_puzzle_split,
)
from repro_control.hashing import (
    atomic_write_json,
    canonical_root,
    file_entry,
    sha256_file,
)
from repro_control.heartbeat import heartbeat, marker
from repro_control.runtime import load_job_manifest, verify_science_spec


def _verify_handoff(handoff: Path, *, verify_checkpoints: bool) -> list[dict]:
    package_manifest = json.loads((handoff / "PACKAGE-MANIFEST.json").read_text())
    for row in package_manifest["entries"]:
        path = handoff / row["path"]
        if path.stat().st_size != row["bytes"] or sha256_file(path) != row["sha256"]:
            raise SystemExit(f"handoff mismatch: {row['path']}")
    imports = json.loads((handoff / "IMPORTS.json").read_text())
    if len(imports) != 9:
        raise SystemExit("IMPORTS.json must contain exactly nine rows")
    validate_smoke_fixture_registry(handoff / "SMOKE-FIXTURES.json")
    if verify_checkpoints:
        for row in imports:
            load_neutral_checkpoint(
                handoff / "checkpoints" / row["neutral_alias"] / "checkpoint.pkl",
                handoff / "configs" / f"{row['neutral_alias']}.json",
                row,
            )
    return imports


def _tree_manifest(root: Path) -> dict:
    entries = [
        file_entry(path, relative_to=root)
        for path in sorted(root.rglob("*"))
        if path.is_file()
    ]
    return {"format": 1, "entries": entries, "root_sha256": canonical_root(entries)}


def _copy_split(source_dataset: Path, split: str, target_dataset: Path) -> None:
    target_dataset.mkdir(parents=True, exist_ok=True)
    shutil.copytree(source_dataset / split, target_dataset / split)
    for name in ("metadata.json", "config.json"):
        source = source_dataset / name
        if source.is_file():
            shutil.copy2(source, target_dataset / name)


def _maze_training_identities(dataset_root: Path) -> frozenset[str]:
    import numpy as np

    inputs = np.load(dataset_root / "train/all__inputs.npy", mmap_mode="r")
    labels = np.load(dataset_root / "train/all__labels.npy", mmap_mode="r")
    return frozenset(
        canonical_maze_identity(input_row, label_row)
        for input_row, label_row in zip(inputs, labels, strict=True)
    )


def _build_registered_inputs(
    handoff: Path,
    workspace: Path,
    *,
    sudoku_revision: str,
) -> dict:
    from sheaf_admm.data.build_maze import MazeConfig
    from sheaf_admm.data.build_maze import build as build_maze
    from sheaf_admm.data.build_mnist import MNISTConfig
    from sheaf_admm.data.build_mnist import build as build_mnist
    from sheaf_admm.data.build_sudoku import SudokuConfig
    from sheaf_admm.data.build_sudoku import build as build_sudoku

    build_root = workspace / "build"
    maze_root = build_root / "maze_std3_19px_10k"
    mnist_root = build_root / "mnist"
    sudoku_root = build_root / "sudoku_easy"

    marker("HEARTBEAT", "CPU_IMPORT:build-maze")
    build_maze(
        MazeConfig(
            height=19,
            width=19,
            train_size=10_000,
            test_size=1_000,
            min_path_length=18,
            train_augment=True,
            test_augment=False,
            seed=0,
            output_dir=maze_root,
        ),
        ood_sizes=False,
    )
    marker("HEARTBEAT", "CPU_IMPORT:build-mnist")
    build_mnist(
        MNISTConfig(
            padding=0,
            seed=0,
            normalize=True,
            gen_robustness=False,
            output_dir=mnist_root,
        )
    )
    marker("HEARTBEAT", "CPU_IMPORT:build-sudoku")
    build_sudoku(
        SudokuConfig(
            dataset_revision=sudoku_revision,
            seed=0,
            train_size=50_000,
            test_size=2_000,
            difficulty_max=2.0,
            train_augment=True,
            test_augment=False,
            output_dir=sudoku_root,
        )
    )

    train_root = workspace / "train"
    eval_root = workspace / "eval"
    imports_root = workspace / "imports"
    _copy_split(maze_root, "train", train_root / "maze_std3_19px_10k")
    _copy_split(mnist_root, "train", train_root / "mnist")
    _copy_split(sudoku_root, "train", train_root / "sudoku_easy")
    _copy_split(maze_root, "test", eval_root / "maze_std3_19px_10k")
    _copy_split(mnist_root, "test", eval_root / "mnist")
    _copy_split(sudoku_root, "test_hard", eval_root / "sudoku_easy")

    test_count = 10_000
    mask_bank = build_mnist_mask_bank(
        f"ml-datasets-0.2.1:{sha256_file(mnist_root / 'test/images.npy')}",
        range(test_count),
    )
    write_mnist_mask_bank(eval_root / "mnist/drop30-mask-bank.json", mask_bank)

    marker("HEARTBEAT", "CPU_IMPORT:build-c5")
    training_identities = _maze_training_identities(maze_root)
    c5_receipts = {}
    for size in C5_SIZES:
        arrays, c5_manifest = generate_c5_size(
            size,
            training_identities=training_identities,
            examples=1000,
        )
        c5_root = eval_root / "c5" / f"{size}x{size}"
        write_puzzle_split(
            c5_root,
            "test",
            arrays["inputs"],
            arrays["labels"],
            height=size,
            width=size,
        )
        atomic_write_json(c5_root / "generation-manifest.json", c5_manifest)
        c5_receipts[str(size)] = c5_manifest

    shutil.copytree(handoff / "checkpoints", imports_root / "checkpoints")
    shutil.copytree(handoff / "configs", imports_root / "configs")
    shutil.copy2(handoff / "IMPORTS.json", imports_root / "IMPORTS.json")
    shutil.copy2(handoff / "RIGHTS.json", imports_root / "RIGHTS.json")
    manifests = {
        "train": _tree_manifest(train_root),
        "eval": _tree_manifest(eval_root),
        "imports": _tree_manifest(imports_root),
    }
    for name, value in manifests.items():
        atomic_write_json(workspace / f"{name}-tree-manifest.json", value)
    return {
        "tree_manifests": manifests,
        "mnist_mask_bank_sha256": mask_bank["bank_sha256"],
        "maze_training_identity_count": len(training_identities),
        "c5": c5_receipts,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--science-spec", type=Path, required=True)
    parser.add_argument("--job-manifest", type=Path, required=True)
    parser.add_argument("--handoff", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--sudoku-revision", required=True)
    parser.add_argument("--verify-checkpoints", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    spec = verify_science_spec(args.science_spec)
    manifest = load_job_manifest(args.job_manifest)
    if manifest["job_class"] != "CPU_IMPORT":
        raise SystemExit("CPU import entrypoint requires a CPU_IMPORT manifest")
    imports = _verify_handoff(args.handoff, verify_checkpoints=args.verify_checkpoints)
    if args.dry_run:
        print(
            json.dumps(
                {
                    "contract_verified": True,
                    "verified_import_rows": len(imports),
                    "outcomes": {},
                }
            )
        )
        return 0
    if args.sudoku_revision != spec["data_rules"]["sudoku"]["revision"]:
        raise SystemExit("Sudoku revision differs from the frozen science spec")
    args.output_dir.mkdir(parents=True, exist_ok=False)
    marker("CPU_READY", "CPU_IMPORT")
    with (
        heartbeat("CPU_IMPORT"),
        tempfile.TemporaryDirectory(prefix="sheaf-import-") as temporary,
    ):
        workspace = Path(temporary)
        build_receipt = _build_registered_inputs(
            args.handoff,
            workspace,
            sudoku_revision=args.sudoku_revision,
        )
        archive_hashes = {}
        for name in ("train", "eval", "imports"):
            archive = args.output_dir / f"{name}-data.tar.gz"
            archive_hashes[name] = create_deterministic_tar_gz(workspace / name, archive)
            shutil.copy2(
                workspace / f"{name}-tree-manifest.json",
                args.output_dir / f"{name}-tree-manifest.json",
            )
    receipt = {
        "format": 1,
        "logical_id": manifest["logical_id"],
        "attempt_id": manifest["attempt_id"],
        "verified_import_rows": len(imports),
        "sudoku_revision": args.sudoku_revision,
        "archives": archive_hashes,
        "build": build_receipt,
        "outcomes": {},
    }
    atomic_write_json(args.output_dir / "import-receipt.json", receipt)
    finalize_attempt(
        args.output_dir,
        logical_id=manifest["logical_id"],
        attempt_id=manifest["attempt_id"],
        expected_outputs=manifest["expected_outputs"],
    )
    verify_read_back(
        args.output_dir,
        logical_id=manifest["logical_id"],
        attempt_id=manifest["attempt_id"],
    )
    marker("DONE", f"{manifest['logical_id']} {manifest['attempt_id']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())