Blame for autojev/train.py

"""Full-weight cross-entropy training with periodic, fixed-fold evaluation."""

import argparse
import copy
import hashlib
import json
import math
import os
from pathlib import Path
import random
import subprocess
import shutil
import time
from collections.abc import Callable, Sequence
from typing import cast

import torch
import torch.nn.functional as F

from autojev.evaluate import (
    Metrics, Prediction, calibration_ok, evaluate_logits, fit_temperature,
    hard_label, label_index, metrics, options, read_predictions, read_rows,
    selection_key, validate_coverage, write_json,
)
from autojev.events import record
from autojev.model import DecisionModel
from autojev.optim import CPUOffloadAdamW
from autojev.sft_pipeline import Receipt, ids_sha256, load_schedule, path_at
from autojev.types import Example, JSONValue


class Arguments(argparse.Namespace):
    train: str
    schedule: str | None
    development: str
    temperature: str
    reference: str
    public: str | None
    guard_data: str | None
    guard_reference: str | None
    quick_development: str | None
    quick_guard_data: str | None
    quick_temperature: str | None
    quick_eval_every: int
    run: str
    output: str
    base_model: str
    revision: str
    cache_dir: str | None
    epochs: int
    batch_size: int
    effective_batch_size: int
    token_budget: int
    max_length: int
    lr: float
    weight_decay: float
    seed: int
    eval_every: int
    public_eval_every: int
    resume_every: int
    stop_after: int | None
    resume: str | None
    initial_checkpoint: str | None
    cpu_threads: int
    allow_shared_families: bool
    keep_only_selected: bool
    no_resume: bool
    resume_unlink_first: bool


KEEP_ONLY_SELECTED = False


def digest(path: str | Path) -> str:
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def initial_identity(path: str, base_model: str, revision: str) -> dict[str, JSONValue]:
    """Bind a decision artifact used to start a new stage, not an optimizer resume."""
    directory = Path(path).resolve()
    saved = cast(dict[str, JSONValue], json.loads((directory / "decision_config.json").read_text()))
    if saved.get("format_version") != 1 or saved.get("base_model") != base_model or saved.get("revision") != revision:
        raise ValueError("Initial checkpoint format/base model/revision differs from the requested model")
    temperature = saved.get("temperature")
    if not isinstance(temperature, (int, float)) or isinstance(temperature, bool) or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Initial checkpoint temperature must be positive and finite")
    index = cast(dict[str, JSONValue], json.loads((directory / "model.safetensors.index.json").read_text()))
    weights = index.get("weight_map")
    if not isinstance(weights, dict) or not weights:
        raise ValueError("Initial checkpoint must have an indexed set of backbone weights")
    names = {"config.json", "decision_config.json", "model.safetensors.index.json", "readout.safetensors",
             "tokenizer.json", "tokenizer_config.json", "processor_config.json", "chat_template.jinja"}
    for name in weights.values():
        if not isinstance(name, str) or Path(name).name != name or not name.endswith(".safetensors"):
            raise ValueError("Checkpoint shard names must be local safetensors filenames")
        names.add(name)
    # team-1: text-only base saves no processor_config.json; hash only files that exist
    hashes: dict[str, JSONValue] = {name: digest(directory / name) for name in sorted(names) if name != "processor_config.json" or (directory / name).exists()}
    result: dict[str, JSONValue] = {
        "path": str(directory), "base_model": base_model, "revision": revision,
        "temperature": temperature, "decision_config": saved, "files_sha256": hashes,
        "artifact_sha256": hashlib.sha256(json.dumps(hashes, sort_keys=True, separators=(",", ":")).encode()).hexdigest(),
    }
    release_manifest = directory / "release-manifest.json"
    if release_manifest.exists():
        result["release_manifest_sha256"] = digest(release_manifest)
    return result


def append(path: Path, value: object) -> None:
    with path.open("a") as stream:
        stream.write(json.dumps(value, ensure_ascii=False) + "\n")


class ScheduledRows:
    """A frozen ID order whose audited text may arrive in immutable tranches."""

    def __init__(self, path: Path, calibration: Path, protected: Sequence[Example], run: str) -> None:
        self.path = path.absolute()
        self.schedule = load_schedule(self.path)
        self.sha256 = digest(self.path)
        self.calibration = calibration
        self.calibration_sha256 = digest(calibration)
        self.run = run
        self.ids = self.schedule["ordered_ids"]
        self.protected_ids = {row["id"] for row in protected}
        self.protected_families = {(str(row["source"].get("dataset", row["suite"])), row["family"]) for row in protected}
        if set(self.ids) & self.protected_ids:
            raise ValueError("Scheduled training IDs overlap a held-out fold")
        for source, checksum in ((self.schedule["plan_path"], self.schedule["plan_sha256"]),
                                 (self.schedule["original_train_path"], self.schedule["original_train_sha256"])):
            if digest(path_at(self.path, source)) != checksum:
                raise ValueError(f"Scheduled source changed: {source}")
        if path_at(self.path, self.schedule["calibration_path"]).resolve() != calibration.resolve():
            raise ValueError("--temperature must be the calibration fold bound by the schedule")
        self.tranche_for_id = {identifier: part["index"] for part in self.schedule["tranches"] for identifier in part["ordered_ids"]}
        self.receipts: list[dict[str, JSONValue]] = []
        self.cache: dict[int, dict[str, Example]] = {}

    def _read(self, index: int) -> tuple[dict[str, JSONValue], list[Example]]:
        if digest(self.path) != self.sha256 or digest(self.calibration) != self.calibration_sha256:
            raise ValueError("The frozen schedule or calibration fold changed")
        part = self.schedule["tranches"][index]
        manifest = path_at(self.path, part["manifest_path"])
        receipt = cast(Receipt, json.loads(manifest.read_text()))
        expected = {"schema_version": 1, "complete": True, "schedule_sha256": self.sha256,
                    "tranche": index, "rows": len(part["ordered_ids"]),
                    "ordered_ids_sha256": ids_sha256(part["ordered_ids"]),
                    "calibration_sha256": self.calibration_sha256,
                    "previous_receipts_sha256": [value["manifest_sha256"] for value in self.receipts[:index]]}
        if any(receipt.get(key) != value for key, value in expected.items()):
            raise ValueError(f"Tranche {index} readiness receipt differs from the frozen schedule")
        data_path = path_at(self.path, part["data_path"])
        audit_path = path_at(self.path, receipt["audit_path"])
        for path, checksum in ((data_path, receipt["data_sha256"]), (audit_path, receipt["audit_sha256"]),
                               (path_at(self.path, part["synthetic_path"]), receipt["synthetic_sha256"]),
                               (path_at(self.path, self.schedule["controls_audit_path"]), receipt["controls_audit_sha256"])):
            if digest(path) != checksum:
                raise ValueError(f"Tranche {index} bound file changed: {path}")
        audit = cast(dict[str, JSONValue], json.loads(audit_path.read_text()))
        if (audit.get("complete") is not True or audit.get("phase") != "tranche"
                or audit.get("schedule_sha256") != self.sha256 or audit.get("tranche") != index
                or audit.get("accepted_rows") != len(part["synthetic_ids"])):
            raise ValueError(f"Tranche {index} has no complete matching synthetic audit")
        audited_train = cast(dict[str, JSONValue], cast(dict[str, JSONValue], audit["outputs"])["train"])
        if audited_train.get("sha256") != receipt["synthetic_sha256"]:
            raise ValueError(f"Tranche {index} synthetic rows differ from their audit")
        controls = cast(dict[str, JSONValue], json.loads(path_at(self.path, self.schedule["controls_audit_path"]).read_text()))
        if controls.get("complete") is not True or controls.get("phase") != "controls":
            raise ValueError("The fixed synthetic controls must pass their complete audit first")
        rows = read_rows(data_path)
        if [row["id"] for row in rows] != part["ordered_ids"]:
            raise ValueError(f"Tranche {index} text IDs/order differ from the frozen schedule")
        families = {(str(row["source"].get("dataset", row["suite"])), row["family"]) for row in rows}
        if families & self.protected_families:
            raise ValueError(f"Tranche {index} has a family overlapping a held-out fold")
        provenance = cast(dict[str, JSONValue], dict(receipt))
        provenance.update(manifest_path=str(manifest), manifest_sha256=digest(manifest), data_path=str(data_path))
        return provenance, rows

    def batch(self, step: int, batch_size: int, before_wait: Callable[[], None] | None = None) -> list[Example]:
        selected_ids = self.ids[step * batch_size:(step + 1) * batch_size]
        if not selected_ids:
            raise ValueError("The training cursor is outside the frozen schedule")
        needed = {self.tranche_for_id[identifier] for identifier in selected_ids}
        for index in sorted(needed):
            if index in self.cache:
                continue
            manifest = path_at(self.path, self.schedule["tranches"][index]["manifest_path"])
            waiting_started = time.monotonic()
            waiting = False
            while not manifest.exists():
                if not waiting and before_wait is not None:
                    before_wait()
                waiting = True
                record("training_waiting_for_data", run=self.run, step=step, next_step=step + 1,
                       tranche=index, schedule_sha256=self.sha256, manifest=str(manifest),
                       waiting_seconds=time.monotonic() - waiting_started)
                time.sleep(20)
            provenance, rows = self._read(index)
            if index == len(self.receipts):
                self.receipts.append(provenance)
                record("training_tranche_loaded", run=self.run, step=step, tranche=index,
                       receipt=provenance, waiting_seconds=time.monotonic() - waiting_started if waiting else 0)
            elif index >= len(self.receipts) or self.receipts[index] != provenance:
                raise ValueError("A previously validated tranche receipt changed, or a tranche was skipped")
            self.cache[index] = {row["id"]: row for row in rows}
        result = [self.cache[self.tranche_for_id[identifier]][identifier] for identifier in selected_ids]
        self.cache = {index: rows for index, rows in self.cache.items() if index in needed}
        return sorted(result, key=length_estimate)

    def snapshot(self, examples_seen: int) -> dict[str, JSONValue]:
        if not 0 <= examples_seen <= len(self.ids):
            raise ValueError("Invalid consumed-row cursor")
        return {"schedule_path": str(self.path), "schedule_sha256": self.sha256,
                "scheduled_rows": len(self.ids), "examples_seen": examples_seen,
                "consumed_ids_sha256": ids_sha256(self.ids[:examples_seen]),
                "validated_tranches": cast(JSONValue, copy.deepcopy(self.receipts))}

    def restore(self, saved: dict[str, JSONValue], examples_seen: int) -> None:
        receipts = saved.get("validated_tranches")
        if not isinstance(receipts, list) or len(receipts) > len(self.schedule["tranches"]):
            raise ValueError("Invalid saved tranche receipt prefix")
        if examples_seen > sum(len(part["ordered_ids"]) for part in self.schedule["tranches"][:len(receipts)]):
            raise ValueError("The saved cursor includes rows without validated tranche receipts")
        for index, previous in enumerate(receipts):
            provenance, _ = self._read(index)
            if provenance != previous:
                raise ValueError(f"Tranche {index} changed since the saved optimizer cursor")
            self.receipts.append(provenance)
        if self.snapshot(examples_seen) != saved:
            raise ValueError("The saved training schedule/cursor differs from this run")


def length_estimate(row: Example) -> int:
    measured = row["source"].get("input_tokens")
    if isinstance(measured, int) and not isinstance(measured, bool):
        return measured + 16
    return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192 + 512 * len(row.get("images", []))


def fixed_subset(path: str | None, full: Sequence[Example]) -> list[Example]:
    """Keep diagnostic panels identical to their frozen parent rows, including option order."""
    if path is None:
        return []
    rows = read_rows(Path(path))
    originals = {row["id"]: json.dumps(row, ensure_ascii=False) for row in full}
    if (not rows or len({row["id"] for row in rows}) != len(rows)
            or any(originals.get(row["id"]) != json.dumps(row, ensure_ascii=False) for row in rows)):
        raise ValueError(f"Quick evaluation file is not an exact, unique subset of its full fold: {path}")
    return rows


def microbatches(rows: Sequence[Example], batch_size: int, token_budget: int) -> list[list[Example]]:
    result: list[list[Example]] = []
    pending: list[Example] = []
    longest = 0
    for row in rows:
        length = length_estimate(row)
        if pending and (len(pending) == batch_size or max(longest, length) * (len(pending) + 1) > token_budget):
            result.append(pending)
            pending, longest = [], 0
        pending.append(row)
        longest = max(longest, length)
    if pending:
        result.append(pending)
    return result


def targets(rows: Sequence[Example], device: torch.device) -> torch.Tensor:
    values = torch.zeros((len(rows), 255), dtype=torch.float32, device=device)
    for index, row in enumerate(rows):
        labels, target = options(row["question"]), row["target"]
        if isinstance(target, list):
            distribution = target
        elif row["question"]["type"] == "noul":
            positive = float(cast(float, target))
            distribution = [1.0 - positive, positive]
        else:
            distribution = [float(label == target) for label in labels]
        if len(distribution) != len(labels) or any(not math.isfinite(p) or p < 0 for p in distribution) or abs(sum(distribution) - 1) > 1e-6:
            raise ValueError(f"Invalid training target: {row['id']}")
        values[index, :len(distribution)] = torch.tensor(distribution, device=device)
    return values


def augment(rows: Sequence[Example], rng: random.Random) -> list[Example]:
    result = copy.deepcopy(list(rows))
    for row in result:
        if row["question"]["type"] == "choice":
            criteria = row["question"]["criteria"]
            target = row["target"]
            weights = dict(zip(criteria, target, strict=True)) if isinstance(target, list) else None
            items = list(criteria.items())
            rng.shuffle(items)
            row["question"]["criteria"] = dict(items)
            if weights is not None:
                row["target"] = [weights[key] for key, _ in items]
    return result


@torch.inference_mode()
def infer(model: DecisionModel, rows: Sequence[Example], args: Arguments) -> list[list[float]]:
    was_training = model.training
    model.eval()
    result: list[list[float]] = []
    for batch in microbatches(rows, args.batch_size, args.token_budget):
        logits = model(model.prepare(batch, max_length=args.max_length)).detach().cpu()
        for row, values in zip(batch, logits, strict=True):
            result.append(cast(list[float], values[:len(options(row["question"]))].tolist()))
    model.train(was_training)
    return result


def save_predictions(path: Path, predictions: Sequence[Prediction]) -> None:
    with path.open("w") as stream:
        for prediction in predictions:
            stream.write(json.dumps(prediction, ensure_ascii=False) + "\n")


def sync_directory(path: Path) -> None:
    descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def select_checkpoint(root: Path, step: int) -> None:
    pending = root / "selected.pending"
    pending.unlink(missing_ok=True)
    pending.symlink_to(f"step-{step:05d}", target_is_directory=True)
    os.replace(pending, root / "selected")
    sync_directory(root)


def save_selected(model: DecisionModel, root: Path, temperature: float, step: int, provenance: dict[str, str],
                  initial_artifact: dict[str, JSONValue] | None = None,
                  training_schedule: dict[str, JSONValue] | None = None) -> None:
    destination = root / f"step-{step:05d}"
    metadata: dict[str, JSONValue] = {"step": step, "provenance": cast(JSONValue, provenance)}
    if initial_artifact is not None:
        metadata["initial_artifact"] = initial_artifact
    if training_schedule is not None:
        metadata["training_schedule"] = training_schedule
    previous = (root / "selected").resolve() if (root / "selected").is_symlink() else None
    model.save(destination, temperature=temperature, **metadata)
    for file in destination.rglob("*"):
        if file.is_file():
            with file.open("rb") as stream:
                os.fsync(stream.fileno())
    sync_directory(destination)
    select_checkpoint(root, step)
    if KEEP_ONLY_SELECTED and previous is not None and previous != destination.resolve() and previous.parent == root.resolve()             and previous.name.startswith("step-") and previous.is_dir():
        shutil.rmtree(previous)
        record("checkpoint_pruned", path=str(previous), selected_step=step)


def archive_interrupted_tail(run: Path, output: Path, step: int) -> None:
    """Keep abandoned work as history while the active trajectory resumes its saved cursor."""
    suffix = str(time.time_ns())
    archive = run / "interrupted" / suffix
    for name in ("training.jsonl", "evaluations.jsonl", "public-evaluations.jsonl", "quick-evaluations.jsonl"):
        path = run / name
        if not path.exists():
            continue
        lines = path.read_text().splitlines()
        tail = [line for line in lines if json.loads(line)["step"] > step]
        if tail:
            archive.mkdir(parents=True, exist_ok=True)
            (archive / name).write_text("\n".join(tail) + "\n")
            path.write_text("\n".join(line for line in lines if json.loads(line)["step"] <= step) + "\n")
    for prefix in ("development", "temperature", "public", "guard", "quick-development", "quick-temperature", "quick-guard"):
        for path in run.glob(f"{prefix}-*.jsonl"):
            suffix_step = path.stem.rsplit("-", 1)[1]
            if suffix_step.isdigit() and int(suffix_step) > step:
                archive.mkdir(parents=True, exist_ok=True)
                path.rename(archive / path.name)
    for path in output.glob("step-*"):
        if path.is_dir() and path.name.rsplit("-", 1)[1].isdigit() and int(path.name.rsplit("-", 1)[1]) > step:
            destination = output / "interrupted" / suffix
            destination.mkdir(parents=True, exist_ok=True)
            path.rename(destination / path.name)
    record("training_resumed", run=run.name, resume_step=step, archived_tail=archive.exists(), attempt=suffix)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    for name, default in [("train", "data/train.jsonl"), ("development", "data/dev.jsonl"),
                          ("temperature", "data/temperature.jsonl"), ("reference", "progress/baselines/jev-development.predictions.jsonl")]:
        parser.add_argument(f"--{name}", default=default)
    parser.add_argument("--schedule", help="Frozen global ID schedule with immutable audited tranches; replaces --train")
    parser.add_argument("--public")
    parser.add_argument("--guard-data", help="Separate development cohort whose calibration also gates selection")
    parser.add_argument("--guard-reference", help="Complete fixed Jev predictions for --guard-data")
    parser.add_argument("--quick-development", help="Frozen subset of development for diagnostic curves only")
    parser.add_argument("--quick-guard-data", help="Frozen subset of the guard cohort for diagnostic curves only")
    parser.add_argument("--quick-temperature", help="Frozen subset of temperature data, used only to calibrate quick curves")
    parser.add_argument("--quick-eval-every", type=int, default=20)
    parser.add_argument("--run", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--base-model", default="Qwen/Qwen3.8-27B")
    parser.add_argument("--revision", default="1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0")
    parser.add_argument("--cache-dir")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--effective-batch-size", type=int, default=256)
    parser.add_argument("--token-budget", type=int, default=8192)
    parser.add_argument("--max-length", type=int, default=8192)
    parser.add_argument("--lr", type=float, default=2e-6)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--seed", type=int, default=20260920)
    parser.add_argument("--eval-every", type=int, default=20)
    parser.add_argument("--public-eval-every", type=int, default=100)
    parser.add_argument("--resume-every", type=int, default=100)
    parser.add_argument("--stop-after", type=int)
    parser.add_argument("--resume")
    parser.add_argument("--initial-checkpoint", help="Start a new full-weight SFT stage from this artifact with a fresh optimizer")
    parser.add_argument("--cpu-threads", type=int, default=32)
    # VIDRAFT team-1 additions (defaults keep upstream behaviour).
    parser.add_argument("--allow-shared-families", action="store_true",
                        help="Held-out folds share (dataset, family) with train; IDs and source.group_id must still be disjoint")
    parser.add_argument("--keep-only-selected", action="store_true",
                        help="Delete the previously selected step-* artifact after a newer one is selected (disk budget)")
    parser.add_argument("--no-resume", action="store_true", help="Never write resume.pt (optimizer state is ~12 bytes/param)")
    parser.add_argument("--resume-unlink-first", action="store_true",
                        help="Remove the old resume.pt before writing the new one (halves peak disk; not atomic)")
    args = parser.parse_args(namespace=Arguments())
    global KEEP_ONLY_SELECTED
    KEEP_ONLY_SELECTED = args.keep_only_selected
    if min(args.epochs, args.batch_size, args.effective_batch_size, args.token_budget, args.eval_every, args.public_eval_every, args.resume_every, args.quick_eval_every) < 1:
        raise ValueError("Batch, epoch and interval settings must be positive")
    if args.stop_after is not None and args.stop_after < 1:
        raise ValueError("A pilot must include at least one training step")
    if bool(args.guard_data) != bool(args.guard_reference):
        raise ValueError("--guard-data and --guard-reference must be supplied together")
    quick_paths = (args.quick_development, args.quick_guard_data, args.quick_temperature)
    if any(quick_paths) and (not all(quick_paths) or not args.guard_data):
        raise ValueError("Quick evaluation needs all three frozen subsets and a full guard cohort/reference")
    if args.initial_checkpoint and not os.environ.get("AUTOJEV_EVENTS"):
        raise ValueError("Set AUTOJEV_EVENTS to the new stage's event log before continuing from a checkpoint")
    if args.schedule and (args.epochs != 1 or not os.environ.get("AUTOJEV_EVENTS")):
        raise ValueError("A streamed schedule requires one epoch and an explicit AUTOJEV_EVENTS path")
    run, output = Path(args.run), Path(args.output)
    if args.initial_checkpoint and any(path.resolve().is_relative_to(Path(args.initial_checkpoint).resolve()) for path in (run, output)):
        raise ValueError("New run/output directories must be outside the initial checkpoint artifact")
    run.mkdir(parents=True, exist_ok=True)
    output.mkdir(parents=True, exist_ok=True)
    if (run / "config.json").exists() and not args.resume:
        raise ValueError("Run already exists; resume explicitly or choose a new run ID")
    train = [] if args.schedule else read_rows(Path(args.train))
    development, temperature_rows = (read_rows(Path(path)) for path in (args.development, args.temperature))
    if (not train and not args.schedule) or not development or not temperature_rows:
        raise ValueError("Training, development and temperature folds must be nonempty")
    reference_predictions = read_predictions(Path(args.reference))
    validate_coverage(development, reference_predictions)
    reference = metrics(reference_predictions)
    public_rows = read_rows(Path(args.public)) if args.public else []
    guard_rows = read_rows(Path(args.guard_data)) if args.guard_data else []
    guard_reference: Metrics | None = None
    guard_predictions: list[Prediction] = []
    if args.guard_data:
        if not guard_rows:
            raise ValueError("The guard cohort must be nonempty")
        guard_predictions = read_predictions(Path(cast(str, args.guard_reference)))
        validate_coverage(guard_rows, guard_predictions)
        guard_reference = metrics(guard_predictions)
    quick_development = fixed_subset(args.quick_development, development)
    quick_guard = fixed_subset(args.quick_guard_data, guard_rows)
    quick_temperature = fixed_subset(args.quick_temperature, temperature_rows)
    quick_ids = {row["id"] for row in quick_development}
    quick_guard_ids = {row["id"] for row in quick_guard}
    quick_reference = metrics([value for value in reference_predictions if value["id"] in quick_ids]) if quick_development else None
    quick_guard_reference = metrics([value for value in guard_predictions if value["id"] in quick_guard_ids]) if quick_guard else None
    partitions = {"train": train, "development": development, "temperature": temperature_rows,
                  "public": public_rows, "guard": guard_rows}
    families: dict[str, set[tuple[str, str]]] = {}
    identifiers: dict[str, set[str]] = {}
    source_groups: dict[str, set[tuple[str, str]]] = {}
    for name, rows in partitions.items():
        identifiers[name] = {row["id"] for row in rows}
        if len(identifiers[name]) != len(rows):
            raise ValueError(f"Duplicate IDs in {name}")
        families[name] = {(str(row["source"].get("dataset", row["suite"])), row["family"]) for row in rows}
        source_groups[name] = {(row["family"], str(row["source"].get("group_id", row["id"]))) for row in rows}
        for other in identifiers:
            if name == other:
                continue
            if identifiers[name] & identifiers[other]:
                raise ValueError(f"Partitions overlap: {name}/{other}")
            if args.allow_shared_families:
                if source_groups[name] & source_groups[other]:
                    raise ValueError(f"Partitions share source groups: {name}/{other}")
            elif families[name] & families[other]:
                raise ValueError(f"Partitions overlap: {name}/{other}")
    scheduled = ScheduledRows(Path(args.schedule), Path(args.temperature), development + temperature_rows + public_rows + guard_rows, run.name) if args.schedule else None
    if scheduled and scheduled.schedule["seed"] != args.seed:
        raise ValueError("The training seed must match the frozen schedule seed")
    train_rows = len(scheduled.ids) if scheduled else len(train)
    hashes = {name: digest(path) for name, path in [("development", args.development), ("temperature", args.temperature), ("reference", args.reference)]}
    hashes["train_schedule" if scheduled else "train"] = scheduled.sha256 if scheduled else digest(args.train)
    if args.public:
        hashes["public"] = digest(args.public)
    if args.guard_data:
        hashes["guard_data"] = digest(args.guard_data)
        hashes["guard_reference"] = digest(cast(str, args.guard_reference))
    for key, path in (("quick_development", args.quick_development), ("quick_guard_data", args.quick_guard_data), ("quick_temperature", args.quick_temperature)):
        if path is not None:
            hashes[key] = digest(path)
    initial_artifact = initial_identity(args.initial_checkpoint, args.base_model, args.revision) if args.initial_checkpoint else None
    revision = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
    package = Path(__file__).resolve().parent
    code_hashes = {name: digest(package / name) for name in ("train.py", "model.py", "optim.py", "evaluate.py", "types.py", "events.py")}
    if scheduled:
        code_hashes["sft_pipeline.py"] = digest(package / "sft_pipeline.py")
    code_hashes["uv.lock"] = digest(package.parents[1] / "uv.lock")
    config = {**vars(args), "data_sha256": hashes, "code_sha256": code_hashes, "git_commit": revision,
              "initial_artifact": initial_artifact, "events_path": os.getenv("AUTOJEV_EVENTS", "progress/events.jsonl"),
              "train_rows": train_rows, "development_rows": len(development), "temperature_rows": len(temperature_rows),
              "guard_rows": len(guard_rows), "quick_development_rows": len(quick_development),
              "quick_guard_rows": len(quick_guard), "quick_temperature_rows": len(quick_temperature)}
    if not args.resume:
        write_json(run / "config.json", config)
        if scheduled:
            scheduled.batch(0, args.effective_batch_size)
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    started = time.monotonic()
    model = DecisionModel(checkpoint=args.initial_checkpoint, train=True, base_model=args.base_model,
                          revision=args.revision, cache_dir=args.cache_dir, gradient_checkpointing=True,
                          cpu_threads=args.cpu_threads)
    optimizer = CPUOffloadAdamW(model.named_parameters(), lr=args.lr, weight_decay=args.weight_decay)
    parameters = [parameter for parameter in model.parameters() if parameter.requires_grad]
    groups: list[list[Example]] = []
    for epoch in range(args.epochs):
        ordered = list(train)
        random.Random(args.seed + epoch).shuffle(ordered)
        for offset in range(0, len(ordered), args.effective_batch_size):
            groups.append(sorted(ordered[offset:offset + args.effective_batch_size], key=length_estimate))
    total_steps = math.ceil(train_rows / args.effective_batch_size) if scheduled else len(groups)
    step, examples_seen = 0, 0
    last_resume_step: int | None = None
    best: Metrics | None = None
    best_guard: Metrics | None = None
    best_step: int | None = None
    selected_temperature: float | None = None
    if args.resume:
        state = torch.load(args.resume, map_location="cpu", weights_only=True)
        if state["data_sha256"] != hashes or state["total_steps"] != total_steps:
            raise ValueError("Resume data/schedule differs from the saved run")
        if state["config"]["code_sha256"] != code_hashes:
            raise ValueError("Training implementation differs from the saved run")
        if state["config"].get("initial_artifact") != initial_artifact:
            raise ValueError("Initial checkpoint content or provenance differs from the saved stage")
        if state["config"].get("events_path") != config["events_path"]:
            raise ValueError("Resume event log differs from the saved stage")
        for key in vars(args):
            if key in {"resume", "stop_after"}:
                continue
            if state["config"][key] != vars(args)[key]:
                raise ValueError(f"Resume configuration differs: {key}")
        step, examples_seen = state["step"], state["examples_seen"]
        if scheduled:
            if examples_seen != min(step * args.effective_batch_size, train_rows):
                raise ValueError("The saved consumed-row cursor differs from its global update")
            scheduled.restore(state["training_schedule"], examples_seen)
        optimizer.load_state_dict(state["optimizer"])
        last_resume_step = step
        if args.stop_after is not None and args.stop_after <= step:
            raise ValueError("The requested stopping step must follow the saved step")
        best, best_step, selected_temperature = state["best"], state["best_step"], state["selected_temperature"]
        best_guard = state.get("best_guard")
        random.setstate(state["python_rng"])
        torch.set_rng_state(state["torch_rng"])
        torch.cuda.set_rng_state_all(state["cuda_rng"])
        del state
        archive_interrupted_tail(run, output, step)
        if best_step is not None:
            selected = output / f"step-{best_step:05d}"
            if not selected.is_dir():
                raise ValueError("The checkpoint selected by the resume state is missing")
            select_checkpoint(output, best_step)
        else:
            (output / "selected").unlink(missing_ok=True)
    record("training_started", run=run.name, git_commit=revision, data_sha256=hashes,
           initial_artifact=initial_artifact, initialization="exact_resume" if args.resume else "checkpoint_fresh_optimizer" if initial_artifact else "base_fresh_optimizer",
           step=step, total_steps=total_steps, training_schedule=scheduled.snapshot(examples_seen) if scheduled else None,
           trainable_parameters=sum(p.numel() for p in parameters), base_model=model.base_model,
           base_revision=model.revision, text_only=model.text_only)

    def evaluate(include_public: bool = False, diagnostic: bool = False) -> None:
        nonlocal best, best_guard, best_step, selected_temperature
        began = time.monotonic()
        temperature_logits = infer(model, temperature_rows, args)
        fitted = fit_temperature(temperature_logits, [label_index(options(row["question"]), hard_label(row)) for row in temperature_rows])
        logits = infer(model, development, args)
        raw_predictions, fitted_predictions = evaluate_logits(development, logits), evaluate_logits(development, logits, fitted)
        raw, calibrated = metrics(raw_predictions), metrics(fitted_predictions)
        guard_raw: Metrics | None = None
        guard_calibrated: Metrics | None = None
        guard_ok = True
        if guard_rows:
            guard_logits = infer(model, guard_rows, args)
            guard_raw = metrics(evaluate_logits(guard_rows, guard_logits))
            guard_predictions = evaluate_logits(guard_rows, guard_logits, fitted)
            guard_calibrated = metrics(guard_predictions)
            guard_ok = calibration_ok(guard_calibrated, cast(Metrics, guard_reference))
            save_predictions(run / f"guard-{step:05d}.jsonl", guard_predictions)
        public_ok = calibration_ok(calibrated, reference)
        eligible = public_ok and guard_ok
        improved = not diagnostic and eligible and (best is None or selection_key(calibrated, step) < selection_key(best, cast(int, best_step)))
        if improved:
            save_selected(model, output, fitted, step, {"run": run.name, "git_commit": revision, **hashes, **code_hashes},
                          initial_artifact, scheduled.snapshot(examples_seen) if scheduled else None)
            best, best_guard, best_step, selected_temperature = calibrated, guard_calibrated, step, fitted
        save_predictions(run / f"development-{step:05d}.jsonl", fitted_predictions)
        save_predictions(run / f"temperature-{step:05d}.jsonl", evaluate_logits(temperature_rows, temperature_logits, fitted))
        guard_target = guard_calibrated is not None and guard_reference is not None and guard_ok and guard_calibrated["accuracy"] > guard_reference["accuracy"]
        value = record("evaluation", run=run.name, step=step, examples_seen=examples_seen, raw=raw, fitted=calibrated,
                       temperature=fitted, jev=reference, eligible=eligible, public_calibration_ok=public_ok,
                       guard_raw=guard_raw, guard_fitted=guard_calibrated, guard_jev=guard_reference,
                       guard_calibration_ok=guard_ok, guard_empirical_target_met=guard_target,
                       development_target_met=eligible and calibrated["accuracy"] > reference["accuracy"],
                       diagnostic=diagnostic, selected_step=best_step, elapsed_seconds=time.monotonic() - started,
                       evaluation_seconds=time.monotonic() - began)
        append(run / "evaluations.jsonl", value)
        print(json.dumps(value), flush=True)
        if include_public and public_rows:
            predictions = evaluate_logits(public_rows, infer(model, public_rows, args), fitted)
            save_predictions(run / f"public-{step:05d}.jsonl", predictions)
            value = record("public_evaluation", run=run.name, step=step, temperature=fitted, metrics=metrics(predictions), status="known_suite_development_check")
            append(run / "public-evaluations.jsonl", value)

    def save_resume() -> None:
        nonlocal last_resume_step
        began = time.monotonic()
        state = {"optimizer": optimizer.state_dict(), "step": step, "examples_seen": examples_seen, "total_steps": total_steps, "data_sha256": hashes, "config": config,
                 "best": best, "best_guard": best_guard, "best_step": best_step, "selected_temperature": selected_temperature,
                 "python_rng": random.getstate(), "torch_rng": torch.get_rng_state(), "cuda_rng": torch.cuda.get_rng_state_all()}
        if scheduled:
            state["training_schedule"] = scheduled.snapshot(examples_seen)
        pending = output / "resume.pt.pending"
        if args.resume_unlink_first:
            (output / "resume.pt").unlink(missing_ok=True)
        torch.save(state, pending)
        with pending.open("rb") as stream:
            os.fsync(stream.fileno())
        os.replace(pending, output / "resume.pt")
        sync_directory(output)
        last_resume_step = step
        record("resume_saved", run=run.name, step=step, bytes=(output / "resume.pt").stat().st_size, seconds=time.monotonic() - began)

    def evaluate_quick() -> None:
        began = time.monotonic()
        temperature_logits = infer(model, quick_temperature, args)
        fitted = fit_temperature(temperature_logits, [label_index(options(row["question"]), hard_label(row)) for row in quick_temperature])
        logits = infer(model, quick_development, args)
        predictions = evaluate_logits(quick_development, logits, fitted)
        guard_logits = infer(model, quick_guard, args)
        guard_predictions = evaluate_logits(quick_guard, guard_logits, fitted)
        save_predictions(run / f"quick-development-{step:05d}.jsonl", predictions)
        save_predictions(run / f"quick-guard-{step:05d}.jsonl", guard_predictions)
        save_predictions(run / f"quick-temperature-{step:05d}.jsonl", evaluate_logits(quick_temperature, temperature_logits, fitted))
        value = record("quick_evaluation", run=run.name, step=step, examples_seen=examples_seen,
                       raw=metrics(evaluate_logits(quick_development, logits)), fitted=metrics(predictions),
                       guard_raw=metrics(evaluate_logits(quick_guard, guard_logits)), guard_fitted=metrics(guard_predictions),
                       temperature=fitted, jev=quick_reference, guard_jev=quick_guard_reference,
                       diagnostic=True, selection_candidate=False, selected_step=best_step,
                       temperature_rows=len(quick_temperature), development_rows=len(quick_development), guard_rows=len(quick_guard),
                       elapsed_seconds=time.monotonic() - started, evaluation_seconds=time.monotonic() - began)
        append(run / "quick-evaluations.jsonl", value)
        print(json.dumps(value), flush=True)

    def save_before_wait() -> None:
        if step > 0 and last_resume_step != step:
            save_resume()

    if step == 0:
        evaluate(include_public=bool(public_rows))
    stop = min(total_steps, args.stop_after) if args.stop_after is not None else total_steps
    while step < stop:
        rows = scheduled.batch(step, args.effective_batch_size, save_before_wait) if scheduled else groups[step]
        began = time.monotonic()
        model.train()
        group = augment(rows, random.Random(args.seed + 100003 * (step + 1)))
        optimizer.zero_grad()
        total_loss = 0.0
        input_tokens = 0
        for batch in microbatches(group, args.batch_size, args.token_budget):
            prepared = model.prepare(batch, max_length=args.max_length)
            logits = model(prepared)
            target = targets(batch, logits.device)
            loss = -(target * F.log_softmax(logits, dim=-1)).sum(-1).mean()
            if not torch.isfinite(loss):
                raise RuntimeError(f"Nonfinite loss at step {step + 1}")
            (loss * len(batch) / len(group)).backward()
            total_loss += float(loss.detach()) * len(batch)
            input_tokens += prepared.input_tokens
            del logits, loss, target, prepared
        gradient_norm = float(torch.nn.utils.clip_grad_norm_(parameters, 1.0))
        if not math.isfinite(gradient_norm):
            raise RuntimeError(f"Nonfinite gradient at step {step + 1}")
        step += 1
        warmup = max(1, int(0.05 * total_steps))
        factor = step / warmup if step <= warmup else 0.1 + 0.45 * (1 + math.cos(math.pi * (step - warmup) / max(1, total_steps - warmup)))
        for group_parameters in optimizer.param_groups:
            group_parameters["lr"] = args.lr * factor
        optimizer_started = time.monotonic()
        optimizer.step()
        torch.cuda.synchronize()
        examples_seen += len(group)
        value = record("training_step", run=run.name, step=step, loss=total_loss / len(group), learning_rate=args.lr * factor,
                       examples_seen=examples_seen, group_examples=len(group), input_tokens=input_tokens, gradient_norm=gradient_norm,
                       step_seconds=time.monotonic() - began, optimizer_seconds=time.monotonic() - optimizer_started,
                       elapsed_seconds=time.monotonic() - started, gpu_peak_gb=torch.cuda.max_memory_allocated() / 1e9)
        append(run / "training.jsonl", value)
        print(json.dumps(value), flush=True)
        if step % args.eval_every == 0 or step == stop:
            evaluate(include_public=step % args.public_eval_every == 0 or step == total_steps,
                     diagnostic=step % args.eval_every != 0 and step != total_steps)
        elif quick_development and step % args.quick_eval_every == 0:
            evaluate_quick()
        if not args.no_resume and (step % args.resume_every == 0 or step == stop):
            save_resume()
    summary = {"run": run.name, "steps": step, "planned_steps": total_steps, "complete": step == total_steps,
               "examples_seen": examples_seen, "best_step": best_step, "selected_temperature": selected_temperature,
               "selected_metrics": best, "selected_guard_metrics": best_guard, "jev": reference, "guard_jev": guard_reference,
               "initial_artifact": initial_artifact, "data_sha256": hashes, "git_commit": revision,
               "elapsed_seconds": time.monotonic() - started, "checkpoint": str(output / "selected") if best else None}
    if scheduled:
        summary["training_schedule"] = scheduled.snapshot(examples_seen)
    write_json(run / "summary.json", summary)
    record("training_finished" if step == total_steps else "training_paused", **summary)
    print(json.dumps(summary, indent=2), flush=True)


if __name__ == "__main__":
    main()