Blame for tests/test_checkpoint_manifest.py

import copy
import json
from pathlib import Path
import subprocess
import sys

import pytest
import yaml
from jsonschema.validators import Draft202012Validator

from dlb.checkpoints import (
    CheckpointResource,
    DigestSpec,
    GDriveSource,
    load_checkpoint_manifest,
    validate_checkpoint_coverage,
)
from dlb.registry import DatasetSupport, load_registry


ROOT = Path(__file__).parents[1]


@pytest.fixture
def registry():
    return load_registry(ROOT / "configs" / "experiments.yaml")


@pytest.fixture
def checkpoints():
    return load_checkpoint_manifest(ROOT / "artifacts" / "checkpoints.yaml")


@pytest.fixture
def per_file_resource():
    return CheckpointResource.model_validate(
        {
            "id": "per_file_fixture",
            "backend": "huggingface",
            "provenance": "official",
            "teacher_family": "masked_mdlm",
            "destination": "official/per-file-fixture",
            "license": "test",
            "terms_url": "https://example.test/terms",
            "digest": {
                "policy": "sha256",
                "per_file_sha256": {
                    "config.json": "1" * 64,
                    "model.safetensors": "2" * 64,
                },
            },
            "required_files": ["config.json", "model.safetensors"],
            "source": {
                "repo_id": "owner/model",
                "revision": "a" * 40,
                "allow_patterns": ["config.json", "model.safetensors"],
            },
        }
    )


def test_every_supported_cell_has_checkpoint_or_training_recipe(registry, checkpoints):
    """Catch supported experiments with no auditable acquisition path."""

    for model_name, model in registry.models.items():
        for dataset, support in model.datasets.items():
            if support.status == "supported":
                assert (model_name, dataset) in checkpoints.coverage or support.train_recipe


def test_runtime_coverage_validation_rejects_a_recipe_for_the_wrong_cell(
    registry, checkpoints
):
    """Catch a registry recipe reference that points at another model/dataset cell."""

    validate_checkpoint_coverage(registry, checkpoints)
    mutated = registry.model_copy(deep=True)
    mutated.models["duo_dcd"].datasets["lm1b"].train_recipe = "candi_owt"

    with pytest.raises(ValueError, match="does not describe duo_dcd/lm1b"):
        validate_checkpoint_coverage(mutated, checkpoints)


def test_checkpoint_sources_are_typed_and_public_resources_are_auditable(
    checkpoints, per_file_resource
):
    """Catch an unimplemented backend or a public resource missing terms/digest policy."""

    assert {resource.backend for resource in checkpoints.resources.values()} <= {
        "huggingface",
        "gdrive",
        "zenodo",
        "direct",
    }
    assert {resource.backend for resource in checkpoints.resources.values()} >= {
        "huggingface",
        "gdrive",
        "zenodo",
    }
    audited_resources = [*checkpoints.resources.values(), per_file_resource]
    for resource in audited_resources:
        assert resource.terms_url.startswith("https://")
        assert resource.digest.policy in {"sha256", "capture_after_download"}
        if resource.digest.policy == "sha256":
            if resource.digest.sha256 is not None:
                assert len(resource.digest.sha256) == 64
                assert not resource.digest.per_file_sha256
            else:
                assert resource.digest.per_file_sha256
                assert set(resource.digest.per_file_sha256) == set(
                    resource.required_files
                )
                assert all(
                    len(digest) == 64
                    for digest in resource.digest.per_file_sha256.values()
                )
        else:
            assert resource.digest.sha256 is None
            assert not resource.digest.per_file_sha256


def test_huggingface_resources_declare_required_payload_inventory(checkpoints):
    """Catch allow filters being mistaken for proof that model weights arrived."""

    for resource in checkpoints.resources.values():
        if resource.backend != "huggingface":
            continue
        assert resource.required_files
        assert "model.safetensors" in resource.required_files
        assert set(resource.required_files) <= set(resource.source.allow_patterns)


def test_multifile_digest_policy_is_typed_and_aggregate_sha_is_rejected():
    """Catch directory backends silently ignoring one aggregate SHA-256."""

    first = "1" * 64
    second = "2" * 64
    digest = DigestSpec(
        policy="sha256",
        per_file_sha256={"config.json": first, "model.safetensors": second},
    )
    assert digest.per_file_sha256 == {
        "config.json": first,
        "model.safetensors": second,
    }

    with pytest.raises(ValueError, match="aggregate SHA-256"):
        CheckpointResource.model_validate(
            {
                "id": "bad_hf",
                "backend": "huggingface",
                "provenance": "official",
                "teacher_family": "masked_mdlm",
                "destination": "official/bad",
                "license": "test",
                "terms_url": "https://example.test/terms",
                "digest": {"policy": "sha256", "sha256": first},
                "required_files": ["config.json", "model.safetensors"],
                "source": {
                    "repo_id": "owner/model",
                    "revision": "a" * 40,
                    "allow_patterns": ["config.json", "model.safetensors"],
                },
            }
        )


def test_gdrive_path_to_file_id_inventory_is_one_to_one():
    """Catch one mutable Drive object being relabeled as multiple declared files."""

    with pytest.raises(ValueError, match="unique"):
        GDriveSource(
            folder_id="folder123",
            expected_files={"model.bin": "sameID", "config.json": "sameID"},
        )


@pytest.mark.parametrize("schema_version", [True, 1.0, "1", None])
def test_manifest_schema_version_requires_the_exact_integer_one(
    tmp_path, schema_version
):
    """Catch bool, float, string, or null schema versions coercing to version 1."""

    document = yaml.safe_load((ROOT / "artifacts" / "checkpoints.yaml").read_text())
    document["schema_version"] = schema_version
    manifest = tmp_path / "checkpoints.yaml"
    manifest.write_text(yaml.safe_dump(document, sort_keys=False), encoding="utf-8")

    with pytest.raises(ValueError):
        load_checkpoint_manifest(manifest)


def test_manifest_schema_version_is_required(tmp_path):
    document = yaml.safe_load((ROOT / "artifacts" / "checkpoints.yaml").read_text())
    document.pop("schema_version")
    manifest = tmp_path / "checkpoints.yaml"
    manifest.write_text(yaml.safe_dump(document, sort_keys=False), encoding="utf-8")

    with pytest.raises(ValueError):
        load_checkpoint_manifest(manifest)


def test_known_primary_sources_and_immutable_hf_revisions_are_recorded(checkpoints):
    """Catch moving Hub references or omission of a required upstream release."""

    expected_hf = {
        "flm_lm1b_hf": ("david3684/FLM-B-LM1B", "cd0ca433d11d9392494de8122f3b63cc996fd4ce"),
        "fmlm_lm1b_hf": ("david3684/FMLM-B-LM1B", "356592ce2d9af4adbf9d7b55e8e00bbfb37eb4c8"),
        "flm_owt_hf": ("david3684/FLM-B-OWT", "624471b934fdd0421757d62290f7e639f32566d3"),
        "fmlm_owt_hf": ("david3684/FMLM-B-OWT", "483ea1b38bba56632cd40dc5a3c70a2340bb4946"),
        "langflow_owt_hf": (
            "Continuous-Rivals-Discrete/langflow-owt",
            "a08f933dd337d52762fec5ef7d60c131896cc341",
        ),
        "langflow_lm1b_hf": (
            "Continuous-Rivals-Discrete/langflow-lm1b",
            "d5d395d9bdf01f697ba8830110e3e78c63d77cb5",
        ),
        "duo_owt_hf": ("s-sahoo/duo", "ef8ad7a25a4bf9624306230913dd7ae5cb2a00bd"),
        "duo_dcd_owt_hf": (
            "s-sahoo/duo-distilled",
            "08ee71faabfc0e2eaaa6cb1833ed3a486931ef2b",
        ),
        "mdlm_owt_hf": (
            "kuleshov-group/mdlm-owt",
            "d0958fa851335ece6c15260ce0025f030673c0fb",
        ),
        "sdtt_owt_hf": ("jdeschena/sdtt", "c84bdf45ce3a80f5a70bb09548093f1c8f2ac7a7"),
    }
    for resource_name, (repo_id, revision) in expected_hf.items():
        source = checkpoints.resources[resource_name].source
        assert source.repo_id == repo_id
        assert source.revision == revision
        assert source.allow_patterns

    flm_drive = checkpoints.resources["flm_lm1b_reproductions"].source
    assert flm_drive.folder_id == (
        "1TJO3aFWqI7ukbmjciZ6krAUFlAak1itl"
    )
    assert flm_drive.expected_files == {
        "lm1b_CANDI.ckpt": "19TbJu1VHWX5EHqE8Si4xWqBhJlTqaGjf",
        "lm1b_Duo.ckpt": "1yHtaYrLX34rS5tu9V1SKF5NvwTVxPWoF",
        "lm1b_MDLM_.ckpt": "1jp6Jka8tQJJ9zljihdLErokFp9TDvlsl",
    }
    rdlm_drive = checkpoints.resources["rdlm_lm1b_drive"].source
    assert rdlm_drive.folder_id == (
        "1aDTZtPIxAxQrkaRSahjuWbkxX1OYq9CC"
    )
    assert rdlm_drive.expected_files["LM1B/checkpoint.pth"] == (
        "1zAqP-DtZigoiChpNfgJgJobKIZ9VOaro"
    )
    assert checkpoints.resources["di4c_mdlm_owt_zenodo"].source.record_id == 15124163
    expected_reproductions = {
        ("duo", "lm1b"): ("lm1b_Duo.ckpt", "uniform_duo"),
        ("mdlm", "lm1b"): ("lm1b_MDLM_.ckpt", "masked_mdlm"),
        ("candi", "lm1b"): ("lm1b_CANDI.ckpt", "hybrid_candi"),
    }
    for cell, (path, teacher_family) in expected_reproductions.items():
        coverage = checkpoints.coverage[cell]
        assert coverage.resource == "flm_lm1b_reproductions"
        assert coverage.path == path
        assert coverage.teacher_family == teacher_family


def test_coverage_preserves_registry_provenance_and_teacher_family(registry, checkpoints):
    """Catch official/reproduction conflation or a distilled cell using the wrong teacher."""

    expected_teachers = {
        "duo": "uniform_duo",
        "duo_dcd": "uniform_duo",
        "duo_di4c": "uniform_duo",
        "mdlm": "masked_mdlm",
        "mdlm_sdtt": "masked_mdlm",
        "mdlm_di4c": "masked_mdlm",
        "flm": "continuous_flm",
        "fmlm": "continuous_flm",
        "langflow": "continuous_langflow",
        "candi": "hybrid_candi",
        "rdlm": "continuous_rdlm",
    }
    for cell, coverage in checkpoints.coverage.items():
        model, dataset = cell
        support = registry.models[model].datasets[dataset]
        resource = checkpoints.resources[coverage.resource]
        assert resource.provenance == support.provenance
        assert (coverage.teacher_family or resource.teacher_family) == expected_teachers[model]

    for model_name, model in registry.models.items():
        for support in model.datasets.values():
            if support.train_recipe:
                recipe = checkpoints.recipes[support.train_recipe]
                assert recipe.teacher_family == expected_teachers[model_name]

    masked_di4c = checkpoints.resources["di4c_mdlm_owt_zenodo"]
    assert masked_di4c.teacher_family == "masked_mdlm"
    assert checkpoints.coverage[("mdlm_di4c", "owt")].resource == masked_di4c.id
    assert all(
        coverage.resource != masked_di4c.id
        for cell, coverage in checkpoints.coverage.items()
        if cell[0] == "duo_di4c"
    )


def test_flm_family_uses_official_lightning_checkpoint_recipes(registry, checkpoints):
    """Catch official FLM/FMLM .ckpt files being replaced by incompatible HF dirs."""

    expected = {
        ("flm", "lm1b"): ("flm_lm1b_official_ckpt", "official/flm_ckpt/lm1b/lm1b_flm.ckpt"),
        ("flm", "owt"): ("flm_owt_official_ckpt", "official/flm_ckpt/owt/owt_flm.ckpt"),
        (
            "fmlm",
            "lm1b",
        ): ("fmlm_lm1b_official_ckpt", "official/fmlm_ckpt/lm1b/lm1b_fmlm.ckpt"),
        ("fmlm", "owt"): ("fmlm_owt_official_ckpt", "official/fmlm_ckpt/owt/owt_fmlm.ckpt"),
    }
    for cell, (recipe_id, path) in expected.items():
        model, dataset = cell
        support = registry.models[model].datasets[dataset]
        recipe = checkpoints.recipes[recipe_id]

        assert cell not in checkpoints.coverage
        assert support.train_recipe == recipe_id
        assert recipe.provenance == "official"
        assert recipe.teacher_family == "continuous_flm"
        assert f"checkpoints/{path}" == f"{recipe.output}/{recipe.sampling_checkpoint}"


def test_training_recipes_are_explicit_and_pinned(checkpoints, registry):
    """Catch a recipe fallback that cannot identify its source revision and command."""

    recipe_ids = {
        support.train_recipe
        for model in registry.models.values()
        for support in model.datasets.values()
        if support.train_recipe
    }
    assert recipe_ids
    assert recipe_ids <= set(checkpoints.recipes)
    for recipe_id in recipe_ids:
        recipe = checkpoints.recipes[recipe_id]
        assert len(recipe.source_commit) == 40
        assert recipe.command.strip()
        if recipe.provenance == "official":
            assert recipe.command.startswith("gdown --folder ")
            assert "cp downloads/flm_official_ckpt/" in recipe.command
        else:
            assert recipe.command.startswith(
                ("bash scripts/train/", "bash scripts/distill/")
            )
        assert recipe.output.startswith("checkpoints/")


def test_candi_and_duo_dcd_use_project_owned_task12_wrappers(checkpoints):
    """Catch nonexistent clone paths and upstream jobs that ignore declared outputs."""

    candi = checkpoints.recipes["candi_owt"].command
    assert candi.startswith("bash scripts/train/candi.sh")
    assert "--source upstreams/candi" in candi
    assert "--dataset owt" in candi
    assert "--output checkpoints/reference_reproduction/candi/owt" in candi
    assert "slurm" not in candi.lower()

    dcd = checkpoints.recipes["duo_dcd_lm1b"].command
    assert dcd.startswith("bash scripts/distill/duo_dcd.sh")
    assert "--source upstreams/duo" in dcd
    assert "--dataset lm1b" in dcd
    assert "--teacher checkpoints/reference_reproduction/flm_baselines/lm1b/lm1b_Duo.ckpt" in dcd
    assert "--output checkpoints/reference_reproduction/duo_dcd/lm1b" in dcd
    assert "--rounds 8" in dcd
    assert "--steps-per-round 10000" in dcd
    assert "--global-batch-size 128" in dcd
    assert "--learning-rate 6e-5" in dcd


def test_distillation_recipes_pin_the_declared_teacher_family(checkpoints):
    """Catch recipes falling back to default MDLM teachers or non-Di4C training."""

    dcd = checkpoints.recipes["duo_dcd_lm1b"].command
    assert dcd.startswith("bash scripts/distill/duo_dcd.sh")
    assert "--teacher " in dcd
    assert "lm1b_Duo.ckpt" in dcd

    for recipe_id in ("duo_di4c_lm1b", "duo_di4c_owt"):
        recipe = checkpoints.recipes[recipe_id]
        command = recipe.command
        assert recipe.teacher_adapter == "uniform_to_absorbing"
        assert command.startswith("bash scripts/distill/di4c.sh")
        assert "is_di4c=true" in command
        assert "--teacher-family uniform_duo" in command
        if recipe_id.endswith("_lm1b"):
            assert "--student-init teacher" in command
        else:
            assert "--student-init hf_small" in command
        assert "Duo.ckpt" in command or "duo.ckpt" in command
        assert "MDLM" not in command
        assert "python src/sdtt/main.py" not in command

    masked_di4c = checkpoints.recipes["mdlm_di4c_lm1b"]
    assert masked_di4c.teacher_adapter == "masked_to_absorbing"
    assert masked_di4c.command.startswith("bash scripts/distill/di4c.sh")
    assert "is_di4c=true" in masked_di4c.command
    assert "--teacher-family masked_mdlm" in masked_di4c.command
    assert "--student-init teacher" in masked_di4c.command
    assert "lm1b_MDLM_.ckpt" in masked_di4c.command
    assert "python src/sdtt/main.py" not in masked_di4c.command

    masked_sdtt = checkpoints.recipes["mdlm_sdtt_lm1b"]
    assert masked_sdtt.teacher_adapter == "masked_to_absorbing"
    assert masked_sdtt.command.startswith("bash scripts/distill/mdlm_sdtt.sh")
    assert "--teacher-family masked_mdlm" in masked_sdtt.command
    assert "lm1b_MDLM_.ckpt" in masked_sdtt.command
    assert "python src/sdtt/main.py" not in masked_sdtt.command

    for recipe_id in (
        "duo_dcd_lm1b",
        "duo_di4c_lm1b",
        "mdlm_sdtt_lm1b",
        "mdlm_di4c_lm1b",
    ):
        command = checkpoints.recipes[recipe_id].command
        assert "lm1b" in command
        assert "--output " in command


def test_train_recipe_is_typed_and_forbidden_on_unsupported_cells():
    """Catch recipes attached to unsupported cells or blank recipe identifiers."""

    supported = DatasetSupport(
        status="supported",
        provenance="reference_reproduction",
        train_recipe="duo_lm1b",
    )
    assert supported.train_recipe == "duo_lm1b"
    with pytest.raises(ValueError, match="unsupported datasets"):
        DatasetSupport(
            status="unsupported",
            reason="No public or reproducible experiment is defined.",
            train_recipe="not_allowed",
        )
    with pytest.raises(ValueError):
        DatasetSupport(
            status="supported",
            provenance="self_trained",
            train_recipe=" ",
        )


def test_json_schema_accepts_recipes_and_rejects_them_on_unsupported_cells():
    """Catch drift between the Pydantic and JSON Schema recipe contracts."""

    document = yaml.safe_load((ROOT / "configs" / "experiments.yaml").read_text())
    schema = json.loads((ROOT / "configs" / "schema.json").read_text())
    validator = Draft202012Validator(schema)
    assert list(validator.iter_errors(document)) == []

    invalid = copy.deepcopy(document)
    invalid["models"]["rdlm"]["datasets"]["owt"]["train_recipe"] = "bad"
    assert list(validator.iter_errors(invalid))


def test_fetch_dry_run_enumerates_public_resources_and_recipes_without_writes(tmp_path):
    """Catch dry-run network/filesystem side effects or omitted acquisition fallbacks."""

    result = subprocess.run(
        [
            sys.executable,
            str(ROOT / "scripts" / "fetch_checkpoints.py"),
            "--root",
            str(tmp_path),
            "--config",
            str(ROOT / "artifacts" / "checkpoints.yaml"),
            "--registry",
            str(ROOT / "configs" / "experiments.yaml"),
            "--dry-run",
            "--all-public",
        ],
        check=False,
        capture_output=True,
        text=True,
    )

    assert result.returncode == 0, result.stderr
    lines = result.stdout.splitlines()
    manifest = load_checkpoint_manifest(ROOT / "artifacts" / "checkpoints.yaml")
    assert sum(line.startswith("RESOURCE ") for line in lines) == len(manifest.resources)
    assert sum(line.startswith("RECIPE ") for line in lines) == len(manifest.recipes)
    assert any("huggingface" in line and "official/flm/lm1b" in line for line in lines)
    assert any("gdrive" in line and "official/rdlm/lm1b" in line for line in lines)
    assert any("zenodo" in line and "official/mdlm_di4c/owt" in line for line in lines)
    assert not (tmp_path / "checkpoints").exists()
    assert not (tmp_path / "artifacts" / "checkpoint_lock.json").exists()