#!/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())