scripts/train_repro.py

#!/usr/bin/env python3
"""Scientific training contract gate for one registered isolated seed."""

from __future__ import annotations

import argparse
import json
import os
import pickle
import subprocess
import sys
from pathlib import Path

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

from repro_control.archives import safe_extract_tar_gz
from repro_control.artifacts import finalize_attempt, verify_read_back
from repro_control.configs import (
    config_sha256,
    load_json_config,
    registered_config_for,
    verify_before_optimizer,
)
from repro_control.hashing import atomic_write_json, sha256_file
from repro_control.heartbeat import heartbeat, marker
from repro_control.interventions import contains_restriction_map
from repro_control.runtime import (
    configure_scientific_runtime,
    load_job_manifest,
    require_control_freeze,
    require_jax_platform,
    verify_science_spec,
)


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("--control-dir", type=Path, required=True)
    parser.add_argument("--freeze-sha256", required=True)
    parser.add_argument("--data-dir", type=Path)
    parser.add_argument("--train-archive", type=Path)
    parser.add_argument("--output-dir", type=Path)
    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, freeze_sha256=args.freeze_sha256)
    if manifest["job_class"] != "SCIENTIFIC_TRAIN":
        raise SystemExit("training entrypoint requires SCIENTIFIC_TRAIN")
    if manifest["logical_id"] not in spec["run_identities"]:
        raise SystemExit("unregistered training logical identity")
    require_control_freeze(args.control_dir, args.freeze_sha256, manifest["science_spec_sha256"])
    configure_scientific_runtime()
    if args.dry_run:
        print(
            json.dumps(
                {
                    "contract_verified": True,
                    "logical_id": manifest["logical_id"],
                    "outcomes": {},
                }
            )
        )
        return 0
    if args.data_dir is None or args.output_dir is None:
        raise SystemExit("--data-dir and --output-dir are required for execution")
    if args.train_archive:
        safe_extract_tar_gz(args.train_archive, args.data_dir)
    require_jax_platform("gpu")
    marker("GPU_READY", manifest["logical_id"])
    logical_id = manifest["logical_id"]
    seed = manifest["seed"]
    expected_path = args.control_dir / "registered-configs" / f"{logical_id}.json"
    expected = load_json_config(expected_path)
    expected_sha256 = config_sha256(expected)
    if expected_sha256 != manifest["hashes"]["config"]:
        raise SystemExit("registered config hash differs from the Job manifest")
    independently_composed = registered_config_for(
        Path(__file__).resolve().parents[1],
        logical_id,
        data_root=str(args.data_dir),
    )
    verify_before_optimizer(
        independently_composed,
        expected,
        expected_sha256=expected_sha256,
    )
    args.output_dir.mkdir(parents=True, exist_ok=False)
    with heartbeat(f"SCIENTIFIC_TRAIN:{logical_id}"):
        if logical_id.startswith("C3-MNIST-CNN-"):
            from repro_control.cnn_training import train_fixed_cnn

            train_fixed_cnn(
                Path(expected["data"]["dir"]),
                args.output_dir,
                seed=seed,
                registered_config=expected,
                registered_config_sha256=expected_sha256,
            )
        else:
            root = Path(__file__).resolve().parents[1]
            if logical_id.startswith("C1-SUD-MPNN225-"):
                experiment = "sudoku_mpnn"
                overrides = [
                    "model.d_v=225",
                    "training.lr=1.7e-3",
                    "training.epochs=10",
                ]
            elif logical_id.startswith("C2-MAZE-MPNN84-"):
                experiment = "maze_mpnn"
                overrides = []
            elif logical_id.startswith("C4-SUD-IDENTITY-"):
                experiment = "sudoku_sheaf"
                overrides = ["model.rm_init=identity", "+model.rm_constant=true"]
            elif logical_id.startswith("C4-MAZE-QUADRATIC-"):
                experiment = "maze_sheaf"
                overrides = ["model.objective_mode=quadratic"]
            else:
                raise SystemExit("no registered training backend")
            command = [
                sys.executable,
                str(root / "scripts" / "train.py"),
                f"+experiment={experiment}",
                f"training.seed={seed}",
                f"data.dir={expected['data']['dir']}",
                "data.val_splits=[]",
                "wandb.mode=disabled",
                f"hydra.run.dir={args.output_dir}",
                *overrides,
            ]
            env = os.environ.copy()
            env.update(
                REPRO_EXPECTED_CONFIG_PATH=str(expected_path),
                REPRO_EXPECTED_CONFIG_SHA256=expected_sha256,
            )
            subprocess.run(command, check=True, env=env)

    checkpoint_path = args.output_dir / "checkpoint.pkl"
    with checkpoint_path.open("rb") as handle:
        checkpoint = pickle.load(handle)
        if handle.read(1):
            raise SystemExit("checkpoint contains trailing bytes")
    if checkpoint.get("ema_params") is None:
        raise SystemExit("training checkpoint lacks EMA parameters")
    import jax

    parameter_count = sum(
        int(leaf.size) for leaf in jax.tree_util.tree_leaves(checkpoint["params"])
    )
    identity_parity = None
    if logical_id.startswith("C4-SUD-IDENTITY-"):
        parity_path = args.output_dir / "identity-parity.json"
        if not parity_path.is_file():
            raise SystemExit("identity training lacks the pre-step parity receipt")
        identity_parity = json.loads(parity_path.read_text())
        if not all(
            identity_parity.get(name) is True
            for name in (
                "common_parameter_leaves_bit_identical",
                "common_optimizer_leaves_bit_identical",
                "counters_bit_identical",
                "checked_before_first_optimizer_step",
            )
        ):
            raise SystemExit("identity pre-step parity receipt is not affirmative")
        if contains_restriction_map(checkpoint["params"]):
            raise SystemExit("identity checkpoint unexpectedly contains restriction-map parameters")
    receipt = {
        "format": 1,
        "logical_id": logical_id,
        "attempt_id": manifest["attempt_id"],
        "seed": seed,
        "registered_config_sha256": expected_sha256,
        "checkpoint_sha256": sha256_file(checkpoint_path),
        "history_sha256": sha256_file(args.output_dir / "history.json"),
        "parameter_count": parameter_count,
        "ema": True,
        "fixed_final_epoch": int(expected["training"]["epochs"]) - 1,
        "identity_common_parity": identity_parity,
        "outcomes": {},
    }
    atomic_write_json(args.output_dir / "training-receipt.json", receipt)
    finalize_attempt(
        args.output_dir,
        logical_id=logical_id,
        attempt_id=manifest["attempt_id"],
        expected_outputs=manifest["expected_outputs"],
    )
    verify_read_back(
        args.output_dir,
        logical_id=logical_id,
        attempt_id=manifest["attempt_id"],
    )
    marker("DONE", f"{logical_id} {manifest['attempt_id']}")
    return 0


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