source/experiment.py

"""Deadline-bounded public-data training, exact artifact reconstruction and eval.

Selection uses a recorded validation-only criterion. Calibration and test/whole-family
holdout evaluation happen once in the separate finalize command. All checkpoint
loads require trusted project artifacts; no weights or optimizer state go in Git.
"""
import argparse
from collections import Counter, defaultdict
from datetime import datetime, timezone
import hashlib
import importlib.metadata
import json
import math
import os
from pathlib import Path
import platform
import random
import signal
import statistics
import subprocess
import time

import torch

from training_model import ADAPTER_VERSION, PROMPT_VERSION, TrainableScorer
from selection import SELECTION_METRIC, validation_selection
from data_transition import verify_train_data_transition

STOP = False


class TrainingValidationInterrupted(RuntimeError):
    """A bounded validation pass yielded to a training stop or deadline."""


def request_stop(*args):
    global STOP
    STOP = True


def sha256(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def write_json(path, value):
    path = Path(path)
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n")
    tmp.replace(path)


def save_torch(path, value):
    path = Path(path)
    tmp = path.with_suffix(".tmp")
    torch.save(value, tmp)
    tmp.replace(path)


def mem_available():
    return int(next(line.split()[1] for line in Path("/proc/meminfo").read_text().splitlines()
                    if line.startswith("MemAvailable:"))) * 1024


def guard_memory(device="cuda"):
    Path("/proc/self/oom_score_adj").write_text("0")
    host_snapshot = subprocess.check_output(['free','-b'],text=True,timeout=10)
    gpu_snapshot = subprocess.check_output(
        ['nvidia-smi','--query-compute-apps=pid,process_name,used_memory','--format=csv'],
        text=True,timeout=10) if device == 'cuda' else None
    print(json.dumps({'event':'memory_preflight','pid':os.getpid(),'free_bytes':host_snapshot,
                      'gpu_processes':gpu_snapshot,'mem_available_bytes':mem_available(),
                      'oom_score_adj':Path('/proc/self/oom_score_adj').read_text().strip()}),flush=True)
    if mem_available() < 24 * 2**30:
        raise RuntimeError("Requires at least 24 GiB currently available unified RAM")
    torch.set_num_threads(8)
    if device == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA unavailable")
        torch.cuda.set_per_process_memory_fraction(16 * 2**30 / torch.cuda.get_device_properties(0).total_memory)
        torch.backends.cuda.matmul.allow_tf32 = False


def sync(scorer):
    if scorer.device.type == "cuda":
        torch.cuda.synchronize()


def load_artifact(path, device="cuda"):
    artifact = torch.load(path, map_location="cpu", weights_only=False)
    if artifact.get("format") != "opensysone-adapter-v1":
        raise ValueError("Unrecognized artifact format")
    config = artifact["config"]
    scorer = TrainableScorer(config["model"], rank=config["rank"], alpha=config["alpha"],
                            adapters=config["adapters"], device=device, max_tokens=config["max_tokens"],
                            branch_batch_size=config["branch_batch_size"])
    if scorer.provenance != artifact["model_provenance"]:
        raise ValueError("Artifact/base provenance mismatch")
    if artifact["prompt_version"] != PROMPT_VERSION or artifact["adapter_version"] != ADAPTER_VERSION:
        raise ValueError("Artifact implementation version mismatch")
    scorer.restore_trainable(artifact["trainable_state"])
    return scorer, artifact


def data_for(scorer, dataset, output):
    dataset = Path(dataset)
    source = json.loads((dataset / "manifest.json").read_text())
    for split, checksum in source["split_sha256"].items():
        if sha256(dataset / f"{split}.jsonl") != checksum:
            raise ValueError(f"Frozen dataset hash mismatch: {split}")
    signature = hashlib.sha256(json.dumps({"data": source["split_sha256"], "model": scorer.provenance,
        "implementation": sha256("training_model.py"), "max_tokens": scorer.max_tokens}, sort_keys=True).encode()).hexdigest()
    cache = dataset / f"tokens-{signature[:16]}.pt"
    if cache.exists():
        saved = torch.load(cache, map_location="cpu", weights_only=False)
        if saved["signature"] != signature:
            raise ValueError("Token cache signature mismatch")
        data, audit = saved["data"], saved["audit"]
    else:
        data, audit = {}, {}
        for split in source["split_sha256"]:
            retained, dropped = [], []
            for line in (dataset / f"{split}.jsonl").read_text().splitlines():
                row = json.loads(line)
                try:
                    row["_sequences"] = scorer.sequences(row)
                except ValueError as error:
                    if "no truncation" not in str(error):
                        raise
                    dropped.append(row["id"])
                    continue
                retained.append(row)
            data[split] = retained
            audit[split] = {"retained": len(retained), "dropped_ids": dropped,
                            "family_counts": dict(Counter(r["family"] for r in retained)),
                            "retained_id_sha256": hashlib.sha256("\n".join(r["id"] for r in retained).encode()).hexdigest(),
                            "max_branch_tokens": max((len(s) for r in retained for s in r["_sequences"]), default=0)}
            print(json.dumps({"event": "tokenized", "split": split, "audit": {k:v for k,v in audit[split].items() if k != "dropped_ids"}}), flush=True)
        save_torch(cache, {"signature": signature, "data": data, "audit": audit})
    write_json(Path(output) / "data_filter.json", audit)
    return data, signature


def validation_cut(rows, per_family=32):
    groups = defaultdict(list)
    for row in rows:
        groups[row["family"]].append(row)
    return [row for family in sorted(groups) for row in groups[family][:per_family]]


@torch.inference_mode()
def predict(scorer, rows, temperature=1.0, token=False, deadline=None):
    scorer.eval()
    result = []
    start = time.perf_counter()
    for index,row in enumerate(rows):
        if deadline is not None and (STOP or time.time() >= deadline):
            raise TrainingValidationInterrupted("Training validation interrupted before completion")
        scores = scorer.scores_token_baseline([row]) if token else scorer.score_examples([row])
        logits = scores[0].float() / temperature
        result.append({"id": row["id"], "group": row["group"], "family": row["family"],
                       "target": row["target"], "choices": row["choices"], "logits": scores[0].float().tolist(),
                       "probabilities": logits.softmax(0).tolist(), "log_probabilities": logits.log_softmax(0).tolist()})
        if (index+1) % 128 == 0:
            print(json.dumps({'event':'prediction_progress','decisions':index+1,'total':len(rows),
                              'elapsed_seconds':time.perf_counter()-start}),flush=True)
    return result


def metrics(predictions):
    if not predictions:
        raise ValueError("Cannot report an empty evaluation")
    bins = [{"count": 0, "confidence_sum": 0.0, "correct_sum": 0.0} for _ in range(10)]
    values = []
    for row in predictions:
        probabilities = row["probabilities"]
        predicted = max(range(len(probabilities)), key=probabilities.__getitem__)
        confidence = probabilities[predicted]
        correct = float(predicted == row["target"])
        nll = -row["log_probabilities"][row["target"]]
        brier = sum((p - float(i == row["target"])) ** 2 for i,p in enumerate(probabilities))
        values.append((correct, nll, brier, confidence))
        bucket = bins[min(9, int(confidence * 10))]
        bucket["count"] += 1
        bucket["confidence_sum"] += confidence
        bucket["correct_sum"] += correct
    n = len(values)
    sorted_values = sorted(values, key=lambda v:v[3], reverse=True)
    coverage = {}
    for fraction in (0.25, 0.5, 0.75, 1.0):
        selected = sorted_values[:max(1, math.ceil(n * fraction))]
        coverage[str(fraction)] = {"n": len(selected), "accuracy": statistics.mean(v[0] for v in selected),
                                  "min_confidence": selected[-1][3]}
    return {"n": n, "accuracy": statistics.mean(v[0] for v in values),
            "nll": statistics.mean(v[1] for v in values), "brier_multiclass_sum": statistics.mean(v[2] for v in values),
            "ece_top_label_10_equal_width_bins": sum(abs(b["confidence_sum"] - b["correct_sum"]) for b in bins) / n,
            "reliability_bins": bins, "accuracy_vs_coverage": coverage}


def report(predictions):
    grouped = defaultdict(list)
    for row in predictions:
        grouped[row["family"]].append(row)
    return {"overall": metrics(predictions), "per_family": {f:metrics(rows) for f,rows in grouped.items()}}


def objective(predictions):
    return statistics.mean(value["nll"] for value in report(predictions)["per_family"].values())


def decision_backward(scorer, row, divisor, two_pass=False):
    """Exact categorical gradient with one candidate graph alive in two-pass mode.

    Requires deterministic, dropout-free forwards. The first pass supplies the
    softmax derivative; recomputation applies that derivative to each scalar score.
    """
    if not two_pass:
        score = scorer.score_examples([row])[0].float()
        loss = -score.log_softmax(0)[row["target"]]
        if not torch.isfinite(loss):
            raise RuntimeError("Non-finite training loss")
        (loss / divisor).backward()
        return loss.item()
    with torch.no_grad():
        score = scorer.score_examples([row])[0].float()
        loss = -score.log_softmax(0)[row["target"]]
        derivative = score.softmax(0)
        derivative[row["target"]] -= 1
    if not torch.isfinite(loss):
        raise RuntimeError("Non-finite training loss")
    for index, sequence in enumerate(row["_sequences"]):
        value = scorer.score_examples([{"_sequences": [sequence]}])[0][0]
        value.backward(gradient=derivative[index] / divisor)
    return loss.item()


@torch.inference_mode()
def correctness(scorer, rows, token_initial=False, allow_stop=False):
    scorer.eval()
    chosen = validation_cut(rows, 2)
    original_size = scorer.branch_batch_size
    def scores(examples, token=False):
        if allow_stop and STOP:
            raise TrainingValidationInterrupted("Final training correctness stopped before completion")
        return scorer.scores_token_baseline(examples) if token else scorer.score_examples(examples)
    base = [s.detach().clone() for s in scores(chosen)]
    def difference(a,b):
        return max((x.float().softmax(0) - y.float().softmax(0)).abs().max().item() for x,y in zip(a,b))
    checks = {}
    try:
        for size in (1, 2, 4):
            scorer.branch_batch_size = size
            checks[f"branch_chunks_{size}_probability_max_abs"] = difference(base, scores(chosen))
        scorer.branch_batch_size = original_size
        single = [scores([row])[0] for row in chosen]
        checks["question_isolation_probability_max_abs"] = difference(base, single)
        reversed_rows = []
        for row in chosen:
            reversed_rows.append({**row, "choices": list(reversed(row["choices"])),
                                  "_sequences": list(reversed(row["_sequences"]))})
        flipped = [s.flip(0) for s in scores(reversed_rows)]
        checks["candidate_permutation_probability_max_abs"] = difference(base, flipped)
        checks["repeat_probability_max_abs"] = difference(base, scores(chosen))
        if token_initial:
            checks["pretrained_readout_probability_max_abs"] = difference(base, scores(chosen, token=True))
    finally:
        scorer.branch_batch_size = original_size
    checks["tolerance_probability_abs"] = 1e-4
    if any(value > 1e-4 for key,value in checks.items() if key.endswith("max_abs")):
        raise RuntimeError(f"FP32 correctness gate failed: {checks}")
    return checks


def restore_warm_start(scorer, saved, config, signature):
    """Restore compatible trained weights while retaining a fresh optimizer/RNG.

    Optimization settings may change. An explicitly requested data expansion
    verifies both datasets and preserves every reserved evaluation split.
    """
    if saved.get("format") != "opensysone-adapter-v1":
        raise ValueError("Unrecognized warm-start artifact format")
    if (saved.get("prompt_version") != PROMPT_VERSION or
            saved.get("adapter_version") != ADAPTER_VERSION):
        raise ValueError("Warm-start implementation version mismatch")
    for key in ("rank", "alpha", "adapters"):
        if saved["config"].get(key) != config[key]:
            raise ValueError(f"Warm-start must preserve {key}")
    if saved["model_provenance"] != scorer.provenance:
        raise ValueError("Warm-start model/data mismatch")
    transition = None
    if config.get("allow_train_data_change", False):
        transition = verify_train_data_transition(saved, config, signature)
    elif saved["data_signature"] != signature:
        raise ValueError("Warm-start model/data mismatch")
    scorer.restore_trainable(saved["trainable_state"])
    return transition


def validation_fits(deadline, measured_seconds, now=None):
    """Do not start a full validation pass that would consume the stop margin."""
    now = time.time() if now is None else now
    return now + max(60, measured_seconds * 1.25 + 30) < deadline


def selection_result(predictions, metric):
    if metric == "raw_nll":
        return {"metric": metric, "score": objective(predictions),
                "raw_macro_nll": objective(predictions)}
    if metric != SELECTION_METRIC:
        raise ValueError("Unknown checkpoint selection metric")
    return validation_selection(predictions)


def reselect_inherited(artifact, predictions, metric):
    """Change selection metadata only; retained weights/provenance stay intact."""
    result = selection_result(predictions, metric)
    updated = dict(artifact)
    updated.update(selection_metric=metric,
                   best_validation_macro_nll=result["raw_macro_nll"],
                   best_validation_selection_score=result["score"],
                   validation_selection=result)
    return updated, result


def train(args):
    if getattr(args, "allow_train_data_change", False) and not args.warm_start:
        raise ValueError("--allow-train-data-change requires --warm-start")
    out = Path(args.output).resolve()
    out.mkdir(parents=True, exist_ok=False)
    guard_memory()
    signal.signal(signal.SIGTERM, request_stop)
    signal.signal(signal.SIGINT, request_stop)
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    config = vars(args).copy()
    config["model"] = str(Path(args.model).resolve())
    config["dataset"] = str(Path(args.dataset).resolve())
    config["adapters"] = not args.head_only
    config["selection_metric"] = getattr(args, "selection_metric", "raw_nll")
    parent_path = args.resume or args.warm_start
    initialization = {"kind": "resume" if args.resume else "warm_start" if args.warm_start else "pretrained",
                      "parent_checkpoint": str(Path(parent_path).resolve()) if parent_path else None,
                      "parent_checkpoint_sha256": sha256(parent_path) if parent_path else None,
                      "restores_optimizer": bool(args.resume), "restores_rng": bool(args.resume)}
    manifest = {"config": config, "pid": os.getpid(), "hostname": platform.node(),
        "started_utc": datetime.now(timezone.utc).isoformat(),
        "git_commit": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(),
        "git_status": subprocess.check_output(["git", "status", "--porcelain"], text=True),
        "source_sha256": {str(p):sha256(p) for p in [*Path('.').glob('*.py'), *Path('scripts').glob('*')] if p.is_file()},
        "packages": {p:importlib.metadata.version(p) for p in ("torch", "transformers", "pyarrow", "numpy")},
        "cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name(),
        "capability": torch.cuda.get_device_capability(), "cuda_cap_bytes": 16 * 2**30,
        "initial_mem_available_bytes": mem_available(), "oom_score_adj": Path('/proc/self/oom_score_adj').read_text().strip(),
        "parent_checkpoint_sha256": initialization["parent_checkpoint_sha256"],
        "initialization": initialization}
    write_json(out / "manifest.json", manifest)
    start = time.perf_counter()
    scorer = TrainableScorer(config["model"], rank=args.rank, alpha=args.alpha,
                            adapters=config["adapters"], max_tokens=args.max_tokens,
                            branch_batch_size=args.branch_batch_size)
    if args.two_pass:
        if args.branch_batch_size != 1:
            raise ValueError("Two-pass training requires branch_batch_size=1 to preserve recomputation shapes")
        if any(isinstance(m, torch.nn.Dropout) and m.p for m in scorer.modules()) or getattr(scorer.lm.config, "attention_dropout", 0):
            raise ValueError("Two-pass gradients require dropout-free forwards")
    data, signature = data_for(scorer, args.dataset, out)
    manifest.update(model_provenance=scorer.provenance, data_signature=signature,
                    total_parameters=sum(p.numel() for p in scorer.parameters()),
                    trainable_parameters=sum(p.numel() for p in scorer.parameters() if p.requires_grad),
                    adapter_modules=scorer.adapter_names, load_and_tokenize_seconds=time.perf_counter() - start)
    write_json(out / "manifest.json", manifest)
    optimizer = torch.optim.AdamW([
        {"params": [p for p in scorer.lm.parameters() if p.requires_grad], "lr": args.lr},
        {"params": scorer.head.parameters(), "lr": args.head_lr}], weight_decay=0.01)
    parameters = [p for p in scorer.parameters() if p.requires_grad]
    maximum_steps = math.ceil(len(data["train"]) * args.epochs / args.effective_batch)
    target_steps = min(args.steps, maximum_steps) if args.steps else maximum_steps
    completed, best, best_selection, stale = 0, math.inf, math.inf, 0
    parent_best = None
    if args.warm_start:
        saved = torch.load(args.warm_start, map_location="cpu", weights_only=False)
        transition = restore_warm_start(scorer, saved, config, signature)
        initialization.update(parent_step=saved["step"], parent_source_commit=saved["source_commit"])
        if transition is not None:
            initialization["data_transition"] = transition
        del saved
    if args.resume:
        saved = torch.load(args.resume, map_location="cpu", weights_only=False)
        initialization.update(parent_step=saved["step"], parent_source_commit=saved["source_commit"],
                              parent_initialization=saved.get("initialization"))
        for key in ("rank", "alpha", "adapters", "max_tokens", "branch_batch_size", "seed", "effective_batch", "lr", "head_lr", "schedule_steps", "epochs", "validation_per_family", "two_pass"):
            if saved["config"].get(key,False) != config[key]:
                raise ValueError(f"Resume must preserve {key}")
        if saved["model_provenance"] != scorer.provenance or saved["data_signature"] != signature:
            raise ValueError("Resume model/data mismatch")
        scorer.restore_trainable(saved["trainable_state"])
        optimizer.load_state_dict(saved["optimizer"])
        random.setstate(saved["random_state"])
        torch.set_rng_state(saved["torch_rng"])
        torch.cuda.set_rng_state_all(saved["cuda_rng"])
        completed, best, stale = saved["step"], saved["best_validation_macro_nll"], saved["stale_evaluations"]
        previous_metric = saved.get("selection_metric", "raw_nll")
        initialization["selection_policy_change"] = {
            "from": previous_metric, "to": config["selection_metric"],
            "optimizer_and_rng_unchanged": True}
        parent_best = Path(args.resume).resolve().parent / "best.pt"
        if not parent_best.exists():
            raise ValueError("Resume requires the parent's validation-selected best.pt")
        inherited = torch.load(parent_best, map_location="cpu", weights_only=False)
        if inherited["data_signature"] != signature:
            raise ValueError("Parent best artifact data mismatch")
        # Legacy campaigns copied an inherited best artifact without its raw
        # prediction file. Its recorded output locates that original evidence.
        evidence_directories = [parent_best.parent]
        if inherited["config"].get("output"):
            evidence_directories.append(Path(inherited["config"]["output"]))
        candidates = []
        for directory in evidence_directories:
            candidates.extend([directory / f"validation_step_{inherited['step']:06d}_predictions.json",
                               directory / "best_validation_predictions.json"])
            if inherited["step"] == 0:
                candidates.append(directory / "initial_validation_predictions.json")
        best_predictions = None
        for candidate in candidates:
            if candidate.exists():
                candidate_predictions = json.loads(candidate.read_text())
                # A canonical file can belong to a different step after interruption.
                if abs(objective(candidate_predictions) - inherited["best_validation_macro_nll"]) > 2e-5:
                    continue
                best_predictions = candidate_predictions
                write_json(out / "best_validation_predictions.json", best_predictions)
                write_json(out / f"validation_step_{inherited['step']:06d}_predictions.json", best_predictions)
                break
        if best_predictions is None:
            raise ValueError("Inherited best requires matching raw validation evidence")
        inherited, inherited_selection = reselect_inherited(inherited, best_predictions, config["selection_metric"])
        best, best_selection = inherited_selection["raw_macro_nll"], inherited_selection["score"]
        if previous_metric != config["selection_metric"]:
            stale = 0
        write_json(out / "inherited_validation_selection.json", inherited_selection)
        write_json(out / "best_validation_selection.json", inherited_selection)
        save_torch(out / "best.pt", inherited)
        del saved, inherited
    write_json(out / "manifest.json", manifest)

    def checkpoint(path, step, resumable=True):
        artifact = {"format": "opensysone-adapter-v1", "step": step, "config": config,
                    "model_provenance": scorer.provenance, "data_signature": signature,
                    "prompt_version": PROMPT_VERSION, "adapter_version": ADAPTER_VERSION,
                    "trainable_state": scorer.trainable_state(), "best_validation_macro_nll": best,
                    "selection_metric": config["selection_metric"],
                    "best_validation_selection_score": best_selection,
                    "stale_evaluations": stale, "source_commit": manifest["git_commit"],
                    "initialization": initialization}
        if resumable:
            artifact.update(optimizer=optimizer.state_dict(), random_state=random.getstate(),
                            torch_rng=torch.get_rng_state(), cuda_rng=torch.cuda.get_rng_state_all())
        save_torch(path, artifact)

    # A durable reconstruction exists before any correctness/evaluation work.
    checkpoint(out / "checkpoint.pt", completed)
    checks = correctness(scorer, data["validation"], token_initial=not parent_path)
    write_json(out / "correctness_initial.json", checks)
    val = validation_cut(data["validation"], args.validation_per_family)
    validation_started = time.monotonic()
    predictions = predict(scorer, val)
    validation_seconds = time.monotonic() - validation_started
    write_json(out / "resumed_initial_predictions.json" if args.resume else out / "initial_validation_predictions.json", predictions)
    initial_report = report(predictions)
    initial_objective = objective(predictions)
    initial_selection = selection_result(predictions, config["selection_metric"])
    write_json(out / "initial_validation_selection.json", initial_selection)
    if args.resume:
        write_json(out / f"validation_step_{completed:06d}_predictions.json", predictions)
    initial_improved = initial_selection["score"] < best_selection - 0.001
    if not args.resume or initial_improved:
        best, best_selection, stale = initial_objective, initial_selection["score"], 0
        write_json(out / "best_validation_predictions.json", predictions)
        write_json(out / "best_validation_selection.json", initial_selection)
        checkpoint(out / "best.pt", completed, False)
    if args.resume:
        print(json.dumps({"event": "resumed_initial_selection", "step": completed,
                          "macro_nll": initial_objective, "best_macro_nll": best,
                          "selection_metric": config["selection_metric"],
                          "selection_score": initial_selection["score"], "best_selection_score": best_selection,
                          "improved": initial_improved}), flush=True)
    # Persist an initial/resumed selection before a stop can interrupt new updates.
    checkpoint(out / "checkpoint.pt", completed)
    print(json.dumps({"event": "ready", "output": str(out), "trainable_parameters": manifest["trainable_parameters"],
                      "validation_macro_nll": initial_objective, "target_steps": target_steps,
                      "correctness": checks}), flush=True)
    deadline = datetime.fromisoformat(args.deadline.replace("Z", "+00:00")).timestamp() if args.deadline else math.inf
    last_save = time.monotonic()
    history = []
    orders = {}
    train_start = time.monotonic()
    status = "completed_step_target" if args.steps and args.steps < maximum_steps else "completed_epochs"
    for step in range(completed + 1, target_steps + 1):
        if STOP or time.time() >= deadline:
            status = "interrupted" if STOP else "training_deadline"
            break
        if mem_available() < 16 * 2**30:
            checkpoint(out / "checkpoint.pt", completed)
            raise RuntimeError("Host availability below 16 GiB; checkpoint saved")
        scorer.train()
        optimizer.zero_grad(set_to_none=True)
        batch = []
        for absolute in range((step - 1) * args.effective_batch,
                              min(step * args.effective_batch, len(data["train"]) * args.epochs)):
            epoch, index = divmod(absolute, len(data["train"]))
            if epoch not in orders:
                order = list(range(len(data["train"])))
                random.Random(args.seed + epoch).shuffle(order)
                orders = {epoch:order}
            batch.append(data["train"][orders[epoch][index]])
        sync(scorer)
        tick = time.perf_counter()
        loss_sum = actual_tokens = padded_tokens = branches = 0
        for row in batch:
            loss_sum += decision_backward(scorer,row,len(batch),args.two_pass)
            sequences = row["_sequences"]
            actual_tokens += sum(map(len, sequences))
            branches += len(sequences)
            for start_index in range(0, len(sequences), args.branch_batch_size):
                chunk = sequences[start_index:start_index + args.branch_batch_size]
                padded_tokens += max(map(len, chunk)) * len(chunk)
        norm = torch.nn.utils.clip_grad_norm_(parameters, 1.0, error_if_nonfinite=True)
        # Fixed schedule length is persisted separately for resumable pilots/campaigns.
        schedule_steps = args.schedule_steps or maximum_steps
        warmup = max(1, min(100, schedule_steps // 20))
        factor = min(1.0, step / warmup) if step <= warmup else max(0.1, 0.5 * (1 + math.cos(math.pi * min(1, (step - warmup) / max(1, schedule_steps - warmup)))))
        for group, lr in zip(optimizer.param_groups, (args.lr, args.head_lr)):
            group["lr"] = lr * factor
        optimizer.step()
        sync(scorer)
        completed = step
        item = {"step": step, "loss": loss_sum / len(batch), "gradient_norm": norm.item(),
                "seconds": time.perf_counter() - tick, "decisions": len(batch), "branches": branches,
                "actual_branch_tokens": actual_tokens, "padded_branch_tokens": padded_tokens,
                "peak_cuda_allocated_bytes": torch.cuda.max_memory_allocated(), "lr_factor": factor}
        history.append(item)
        with (out / "training.jsonl").open("a") as handle:
            handle.write(json.dumps(item) + "\n")
        if step % 10 == 0 or step == target_steps:
            print(json.dumps({"event": "step", **item}), flush=True)
        if step % args.save_steps == 0 or time.monotonic() - last_save >= args.save_seconds:
            checkpoint(out / "checkpoint.pt", completed)
            last_save = time.monotonic()
        if step % args.eval_steps == 0 or step == target_steps:
            checkpoint(out / "checkpoint.pt", completed)
            if not validation_fits(deadline, validation_seconds):
                print(json.dumps({"event": "validation_skipped_for_deadline", "step": completed,
                                  "measured_validation_seconds": validation_seconds}), flush=True)
                continue
            validation_started = time.monotonic()
            try:
                predictions = predict(scorer, val, deadline=deadline)
            except TrainingValidationInterrupted:
                status = "interrupted" if STOP else "training_deadline"
                print(json.dumps({"event": "validation_interrupted", "step": completed,
                                  "status": status}), flush=True)
                break
            validation_seconds = max(validation_seconds, time.monotonic() - validation_started)
            score = objective(predictions)
            selected_score = selection_result(predictions, config["selection_metric"])
            improved = selected_score["score"] < best_selection - 0.001
            # Save evidence before publishing a newly selected checkpoint. Fleet
            # selection can always retrieve predictions for the durable best step.
            write_json(out / f"validation_step_{completed:06d}_predictions.json", predictions)
            if improved:
                best, best_selection, stale = score, selected_score["score"], 0
                write_json(out / "best_validation_predictions.json", predictions)
                write_json(out / "best_validation_selection.json", selected_score)
                checkpoint(out / "best.pt", completed, False)
            else:
                stale += 1
            validation = {"step": completed, "macro_nll": score, "best_macro_nll": best,
                          "selection": selected_score, "best_selection_score": best_selection,
                          "improved": improved, "metrics": report(predictions)}
            with (out / "validation.jsonl").open("a") as handle:
                handle.write(json.dumps(validation) + "\n")
            print(json.dumps({"event": "validation", "step": completed, "macro_nll": score, "best": best}), flush=True)
            checkpoint(out / "checkpoint.pt", completed)
            last_save = time.monotonic()
            if stale >= args.patience:
                status = "validation_early_stop"
                break
    checkpoint(out / "checkpoint.pt", completed)
    try:
        checks = correctness(scorer, data["validation"], allow_stop=True)
    except TrainingValidationInterrupted:
        checks = None
        final_correctness_status = "skipped_on_stop"
        status = "interrupted"
    else:
        final_correctness_status = "passed"
        write_json(out / "correctness_final.json", checks)
    summary = {"status": status, "completed_steps": completed, "target_steps": target_steps,
               "initial_validation": initial_report, "best_validation_macro_nll": best,
               "selection_metric": config["selection_metric"], "best_validation_selection_score": best_selection,
               "training_seconds": time.monotonic() - train_start,
               "median_step_seconds": statistics.median(item["seconds"] for item in history) if history else None,
               "processed_decisions": sum(item["decisions"] for item in history),
               "actual_branch_tokens": sum(item["actual_branch_tokens"] for item in history),
               "padded_branch_tokens": sum(item["padded_branch_tokens"] for item in history),
               "peak_cuda_allocated_bytes": torch.cuda.max_memory_allocated(),
               "peak_cuda_reserved_bytes": torch.cuda.max_memory_reserved(),
               "final_mem_available_bytes": mem_available(), "correctness": checks,
               "final_correctness_status": final_correctness_status,
               "checkpoint_sha256": sha256(out / "checkpoint.pt"), "best_sha256": sha256(out / "best.pt")}
    write_json(out / "summary.json", summary)
    print(json.dumps({"event": "training_complete", **summary}), flush=True)


def fit_temperature(predictions):
    candidates = torch.logspace(-1, 1.3, 101).tolist()
    return min(candidates, key=lambda t: statistics.mean(
        -torch.tensor(row["logits"]).div(t).log_softmax(0)[row["target"]].item() for row in predictions))


def with_temperature(predictions, temperature):
    return [{**row, "probabilities": (torch.tensor(row["logits"]) / temperature).softmax(0).tolist(),
             "log_probabilities": (torch.tensor(row["logits"]) / temperature).log_softmax(0).tolist()}
            for row in predictions]


def bootstrap_difference(base, tuned, repetitions=400):
    import numpy as np
    if [row["id"] for row in base] != [row["id"] for row in tuned]:
        raise ValueError("Bootstrap needs matched predictions")
    groups = defaultdict(list)
    def losses(row):
        p = row["probabilities"]
        return [float(max(range(len(p)), key=p.__getitem__) == row["target"]),
                -row["log_probabilities"][row["target"]],
                sum((v - float(i == row["target"])) ** 2 for i,v in enumerate(p))]
    for a,b in zip(base,tuned):
        groups[(a["family"],a["group"])].append(np.asarray(losses(b)) - np.asarray(losses(a)))
    by_family = defaultdict(list)
    for (family,_),rows in groups.items():
        by_family[family].append((np.sum(rows,axis=0),len(rows)))
    arrays = [(np.asarray([r[0] for r in rows]),np.asarray([r[1] for r in rows]))
              for rows in by_family.values()]
    rng = np.random.default_rng(907)
    samples = []
    for _ in range(repetitions):
        total, count = np.zeros(3), 0
        for sums,counts in arrays:
            selected = rng.integers(0,len(sums),len(sums))
            total += sums[selected].sum(axis=0)
            count += counts[selected].sum()
        samples.append(total/count)
    lo, hi = np.quantile(samples,[0.025,0.975],axis=0)
    point = sum((sums.sum(axis=0) for sums,_ in arrays),np.zeros(3))/len(base)
    return {"method": f"{repetitions} stratified source-group bootstrap resamples; decision-weighted tuned minus base",
            "point_delta":dict(zip(('accuracy','nll','brier'),map(float,point))),
            "accuracy": [float(lo[0]),float(hi[0])], "nll": [float(lo[1]),float(hi[1])],
            "brier": [float(lo[2]),float(hi[2])]}


def finalize(args):
    out = Path(args.output).resolve()
    out.mkdir(parents=True, exist_ok=False)
    guard_memory()
    scorer, selected = load_artifact(args.checkpoint)
    data, signature = data_for(scorer, args.dataset, out)
    if signature != selected["data_signature"]:
        raise ValueError("Evaluation dataset differs from training")
    write_json(out / "manifest.json", {"pid": os.getpid(), "source_commit": subprocess.check_output(
        ['git','rev-parse','HEAD'],text=True).strip(), "checkpoint_sha256": sha256(args.checkpoint),
        "selected_step": selected["step"], "data_signature": signature,
        "model_provenance": scorer.provenance, "training_config": selected["config"],
        "training_source_commit": selected["source_commit"],
        "source_sha256": {str(p):sha256(p) for p in [*Path('.').glob('*.py'), *Path('scripts').glob('*')] if p.is_file()},
        "packages": {p:importlib.metadata.version(p) for p in ("torch", "transformers", "pyarrow", "numpy")},
        "hostname": platform.node(), "cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name(),
        "cuda_cap_bytes": 16 * 2**30, "initial_mem_available_bytes": mem_available(),
        "oom_score_adj": Path('/proc/self/oom_score_adj').read_text().strip(),
        "selection": {"metric": selected.get("selection_metric", "raw_nll"),
                      "score": selected.get("best_validation_selection_score", selected["best_validation_macro_nll"]),
                      "raw_macro_nll": selected["best_validation_macro_nll"],
                      "scope": "validation only; reserved calibration/test/holdout not used"},
        "started_utc": datetime.now(timezone.utc).isoformat()})
    write_json(out / "correctness.json", correctness(scorer,data["validation"]))
    calibration = predict(scorer,data["calibration"])
    write_json(out / "calibration_predictions.json",calibration)
    temperature = fit_temperature(calibration)
    selected["temperature"] = temperature
    selected["calibration_status"] = "one global temperature fitted on separate known-family calibration data; whole-family calibration unproven"
    selected.pop("optimizer",None)
    # The deployable artifact is durable before untouched evaluation starts.
    save_torch(out / "model.pt",selected)
    result = {"temperature": temperature, "selected_step":selected["step"], "status":"pending_test",
              "claim_scope":"Public decision benchmark; no claim of Jev-level intelligence or general calibration"}
    all_tuned = {}
    for split in ("test","holdout"):
        predictions = predict(scorer,data[split])
        calibrated = with_temperature(predictions,temperature)
        all_tuned[split] = calibrated
        write_json(out / f"{split}_trained_predictions.json",predictions)
        write_json(out / f"{split}_calibrated_predictions.json",calibrated)
        result[split] = {"trained":report(predictions), "calibrated":report(calibrated)}
        write_json(out / "metrics.json",result)
        print(json.dumps({"event":"evaluation","split":split,"metrics":result[split]}),flush=True)
    # Reconstruct the unchanged pretrained readout in the same model allocation.
    # Reset adapters without creating a second backbone on the GPU.
    with torch.no_grad():
        for name,p in scorer.named_parameters():
            if name.endswith("adapter_b"):
                p.zero_()
        scorer.head.weight.copy_((scorer.lm.lm_head.weight[scorer.yes_no_ids[0]] -
                                 scorer.lm.lm_head.weight[scorer.yes_no_ids[1]]).unsqueeze(0))
        bias = scorer.lm.lm_head.bias
        scorer.head.bias.fill_(0 if bias is None else bias[scorer.yes_no_ids[0]]-bias[scorer.yes_no_ids[1]])
    base_calibration = predict(scorer,data["calibration"])
    base_temperature = fit_temperature(base_calibration)
    write_json(out / "base_calibration_predictions.json",base_calibration)
    result["base_temperature"] = base_temperature
    for split in ("test","holdout"):
        predictions = predict(scorer,data[split])
        calibrated = with_temperature(predictions,base_temperature)
        write_json(out / f"{split}_base_predictions.json",predictions)
        write_json(out / f"{split}_base_calibrated_predictions.json",calibrated)
        result[split].update(base=report(predictions),base_calibrated=report(calibrated),
                             calibrated_difference_95pct=bootstrap_difference(calibrated,all_tuned[split]))
    result.update(status="complete",model_sha256=sha256(out / "model.pt"),
                  peak_cuda_allocated_bytes=torch.cuda.max_memory_allocated(),
                  peak_cuda_reserved_bytes=torch.cuda.max_memory_reserved(),final_mem_available_bytes=mem_available())
    write_json(out / "metrics.json",result)
    print(json.dumps({"event":"finalized","output":str(out),"model_sha256":result["model_sha256"]}),flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command",required=True)
    training = sub.add_parser("train")
    training.add_argument("--model",required=True)
    training.add_argument("--dataset",required=True)
    training.add_argument("--output",required=True)
    initialization = training.add_mutually_exclusive_group()
    initialization.add_argument("--resume", help="Restore weights, optimizer and RNG without changing training configuration")
    initialization.add_argument("--warm-start", help="Initialize compatible trained weights with a fresh optimizer and training configuration")
    training.add_argument("--allow-train-data-change", action="store_true",
                          help="Explicit warm-start expansion; verify parent lineage and preserve all reserved data")
    training.add_argument("--steps",type=int)
    training.add_argument("--epochs",type=int,default=3)
    training.add_argument("--rank",type=int,default=16)
    training.add_argument("--alpha",type=float,default=32)
    training.add_argument("--head-only",action="store_true")
    training.add_argument("--two-pass",action="store_true")
    training.add_argument("--lr",type=float,default=1e-4)
    training.add_argument("--head-lr",type=float,default=1e-4)
    training.add_argument("--seed",type=int,default=431)
    training.add_argument("--effective-batch",type=int,default=4)
    training.add_argument("--branch-batch-size",type=int,default=2)
    training.add_argument("--max-tokens",type=int,default=768)
    training.add_argument("--save-steps",type=int,default=250)
    training.add_argument("--save-seconds",type=int,default=900)
    training.add_argument("--eval-steps",type=int,default=500)
    training.add_argument("--validation-per-family",type=int,default=32)
    training.add_argument("--patience",type=int,default=8)
    training.add_argument("--deadline")
    training.add_argument("--schedule-steps",type=int)
    training.add_argument("--selection-metric", choices=("raw_nll", SELECTION_METRIC), default="raw_nll",
                          help="Validation-only checkpoint selection; changing it preserves optimizer/RNG but re-scores inherited evidence")
    evaluation = sub.add_parser("finalize")
    evaluation.add_argument("--checkpoint",required=True)
    evaluation.add_argument("--dataset",required=True)
    evaluation.add_argument("--output",required=True)
    args = parser.parse_args()
    if args.command == "train":
        if any(getattr(args,key) <= 0 for key in ('epochs','rank','alpha','lr','head_lr','effective_batch','branch_batch_size','max_tokens','save_steps','save_seconds','eval_steps','validation_per_family','patience')):
            parser.error("training sizes, rates and cadences must be positive")
        if args.steps is not None and args.steps <= 0:
            parser.error("--steps must be positive")
        if args.schedule_steps is not None and args.schedule_steps <= 0:
            parser.error("--schedule-steps must be positive")
        train(args)
    else:
        finalize(args)


if __name__ == "__main__":
    main()