Blame for tests/test_teacher_adapters.py

from dataclasses import replace
import hashlib
import json
from pathlib import Path
import sys
from types import SimpleNamespace

import pytest
import yaml

import dlb.adapters.capture as capture_module
from dlb.adapters.base import AdapterError
from dlb.adapters.candi import CANDIAdapter
from dlb.adapters.duo import DuoAdapter
from dlb.adapters.flm import FLMAdapter
from dlb.adapters.mdlm import MDLMAdapter
from dlb.command import main as command_main
from dlb.checkpoints import load_checkpoint_manifest
from dlb.registry import load_registry
from dlb.runner import RunRequest, _resolve_checkpoint_provenance


ROOT = Path(__file__).parents[1]


def request(model: str, dataset: str, steps: int = 32, samples: int = 17) -> RunRequest:
    return RunRequest(
        run_id=f"{model}-{dataset}-steps-{steps}",
        model_id=model,
        dataset_id=dataset,
        step_count=steps,
        seed=7,
        sample_count=samples,
    )


def run_dir(root: Path, item: RunRequest) -> Path:
    return (
        root
        / "results"
        / "samples"
        / item.dataset_id
        / item.model_id
        / f"steps_{item.step_count}"
    )


def override(command: list[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_hf_masked_lm_adapter_only_truncates_duo_extra_vocab_class() -> None:
    """Catch Duo's HF checkpoint leaking its non-tokenizer runtime class into sampling."""

    class FakeLogits:
        def __init__(self, last_dim: int):
            self.shape = (2, 3, last_dim)

        def __getitem__(self, key):
            assert key == (Ellipsis, slice(None, 50_257, None))
            return FakeLogits(50_257)

    class FakeMaskedLM:
        def __init__(self, *, model_type: str, bare_logits: bool = False):
            self.config = SimpleNamespace(model_type=model_type, vocab_size=50_258)
            self.bare_logits = bare_logits

        def forward(self, input_ids, timesteps):
            del input_ids, timesteps
            logits = FakeLogits(50_258)
            if self.bare_logits:
                return logits
            return SimpleNamespace(logits=logits)

    duo = capture_module._adapt_hf_masked_lm_backbone(
        FakeMaskedLM(model_type="DUO")
    )
    duo_bare = capture_module._adapt_hf_masked_lm_backbone(
        FakeMaskedLM(model_type="DUO", bare_logits=True)
    )
    other = capture_module._adapt_hf_masked_lm_backbone(
        FakeMaskedLM(model_type="MDLM")
    )

    duo_logits = duo(
        x=[[0, 0, 0], [0, 0, 0]],
        sigma=[1, 1],
    )
    duo_bare_logits = duo_bare(
        x=[[0, 0, 0], [0, 0, 0]],
        sigma=[1, 1],
    )
    other_logits = other(
        x=[[0, 0, 0], [0, 0, 0]],
        sigma=[1, 1],
    )

    assert duo_logits.shape == (2, 3, 50_257)
    assert duo_bare_logits.shape == (2, 3, 50_257)
    assert other_logits.shape == (2, 3, 50_258)


def test_duo_sampler_vocab_uses_canonical_gpt2_size_when_tokenizer_adds_pad() -> None:
    owner = SimpleNamespace(
        vocab_size=50_258,
        backbone=SimpleNamespace(
            config=SimpleNamespace(model_type="DUO", vocab_size=50_258)
        ),
    )

    assert capture_module._duo_canonical_sampler_vocab_size(owner) == 50_257

    owner.vocab_size = 50_257
    assert capture_module._duo_canonical_sampler_vocab_size(owner) is None


def prepare_conversion_root(tmp_path: Path) -> Path:
    (tmp_path / "artifacts").mkdir()
    (tmp_path / "configs").mkdir()
    (tmp_path / "configs" / "experiments.yaml").write_text(
        (ROOT / "configs" / "experiments.yaml").read_text(encoding="utf-8"),
        encoding="utf-8",
    )
    (tmp_path / "artifacts" / "data.yaml").write_text(
        (ROOT / "artifacts" / "data.yaml").read_text(encoding="utf-8"),
        encoding="utf-8",
    )
    (tmp_path / "artifacts" / "checkpoints.yaml").write_text(
        (ROOT / "artifacts" / "checkpoints.yaml").read_text(encoding="utf-8"),
        encoding="utf-8",
    )
    return tmp_path


def canonical_conversion_request(root: Path, item: RunRequest) -> RunRequest:
    registry = load_registry(root / "configs" / "experiments.yaml")
    manifest_model = load_checkpoint_manifest(root / "artifacts" / "checkpoints.yaml")
    support = registry.models[item.model_id].datasets[item.dataset_id]
    manifest = root / "artifacts" / "checkpoints.yaml"
    manifest_sha256 = hashlib.sha256(manifest.read_bytes()).hexdigest()
    if support.train_recipe:
        recipe = manifest_model.recipes[support.train_recipe]
        checkpoint = root / recipe.output / str(recipe.sampling_checkpoint)
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        checkpoint.write_bytes(b"ckp")
        provenance = _resolve_checkpoint_provenance(root, item, support.train_recipe)
        return replace(
            item,
            checkpoint_sha256=provenance.sha256,
            checkpoint_lock_id=provenance.lock_id,
            checkpoint_selection=provenance.selection,
            checkpoint_teacher_family=provenance.teacher_family,
        )

    coverage = manifest_model.coverage[(item.model_id, item.dataset_id)]
    resource = coverage.resource
    selected_path = coverage.path or "model.safetensors"
    files = [
        {
            "path": f"checkpoints/fixture/{resource}/{selected_path}",
            "size_bytes": 3,
            "sha256": "b" * 64,
        }
    ]
    (root / "artifacts" / "checkpoint_lock.json").write_text(
        json.dumps(
            {
                "manifest_sha256": manifest_sha256,
                "resources": {resource: {"status": "downloaded", "files": files}},
            }
        ),
        encoding="utf-8",
    )
    provenance = _resolve_checkpoint_provenance(root, item, support.train_recipe)
    return replace(
        item,
        checkpoint_sha256=provenance.sha256,
        checkpoint_lock_id=provenance.lock_id,
        checkpoint_selection=provenance.selection,
        checkpoint_teacher_family=provenance.teacher_family,
    )


def write_standard_output(path: Path, texts: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps({"generative_ppl": 1.0, "entropy": 2.0, "generated_seqs": texts}),
        encoding="utf-8",
    )


def write_capture(
    path: Path, count: int, token_length: int, *, duplicate_id: bool = False
) -> None:
    samples = [
        {
            "sample_id": 0 if duplicate_id and index == 1 else index,
            "text": f"sample {index}",
            "token_ids": [index % 100 + 1] * token_length,
        }
        for index in range(count)
    ]
    path.write_text(
        json.dumps({"schema": "dlb-upstream-token-capture-v1", "samples": samples}),
        encoding="utf-8",
    )


@pytest.mark.parametrize(
    ("adapter", "model", "upstream"),
    [
        (FLMAdapter(), "flm", "flm"),
        (DuoAdapter(), "duo", "duo"),
        (MDLMAdapter(), "mdlm", "mdlm"),
        (CANDIAdapter(), "candi", "candi"),
    ],
)
def test_commands_reference_real_pinned_entrypoints_and_actual_hydra_keys(
    adapter, model: str, upstream: str
) -> None:
    """Catch adapters drifting to a nonexistent entrypoint or invented override."""

    item = request(model, "owt")
    command = adapter.render_command(item, run_dir(ROOT, item), dry_run=True)
    entrypoint = Path(override(command, "--upstream-entrypoint"))

    assert entrypoint == ROOT / "upstreams" / upstream / "main.py"
    assert entrypoint.is_file()
    config = yaml.safe_load((entrypoint.parent / "configs" / "config.yaml").read_text())
    data_config = override(command, "data")
    assert (entrypoint.parent / "configs" / "data" / f"{data_config}.yaml").is_file()
    for key in ("sampling.steps", "sampling.num_sample_batches", "loader.eval_batch_size"):
        section, child = key.split(".")
        assert child in config[section]
        assert override(command, key)


def test_flm_and_fmlm_render_exact_paper_sampling_semantics() -> None:
    """Catch FLM losing Euler or FMLM losing its dataset-specific paper gamma."""

    flm = request("flm", "owt")
    fmlm_lm1b = request("fmlm", "lm1b", steps=8)
    fmlm_owt = request("fmlm", "owt", steps=8)

    flm_command = FLMAdapter().render_command(flm, run_dir(ROOT, flm), dry_run=True)
    lm1b_command = FLMAdapter().render_command(
        fmlm_lm1b, run_dir(ROOT, fmlm_lm1b), dry_run=True
    )
    owt_command = FLMAdapter().render_command(
        fmlm_owt, run_dir(ROOT, fmlm_owt), dry_run=True
    )

    assert override(flm_command, "sampling.solver") == "euler"
    assert override(flm_command, "sampling.steps") == "32"
    assert override(lm1b_command, "algo") == "fmlm"
    assert override(lm1b_command, "sampling.gamma") == "0.8"
    assert override(owt_command, "sampling.gamma") == "1.0"


def test_flm_family_commands_use_official_lightning_checkpoints(tmp_path: Path) -> None:
    """Catch FLM/FMLM official ckpts being routed through the HF AutoModel branch."""

    cells = (
        ("flm", "lm1b", "checkpoints/official/flm_ckpt/lm1b/lm1b_flm.ckpt"),
        ("flm", "owt", "checkpoints/official/flm_ckpt/owt/owt_flm.ckpt"),
        ("fmlm", "lm1b", "checkpoints/official/fmlm_ckpt/lm1b/lm1b_fmlm.ckpt"),
        ("fmlm", "owt", "checkpoints/official/fmlm_ckpt/owt/owt_fmlm.ckpt"),
    )
    adapter = FLMAdapter()
    root = prepare_conversion_root(tmp_path)
    (root / "upstreams" / "flm").mkdir(parents=True)
    (root / "upstreams" / "flm" / "main.py").write_text("# fixture\n", encoding="utf-8")

    for model, dataset, checkpoint in cells:
        item = request(model, dataset)
        command = adapter.render_command(item, run_dir(root, item), dry_run=True)

        assert override(command, "algo.backbone") == "dit"
        assert Path(override(command, "eval.checkpoint_path")) == root / checkpoint


@pytest.mark.parametrize(
    ("adapter", "model", "sampler_key", "sampler_value", "noise_value"),
    [
        (DuoAdapter(), "duo", "sampling.predictor", "ancestral", "ancestral"),
        (DuoAdapter(), "duo_dcd", "sampling.predictor", "ancestral", "ancestral"),
        (MDLMAdapter(), "mdlm", "sampling.predictor", "ddpm", "True"),
    ],
)
def test_discrete_commands_use_actual_ancestral_keys_without_invented_temperature(
    adapter, model: str, sampler_key: str, sampler_value: str, noise_value: str
) -> None:
    """Catch a non-ancestral sampler or a Hydra override absent from pinned configs."""

    item = request(model, "owt", steps=8)
    command = adapter.render_command(item, run_dir(ROOT, item), dry_run=True)

    assert override(command, sampler_key) == sampler_value
    assert override(command, "sampling.noise_removal") == noise_value
    assert not any(argument.startswith("sampling.temperature=") for argument in command)


def test_candi_command_uses_the_pinned_hybrid_sampler() -> None:
    """Catch CANDI being reduced to a discrete-only or continuous-only sampler."""

    item = request("candi", "owt", steps=16)
    command = CANDIAdapter().render_command(item, run_dir(ROOT, item), dry_run=True)

    assert override(command, "algo") == "candi"
    assert override(command, "algo.sampler") == "cached"
    assert override(command, "algo.mixed_coeff") == "0.5"
    assert override(command, "algo.step_size") == "1.0"
    assert override(command, "algo.temp") == "1.0"
    assert override(command, "sampling.steps") == "16"


def test_commands_accept_smoke_results_root_under_project_results(tmp_path: Path) -> None:
    """Catch adapters rejecting runner requests with --results-root results/smoke."""

    root = prepare_conversion_root(tmp_path)
    (root / "upstreams" / "candi").mkdir(parents=True)
    (root / "upstreams" / "candi" / "main.py").write_text("# fixture\n", encoding="utf-8")
    results_root = root / "results" / "smoke"
    item = replace(request("candi", "lm1b", steps=1, samples=1), results_root=str(results_root))
    output_dir = (
        results_root
        / "samples"
        / item.dataset_id
        / item.model_id
        / f"steps_{item.step_count}"
    )

    command = CANDIAdapter().render_command(item, output_dir, dry_run=True)

    assert Path(override(command, "--capture-path")).parent == output_dir


@pytest.mark.parametrize(
    ("adapter", "model", "dataset", "batch_size"),
    [
        (FLMAdapter(), "flm", "lm1b", 32),
        (FLMAdapter(), "fmlm", "owt", 16),
        (DuoAdapter(), "duo", "owt", 8),
        (MDLMAdapter(), "mdlm", "owt", 16),
        (CANDIAdapter(), "candi", "owt", 2),
    ],
)
def test_sample_batches_are_ceiled_from_the_configured_eval_batch(
    adapter, model: str, dataset: str, batch_size: int
) -> None:
    """Catch floor division silently generating fewer samples than requested."""

    item = request(model, dataset, samples=batch_size + 1)
    command = adapter.render_command(item, run_dir(ROOT, item), dry_run=True)

    assert override(command, "loader.eval_batch_size") == str(batch_size)
    assert override(command, "sampling.num_sample_batches") == "2"
    assert override(command, "trainer.devices") == "1"


def test_rendered_commands_have_only_concrete_nonempty_arguments() -> None:
    """Catch unresolved templates or sentinel paths reaching subprocess argv."""

    adapters = {
        "flm": FLMAdapter(),
        "fmlm": FLMAdapter(),
        "duo": DuoAdapter(),
        "duo_dcd": DuoAdapter(),
        "mdlm": MDLMAdapter(),
        "candi": CANDIAdapter(),
    }
    for model, adapter in adapters.items():
        for dataset in ("lm1b", "owt"):
            item = request(model, dataset)
            command = adapter.render_command(item, run_dir(ROOT, item), dry_run=True)
            entrypoint = Path(override(command, "--upstream-entrypoint"))
            data_config = override(command, "data")
            assert (entrypoint.parent / "configs" / "data" / f"{data_config}.yaml").is_file()
            assert command
            assert all(argument and "${" not in argument for argument in command)
            assert all("{" not in argument and "}" not in argument for argument in command)
            assert all("None" not in argument for argument in command)


def test_real_build_rejects_a_missing_checkpoint() -> None:
    """Catch non-dry execution silently relying on a future or wrong checkpoint."""

    item = request("flm", "owt")
    with pytest.raises(AdapterError, match="checkpoint"):
        FLMAdapter().build_command(item, run_dir(ROOT, item))


def test_conversion_prefers_captured_token_ids_and_trims_only_ceiling_excess(
    tmp_path: Path,
) -> None:
    """Catch re-tokenization despite IDs being available or nondeterministic trimming."""

    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("flm", "owt", samples=17))
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)
    texts = [f"sample {index}" for index in range(32)]
    write_standard_output(output_dir / "upstream_samples.json", texts)
    write_capture(output_dir / "upstream_token_ids.json", 32, 1024)

    records = list(FLMAdapter().convert_outputs(item, output_dir))

    assert len(records) == 17
    assert [record.sample_id for record in records] == list(range(17))
    assert records[1].token_ids == [2] * 1024
    metadata = json.loads((output_dir / "conversion_metadata.json").read_text())
    assert metadata["token_ids_source"] == "upstream"
    assert metadata["trimmed_samples"] == 15
    assert metadata["generation_seconds_source"] == "unavailable_excluded_sentinel"
    assert metadata["checkpoint_sha256"] == item.checkpoint_sha256
    assert metadata["checkpoint_lock_id"] == item.checkpoint_lock_id
    assert metadata["checkpoint_selection"] == item.checkpoint_selection
    assert metadata["teacher_family"] == "continuous_flm"
    assert all(record.generation_seconds == 0.0 for record in records)


def test_conversion_rejects_missing_runner_checkpoint_provenance_before_reading_outputs(
    tmp_path: Path,
) -> None:
    """Catch direct conversion bypassing the runner's canonical checkpoint resolution."""

    root = prepare_conversion_root(tmp_path)
    item = request("flm", "owt", samples=1)
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)

    with pytest.raises(AdapterError, match="runner-resolved checkpoint provenance is required"):
        list(FLMAdapter().convert_outputs(item, output_dir))


def test_conversion_rejects_wrong_teacher_family_before_reading_outputs(
    tmp_path: Path,
) -> None:
    """Catch a valid checkpoint identity being relabeled as another teacher family."""

    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("flm", "owt", samples=1))
    item = replace(item, checkpoint_teacher_family="continuous_rdlm")
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)

    with pytest.raises(AdapterError, match="teacher family"):
        list(FLMAdapter().convert_outputs(item, output_dir))


def test_conversion_rejects_extra_missing_empty_duplicate_and_invalid_samples(
    tmp_path: Path,
) -> None:
    """Catch malformed upstream artifacts being normalized into plausible samples."""

    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("flm", "owt", samples=17))
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)
    write_standard_output(output_dir / "upstream_samples.json", [f"sample {i}" for i in range(33)])
    write_capture(output_dir / "upstream_token_ids.json", 33, 1024)
    with pytest.raises(AdapterError, match="expected 32 generated samples, found 33"):
        list(FLMAdapter().convert_outputs(item, output_dir))

    write_standard_output(output_dir / "upstream_samples.json", [f"sample {i}" for i in range(32)])
    write_capture(output_dir / "upstream_token_ids.json", 32, 1024, duplicate_id=True)
    with pytest.raises(AdapterError, match="duplicate sample_id"):
        list(FLMAdapter().convert_outputs(item, output_dir))

    write_capture(output_dir / "upstream_token_ids.json", 32, 1024)
    capture = json.loads((output_dir / "upstream_token_ids.json").read_text())
    capture["samples"][0]["token_ids"] = [50257] + [1] * 1023
    (output_dir / "upstream_token_ids.json").write_text(json.dumps(capture))
    with pytest.raises(AdapterError, match="invalid token"):
        list(FLMAdapter().convert_outputs(item, output_dir))

    write_capture(output_dir / "upstream_token_ids.json", 32, 1024)
    actual = json.loads((output_dir / "upstream_samples.json").read_text())
    actual["generated_seqs"][0] = "   "
    (output_dir / "upstream_samples.json").write_text(json.dumps(actual))
    with pytest.raises(AdapterError, match="empty text"):
        list(FLMAdapter().convert_outputs(item, output_dir))


def test_mdlm_conversion_accepts_only_the_documented_capture_format(tmp_path: Path) -> None:
    """Catch parsing MDLM's lossy final-batch stdout representation."""

    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("mdlm", "lm1b", samples=17))
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)
    write_capture(output_dir / "upstream_token_ids.json", 32, 128)

    records = list(MDLMAdapter().convert_outputs(item, output_dir))
    assert len(records) == 17
    assert records[0].token_ids == [1] * 128

    (output_dir / "upstream_token_ids.json").write_text(json.dumps(["not", "supported"]))
    with pytest.raises(AdapterError, match="unexpected capture format"):
        list(MDLMAdapter().convert_outputs(item, output_dir))


@pytest.mark.parametrize(
    ("dataset", "length", "batch_size"), [("lm1b", 128, 32), ("owt", 1024, 16)]
)
def test_conversion_requires_exact_canonical_token_length_for_each_dataset(
    tmp_path: Path, dataset: str, length: int, batch_size: int
) -> None:
    """Catch short or long token sequences passing canonical conversion validation."""

    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("flm", dataset, samples=1))
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)
    texts = [f"sample {index}" for index in range(batch_size)]
    write_standard_output(output_dir / "upstream_samples.json", texts)

    write_capture(output_dir / "upstream_token_ids.json", batch_size, length)
    records = list(FLMAdapter().convert_outputs(item, output_dir))
    assert len(records[0].token_ids) == length

    write_capture(output_dir / "upstream_token_ids.json", batch_size, length - 1)
    with pytest.raises(AdapterError, match=f"expected {length} tokens"):
        list(FLMAdapter().convert_outputs(item, output_dir))

    write_capture(output_dir / "upstream_token_ids.json", batch_size, length + 1)
    with pytest.raises(AdapterError, match=f"expected {length} tokens"):
        list(FLMAdapter().convert_outputs(item, output_dir))


def test_text_only_conversion_pads_and_truncates_with_the_pinned_offline_tokenizer(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Catch text fallback emitting variable-length IDs or hiding its transformation."""

    class FakeTokenizer:
        eos_token = "<eos>"
        pad_token = None
        pad_token_id = 50256

        def __call__(self, texts, **settings):
            assert settings == {
                "add_special_tokens": False,
                "padding": "max_length",
                "truncation": True,
                "max_length": 1024,
                "return_attention_mask": False,
            }
            return {"input_ids": [[11] + [50256] * 1023 for _ in texts]}

    class FakeAutoTokenizer:
        @staticmethod
        def from_pretrained(name, **settings):
            assert name == "gpt2"
            assert settings == {
                "revision": "607a30d783dfa663caf39e06633721c8d4cfcd7e",
                "local_files_only": True,
            }
            return FakeTokenizer()

    monkeypatch.setitem(
        sys.modules, "transformers", SimpleNamespace(AutoTokenizer=FakeAutoTokenizer)
    )
    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("flm", "owt", samples=1))
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)
    write_standard_output(
        output_dir / "upstream_samples.json", [f"sample {index}" for index in range(16)]
    )

    records = list(FLMAdapter().convert_outputs(item, output_dir))
    metadata = json.loads((output_dir / "conversion_metadata.json").read_text())

    assert records[0].token_ids == [11] + [50256] * 1023
    assert metadata["token_ids_transformation"] == {
        "add_special_tokens": False,
        "max_length": 1024,
        "padding": "max_length",
        "pad_token_id": 50256,
        "truncation": True,
    }


def test_command_cli_dry_renders_six_models_on_both_datasets_without_writes(
    capsys: pytest.CaptureFixture[str],
) -> None:
    """Catch dry-run omitting a cell or accidentally trying to acquire artifacts."""

    result = command_main(
        [
            "--root",
            str(ROOT),
            "--models",
            "flm,fmlm,duo,duo_dcd,mdlm,candi",
            "--datasets",
            "lm1b,owt",
            "--dry-run",
        ]
    )

    records = [json.loads(line) for line in capsys.readouterr().out.splitlines()]
    assert result == 0
    assert len(records) == 12
    assert {(record["model"], record["dataset"]) for record in records} == {
        (model, dataset)
        for model in ("flm", "fmlm", "duo", "duo_dcd", "mdlm", "candi")
        for dataset in ("lm1b", "owt")
    }
    assert all(record["status"] == "supported" and record["command"] for record in records)


def test_command_cli_emits_a_structured_unsupported_record(
    capsys: pytest.CaptureFixture[str],
) -> None:
    """Catch unsupported combinations terminating the whole render matrix."""

    assert command_main(["--root", str(ROOT), "--models", "rdlm", "--datasets", "owt", "--dry-run"]) == 0
    record = json.loads(capsys.readouterr().out)
    assert record["status"] == "unsupported"
    assert record["model"] == "rdlm"
    assert record["dataset"] == "owt"
    assert record["reason"]