| 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() |
|
|
