Blame for tests/test_training_recipes.py

from __future__ import annotations

import json
import os
from pathlib import Path
import subprocess
import sys
import types
import importlib.util

import pytest

import dlb.recipes as recipes_module
from dlb.recipes import (
    BACKBONE_TENSOR_KEYS,
    RecipeError,
    build_launch,
    expected_backbone_keys,
    load_recipe,
    masked_to_absorbing,
    uniform_to_absorbing,
    validate_effective_batch,
)


ROOT = Path(__file__).parents[1]


def override(command: tuple[str, ...], key: str) -> str:
    prefix = key + "="
    values = [argument[len(prefix) :] for argument in command if argument.startswith(prefix)]
    assert len(values) == 1, (key, command)
    return values[0]


def test_flm_recipe_matches_paper_v3_reference() -> None:
    recipe = load_recipe("flm", "lm1b")

    assert recipe.max_steps == 1_000_000
    assert recipe.global_batch_size == 512
    assert recipe.learning_rate == pytest.approx(3e-4)
    assert recipe.warmup_steps == 2_500
    assert recipe.optimizer == "adam"
    assert recipe.sequence_length == 128
    assert recipe.evidence == "paper_v3_main"


def test_fmlm_recipe_uses_main_text_ce_schedule_not_release_drift() -> None:
    recipe = load_recipe("fmlm", "lm1b")

    assert recipe.max_steps == 100_000
    assert recipe.global_batch_size == 512
    assert recipe.learning_rate == pytest.approx(3e-4)
    assert recipe.warmup_steps == 2_500
    assert recipe.distillation_loss == "cross_entropy"
    assert recipe.release_script_max_steps == 1_000_000


def test_flm_release_script_drift_is_explicit() -> None:
    recipe = load_recipe("flm", "owt")

    assert recipe.max_steps == 1_000_000
    assert recipe.release_script_max_steps == 1_500_000
    assert recipe.release_evidence == "official_release_script"
    assert "trainer.max_steps=1500000" in (
        ROOT / "upstreams" / "flm" / "scripts" / "train_lm1b_flm.sh"
    ).read_text()
    assert "trainer.max_steps=1000000" in (
        ROOT / "upstreams" / "flm" / "scripts" / "train_lm1b_fmlm_denoiser.sh"
    ).read_text()


def test_sdtt_and_dcd_round_schedule() -> None:
    for name in ("mdlm_sdtt", "duo_dcd"):
        recipe = load_recipe(name, "lm1b")
        assert recipe.rounds == 8
        assert recipe.steps_per_round == 10_000
        assert recipe.max_steps == 80_000
        assert recipe.global_batch_size == 128
        assert recipe.learning_rate == pytest.approx(6e-5)
        assert recipe.warmup_steps == 2_500


@pytest.mark.parametrize(
    ("dataset", "cache_name"),
    [
        ("lm1b", "bert-base-uncased.pkl"),
        ("owt", "gpt2.pkl"),
    ],
)
def test_duo_dcd_uses_repo_integral_cache_not_tokenizer_path(
    dataset: str, cache_name: str
) -> None:
    launch = build_launch(
        load_recipe("duo_dcd", dataset),
        root=ROOT,
        source=ROOT / "upstreams" / "duo",
        output=ROOT / "checkpoints" / "self_trained" / dataset / "duo_dcd",
        teacher=ROOT / "checkpoints" / "fixture" / "uniform_duo.ckpt",
        devices=8,
        nodes=1,
        per_device_batch_size=16,
        seed=42,
        resume=True,
    )

    integral_cache = override(launch.command, "algo.integral_cache_path")
    curriculum_integral_cache = override(
        launch.command, "+algo.curriculum.integral_cache_path"
    )
    assert integral_cache == str(
        (ROOT / "upstreams" / "duo" / "integral" / cache_name).absolute()
    )
    assert curriculum_integral_cache == integral_cache
    assert "snapshots" not in integral_cache
    assert override(launch.command, "+algo.curriculum.mode") == "simple"
    assert override(launch.command, "+algo.curriculum.gamma_min") == "-4"
    assert override(launch.command, "+algo.curriculum.gamma_max") == "-1"
    assert Path(launch.command[2]).name == "train_duo_dcd.py"


def test_duo_dcd_wrapper_accepts_trainer_base_nll_signature(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    spec = importlib.util.spec_from_file_location(
        "train_duo_dcd_for_test", ROOT / "adapters" / "train_duo_dcd.py"
    )
    assert spec is not None and spec.loader is not None
    train_duo_dcd = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(train_duo_dcd)

    calls = []

    class Distillation:
        def nll(
            self,
            x0,
            output_tokens,
            current_accumulation_step=None,
            train_mode=None,
        ):
            calls.append(
                (x0, output_tokens, current_accumulation_step, train_mode)
            )
            return "loss"

    algo = type(sys)("algo")
    algo.Distillation = Distillation
    monkeypatch.setitem(sys.modules, "algo", algo)

    train_duo_dcd._patch_distillation_nll()

    assert Distillation().nll("x0", None, "output", "accum", True) == "loss"
    assert calls == [("x0", "output", "accum", True)]


def test_di4c_sampling_checkpoints_match_reference_selection() -> None:
    assert load_recipe("duo_di4c", "lm1b").sampling_step == 20_000
    assert load_recipe("mdlm_di4c", "lm1b").sampling_step == 20_000
    assert load_recipe("duo_di4c", "owt").sampling_step == 50_000


def test_lm1b_di4c_uses_project_wrapper_and_teacher_student_init() -> None:
    recipe = load_recipe("duo_di4c", "lm1b")
    launch = build_launch(
        recipe,
        root=ROOT,
        source=ROOT / "upstreams" / "di4c",
        output=ROOT / "checkpoints" / "reference_reproduction" / "duo_di4c" / "lm1b",
        teacher=ROOT / "checkpoints" / "fixture" / "uniform_duo.ckpt",
        devices=2,
        nodes=1,
        per_device_batch_size=1,
        seed=42,
        resume=True,
    )

    assert Path(launch.command[2]).name == "train_di4c.py"
    assert override(launch.command, "+dlb_student_init") == "teacher"


def test_di4c_student_init_can_be_overridden_to_scratch() -> None:
    recipe = load_recipe("mdlm_di4c", "lm1b")
    launch = build_launch(
        recipe,
        root=ROOT,
        source=ROOT / "upstreams" / "di4c",
        output=ROOT / "checkpoints" / "reference_reproduction" / "mdlm_di4c" / "lm1b",
        teacher=ROOT / "checkpoints" / "fixture" / "masked_mdlm.ckpt",
        devices=2,
        nodes=1,
        per_device_batch_size=1,
        seed=42,
        resume=True,
        di4c_student_init="scratch",
    )

    assert Path(launch.command[2]).name == "train_di4c.py"
    assert override(launch.command, "+dlb_student_init") == "scratch"


def test_di4c_wrapper_registers_hydra_resolvers_before_composition(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from adapters import train_di4c

    registered = {}

    class FakeOmegaConf:
        @staticmethod
        def register_new_resolver(name, resolver, *, replace=False):
            registered[name] = (resolver, replace)

    monkeypatch.setitem(
        sys.modules,
        "omegaconf",
        types.SimpleNamespace(OmegaConf=FakeOmegaConf),
    )
    train_di4c._register_omegaconf_resolvers()

    assert set(registered) == {"cwd", "device_count", "div_up", "eval"}
    assert all(replace for _, replace in registered.values())
    assert registered["div_up"][0](5, 2) == 3
    assert registered["eval"][0]("1 + 1") == 2
    assert registered["eval"][0]("len([1, 2, 3])") == 3
    assert registered["eval"][0]("__import__('math').ceil(1.2)") == 2
    assert isinstance(registered["cwd"][0](), str)
    assert isinstance(registered["device_count"][0](), int)


def test_di4c_wrapper_adds_original_script_directory_to_import_path(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from adapters import train_di4c

    source = tmp_path / "upstreams" / "di4c"
    package_root = source / "sdtt" / "src" / "sdtt"
    package_root.mkdir(parents=True)
    (package_root / "main.py").write_text("# fixture\n", encoding="utf-8")
    monkeypatch.chdir(source)
    before = list(sys.path)

    try:
        train_di4c._install_upstream_path()

        assert str(package_root) in sys.path
        assert sys.path.index(str(package_root)) < sys.path.index(
            str(source / "sdtt" / "src")
        )
    finally:
        sys.path[:] = before


def test_di4c_hf_small_student_init_preserves_configured_teacher(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from adapters import train_di4c

    fake_sdtt = types.ModuleType("sdtt")
    fake_sdtt.load_small_student = lambda **kwargs: {
        "source": "upstream-default-loader",
        "kwargs": kwargs,
    }
    fake_core = types.ModuleType("sdtt.core")
    fake_distill = types.ModuleType("sdtt.core.distill")
    fake_multi_round = types.ModuleType("sdtt.core.distill.multi_round_sdtt")
    calls = []

    class FakeMultiRoundSDTT:
        @classmethod
        def from_pretrained(
            cls,
            repo,
            revision,
            *,
            config=None,
            student_as_teacher=True,
        ):
            calls.append(
                {
                    "repo": repo,
                    "revision": revision,
                    "config": config,
                    "student_as_teacher": student_as_teacher,
                }
            )
            return {"source": "direct-hf-small-loader"}

    fake_multi_round.MultiRoundSDTT = FakeMultiRoundSDTT
    monkeypatch.setitem(sys.modules, "sdtt", fake_sdtt)
    monkeypatch.setitem(sys.modules, "sdtt.core", fake_core)
    monkeypatch.setitem(sys.modules, "sdtt.core.distill", fake_distill)
    monkeypatch.setitem(
        sys.modules,
        "sdtt.core.distill.multi_round_sdtt",
        fake_multi_round,
    )

    train_di4c._patch_student_loader()
    config = {"dlb_student_init": "hf_small"}

    result = fake_sdtt.load_small_student(
        loss="kld",
        round=7,
        config=config,
        student_as_teacher=True,
    )

    assert result == {"source": "direct-hf-small-loader"}
    assert calls == [
        {
            "repo": "jdeschena/sdtt",
            "revision": "baselines_kld_step_70000",
            "config": config,
            "student_as_teacher": False,
        }
    ]


@pytest.mark.parametrize(
    ("model", "dataset"),
    [
        ("flm", "lm1b"),
        ("duo", "lm1b"),
        ("mdlm", "lm1b"),
        ("candi", "owt"),
        ("rdlm", "lm1b"),
        ("fmlm", "lm1b"),
        ("duo_dcd", "lm1b"),
        ("mdlm_sdtt", "lm1b"),
        ("duo_di4c", "lm1b"),
        ("mdlm_di4c", "lm1b"),
    ],
)
def test_rendered_commands_use_real_entrypoints_and_absolute_project_paths(
    model: str, dataset: str
) -> None:
    recipe = load_recipe(model, dataset)
    teacher = ROOT / "checkpoints" / "fixture" / "teacher.ckpt" if recipe.teacher_family else None
    devices = 2 if model.endswith("_di4c") else 8
    launch = build_launch(
        recipe,
        root=ROOT,
        source=ROOT / recipe.source_path,
        output=ROOT / "checkpoints" / "self_trained" / dataset / model,
        teacher=teacher,
        devices=devices,
        nodes=1,
        per_device_batch_size=recipe.default_per_device_batch_size,
        seed=42,
        resume=True,
    )

    assert Path(launch.command[0]).is_absolute()
    assert Path(launch.entrypoint).is_absolute()
    assert Path(launch.cwd).is_absolute()
    assert Path(launch.data_path).is_absolute()
    assert Path(launch.output).is_absolute()
    assert all("${" not in argument for argument in launch.command)
    assert override(launch.command, recipe.output_override) == str(launch.output)
    assert override(launch.command, recipe.max_steps_override) == str(recipe.max_steps)
    assert override(launch.command, recipe.learning_rate_override) == str(
        recipe.learning_rate
    )
    assert launch.effective_global_batch_size == recipe.global_batch_size


def test_rendered_hydra_overrides_exist_in_pinned_upstream_configs() -> None:
    flm = build_launch(
        load_recipe("flm", "lm1b"),
        root=ROOT,
        source=ROOT / "upstreams" / "flm",
        output=ROOT / "checkpoints" / "self_trained" / "lm1b" / "flm",
        teacher=None,
        devices=8,
        nodes=1,
        per_device_batch_size=64,
        seed=42,
        resume=True,
    )
    assert override(flm.command, "data") == "lm1b-wrap"
    assert override(flm.command, "algo") == "flm"
    assert override(flm.command, "model.length") == "128"
    assert override(flm.command, "lr_scheduler.num_warmup_steps") == "2500"
    assert override(flm.command, "trainer.accumulate_grad_batches") == "1"
    assert override(flm.command, "wandb") == "null"


def test_all_manual_recipes_render_unique_site_independent_overrides() -> None:
    models = (
        "flm",
        "duo",
        "mdlm",
        "candi",
        "fmlm",
        "duo_dcd",
        "mdlm_sdtt",
        "duo_di4c",
        "mdlm_di4c",
    )
    for model in models:
        for dataset in ("lm1b", "owt"):
            recipe = load_recipe(model, dataset)
            devices = 2 if model.endswith("_di4c") else 8
            teacher = (
                ROOT / "checkpoints" / "fixture" / f"{recipe.teacher_family}.ckpt"
                if recipe.teacher_family
                else None
            )
            launch = build_launch(
                recipe,
                root=ROOT,
                source=ROOT / recipe.source_path,
                output=ROOT / "checkpoints" / "self_trained" / dataset / model,
                teacher=teacher,
                devices=devices,
                nodes=1,
                per_device_batch_size=recipe.default_per_device_batch_size,
                seed=42,
                resume=True,
            )
            keys = [argument.split("=", 1)[0].lstrip("+") for argument in launch.command[3:]]
            assert len(keys) == len(set(keys)), (model, dataset, keys)
            rendered = "\n".join(launch.command)
            assert "/home/" not in rendered
            assert "/share/" not in rendered
            assert "REDACTED" not in rendered
            assert launch.entrypoint.is_file()

    rdlm = load_recipe("rdlm", "lm1b")
    launch = build_launch(
        rdlm,
        root=ROOT,
        source=ROOT / rdlm.source_path,
        output=ROOT / "checkpoints" / "self_trained" / "lm1b" / "rdlm",
        teacher=None,
        devices=8,
        nodes=1,
        per_device_batch_size=64,
        seed=42,
        resume=True,
    )
    keys = [argument.split("=", 1)[0].lstrip("+") for argument in launch.command[3:]]
    assert len(keys) == len(set(keys))
    assert "/home/" not in "\n".join(launch.command)


def test_effective_global_batch_must_be_exact_and_divisible() -> None:
    assert validate_effective_batch(512, devices=8, nodes=1, per_device_batch_size=16) == 4
    with pytest.raises(RecipeError, match="divisible"):
        validate_effective_batch(512, devices=6, nodes=1, per_device_batch_size=32)
    with pytest.raises(RecipeError, match="positive"):
        validate_effective_batch(512, devices=0, nodes=1, per_device_batch_size=32)


def tensor_state(vocab: int = 3, hidden: int = 2) -> dict[str, object]:
    return {
        "vocab_embed.embedding": [
            [10 * row + column + 1 for column in range(hidden)] for row in range(vocab)
        ],
        "output_layer.linear.weight": [
            [100 + 10 * row + column for column in range(hidden)] for row in range(vocab)
        ],
        "output_layer.linear.bias": [200 + row for row in range(vocab)],
    }


def test_uniform_teacher_appends_zero_absorbing_rows() -> None:
    state = tensor_state()
    adapted = uniform_to_absorbing(state, expected_keys=set(BACKBONE_TENSOR_KEYS))

    assert adapted["vocab_embed.embedding"] == [[1, 2], [11, 12], [21, 22], [0, 0]]
    assert adapted["output_layer.linear.weight"][-1] == [0, 0]
    assert adapted["output_layer.linear.bias"] == [200, 201, 202, 0]
    assert state["vocab_embed.embedding"] == [[1, 2], [11, 12], [21, 22]]


def test_uniform_teacher_zeroes_existing_absorbing_row() -> None:
    state = tensor_state(vocab=4)

    adapted = uniform_to_absorbing(
        state,
        target_vocab_size=4,
        expected_keys=set(BACKBONE_TENSOR_KEYS),
    )

    assert adapted["vocab_embed.embedding"] == [[1, 2], [11, 12], [21, 22], [0, 0]]
    assert adapted["output_layer.linear.weight"][-1] == [0, 0]
    assert adapted["output_layer.linear.bias"] == [200, 201, 202, 0]
    assert state["vocab_embed.embedding"][-1] == [31, 32]


def test_masked_teacher_moves_existing_mask_row_to_final_absorbing_state() -> None:
    state = tensor_state()
    adapted = masked_to_absorbing(
        state,
        source_mask_index=1,
        target_vocab_size=4,
        expected_keys=set(BACKBONE_TENSOR_KEYS),
    )

    assert adapted["vocab_embed.embedding"] == [[1, 2], [0, 0], [21, 22], [11, 12]]
    assert adapted["output_layer.linear.weight"][1] == [0, 0]
    assert adapted["output_layer.linear.weight"][-1] == [110, 111]
    assert adapted["output_layer.linear.bias"] == [200, 0, 202, 201]


def test_masked_teacher_does_not_duplicate_an_already_final_mask() -> None:
    state = tensor_state(vocab=4)
    adapted = masked_to_absorbing(
        state,
        source_mask_index=3,
        target_vocab_size=4,
        expected_keys=set(BACKBONE_TENSOR_KEYS),
    )
    assert adapted == state
    assert adapted is not state


def test_teacher_adapters_allow_rotary_buffer_from_pinned_checkpoints() -> None:
    state = tensor_state()
    state["rotary_emb.inv_freq"] = [1, 2]

    adapted = masked_to_absorbing(
        state,
        source_mask_index=1,
        target_vocab_size=4,
        expected_keys=set(BACKBONE_TENSOR_KEYS),
    )

    assert adapted["rotary_emb.inv_freq"] == [1, 2]


def test_adapted_teacher_checkpoint_drops_lightning_metadata_for_weights_only_load(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe = load_recipe("mdlm_di4c", "lm1b")
    launch = build_launch(
        recipe,
        root=tmp_path,
        source=tmp_path / recipe.source_path,
        output=tmp_path / "checkpoints" / "reference_reproduction" / "mdlm_di4c",
        teacher=tmp_path / "teacher.ckpt",
        devices=1,
        nodes=1,
        per_device_batch_size=1,
        seed=42,
        resume=True,
    )
    launch.teacher.write_bytes(b"source checkpoint")
    state = tensor_state()
    for key in expected_backbone_keys() - set(state):
        state[key] = [0]
    source_payload = {
        "state_dict": state,
        "hyper_parameters": {"config": object()},
        "callbacks": {"custom": object()},
    }
    saved_payloads: list[dict[str, object]] = []

    class FakeTorch:
        @staticmethod
        def load(path, **kwargs):
            return source_payload

        @staticmethod
        def save(payload, path):
            saved_payloads.append(payload)
            Path(path).write_bytes(b"adapted checkpoint")

    monkeypatch.setitem(sys.modules, "torch", FakeTorch)
    monkeypatch.setitem(recipes_module.BASE_VOCAB_SIZES, "lm1b", 3)

    recipes_module.adapt_teacher_checkpoint(launch, source_mask_index=1)

    assert len(saved_payloads) == 1
    assert set(saved_payloads[0]) == {"state_dict"}
    assert "hyper_parameters" not in saved_payloads[0]


@pytest.mark.parametrize(
    "mutation,match",
    [
        (lambda state: state.pop("output_layer.linear.bias"), "keys"),
        (lambda state: state.update({"unexpected": [1]}), "keys"),
        (
            lambda state: state.update({"output_layer.linear.weight": [[1], [2], [3]]}),
            "hidden",
        ),
    ],
)
def test_teacher_adapters_fail_closed_on_key_or_shape_mismatch(mutation, match: str) -> None:
    state = tensor_state()
    mutation(state)
    with pytest.raises(RecipeError, match=match):
        uniform_to_absorbing(state, expected_keys=set(BACKBONE_TENSOR_KEYS))


@pytest.mark.parametrize(
    "wrapper",
    [
        "scripts/train/flm.sh",
        "scripts/train/duo.sh",
        "scripts/train/mdlm.sh",
        "scripts/train/candi.sh",
        "scripts/train/rdlm.sh",
        "scripts/distill/fmlm.sh",
        "scripts/distill/duo_dcd.sh",
        "scripts/distill/mdlm_sdtt.sh",
        "scripts/distill/di4c.sh",
    ],
)
def test_every_wrapper_supports_help(wrapper: str) -> None:
    completed = subprocess.run(
        ["bash", str(ROOT / wrapper), "--help"],
        cwd=ROOT,
        env={**os.environ, "DLB_PYTHON": sys.executable},
        text=True,
        capture_output=True,
        check=False,
    )
    assert completed.returncode == 0, completed.stderr
    assert "--dry-run" in completed.stdout


def test_manifest_prerequisite_commands_dry_run_without_external_assets() -> None:
    commands = [
        [
            "bash",
            "scripts/train/candi.sh",
            "--source",
            "upstreams/candi",
            "--dataset",
            "owt",
            "--output",
            "checkpoints/reference_reproduction/candi/owt",
            "--dry-run",
        ],
        [
            "bash",
            "scripts/distill/duo_dcd.sh",
            "--source",
            "upstreams/duo",
            "--dataset",
            "lm1b",
            "--teacher",
            "checkpoints/reference_reproduction/flm_baselines/lm1b/lm1b_Duo.ckpt",
            "--output",
            "checkpoints/reference_reproduction/duo_dcd/lm1b",
            "--rounds",
            "8",
            "--steps-per-round",
            "10000",
            "--global-batch-size",
            "128",
            "--learning-rate",
            "6e-5",
            "--dry-run",
        ],
    ]
    for command in commands:
        completed = subprocess.run(
            command,
            cwd=ROOT,
            env={**os.environ, "DLB_PYTHON": sys.executable},
            text=True,
            capture_output=True,
            check=False,
        )
        assert completed.returncode == 0, completed.stderr
        record = json.loads(completed.stdout)
        assert record["dry_run"] is True
        assert record["command"]
        assert Path(record["output"]).is_absolute()


def test_fake_launch_publishes_before_run_and_resumes_exact_completion(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe = load_recipe("flm", "lm1b")
    launch = build_launch(
        recipe,
        root=tmp_path,
        source=tmp_path / "upstreams" / "flm",
        output=tmp_path / "checkpoints" / "self_trained" / "lm1b" / "flm",
        teacher=None,
        devices=8,
        nodes=1,
        per_device_batch_size=64,
        seed=42,
        resume=True,
    )
    monkeypatch.setattr(recipes_module, "_verify_source", lambda item: None)
    monkeypatch.setattr(
        recipes_module,
        "_prepare_training_cache",
        lambda item, root: {
            "mode": "fake",
            "processed_manifest_sha256": "a" * 64,
        },
    )
    monkeypatch.setattr(
        recipes_module, "_compose_config", lambda item, environment: "mode: train\n"
    )
    calls = 0

    def fake_run(command, **kwargs):
        nonlocal calls
        calls += 1
        assert (launch.output / "recipe.json").is_file()
        assert (launch.output / "launch_argv.json").is_file()
        assert (launch.output / "provenance.json").is_file()
        checkpoint = launch.output / "checkpoints" / "last.ckpt"
        checkpoint.parent.mkdir(parents=True)
        checkpoint.write_bytes(b"fake checkpoint")
        return subprocess.CompletedProcess(command, 0)

    monkeypatch.setattr(recipes_module.subprocess, "run", fake_run)

    assert recipes_module.execute_launch(launch, root=tmp_path) == "completed"
    assert recipes_module.execute_launch(launch, root=tmp_path) == "skipped"
    assert calls == 1
    completed = json.loads((launch.output / "completed.json").read_text())
    assert completed["checkpoint"] == "model.ckpt"
    assert (launch.output / "model.ckpt").read_bytes() == b"fake checkpoint"


def test_distilled_launch_canonicalizes_lightning_step_checkpoint(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe = load_recipe("duo_di4c", "owt")
    teacher = tmp_path / "checkpoints" / "fixture" / "uniform_duo.ckpt"
    teacher.parent.mkdir(parents=True)
    teacher.write_bytes(b"fake teacher")
    launch = build_launch(
        recipe,
        root=tmp_path,
        source=tmp_path / "upstreams" / "di4c",
        output=tmp_path / "checkpoints" / "reference_reproduction" / "duo_di4c" / "owt",
        teacher=teacher,
        devices=2,
        nodes=1,
        per_device_batch_size=1,
        seed=42,
        resume=True,
    )
    monkeypatch.setattr(recipes_module, "_verify_source", lambda item: None)
    monkeypatch.setattr(
        recipes_module,
        "_prepare_training_cache",
        lambda item, root: {
            "mode": "fake",
            "processed_manifest_sha256": "a" * 64,
        },
    )
    monkeypatch.setattr(
        recipes_module, "_compose_config", lambda item, environment: "mode: train\n"
    )
    monkeypatch.setattr(
        recipes_module,
        "adapt_teacher_checkpoint",
        lambda item: {"mode": "fake_teacher_adapter"},
    )

    def fake_run(command, **kwargs):
        checkpoint = launch.output / "checkpoints" / "0-50000.ckpt"
        checkpoint.parent.mkdir(parents=True)
        checkpoint.write_bytes(b"fake distilled checkpoint")
        return subprocess.CompletedProcess(command, 0)

    monkeypatch.setattr(recipes_module.subprocess, "run", fake_run)

    assert recipes_module.execute_launch(launch, root=tmp_path) == "completed"
    completed = json.loads((launch.output / "completed.json").read_text())
    assert completed["checkpoint"] == "student_checkpoints/50000.ckpt"
    assert (launch.output / completed["checkpoint"]).read_bytes() == (
        b"fake distilled checkpoint"
    )


def test_training_cache_bridge_adds_attention_mask_for_duo_text_batches(
    tmp_path: Path,
) -> None:
    datasets = pytest.importorskip("datasets")
    Dataset = datasets.Dataset
    load_from_disk = datasets.load_from_disk

    processed = tmp_path / "data" / "processed" / "lm1b-bert-128"
    (tmp_path / "data" / "manifests").mkdir(parents=True)
    for split in ("train", "validation"):
        Dataset.from_dict(
            {"input_ids": [[101] + [7] * 126 + [102]]}
        ).save_to_disk(str(processed / split))
    (tmp_path / "data" / "manifests" / "lm1b.json").write_text(
        json.dumps({"dataset": "lm1b", "sequence_length": 128}),
        encoding="utf-8",
    )

    launch = build_launch(
        load_recipe("duo_dcd", "lm1b"),
        root=tmp_path,
        source=tmp_path / "upstreams" / "duo",
        output=tmp_path / "checkpoints" / "reference_reproduction" / "duo_dcd" / "lm1b",
        teacher=tmp_path / "checkpoints" / "fixture" / "uniform_duo.ckpt",
        devices=1,
        nodes=1,
        per_device_batch_size=128,
        seed=42,
        resume=True,
    )

    record = recipes_module._prepare_training_cache(launch, tmp_path)
    cached = load_from_disk(
        str(launch.data_path / "lm1b_train_bs128_wrapped.dat")
    )

    assert record["format"] == "hf_dataset_with_attention_mask_v1"
    assert set(cached.column_names) == {"input_ids", "attention_mask"}
    assert cached[0]["attention_mask"] == [1] * 128


def test_source_verification_ignores_python_bytecode_cache_status(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe = load_recipe("mdlm_di4c", "lm1b")
    source = tmp_path / "upstreams" / "di4c"
    entrypoint = source / recipe.entrypoint
    entrypoint.parent.mkdir(parents=True)
    entrypoint.write_text("print('placeholder')\n", encoding="utf-8")
    launch = build_launch(
        recipe,
        root=tmp_path,
        source=source,
        output=tmp_path / "checkpoints" / "reference_reproduction" / "mdlm_di4c",
        teacher=tmp_path / "teacher.ckpt",
        devices=1,
        nodes=1,
        per_device_batch_size=1,
        seed=42,
        resume=True,
    )

    def fake_run(command, **kwargs):
        if command[3] == "rev-parse":
            return subprocess.CompletedProcess(
                command,
                0,
                stdout=recipe.source_commit + "\n",
                stderr="",
            )
        if command[3] == "status":
            return subprocess.CompletedProcess(
                command,
                0,
                stdout="?? sdtt/src/sdtt/__pycache__/\n?? sdtt/src/sdtt/core/__pycache__/\n",
                stderr="",
            )
        raise AssertionError(command)

    monkeypatch.setattr(recipes_module.subprocess, "run", fake_run)

    recipes_module._verify_source(launch)


def test_training_launch_disables_python_bytecode_writes(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    recipe = load_recipe("flm", "lm1b")
    launch = build_launch(
        recipe,
        root=tmp_path,
        source=tmp_path / "upstreams" / "flm",
        output=tmp_path / "checkpoints" / "self_trained" / "lm1b" / "flm",
        teacher=None,
        devices=8,
        nodes=1,
        per_device_batch_size=64,
        seed=42,
        resume=True,
    )
    monkeypatch.setattr(recipes_module, "_verify_source", lambda item: None)
    monkeypatch.setattr(
        recipes_module,
        "_prepare_training_cache",
        lambda item, root: {
            "mode": "fake",
            "processed_manifest_sha256": "a" * 64,
        },
    )

    def fake_compose(item, environment):
        assert environment["PYTHONDONTWRITEBYTECODE"] == "1"
        return "mode: train\n"

    def fake_run(command, **kwargs):
        assert kwargs["env"]["PYTHONDONTWRITEBYTECODE"] == "1"
        checkpoint = launch.output / "checkpoints" / "last.ckpt"
        checkpoint.parent.mkdir(parents=True)
        checkpoint.write_bytes(b"fake checkpoint")
        return subprocess.CompletedProcess(command, 0)

    monkeypatch.setattr(recipes_module, "_compose_config", fake_compose)
    monkeypatch.setattr(recipes_module.subprocess, "run", fake_run)

    assert recipes_module.execute_launch(launch, root=tmp_path) == "completed"