Blame for src/dlb/recipes.py

"""Paper-locked server recipes for baseline training and distillation.

This module is intentionally importable without PyTorch.  Checkpoint tensor
loading happens lazily only when a real server distillation run is launched.
"""

from __future__ import annotations

import argparse
from copy import deepcopy
from dataclasses import asdict, dataclass, replace
from datetime import datetime, timezone
import hashlib
import inspect
import json
import os
from pathlib import Path, PurePosixPath
import shutil
import stat
import subprocess
import sys
from typing import Collection, Mapping, MutableMapping, Sequence

from dlb.io import atomic_json_write, ensure_safe_directory, sha256_file


class RecipeError(ValueError):
    """A requested training run does not satisfy the locked recipe contract."""


class RecipeProcessError(RecipeError):
    """An upstream process failed and its exit status must be preserved."""

    def __init__(self, message: str, exit_status: int) -> None:
        super().__init__(message)
        self.exit_status = exit_status


BACKBONE_TENSOR_KEYS = (
    "vocab_embed.embedding",
    "output_layer.linear.weight",
    "output_layer.linear.bias",
)
OPTIONAL_BACKBONE_STATE_KEYS = frozenset({"rotary_emb.inv_freq"})

SOURCE_COMMITS = {
    "flm": "a1918d5164e5038e37d0b7a4fb2010ce75b863b3",
    "duo": "7c9b498f5b717de064d6fad7e2509c866e6cb620",
    "mdlm": "c112c526d193436838c98d81455ee51f90309470",
    "candi": "cd57ae9eec98d6ac71cd52bdc50eeec8dfd70f91",
    "rdlm": "67443aa6a2d0fa981eb7c6105f9cbc563e59c5c1",
    "sdtt": "1150985e90b8f2d5749e4469d5154eff9ec922c4",
    "di4c": "ac61ff9fe8e85120f9e1d2a8c5a332f8b8353dd3",
}

TOKENIZER_REVISIONS = {
    "lm1b": ("bert-base-uncased", "86b5e0934494bd15c9632b12f734a8a67f723594"),
    "owt": ("gpt2", "607a30d783dfa663caf39e06633721c8d4cfcd7e"),
}
SEQUENCE_LENGTHS = {"lm1b": 128, "owt": 1024}
BASE_VOCAB_SIZES = {"lm1b": 30_522, "owt": 50_257}


@dataclass(frozen=True)
class TrainingRecipe:
    model: str
    dataset: str
    source: str
    source_path: str
    source_commit: str
    entrypoint: str
    max_steps: int
    global_batch_size: int
    learning_rate: float
    warmup_steps: int
    optimizer: str
    precision: str
    sequence_length: int
    checkpoint_every: int
    default_per_device_batch_size: int
    evidence: str
    max_steps_override: str
    learning_rate_override: str
    output_override: str
    algorithm: str | None = None
    rounds: int | None = None
    steps_per_round: int | None = None
    teacher_family: str | None = None
    teacher_adapter: str | None = None
    distillation_loss: str | None = None
    sampling_step: int | None = None
    release_script_max_steps: int | None = None
    release_evidence: str | None = None


@dataclass(frozen=True)
class LaunchSpec:
    recipe: TrainingRecipe
    command: tuple[str, ...]
    cwd: Path
    entrypoint: Path
    data_path: Path
    output: Path
    teacher: Path | None
    adapted_teacher: Path | None
    devices: int
    nodes: int
    per_device_batch_size: int
    gradient_accumulation: int
    effective_global_batch_size: int
    seed: int
    resume: bool


def _base_recipe(
    model: str,
    dataset: str,
    *,
    source: str,
    entrypoint: str = "main.py",
    max_steps: int = 1_000_000,
    global_batch_size: int = 512,
    learning_rate: float = 3e-4,
    warmup_steps: int = 2_500,
    optimizer: str = "adam",
    precision: str = "bf16",
    checkpoint_every: int = 20_000,
    default_per_device_batch_size: int | None = None,
    evidence: str = "paper_v3_main",
    max_steps_override: str = "trainer.max_steps",
    learning_rate_override: str = "optim.lr",
    output_override: str = "checkpointing.save_dir",
    **values: object,
) -> TrainingRecipe:
    if dataset not in SEQUENCE_LENGTHS:
        raise RecipeError(f"unsupported dataset {dataset!r}")
    per_device = default_per_device_batch_size
    if per_device is None:
        per_device = 64 if dataset == "lm1b" else 16
    return TrainingRecipe(
        model=model,
        dataset=dataset,
        source=source,
        source_path=f"upstreams/{source}",
        source_commit=SOURCE_COMMITS[source],
        entrypoint=entrypoint,
        max_steps=max_steps,
        global_batch_size=global_batch_size,
        learning_rate=learning_rate,
        warmup_steps=warmup_steps,
        optimizer=optimizer,
        precision=precision,
        sequence_length=SEQUENCE_LENGTHS[dataset],
        checkpoint_every=checkpoint_every,
        default_per_device_batch_size=per_device,
        evidence=evidence,
        max_steps_override=max_steps_override,
        learning_rate_override=learning_rate_override,
        output_override=output_override,
        **values,
    )


def load_recipe(model: str, dataset: str) -> TrainingRecipe:
    """Return the immutable paper/reference recipe for one training route."""

    if dataset not in SEQUENCE_LENGTHS:
        raise RecipeError(f"unsupported dataset {dataset!r}")
    if model == "flm":
        return _base_recipe(
            model,
            dataset,
            source="flm",
            algorithm="flm",
            release_script_max_steps=1_500_000,
            release_evidence="official_release_script",
        )
    if model == "duo":
        return _base_recipe(
            model,
            dataset,
            source="duo",
            algorithm="duo",
            evidence="paper_appendix_d" if dataset == "lm1b" else "official_release_script",
        )
    if model == "mdlm":
        return _base_recipe(
            model,
            dataset,
            source="mdlm",
            algorithm="mdlm",
            evidence="paper_appendix_d" if dataset == "lm1b" else "official_release_script",
        )
    if model == "candi":
        return _base_recipe(
            model,
            dataset,
            source="candi",
            algorithm="candi",
            evidence="paper_appendix_d" if dataset == "lm1b" else "project_fallback",
        )
    if model == "rdlm":
        if dataset != "lm1b":
            raise RecipeError("RDLM/OWT is unsupported by the official release")
        return _base_recipe(
            model,
            dataset,
            source="rdlm",
            algorithm="rdlm",
            max_steps_override="training.n_iters",
            learning_rate_override="optim.lr",
            output_override="hydra.run.dir",
            evidence="official_release_script",
        )
    if model == "fmlm":
        return _base_recipe(
            model,
            dataset,
            source="flm",
            max_steps=100_000,
            algorithm="fmlm",
            teacher_family="continuous_flm",
            distillation_loss="cross_entropy",
            release_script_max_steps=1_000_000,
            release_evidence="official_release_script",
        )
    if model == "duo_dcd":
        return _base_recipe(
            model,
            dataset,
            source="duo",
            entrypoint="../../adapters/train_duo_dcd.py",
            max_steps=80_000,
            global_batch_size=128,
            learning_rate=6e-5,
            checkpoint_every=10_000,
            default_per_device_batch_size=16,
            algorithm="distillation",
            rounds=8,
            steps_per_round=10_000,
            teacher_family="uniform_duo",
            distillation_loss="kl_backward",
            evidence="paper_appendix_d",
        )
    if model == "mdlm_sdtt":
        return _base_recipe(
            model,
            dataset,
            source="sdtt",
            entrypoint="src/sdtt/main.py",
            max_steps=80_000,
            global_batch_size=128,
            learning_rate=6e-5,
            checkpoint_every=10_000,
            default_per_device_batch_size=16,
            algorithm="multi-round-sdtt",
            rounds=8,
            steps_per_round=10_000,
            teacher_family="masked_mdlm",
            teacher_adapter="masked_to_absorbing",
            distillation_loss="kl_backward",
            sampling_step=70_000,
            evidence="paper_appendix_d",
        )
    if model in {"duo_di4c", "mdlm_di4c"}:
        teacher_family = "uniform_duo" if model == "duo_di4c" else "masked_mdlm"
        adapter = (
            "uniform_to_absorbing"
            if teacher_family == "uniform_duo"
            else "masked_to_absorbing"
        )
        sampling_step = 20_000 if dataset == "lm1b" else 50_000
        return _base_recipe(
            model,
            dataset,
            source="di4c",
            entrypoint="../../adapters/train_di4c.py",
            max_steps=sampling_step,
            global_batch_size=2,
            learning_rate=3e-5,
            checkpoint_every=10_000,
            default_per_device_batch_size=1,
            algorithm="di4c",
            teacher_family=teacher_family,
            teacher_adapter=adapter,
            distillation_loss="dimensional_correlations",
            sampling_step=sampling_step,
            evidence="paper_appendix_d",
        )
    if model == "di4c":
        raise RecipeError("Di4C requires --model duo_di4c or mdlm_di4c")
    raise RecipeError(f"unknown training recipe {model!r}")


def validate_effective_batch(
    global_batch_size: int,
    *,
    devices: int,
    nodes: int,
    per_device_batch_size: int,
) -> int:
    """Return exact accumulation or reject a reduced/non-divisible paper batch."""

    values = (global_batch_size, devices, nodes, per_device_batch_size)
    if any(value <= 0 for value in values):
        raise RecipeError("batch size, devices, nodes, and per-device batch must be positive")
    denominator = devices * nodes * per_device_batch_size
    if global_batch_size % denominator:
        raise RecipeError(
            f"global batch {global_batch_size} is not divisible by "
            f"{devices} devices * {nodes} nodes * {per_device_batch_size} per device"
        )
    accumulation = global_batch_size // denominator
    if devices * nodes * per_device_batch_size * accumulation != global_batch_size:
        raise RecipeError("effective global batch differs from the locked recipe")
    return accumulation


def _tokenizer_snapshot(root: Path, dataset: str) -> Path:
    name, revision = TOKENIZER_REVISIONS[dataset]
    hub_name = "models--" + name.replace("/", "--")
    return (root / "data" / "raw" / "huggingface" / "hub" / hub_name / "snapshots" / revision).absolute()


def _data_cache(root: Path, source: str, dataset: str) -> Path:
    if source == "rdlm":
        return (root / "data" / "raw" / "huggingface" / "datasets").absolute()
    return (root / "data" / "training_cache" / source / dataset).absolute()


def _duo_integral_cache(root: Path, dataset: str) -> Path:
    filename = "bert-base-uncased.pkl" if dataset == "lm1b" else "gpt2.pkl"
    return (root / "upstreams" / "duo" / "integral" / filename).absolute()


def _float_text(value: float) -> str:
    return format(value, ".12g")


def _common_lightning_overrides(
    recipe: TrainingRecipe,
    *,
    root: Path,
    output: Path,
    data_cache: Path,
    tokenizer: Path,
    devices: int,
    nodes: int,
    per_device_batch_size: int,
    accumulation: int,
    seed: int,
    resume: bool,
) -> list[str]:
    if recipe.dataset == "lm1b":
        data_group = "lm1b" if recipe.source == "mdlm" else "lm1b-wrap"
    else:
        data_group = "openwebtext-split"
    arguments = [
        "mode=train",
        f"data={data_group}",
        f"data.cache_dir={data_cache}",
        f"data.tokenizer_name_or_path={tokenizer}",
        "model=small",
        f"model.length={recipe.sequence_length}",
        f"loader.global_batch_size={recipe.global_batch_size}",
        f"loader.eval_global_batch_size={recipe.global_batch_size}",
        f"loader.batch_size={per_device_batch_size}",
        f"loader.eval_batch_size={per_device_batch_size}",
        f"trainer.devices={devices}",
        f"trainer.num_nodes={nodes}",
        f"trainer.accumulate_grad_batches={accumulation}",
        f"trainer.max_steps={recipe.max_steps}",
        f"trainer.precision={recipe.precision}",
        f"optim.lr={_float_text(recipe.learning_rate)}",
        "lr_scheduler=constant_warmup",
        f"lr_scheduler.num_warmup_steps={recipe.warmup_steps}",
        (
            "callbacks.checkpoint_every_n_steps.every_n_train_steps="
            f"{recipe.checkpoint_every}"
        ),
        f"checkpointing.save_dir={output}",
        f"checkpointing.resume_from_ckpt={str(resume).lower()}",
        f"checkpointing.resume_ckpt_path={output / 'checkpoints' / 'last.ckpt'}",
        f"hydra.run.dir={output}",
        "hydra.job.chdir=true",
        "wandb=null",
        f"seed={seed}",
    ]
    if recipe.source == "mdlm" and recipe.dataset == "lm1b":
        arguments.extend(
            [
                "data.wrap=true",
                "+data.insert_train_eos=true",
                "+data.insert_valid_eos=true",
            ]
        )
    return arguments


def _standard_command(
    recipe: TrainingRecipe,
    *,
    python: Path,
    entrypoint: Path,
    root: Path,
    output: Path,
    data_cache: Path,
    tokenizer: Path,
    teacher: Path | None,
    devices: int,
    nodes: int,
    per_device_batch_size: int,
    accumulation: int,
    seed: int,
    resume: bool,
) -> list[str]:
    command = [str(python), "-u", str(entrypoint)]
    command.extend(
        _common_lightning_overrides(
            recipe,
            root=root,
            output=output,
            data_cache=data_cache,
            tokenizer=tokenizer,
            devices=devices,
            nodes=nodes,
            per_device_batch_size=per_device_batch_size,
            accumulation=accumulation,
            seed=seed,
            resume=resume,
        )
    )
    if recipe.model == "flm":
        command.extend(["algo=flm", "algo.double_temb=false"])
    elif recipe.model == "fmlm":
        if teacher is None:
            raise RecipeError("FMLM requires an FLM teacher checkpoint")
        command.extend(
            [
                "algo=fmlm",
                "algo.double_temb=true",
                "algo.distillation_method=PSD",
                "algo.use_mse_loss_psd=false",
                "algo.learnable_loss_weighting=false",
                "algo.initialize_student_from_teacher=true",
                f"algo.teacher_path={teacher}",
            ]
        )
    elif recipe.model == "duo":
        command.extend(
            [
                "algo=duo",
                "algo.curriculum.mode=simple",
                "algo.curriculum.gumbel_tau_log10_start=-3.0",
                "algo.curriculum.gumbel_tau_log10_end=-3.0",
                "algo.curriculum.gamma_min=-3.5",
                "algo.curriculum.gamma_max=-1.75",
                "algo.curriculum.start=0",
                "algo.curriculum.end=500000",
            ]
        )
    elif recipe.model == "duo_dcd":
        if teacher is None:
            raise RecipeError("Duo+DCD requires a uniform Duo teacher checkpoint")
        command.extend(
            [
                "algo=distillation",
                f"training.finetune_path={teacher}",
                "training.ema=0.999",
                "algo.T=512",
                f"algo.integral_cache_path={_duo_integral_cache(root, recipe.dataset)}",
                "+algo.curriculum.mode=simple",
                "+algo.curriculum.gumbel_tau_log10_start=-1",
                "+algo.curriculum.gumbel_tau_log10_end=-1",
                "+algo.curriculum.start=-1",
                "+algo.curriculum.end=-1",
                "+algo.curriculum.gamma_min=-4",
                "+algo.curriculum.gamma_max=-1",
                f"+algo.curriculum.integral_cache_path={_duo_integral_cache(root, recipe.dataset)}",
                f"algo.update_teacher_every={recipe.steps_per_round}",
                "algo.teacher_ema=false",
                "algo.linear_growth_dt=false",
            ]
        )
    elif recipe.model == "mdlm":
        command.extend(["parameterization=subs", "diffusion=absorbing_state"])
    elif recipe.model == "candi":
        command.extend(
            [
                f"scratch_dir={root / 'work' / 'candi'}",
                "algo=candi",
                "algo.sampler=cached",
                "algo.mixed_coeff=0.5",
                "algo.step_size=1.0",
                "algo.temp=1.0",
            ]
        )
    else:
        raise RecipeError(f"no standard command renderer for {recipe.model}")
    return command


def _rdlm_command(
    recipe: TrainingRecipe,
    *,
    python: Path,
    entrypoint: Path,
    output: Path,
    data_cache: Path,
    devices: int,
    accumulation: int,
    seed: int,
) -> list[str]:
    return [
        str(python),
        "-u",
        str(entrypoint),
        "run_mode=train",
        "exp=lm1b",
        "optim=adam",
        f"ngpus={devices}",
        f"training.batch_size={recipe.global_batch_size}",
        f"training.accum={accumulation}",
        f"training.n_iters={recipe.max_steps}",
        f"training.snapshot_freq={recipe.checkpoint_every}",
        "training.snapshot_freq_for_preemption=5000",
        "training.snapshot_sampling=false",
        f"data.cache_dir={data_cache}",
        f"optim.lr={_float_text(recipe.learning_rate)}",
        f"optim.warmup={recipe.warmup_steps}",
        f"hydra.run.dir={output}",
        "use_wandb=false",
        f"seed={seed}",
    ]


def _sdtt_command(
    recipe: TrainingRecipe,
    *,
    python: Path,
    entrypoint: Path,
    output: Path,
    data_cache: Path,
    tokenizer: Path,
    adapted_teacher: Path,
    devices: int,
    nodes: int,
    per_device_batch_size: int,
    accumulation: int,
    seed: int,
    resume: bool,
    di4c_student_init: str | None = None,
) -> list[str]:
    train_name = "lm1b" if recipe.dataset == "lm1b" else "openwebtext-train"
    valid_name = "lm1b" if recipe.dataset == "lm1b" else "openwebtext-valid"
    command = [
        str(python),
        "-u",
        str(entrypoint),
        "mode=train",
        f"data.train={train_name}",
        f"data.valid={valid_name}",
        f"tokenizer.name={tokenizer}",
        f"data_preprocess.data_cache={data_cache}",
        f"data_preprocess.seq_len={recipe.sequence_length}",
        "data_preprocess.group_text=true",
        "data_preprocess.remove_text=true",
        "data_preprocess.add_bos=false",
        "data_preprocess.add_eos=true",
        "data_preprocess.legacy_start_end_bos=true",
        "model=dit-orig-small",
        f"model.length={recipe.sequence_length}",
        "parameterization=multi-round-sdtt",
        f"parameterization.checkpoint_path={adapted_teacher}",
        "parameterization.start_from_hf=false",
        f"parameterization.grow_dt_every={recipe.steps_per_round or 10000}",
        "parameterization.orig_num_sampling_steps=1024",
        f"loader.global_batch_size={recipe.global_batch_size}",
        f"loader.eval_global_batch_size={recipe.global_batch_size}",
        f"loader.batch_size={per_device_batch_size}",
        f"loader.eval_batch_size={per_device_batch_size}",
        f"trainer.devices={devices}",
        f"trainer.num_nodes={nodes}",
        f"trainer.accumulate_grad_batches={accumulation}",
        f"trainer.max_steps={recipe.max_steps}",
        "trainer.precision=bf16-mixed",
        f"optim.lr={_float_text(recipe.learning_rate)}",
        f"lr_scheduler.num_warmup_steps={recipe.warmup_steps}",
        (
            "callbacks.checkpoint_every_n_steps.every_n_train_steps="
            f"{recipe.checkpoint_every}"
        ),
        f"checkpointing.save_dir={output}",
        f"checkpointing.resume_from_ckpt={str(resume).lower()}",
        f"checkpointing.resume_ckpt_path={output / 'checkpoints' / 'last.ckpt'}",
        f"hydra.run.dir={output}",
        "hydra.job.chdir=true",
        "wandb=null",
        f"+seed={seed}",
    ]
    if recipe.model.endswith("_di4c"):
        init_mode = (
            di4c_student_init
            if di4c_student_init is not None
            else ("teacher" if recipe.dataset == "lm1b" else "hf_small")
        )
        if init_mode not in {"hf_small", "teacher", "scratch"}:
            raise RecipeError(f"unsupported Di4C student init mode: {init_mode!r}")
        command.extend(
            [
                "is_di4c=true",
                "is_teacher_di4c=false",
                "T=1024",
                "latent_bsize=16",
                "round=7",
                f"+dlb_student_init={init_mode}",
            ]
        )
    return command


def build_launch(
    recipe: TrainingRecipe,
    *,
    root: Path,
    source: Path,
    output: Path,
    teacher: Path | None,
    devices: int,
    nodes: int,
    per_device_batch_size: int,
    seed: int,
    resume: bool,
    di4c_student_init: str | None = None,
) -> LaunchSpec:
    """Render one real upstream command without touching data, checkpoints, or GPUs."""

    root = root.absolute()
    source = source.absolute()
    output = output.absolute()
    teacher = teacher.absolute() if teacher is not None else None
    accumulation = validate_effective_batch(
        recipe.global_batch_size,
        devices=devices,
        nodes=nodes,
        per_device_batch_size=per_device_batch_size,
    )
    python = Path(sys.executable).absolute()
    entrypoint = (source / recipe.entrypoint).absolute()
    data_cache = _data_cache(root, recipe.source, recipe.dataset)
    tokenizer = _tokenizer_snapshot(root, recipe.dataset)
    adapted_teacher = (
        output / "inputs" / "teacher_adapted.ckpt"
        if recipe.teacher_adapter is not None
        else None
    )
    if recipe.teacher_family is not None and teacher is None:
        raise RecipeError(f"{recipe.model} requires --teacher or --teacher-checkpoint")
    if recipe.source == "rdlm":
        command = _rdlm_command(
            recipe,
            python=python,
            entrypoint=entrypoint,
            output=output,
            data_cache=data_cache,
            devices=devices,
            accumulation=accumulation,
            seed=seed,
        )
    elif recipe.source in {"sdtt", "di4c"}:
        if adapted_teacher is None:
            raise RecipeError("SDTT/Di4C recipes require a teacher adapter")
        command = _sdtt_command(
            recipe,
            python=python,
            entrypoint=entrypoint,
            output=output,
            data_cache=data_cache,
            tokenizer=tokenizer,
            adapted_teacher=adapted_teacher,
            devices=devices,
            nodes=nodes,
            per_device_batch_size=per_device_batch_size,
            accumulation=accumulation,
            seed=seed,
            resume=resume,
            di4c_student_init=di4c_student_init,
        )
    else:
        command = _standard_command(
            recipe,
            python=python,
            entrypoint=entrypoint,
            root=root,
            output=output,
            data_cache=data_cache,
            tokenizer=tokenizer,
            teacher=teacher,
            devices=devices,
            nodes=nodes,
            per_device_batch_size=per_device_batch_size,
            accumulation=accumulation,
            seed=seed,
            resume=resume,
        )
    return LaunchSpec(
        recipe=recipe,
        command=tuple(command),
        cwd=source,
        entrypoint=entrypoint,
        data_path=data_cache,
        output=output,
        teacher=teacher,
        adapted_teacher=adapted_teacher,
        devices=devices,
        nodes=nodes,
        per_device_batch_size=per_device_batch_size,
        gradient_accumulation=accumulation,
        effective_global_batch_size=(
            devices * nodes * per_device_batch_size * accumulation
        ),
        seed=seed,
        resume=resume,
    )


def _shape(value: object, name: str) -> tuple[int, ...]:
    raw = getattr(value, "shape", None)
    if raw is not None:
        try:
            return tuple(int(item) for item in raw)
        except (TypeError, ValueError) as error:
            raise RecipeError(f"{name} has an invalid shape") from error
    if isinstance(value, list):
        if not value:
            return (0,)
        children = [_shape(item, name) for item in value]
        if any(child != children[0] for child in children):
            raise RecipeError(f"{name} is not rectangular")
        return (len(value), *children[0])
    return ()


def _clone(value: object) -> object:
    clone = getattr(value, "detach", None)
    if callable(clone):
        return clone().clone()
    return deepcopy(value)


def _zero_row_like(value: object, row: int) -> object:
    if isinstance(value, list):
        selected = value[row]
        if isinstance(selected, list):
            return [0 for _ in selected]
        return 0
    selected = value[row]
    zeros_like = getattr(selected, "new_zeros", None)
    if callable(zeros_like):
        return zeros_like(selected.shape)
    raise RecipeError("tensor object does not provide new_zeros")


def _append_or_move_row(
    value: object,
    *,
    source_row: int | None,
    target_size: int,
    zero_source: bool,
    name: str,
) -> object:
    shape = _shape(value, name)
    if not shape or shape[0] <= 0:
        raise RecipeError(f"{name} must have a non-empty vocabulary dimension")
    source_size = shape[0]
    if target_size not in {source_size, source_size + 1}:
        raise RecipeError(
            f"{name} vocabulary {source_size} cannot map to target {target_size}"
        )
    if source_row is not None and not 0 <= source_row < source_size:
        raise RecipeError(f"{name} source mask row is out of range")
    if isinstance(value, list):
        result = deepcopy(value)
        selected = (
            deepcopy(value[source_row])
            if source_row is not None
            else _zero_row_like(value, 0)
        )
        if target_size == source_size + 1:
            result.append(selected)
        elif source_row is not None and source_row != target_size - 1:
            result[target_size - 1] = selected
        if zero_source and source_row is not None and source_row != target_size - 1:
            result[source_row] = _zero_row_like(value, source_row)
        return result

    result = value.new_zeros((target_size, *shape[1:]))
    result[:source_size].copy_(value)
    if source_row is not None:
        result[target_size - 1].copy_(value[source_row])
        if zero_source and source_row != target_size - 1:
            result[source_row].zero_()
    return result


def _validate_state(
    state: Mapping[str, object], expected_keys: Collection[str] | None
) -> tuple[int, int]:
    if expected_keys is not None:
        expected = set(expected_keys)
        observed = set(state)
        missing = sorted(expected - observed)
        unexpected = sorted(observed - expected - OPTIONAL_BACKBONE_STATE_KEYS)
        if missing or unexpected:
            raise RecipeError(
                f"teacher state keys differ: missing={missing}, unexpected={unexpected}"
            )
    missing_tensors = sorted(set(BACKBONE_TENSOR_KEYS) - set(state))
    if missing_tensors:
        raise RecipeError(f"teacher state keys are missing {missing_tensors}")
    embedding_shape = _shape(state[BACKBONE_TENSOR_KEYS[0]], BACKBONE_TENSOR_KEYS[0])
    weight_shape = _shape(state[BACKBONE_TENSOR_KEYS[1]], BACKBONE_TENSOR_KEYS[1])
    bias_shape = _shape(state[BACKBONE_TENSOR_KEYS[2]], BACKBONE_TENSOR_KEYS[2])
    if len(embedding_shape) != 2 or len(weight_shape) != 2 or len(bias_shape) != 1:
        raise RecipeError("teacher embedding/head ranks must be 2, 2, and 1")
    vocab, hidden = embedding_shape
    if weight_shape[0] != vocab or bias_shape[0] != vocab:
        raise RecipeError("teacher embedding/head vocabulary dimensions differ")
    if weight_shape[1] != hidden:
        raise RecipeError("teacher embedding/head hidden dimensions differ")
    if vocab <= 0 or hidden <= 0:
        raise RecipeError("teacher vocabulary and hidden dimensions must be positive")
    return vocab, hidden


def uniform_to_absorbing(
    state: Mapping[str, object],
    *,
    target_vocab_size: int | None = None,
    expected_keys: Collection[str] | None = None,
) -> dict[str, object]:
    """Append an all-zero absorbing state to a uniform-Duo backbone."""

    vocab, _ = _validate_state(state, expected_keys)
    target_size = target_vocab_size or vocab + 1
    if target_size not in {vocab, vocab + 1}:
        raise RecipeError(
            f"uniform teacher vocabulary {vocab} cannot map to target {target_size}"
        )
    result = {key: _clone(value) for key, value in state.items()}
    if target_size == vocab:
        for key in BACKBONE_TENSOR_KEYS:
            result[key] = _zero_vocab_row(result[key], target_size - 1)
        _validate_state(result, set(state))
        return result
    for key in BACKBONE_TENSOR_KEYS:
        result[key] = _append_or_move_row(
            state[key],
            source_row=None,
            target_size=target_size,
            zero_source=False,
            name=key,
        )
    _validate_state(result, set(state))
    return result


def _zero_vocab_row(value: object, row: int) -> object:
    if isinstance(value, list):
        result = deepcopy(value)
        result[row] = _zero_row_like(value, row)
        return result
    result = value.clone()
    result[row].zero_()
    return result


def masked_to_absorbing(
    state: Mapping[str, object],
    *,
    source_mask_index: int,
    target_vocab_size: int,
    expected_keys: Collection[str] | None = None,
) -> dict[str, object]:
    """Move an existing masked-model state to the final absorbing-mask row."""

    vocab, _ = _validate_state(state, expected_keys)
    if target_vocab_size not in {vocab, vocab + 1}:
        raise RecipeError(
            f"masked teacher vocabulary {vocab} cannot map to {target_vocab_size}"
        )
    result = {key: _clone(value) for key, value in state.items()}
    if source_mask_index == target_vocab_size - 1 and target_vocab_size == vocab:
        return result
    for key in BACKBONE_TENSOR_KEYS:
        result[key] = _append_or_move_row(
            state[key],
            source_row=source_mask_index,
            target_size=target_vocab_size,
            zero_source=True,
            name=key,
        )
    _validate_state(result, set(state))
    return result


def expected_backbone_keys(n_blocks: int = 12) -> set[str]:
    if n_blocks <= 0:
        raise RecipeError("backbone must contain at least one transformer block")
    keys = {
        "vocab_embed.embedding",
        "sigma_map.mlp.0.weight",
        "sigma_map.mlp.0.bias",
        "sigma_map.mlp.2.weight",
        "sigma_map.mlp.2.bias",
        "output_layer.norm_final.weight",
        "output_layer.linear.weight",
        "output_layer.linear.bias",
        "output_layer.adaLN_modulation.weight",
        "output_layer.adaLN_modulation.bias",
    }
    block_suffixes = {
        "norm1.weight",
        "attn_qkv.weight",
        "attn_out.weight",
        "norm2.weight",
        "mlp.0.weight",
        "mlp.0.bias",
        "mlp.2.weight",
        "mlp.2.bias",
        "adaLN_modulation.weight",
        "adaLN_modulation.bias",
    }
    for index in range(n_blocks):
        keys.update(f"blocks.{index}.{suffix}" for suffix in block_suffixes)
    return keys


def _normalize_backbone_state(state: Mapping[str, object]) -> dict[str, object]:
    normalized: dict[str, object] = {}
    for raw_key, value in state.items():
        key = str(raw_key)
        while key.startswith("_orig_mod."):
            key = key[len("_orig_mod.") :]
        if key.startswith("backbone."):
            key = key[len("backbone.") :]
        if key in normalized:
            raise RecipeError(f"teacher key normalization collision at {key!r}")
        normalized[key] = value
    return normalized


def _callable_sha256(function: object) -> str:
    return hashlib.sha256(inspect.getsource(function).encode("utf-8")).hexdigest()


def adapt_teacher_checkpoint(
    launch: LaunchSpec,
    *,
    source_mask_index: int | None = None,
) -> dict[str, object]:
    """Adapt and atomically publish one server-side teacher checkpoint."""

    if launch.teacher is None or launch.adapted_teacher is None:
        raise RecipeError("teacher adaptation requires source and destination paths")
    _require_regular_file(launch.teacher, "teacher checkpoint")
    try:
        import torch
    except ImportError as error:
        raise RecipeError(
            "PyTorch is required only for real server teacher adaptation"
        ) from error
    try:
        payload = torch.load(
            str(launch.teacher), map_location="cpu", weights_only=False
        )
    except Exception as error:
        raise RecipeError(f"could not load teacher checkpoint: {error}") from error
    if not isinstance(payload, MutableMapping):
        raise RecipeError("teacher checkpoint must contain a mapping")
    raw_state = payload.get("state_dict", payload)
    if not isinstance(raw_state, Mapping):
        raise RecipeError("teacher checkpoint state_dict must be a mapping")
    state = _normalize_backbone_state(raw_state)
    expected = expected_backbone_keys()
    _validate_state(state, expected)
    base_vocab = BASE_VOCAB_SIZES[launch.recipe.dataset]
    observed_vocab, hidden = _validate_state(state, expected)
    if launch.recipe.teacher_adapter == "uniform_to_absorbing":
        target_vocab = base_vocab + 1
        if observed_vocab not in {base_vocab, target_vocab}:
            raise RecipeError(
                "uniform teacher vocabulary is "
                f"{observed_vocab}, expected {base_vocab} or {target_vocab}"
            )
        adapted = uniform_to_absorbing(
            state,
            target_vocab_size=target_vocab,
            expected_keys=expected,
        )
        transform = uniform_to_absorbing
        mask_index = target_vocab - 1
    elif launch.recipe.teacher_adapter == "masked_to_absorbing":
        target_vocab = base_vocab + 1
        if source_mask_index is None:
            source_mask_index = (
                target_vocab - 1 if observed_vocab == target_vocab else 103
            )
        adapted = masked_to_absorbing(
            state,
            source_mask_index=source_mask_index,
            target_vocab_size=target_vocab,
            expected_keys=expected,
        )
        transform = masked_to_absorbing
        mask_index = target_vocab - 1
    else:
        raise RecipeError(f"unknown teacher adapter {launch.recipe.teacher_adapter!r}")
    adapted_vocab, adapted_hidden = _validate_state(adapted, expected)
    if adapted_vocab != base_vocab + 1 or adapted_hidden != hidden:
        raise RecipeError("adapted teacher shape does not match absorbing backbone")

    destination = launch.adapted_teacher
    ensure_safe_directory(destination.parent)
    _require_regular_or_missing(destination, "adapted teacher checkpoint")
    adapted_payload = {"state_dict": adapted}
    temporary = destination.with_name(f".{destination.name}.{os.getpid()}.partial")
    _require_regular_or_missing(temporary, "temporary adapted checkpoint")
    try:
        torch.save(adapted_payload, temporary)
        os.replace(temporary, destination)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass
    record = {
        "schema_version": 1,
        "source_path": str(launch.teacher),
        "source_sha256": sha256_file(launch.teacher),
        "output_path": str(destination),
        "output_sha256": sha256_file(destination),
        "teacher_family": launch.recipe.teacher_family,
        "transformation": launch.recipe.teacher_adapter,
        "transformation_sha256": _callable_sha256(transform),
        "source_mask_index": source_mask_index,
        "target_mask_index": mask_index,
        "source_vocab_size": observed_vocab,
        "target_vocab_size": adapted_vocab,
        "hidden_size": hidden,
        "state_keys_sha256": hashlib.sha256(
            "\n".join(sorted(adapted)).encode("utf-8")
        ).hexdigest(),
    }
    atomic_json_write(destination.with_suffix(".provenance.json"), record)
    return record


def _path_components(path: Path) -> list[Path]:
    components: list[Path] = []
    current = path
    while True:
        components.append(current)
        if current.parent == current:
            return list(reversed(components))
        current = current.parent


def _reject_symlink_components(path: Path, label: str) -> None:
    for component in _path_components(path):
        try:
            metadata = component.lstat()
        except FileNotFoundError:
            continue
        if stat.S_ISLNK(metadata.st_mode):
            raise RecipeError(f"{label} contains a symlink component: {component}")


def _project_path(root: Path, value: Path, label: str) -> Path:
    candidate = value if value.is_absolute() else root / value
    candidate = candidate.absolute()
    try:
        candidate.relative_to(root)
    except ValueError as error:
        raise RecipeError(f"{label} must remain inside project root") from error
    _reject_symlink_components(candidate, label)
    return candidate


def _require_regular_file(path: Path, label: str) -> None:
    try:
        metadata = path.lstat()
    except FileNotFoundError as error:
        raise RecipeError(f"{label} is missing: {path}") from error
    if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode):
        raise RecipeError(f"{label} is not a safe regular file: {path}")


def _require_regular_or_missing(path: Path, label: str) -> None:
    try:
        metadata = path.lstat()
    except FileNotFoundError:
        return
    if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode):
        raise RecipeError(f"{label} is not a safe regular file: {path}")


def _verify_source(launch: LaunchSpec) -> None:
    _require_regular_file(launch.entrypoint, "upstream entrypoint")
    completed = subprocess.run(
        ["git", "-C", str(launch.cwd), "rev-parse", "HEAD"],
        text=True,
        capture_output=True,
        check=False,
    )
    if completed.returncode or completed.stdout.strip() != launch.recipe.source_commit:
        raise RecipeError(
            f"{launch.recipe.source} source is not pinned commit "
            f"{launch.recipe.source_commit}"
        )
    dirty = subprocess.run(
        ["git", "-C", str(launch.cwd), "status", "--porcelain"],
        text=True,
        capture_output=True,
        check=False,
    )
    dirty_lines = _source_dirty_lines(dirty.stdout)
    if dirty.returncode or dirty_lines:
        raise RecipeError(f"{launch.recipe.source} source checkout is dirty")


def _source_dirty_lines(status_output: str) -> list[str]:
    """Return source status lines excluding Python bytecode cache byproducts."""

    return [
        line
        for line in status_output.splitlines()
        if line.strip() and not _is_untracked_python_bytecode_cache(line)
    ]


def _is_untracked_python_bytecode_cache(status_line: str) -> bool:
    if not status_line.startswith("?? "):
        return False
    path = status_line[3:].strip()
    parts = PurePosixPath(path).parts
    return "__pycache__" in parts or path.endswith((".pyc", ".pyo"))


def _processed_dataset(root: Path, dataset: str) -> tuple[Path, Path]:
    name = "lm1b-bert-128" if dataset == "lm1b" else "owt-gpt2-1024"
    processed = root / "data" / "processed" / name
    manifest = root / "data" / "manifests" / f"{dataset}.json"
    if processed.is_symlink() or not processed.is_dir():
        raise RecipeError(f"processed dataset is missing: {processed}")
    _require_regular_file(manifest, "processed dataset manifest")
    value = json.loads(manifest.read_text(encoding="utf-8"))
    if value.get("dataset") != dataset:
        raise RecipeError("processed dataset manifest identifies another dataset")
    if int(value.get("sequence_length", -1)) != SEQUENCE_LENGTHS[dataset]:
        raise RecipeError("processed dataset sequence length differs from recipe")
    return processed, manifest


def _link_tree(source: Path, destination: Path) -> None:
    if destination.exists():
        if destination.is_symlink() or not destination.is_dir():
            raise RecipeError(f"training cache path is unsafe: {destination}")
        source_files = {
            path.relative_to(source): path
            for path in source.rglob("*")
            if path.is_file() and not path.is_symlink()
        }
        destination_files = {
            path.relative_to(destination): path
            for path in destination.rglob("*")
            if path.is_file() and not path.is_symlink()
        }
        if set(source_files) != set(destination_files):
            raise RecipeError(f"training cache inventory differs: {destination}")
        for relative, source_file in source_files.items():
            destination_file = destination_files[relative]
            source_stat = source_file.stat()
            destination_stat = destination_file.stat()
            if source_stat.st_size != destination_stat.st_size:
                raise RecipeError(f"training cache file size differs: {destination_file}")
            if source_stat.st_dev != destination_stat.st_dev or source_stat.st_ino != destination_stat.st_ino:
                if sha256_file(source_file) != sha256_file(destination_file):
                    raise RecipeError(
                        f"training cache file digest differs: {destination_file}"
                    )
        return
    ensure_safe_directory(destination)
    for path in sorted(source.rglob("*")):
        relative = path.relative_to(source)
        target = destination / relative
        if path.is_symlink():
            raise RecipeError(f"processed dataset contains a symlink: {path}")
        if path.is_dir():
            ensure_safe_directory(target)
        elif path.is_file():
            ensure_safe_directory(target.parent)
            try:
                os.link(path, target)
            except OSError:
                shutil.copyfile(path, target)
        else:
            raise RecipeError(f"processed dataset contains a special file: {path}")


def _remove_safe_tree(path: Path) -> None:
    if not path.exists():
        return
    if path.is_symlink() or not path.is_dir():
        raise RecipeError(f"training cache path is unsafe: {path}")
    for child in path.rglob("*"):
        if child.is_symlink():
            raise RecipeError(f"training cache contains a symlink: {child}")
    shutil.rmtree(path)


def _dataset_columns(path: Path) -> set[str]:
    from datasets import load_from_disk

    dataset = load_from_disk(str(path))
    return set(dataset.column_names)


def _materialize_text_training_split(
    source: Path,
    destination: Path,
    *,
    sequence_length: int,
    cache_is_current: bool,
) -> None:
    if destination.exists():
        if destination.is_symlink() or not destination.is_dir():
            raise RecipeError(f"training cache path is unsafe: {destination}")
        if cache_is_current and {"input_ids", "attention_mask"} <= _dataset_columns(
            destination
        ):
            return
        _remove_safe_tree(destination)

    from datasets import Sequence as DatasetSequence
    from datasets import Value, load_from_disk

    dataset = load_from_disk(str(source))
    if "attention_mask" not in dataset.column_names:
        features = deepcopy(dataset.features)
        features["attention_mask"] = DatasetSequence(
            Value("uint8"), length=sequence_length
        )

        def add_attention_mask(batch):
            return {
                "attention_mask": [
                    [1] * len(input_ids) for input_ids in batch["input_ids"]
                ]
            }

        dataset = dataset.map(
            add_attention_mask,
            batched=True,
            features=features,
            load_from_cache_file=False,
            desc="Adding all-one attention masks for upstream training",
        )

    staged = destination.with_name(f".{destination.name}.{os.getpid()}.partial")
    _remove_safe_tree(staged)
    ensure_safe_directory(destination.parent)
    try:
        dataset.save_to_disk(str(staged))
        if destination.exists():
            _remove_safe_tree(destination)
        os.replace(staged, destination)
    finally:
        _remove_safe_tree(staged)


def _prepare_training_cache(launch: LaunchSpec, root: Path) -> dict[str, object]:
    processed, manifest = _processed_dataset(root, launch.recipe.dataset)
    processed_manifest_sha = sha256_file(manifest)
    if launch.recipe.source == "rdlm":
        if launch.data_path.is_symlink() or not launch.data_path.is_dir():
            raise RecipeError(
                "RDLM requires the pinned raw Hugging Face dataset cache; "
                "run scripts/fetch_data.py first"
            )
        return {
            "mode": "upstream_raw_retokenization",
            "processed_preflight": str(processed),
            "processed_manifest_sha256": processed_manifest_sha,
            "cache": str(launch.data_path),
        }

    ensure_safe_directory(launch.data_path)
    cache_record = launch.data_path / "dlb_training_cache.json"
    cache_is_current = False
    if cache_record.is_file() and not cache_record.is_symlink():
        try:
            current = json.loads(cache_record.read_text(encoding="utf-8"))
            cache_is_current = (
                current.get("processed_manifest_sha256") == processed_manifest_sha
                and current.get("format") == "hf_dataset_with_attention_mask_v1"
            )
        except (OSError, json.JSONDecodeError):
            cache_is_current = False
    if launch.recipe.dataset == "lm1b":
        aliases = {
            "lm1b_train_bs128_wrapped.dat": processed / "train",
            "lm1b_test_bs128_wrapped.dat": processed / "validation",
        }
    else:
        aliases = {
            "openwebtext-train_train_bs1024_wrapped.dat": processed / "train",
            "openwebtext-valid_validation_bs1024_wrapped.dat": processed
            / "validation",
        }
    for name, source in aliases.items():
        if not source.is_dir():
            raise RecipeError(f"processed split is missing: {source}")
        _materialize_text_training_split(
            source,
            launch.data_path / name,
            sequence_length=launch.recipe.sequence_length,
            cache_is_current=cache_is_current,
        )
    record = {
        "mode": "materialized_bridge",
        "format": "hf_dataset_with_attention_mask_v1",
        "processed_path": str(processed),
        "processed_manifest_sha256": processed_manifest_sha,
        "cache": str(launch.data_path),
        "aliases": {name: str(path) for name, path in aliases.items()},
    }
    atomic_json_write(launch.data_path / "dlb_training_cache.json", record)
    return record


def _launch_identity(
    launch: LaunchSpec,
    data_record: Mapping[str, object],
    *,
    composed_config_sha256: str,
) -> dict[str, object]:
    teacher_sha = (
        sha256_file(launch.teacher)
        if launch.teacher is not None and launch.teacher.is_file()
        else None
    )
    recipe_value = asdict(launch.recipe)
    value = {
        "schema_version": 1,
        "recipe": recipe_value,
        "command": list(launch.command),
        "cwd": str(launch.cwd),
        "source_commit": launch.recipe.source_commit,
        "data": dict(data_record),
        "teacher_sha256": teacher_sha,
        "teacher_transformation_sha256": (
            _callable_sha256(
                uniform_to_absorbing
                if launch.recipe.teacher_adapter == "uniform_to_absorbing"
                else masked_to_absorbing
            )
            if launch.recipe.teacher_adapter is not None
            else None
        ),
        "composed_config_sha256": composed_config_sha256,
        "seed": launch.seed,
        "devices": launch.devices,
        "nodes": launch.nodes,
        "per_device_batch_size": launch.per_device_batch_size,
        "gradient_accumulation": launch.gradient_accumulation,
        "effective_global_batch_size": launch.effective_global_batch_size,
    }
    value["identity_sha256"] = hashlib.sha256(
        json.dumps(
            value, sort_keys=True, separators=(",", ":"), ensure_ascii=False
        ).encode("utf-8")
    ).hexdigest()
    return value


def _completed_is_reusable(output: Path, identity: Mapping[str, object]) -> bool:
    marker = output / "completed.json"
    if marker.is_symlink() or not marker.is_file():
        return False
    try:
        value = json.loads(marker.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return False
    if value.get("identity_sha256") != identity.get("identity_sha256"):
        raise RecipeError("completed output belongs to a different recipe identity")
    checkpoint = output / str(value.get("checkpoint", ""))
    if checkpoint.is_symlink() or not checkpoint.is_file():
        raise RecipeError("completed output checkpoint is missing or unsafe")
    if sha256_file(checkpoint) != value.get("checkpoint_sha256"):
        raise RecipeError("completed output checkpoint digest differs")
    config = output / str(value.get("config", ""))
    if config.is_symlink() or not config.is_file():
        raise RecipeError("completed output config is missing or unsafe")
    if sha256_file(config) != value.get("config_sha256"):
        raise RecipeError("completed output config digest differs")
    if value.get("config_sha256") != identity.get("composed_config_sha256"):
        raise RecipeError("completed output config differs from current composition")
    return True


def _compose_config(launch: LaunchSpec, environment: Mapping[str, str]) -> str:
    compose = list(launch.command[:3]) + ["--cfg", "job", "--resolve"] + list(
        launch.command[3:]
    )
    completed = subprocess.run(
        compose,
        cwd=launch.cwd,
        env=dict(environment),
        text=True,
        capture_output=True,
        check=False,
    )
    if completed.returncode:
        raise RecipeError(
            "Hydra config composition failed before launch: "
            + completed.stderr[-4000:]
        )
    if not completed.stdout.strip():
        raise RecipeError("Hydra produced an empty composed config")
    return completed.stdout


def _write_text_atomic(path: Path, payload: str) -> None:
    _require_regular_or_missing(path, "text output")
    ensure_safe_directory(path.parent)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.partial")
    _require_regular_or_missing(temporary, "temporary text output")
    try:
        with temporary.open("x", encoding="utf-8", newline="\n") as output:
            output.write(payload)
            output.flush()
            os.fsync(output.fileno())
        _require_regular_or_missing(path, "text output")
        os.replace(temporary, path)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass


def _open_safe_append(path: Path):
    ensure_safe_directory(path.parent)
    _require_regular_or_missing(path, "log output")
    descriptor = os.open(
        path,
        os.O_WRONLY
        | os.O_CREAT
        | os.O_APPEND
        | getattr(os, "O_NOFOLLOW", 0),
        0o600,
    )
    return os.fdopen(descriptor, "ab", buffering=0)


def _copy_atomic(source: Path, destination: Path) -> None:
    _require_regular_file(source, "produced checkpoint")
    _require_regular_or_missing(destination, "canonical checkpoint")
    ensure_safe_directory(destination.parent)
    temporary = destination.with_name(f".{destination.name}.{os.getpid()}.partial")
    _require_regular_or_missing(temporary, "temporary canonical checkpoint")
    try:
        shutil.copyfile(source, temporary)
        os.replace(temporary, destination)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass


def _expected_checkpoint(launch: LaunchSpec) -> tuple[Path, str]:
    recipe = launch.recipe
    if recipe.source == "rdlm":
        index = recipe.max_steps // recipe.checkpoint_every
        return launch.output / "checkpoints" / f"checkpoint_{index}.pth", (
            f"checkpoints/checkpoint_{index}.pth"
        )
    if recipe.model in {"mdlm_sdtt", "duo_di4c", "mdlm_di4c"}:
        if recipe.sampling_step is None:
            raise RecipeError("distilled recipe has no selected sampling step")
        relative = f"student_checkpoints/{recipe.sampling_step}.ckpt"
        candidates = (
            launch.output / relative,
            launch.output / "checkpoints" / f"0-{recipe.sampling_step}.ckpt",
            launch.output / "checkpoints" / "last.ckpt",
        )
        for candidate in candidates:
            if candidate.is_file() and not candidate.is_symlink():
                return candidate, relative
        return candidates[0], relative
    return launch.output / "checkpoints" / "last.ckpt", "model.ckpt"


def execute_launch(launch: LaunchSpec, *, root: Path) -> str:
    """Run one prepared recipe on a server, returning ``completed`` or ``skipped``."""

    _verify_source(launch)
    data_record = _prepare_training_cache(launch, root)
    if launch.teacher is not None:
        _require_regular_file(launch.teacher, "teacher checkpoint")
    ensure_safe_directory(launch.output)
    environment = dict(os.environ)
    environment.update(
        {
            "DLB_ROOT": str(root),
            "HF_HOME": str(root / "data" / "raw" / "huggingface"),
            "HF_DATASETS_CACHE": str(
                root / "data" / "raw" / "huggingface" / "datasets"
            ),
            "HF_HUB_OFFLINE": "1",
            "TRANSFORMERS_OFFLINE": "1",
            "WANDB_MODE": "disabled",
            "PYTHONDONTWRITEBYTECODE": "1",
            "PYTHONUNBUFFERED": "1",
        }
    )
    upstream_python_path = str(launch.entrypoint.parents[1])
    existing_python_path = environment.get("PYTHONPATH")
    environment["PYTHONPATH"] = (
        upstream_python_path
        if not existing_python_path
        else upstream_python_path + os.pathsep + existing_python_path
    )
    composed = _compose_config(launch, environment)
    composed_sha256 = hashlib.sha256(composed.encode("utf-8")).hexdigest()
    identity = _launch_identity(
        launch, data_record, composed_config_sha256=composed_sha256
    )
    if _completed_is_reusable(launch.output, identity):
        return "skipped"
    complete_marker = launch.output / "completed.json"
    if complete_marker.exists():
        raise RecipeError("refusing to overwrite a non-reusable completed run")

    adaptation = None
    if launch.recipe.teacher_adapter is not None:
        adaptation = adapt_teacher_checkpoint(launch)
    atomic_json_write(launch.output / "recipe.json", asdict(launch.recipe))
    atomic_json_write(launch.output / "launch_argv.json", list(launch.command))
    atomic_json_write(
        launch.output / "provenance.json",
        {
            **identity,
            "teacher_adaptation": adaptation,
            "prepared_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        },
    )
    config_path = launch.output / "config.yaml"
    _write_text_atomic(config_path, composed)
    with _open_safe_append(launch.output / "stdout.log") as stdout_file, _open_safe_append(
        launch.output / "stderr.log"
    ) as stderr_file:
        completed = subprocess.run(
            launch.command,
            cwd=launch.cwd,
            env=environment,
            stdout=stdout_file,
            stderr=stderr_file,
            check=False,
        )
    if completed.returncode:
        if completed.returncode < 0:
            return_code = 128 + abs(completed.returncode)
        else:
            return_code = completed.returncode
        raise RecipeProcessError(
            f"upstream training exited with status {return_code}", return_code
        )
    produced, relative = _expected_checkpoint(launch)
    _require_regular_file(produced, "selected training checkpoint")
    canonical = launch.output / relative
    if produced != canonical:
        _copy_atomic(produced, canonical)
    marker = {
        "schema_version": 1,
        "identity_sha256": identity["identity_sha256"],
        "checkpoint": relative,
        "checkpoint_sha256": sha256_file(canonical),
        "config": "config.yaml",
        "config_sha256": sha256_file(config_path),
        "completed_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    }
    atomic_json_write(complete_marker, marker)
    return "completed"


def _dry_record(launch: LaunchSpec) -> dict[str, object]:
    return {
        "schema_version": 1,
        "dry_run": True,
        "model": launch.recipe.model,
        "dataset": launch.recipe.dataset,
        "source": str(launch.cwd),
        "source_commit": launch.recipe.source_commit,
        "entrypoint": str(launch.entrypoint),
        "data": str(launch.data_path),
        "teacher": str(launch.teacher) if launch.teacher is not None else None,
        "adapted_teacher": (
            str(launch.adapted_teacher)
            if launch.adapted_teacher is not None
            else None
        ),
        "output": str(launch.output),
        "command": list(launch.command),
        "effective_global_batch_size": launch.effective_global_batch_size,
        "gradient_accumulation": launch.gradient_accumulation,
        "recipe": asdict(launch.recipe),
    }


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--recipe",
        required=True,
        choices=("flm", "duo", "mdlm", "candi", "rdlm", "fmlm", "duo_dcd", "mdlm_sdtt", "di4c"),
    )
    parser.add_argument("--root", type=Path, default=Path.cwd())
    parser.add_argument("--dataset", choices=("lm1b", "owt"), required=True)
    parser.add_argument("--model", choices=("duo_di4c", "mdlm_di4c"))
    parser.add_argument("--source", type=Path)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--teacher", "--teacher-checkpoint", dest="teacher", type=Path)
    parser.add_argument(
        "--teacher-family", choices=("continuous_flm", "uniform_duo", "masked_mdlm")
    )
    parser.add_argument("--devices", type=int)
    parser.add_argument("--nodes", type=int, default=int(os.getenv("DLB_TRAIN_NODES", "1")))
    parser.add_argument("--per-device-batch-size", type=int)
    parser.add_argument("--global-batch-size", type=int)
    parser.add_argument("--learning-rate", type=float)
    parser.add_argument("--rounds", type=int)
    parser.add_argument("--steps-per-round", type=int)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--resume", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--upstream-override", action="append", default=[])
    parser.add_argument("--student-init", choices=("hf_small", "teacher", "scratch"))
    parser.add_argument("--dry-run", action="store_true")
    return parser


def _assert_locked(name: str, supplied: object, expected: object) -> None:
    if supplied is not None and supplied != expected:
        raise RecipeError(
            f"{name}={supplied!r} differs from locked recipe value {expected!r}"
        )


def main(argv: Sequence[str] | None = None) -> int:
    arguments = _parser().parse_args(argv)
    root = arguments.root.absolute()
    _reject_symlink_components(root, "project root")
    model = arguments.model if arguments.recipe == "di4c" else arguments.recipe
    if model is None:
        raise RecipeError("--model is required by the Di4C wrapper")
    if arguments.recipe != "di4c" and arguments.model is not None:
        raise RecipeError("--model is valid only for the Di4C wrapper")
    recipe = load_recipe(model, arguments.dataset)
    if arguments.student_init is not None and not recipe.model.endswith("_di4c"):
        raise RecipeError("--student-init is valid only for Di4C recipes")
    _assert_locked("global batch size", arguments.global_batch_size, recipe.global_batch_size)
    _assert_locked("learning rate", arguments.learning_rate, recipe.learning_rate)
    _assert_locked("rounds", arguments.rounds, recipe.rounds)
    _assert_locked("steps per round", arguments.steps_per_round, recipe.steps_per_round)
    _assert_locked("teacher family", arguments.teacher_family, recipe.teacher_family)
    allowed_overrides = {
        "is_di4c=true",
    }
    unknown = sorted(set(arguments.upstream_override) - allowed_overrides)
    if unknown:
        raise RecipeError(f"unsupported or unsafe upstream overrides: {unknown}")
    source_value = arguments.source or Path(recipe.source_path)
    output_value = arguments.output or Path(
        "checkpoints", "self_trained", recipe.dataset, recipe.model
    )
    source = _project_path(root, source_value, "source")
    output = _project_path(root, output_value, "output")
    teacher = (
        _project_path(root, arguments.teacher, "teacher")
        if arguments.teacher is not None
        else None
    )
    per_device = (
        arguments.per_device_batch_size or recipe.default_per_device_batch_size
    )
    devices = arguments.devices
    if devices is None:
        default_devices = "2" if recipe.model.endswith("_di4c") else "8"
        devices = int(os.getenv("DLB_TRAIN_DEVICES", default_devices))
    launch = build_launch(
        recipe,
        root=root,
        source=source,
        output=output,
        teacher=teacher,
        devices=devices,
        nodes=arguments.nodes,
        per_device_batch_size=per_device,
        seed=arguments.seed,
        resume=arguments.resume,
        di4c_student_init=arguments.student_init,
    )
    if arguments.dry_run:
        print(json.dumps(_dry_record(launch), sort_keys=True))
        return 0
    status = execute_launch(launch, root=root)
    print(json.dumps({"status": status, "output": str(output)}, sort_keys=True))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except RecipeProcessError as error:
        print(f"ERROR: {error}", file=sys.stderr)
        raise SystemExit(error.exit_status)
    except RecipeError as error:
        print(f"ERROR: {error}", file=sys.stderr)
        raise SystemExit(2)