"""Contracts for the pinned SDTT and Di4C distilled language samplers.""" from dataclasses import replace import hashlib import importlib.util import json from pathlib import Path import re import sys from types import SimpleNamespace import pytest import yaml from dlb.adapters.base import AdapterError from dlb.adapters.di4c import Di4CAdapter from dlb.adapters.sdtt import SDTTAdapter from dlb.command import main as command_main from dlb.runner import RunRequest, _resolve_checkpoint_provenance ROOT = Path(__file__).parents[1] def load_runtime_module(): path = ROOT / "adapters/_distilled_runtime.py" specification = importlib.util.spec_from_file_location("dlb_distilled_runtime", path) assert specification is not None and specification.loader is not None module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) return module def request(model: str, dataset: str, steps: int = 8, samples: int = 17) -> RunRequest: return RunRequest( run_id=f"{model}-{dataset}-steps-{steps}", model_id=model, dataset_id=dataset, step_count=steps, seed=42, sample_count=samples, ) def run_dir(item: RunRequest) -> Path: return ( ROOT / "results" / "samples" / item.dataset_id / item.model_id / f"steps_{item.step_count}" ) def run_dir_at(root: Path, item: RunRequest) -> Path: return ( root / "results" / "samples" / item.dataset_id / item.model_id / f"steps_{item.step_count}" ) def option(command: list[str], key: str) -> str: index = command.index(key) return command[index + 1] def prepare_tokenizer_binding( root: Path, dataset: str ) -> tuple[Path, Path, Path]: """Create the minimal immutable tokenizer lock consumed by wrapper tests.""" data_config = root / "artifacts" / "data.yaml" data_config.parent.mkdir(parents=True, exist_ok=True) data_config.write_text( (ROOT / "artifacts/data.yaml").read_text(encoding="utf-8"), encoding="utf-8", ) document = yaml.safe_load(data_config.read_text(encoding="utf-8")) tokenizer_id = document["datasets"][dataset]["tokenizer"] revision = document["models"][tokenizer_id] snapshot = ( root / "locked-tokenizers" / f"models--{tokenizer_id}" / "snapshots" / revision ) snapshot.mkdir(parents=True) downloads = root / "data/manifests/downloads.json" downloads.parent.mkdir(parents=True) downloads.write_text( json.dumps( { "schema_version": 1, "models": { tokenizer_id: { "repo_id": tokenizer_id, "revision": revision, "snapshot_path": snapshot.relative_to(root).as_posix(), "size_bytes": 1, } }, } ), encoding="utf-8", ) return data_config, downloads, snapshot def prepare_distilled_render_root(tmp_path: Path) -> Path: for relative in ( Path("artifacts/data.yaml"), Path("artifacts/checkpoints.yaml"), Path("configs/experiments.yaml"), Path("configs/sampling/di4c_mdlm_owt.yaml"), Path("adapters/sample_di4c.py"), ): target = tmp_path / relative target.parent.mkdir(parents=True, exist_ok=True) target.write_text((ROOT / relative).read_text(encoding="utf-8"), encoding="utf-8") entrypoint = tmp_path / "upstreams/di4c/sdtt/src/sdtt/main.py" entrypoint.parent.mkdir(parents=True, exist_ok=True) entrypoint.write_text("# fixture\n", encoding="utf-8") return tmp_path @pytest.mark.parametrize( ("model", "dataset", "adapter_type"), [ ("mdlm_sdtt", "lm1b", SDTTAdapter), ("mdlm_sdtt", "owt", SDTTAdapter), ("mdlm_di4c", "lm1b", Di4CAdapter), ("mdlm_di4c", "owt", Di4CAdapter), ("duo_di4c", "lm1b", Di4CAdapter), ("duo_di4c", "owt", Di4CAdapter), ], ) def test_all_six_distilled_cells_render_real_project_wrappers( model: str, dataset: str, adapter_type: type ) -> None: item = request(model, dataset) adapter = adapter_type() if adapter_type is SDTTAdapter else adapter_type(model.split("_", 1)[0]) command = adapter.render_command(item, run_dir(item), dry_run=True) wrapper = ROOT / "adapters" / ("sample_sdtt.py" if model == "mdlm_sdtt" else "sample_di4c.py") assert command[:4] == [command[0], "-B", "-u", str(wrapper)] assert Path(option(command, "--upstream-root")).is_dir() assert option(command, "--sample-count") == "17" assert option(command, "--num-steps") == "8" assert option(command, "--seed") == "42" assert option(command, "--sampler") == "ancestral" assert option(command, "--offline") == "true" assert Path(option(command, "--data-config")) == ROOT / "artifacts/data.yaml" assert Path(option(command, "--downloads-manifest")) == ( ROOT / "data/manifests/downloads.json" ) assert option(command, "--dataset") == dataset assert option(command, "--checkpoint-sha256") == "dry-run-unverified" assert option(command, "--config-sha256") == "dry-run-unverified" assert Path(option(command, "--output")) == run_dir(item) / "upstream_token_ids.json" def test_sdtt_uses_kld_round_seven_and_the_official_local_snapshot() -> None: item = request("mdlm_sdtt", "owt", steps=4) command = SDTTAdapter().render_command(item, run_dir(item), dry_run=True) assert option(command, "--loss") == "kld" assert option(command, "--round") == "7" assert Path(option(command, "--checkpoint")) == ( ROOT / "checkpoints/official/mdlm_sdtt/owt/model.safetensors" ) assert Path(option(command, "--config")) == ( ROOT / "checkpoints/official/mdlm_sdtt/owt/config.json" ) assert Path(option(command, "--tokenizer-snapshot")) == ( ROOT / "data/raw/huggingface/hub/models--gpt2/snapshots" / "607a30d783dfa663caf39e06633721c8d4cfcd7e" ) @pytest.mark.parametrize( ("model", "dataset", "family", "suffix"), [ ("mdlm_di4c", "lm1b", "masked_mdlm", "student_checkpoints/20000.ckpt"), ("mdlm_di4c", "owt", "masked_mdlm", "sdtt7-di4c2.ckpt"), ("duo_di4c", "lm1b", "uniform_duo", "student_checkpoints/20000.ckpt"), ("duo_di4c", "owt", "uniform_duo", "student_checkpoints/50000.ckpt"), ], ) def test_di4c_binds_teacher_family_and_dataset_intermediate_checkpoint( model: str, dataset: str, family: str, suffix: str ) -> None: item = request(model, dataset) command = Di4CAdapter(model.split("_", 1)[0]).render_command( item, run_dir(item), dry_run=True ) assert option(command, "--teacher-family") == family assert Path(option(command, "--checkpoint")).as_posix().endswith(suffix) assert Path(option(command, "--config")).name in { "config.yaml", "di4c_mdlm_owt.yaml", } if family == "uniform_duo": assert "sdtt7-di4c2.ckpt" not in " ".join(command) assert "official/mdlm_di4c" not in " ".join(command) def test_di4c_checkpoints_allow_missing_embedded_config( tmp_path: Path, ) -> None: root = prepare_distilled_render_root(tmp_path) official_owt = request("mdlm_di4c", "owt") reproduced_owt = request("duo_di4c", "owt") recipe_lm1b = request("mdlm_di4c", "lm1b") official_owt_command = Di4CAdapter("mdlm").render_command( official_owt, run_dir_at(root, official_owt), dry_run=True ) reproduced_owt_command = Di4CAdapter("duo").render_command( reproduced_owt, run_dir_at(root, reproduced_owt), dry_run=True ) recipe_lm1b_command = Di4CAdapter("mdlm").render_command( recipe_lm1b, run_dir_at(root, recipe_lm1b), dry_run=True ) assert option(official_owt_command, "--allow-missing-embedded-config") == "true" assert option(reproduced_owt_command, "--allow-missing-embedded-config") == "true" assert option(recipe_lm1b_command, "--allow-missing-embedded-config") == "true" def test_di4c_rejects_cross_family_model_identity_before_checkpoint_resolution() -> None: item = request("duo_di4c", "owt") with pytest.raises(AdapterError, match="does not support"): Di4CAdapter("mdlm").render_command(item, run_dir(item), dry_run=True) def test_distilled_dry_run_matrix_has_six_supported_cells(capsys) -> None: exit_code = command_main( [ "--root", str(ROOT), "--models", "mdlm_sdtt,mdlm_di4c,duo_di4c", "--datasets", "lm1b,owt", "--steps", "8", "--num-samples", "17", "--seed", "42", "--dry-run", ] ) lines = capsys.readouterr().out.splitlines() assert exit_code == 0 assert len(lines) == 6 assert all('"status": "supported"' in line for line in lines) def test_wrapper_files_are_source_auditable_and_upstreams_stay_unmodified() -> None: assert (ROOT / "adapters/sample_sdtt.py").is_file() assert (ROOT / "adapters/sample_di4c.py").is_file() assert (ROOT / "patches/sdtt/README.md").is_file() assert (ROOT / "patches/di4c/README.md").is_file() class FakeTensor: def __init__(self, rows: list[list[int]]) -> None: self.rows = rows def detach(self): return self def cpu(self): return self def tolist(self): return self.rows class FakeModel: def __init__(self, length: int, fail_on_call: int | None = None) -> None: self.length = length self.fail_on_call = fail_on_call self.calls: list[dict[str, object]] = [] def sample(self, **kwargs): self.calls.append(kwargs) if self.fail_on_call == len(self.calls): raise RuntimeError("fake sampler failure") start = sum(int(call["n_samples"]) for call in self.calls[:-1]) return FakeTensor( [[start + index + 1] * self.length for index in range(kwargs["n_samples"])] ) class FakeTokenizer: def batch_decode(self, rows): return [f"decoded {row[0]}" for row in rows] class ResolverGuardedDict(dict): """Model an OmegaConf tree whose unrelated custom resolvers cannot run.""" def items(self): raise RuntimeError("unregistered resolver device_count must not be evaluated") def test_sampling_runtime_streams_exact_batches_and_atomically_publishes(tmp_path: Path) -> None: runtime = load_runtime_module() model = FakeModel(length=8) output = tmp_path / "capture.json" runtime.write_capture_atomic( output, model=model, tokenizer=FakeTokenizer(), sample_count=17, batch_size=16, num_steps=4, seq_len=8, sampler="ancestral", ) document = json.loads(output.read_text(encoding="utf-8")) assert [call["n_samples"] for call in model.calls] == [16, 1] assert all(call["num_steps"] == 4 for call in model.calls) assert [item["sample_id"] for item in document["samples"]] == list(range(17)) assert all(len(item["token_ids"]) == 8 for item in document["samples"]) assert not any(path.name.endswith(".partial") for path in tmp_path.iterdir()) def test_sampling_runtime_failure_preserves_existing_output(tmp_path: Path) -> None: runtime = load_runtime_module() output = tmp_path / "capture.json" output.write_text("old-complete-output", encoding="utf-8") with pytest.raises(RuntimeError, match="fake sampler failure"): runtime.write_capture_atomic( output, model=FakeModel(length=8, fail_on_call=2), tokenizer=FakeTokenizer(), sample_count=17, batch_size=16, num_steps=4, seq_len=8, sampler="ancestral", ) assert output.read_text(encoding="utf-8") == "old-complete-output" assert not any(path.name.endswith(".partial") for path in tmp_path.iterdir()) def test_sampling_runtime_rejects_a_symlink_output(tmp_path: Path) -> None: runtime = load_runtime_module() target = tmp_path / "target.json" target.write_text("do not overwrite", encoding="utf-8") output = tmp_path / "capture.json" output.symlink_to(target) with pytest.raises(ValueError, match="unsafe"): runtime.write_capture_atomic( output, model=FakeModel(length=8), tokenizer=FakeTokenizer(), sample_count=1, batch_size=1, num_steps=1, seq_len=8, sampler="ancestral", ) assert target.read_text(encoding="utf-8") == "do not overwrite" def test_lightning_dictconfig_checkpoint_uses_verified_full_load( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """Catch PyTorch 2.5 weights-only rejecting Lightning's OmegaConf DictConfig.""" runtime = load_runtime_module() checkpoint = tmp_path / "student.ckpt" checkpoint.write_bytes(b"manifest-verified-lightning-checkpoint") expected = hashlib.sha256(checkpoint.read_bytes()).hexdigest() embedded = {"model": {"type": "ddit-orig"}} calls: list[dict[str, object]] = [] def load(path, **kwargs): calls.append({"path": path, **kwargs}) if kwargs.get("weights_only") is not False: raise RuntimeError("Weights only load failed for omegaconf.dictconfig.DictConfig") return { "state_dict": {"backbone.weight": object()}, "hyper_parameters": {"config": embedded}, } monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(load=load)) state, observed = runtime.checkpoint_state(checkpoint, expected) assert calls == [ {"path": str(checkpoint), "map_location": "cpu", "weights_only": False} ] assert state == {"backbone.weight": state["backbone.weight"]} assert observed is embedded def test_checkpoint_loader_rejects_unverified_pickle_before_torch_load( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: runtime = load_runtime_module() checkpoint = tmp_path / "student.ckpt" checkpoint.write_bytes(b"checkpoint") called = False def load(*args, **kwargs): nonlocal called called = True raise AssertionError("must not deserialize unverified bytes") monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(load=load)) with pytest.raises(ValueError, match="SHA-256"): runtime.checkpoint_state(checkpoint, "0" * 64) assert not called def test_real_lightning_omegaconf_fixture_when_runtime_is_available(tmp_path: Path) -> None: torch = pytest.importorskip("torch") OmegaConf = pytest.importorskip("omegaconf").OmegaConf runtime = load_runtime_module() checkpoint = tmp_path / "student.ckpt" config = OmegaConf.create({"model": {"type": "ddit-orig"}}) torch.save( { "state_dict": {"backbone.weight": torch.tensor([1.0])}, "hyper_parameters": {"config": config}, }, checkpoint, ) digest = hashlib.sha256(checkpoint.read_bytes()).hexdigest() state, embedded = runtime.checkpoint_state(checkpoint, digest) assert state["backbone.weight"].tolist() == [1.0] assert OmegaConf.to_container(embedded) == {"model": {"type": "ddit-orig"}} def test_embedded_config_must_match_manifest_selected_architecture() -> None: runtime = load_runtime_module() authoritative = { "model": {"type": "ddit-orig", "hidden_size": 768, "n_blocks": 12}, "parameterization": {"name": "multi-round-sdtt", "sampling_mode": "ancestral"}, "training": {"sampling_eps": 1e-5}, "T": 1024, "time_conditioning": False, } mismatched = json.loads(json.dumps(authoritative)) mismatched["model"]["n_blocks"] = 24 with pytest.raises(ValueError, match="model.n_blocks"): runtime.validate_embedded_config(authoritative, mismatched) def test_missing_embedded_config_is_allowed_only_for_explicit_release_escape() -> None: runtime = load_runtime_module() authoritative = complete_di4c_sampling_config() with pytest.raises(ValueError, match="does not embed"): runtime.validate_embedded_config(authoritative, None) runtime.validate_embedded_config( authoritative, None, allow_missing_embedded_config=True ) def complete_di4c_sampling_config() -> dict[str, object]: config = yaml.safe_load( (ROOT / "configs/sampling/di4c_mdlm_owt.yaml").read_text(encoding="utf-8") ) assert isinstance(config, dict) return config def test_config_validation_selects_only_sampling_fields_and_standard_references() -> None: runtime = load_runtime_module() config = complete_di4c_sampling_config() config["data_preprocess"]["seq_len"] = "${model.length}" config["trainer"] = {"devices": "${device_count:}"} config["loader"] = { "batch_size": "${div_up:${loader.global_batch_size},${trainer.devices}}" } guarded = ResolverGuardedDict(config) binding = runtime.TokenizerBinding("gpt2", "a" * 40, Path("/locked/gpt2")) runtime.validate_sampling_config( guarded, binding=binding, sequence_length=1024, require_di4c=True, ) runtime.validate_embedded_config(guarded, ResolverGuardedDict(config)) def test_config_validation_accepts_locked_tokenizer_snapshot_path() -> None: runtime = load_runtime_module() config = complete_di4c_sampling_config() snapshot = Path("/locked/hub/models--gpt2/snapshots/" + "a" * 40) config["tokenizer"]["name"] = str(snapshot) binding = runtime.TokenizerBinding("gpt2", "a" * 40, snapshot) runtime.validate_sampling_config( config, binding=binding, sequence_length=1024, require_di4c=True, ) def test_real_omegaconf_validation_ignores_unregistered_unrelated_resolvers() -> None: OmegaConf = pytest.importorskip("omegaconf").OmegaConf runtime = load_runtime_module() for name in ("device_count", "div_up", "eval", "cwd"): OmegaConf.clear_resolver(name) config = complete_di4c_sampling_config() config["data_preprocess"]["seq_len"] = "${model.length}" config["trainer"] = {"devices": "${device_count:}"} config["loader"] = { "global_batch_size": 16, "batch_size": "${div_up:${loader.global_batch_size},${trainer.devices}}", } config["checkpointing"] = {"save_dir": "${cwd:}"} config["eval"] = {"batches": "${eval:1 + 1}"} authoritative = OmegaConf.create(config) embedded = OmegaConf.create(config) binding = runtime.TokenizerBinding("gpt2", "a" * 40, Path("/locked/gpt2")) runtime.validate_sampling_config( authoritative, binding=binding, sequence_length=1024, require_di4c=True, ) runtime.validate_embedded_config(authoritative, embedded) @pytest.mark.parametrize( "field", [ "T", "time_conditioning", "model.type", "model.hidden_size", "model.cond_dim", "model.length", "model.n_blocks", "model.n_heads", "model.scale_by_sigma", "model.dropout", "model.tie_word_embeddings", "model.causal", "noise.type", "noise.sigma_min", "noise.sigma_max", "training.ema", "training.antithetic_sampling", "training.importance_sampling", "training.sampling_eps", "training.change_of_variables", "tokenizer.name", "parameterization.name", "parameterization.log_loss_buckets", "parameterization.start_from_hf", "parameterization.distill_mode", "parameterization.num_distill_steps", "parameterization.min_num_sampling_steps", "parameterization.grow_dt_every", "parameterization.orig_num_sampling_steps", "parameterization.sampling_mode", "parameterization.loss_precision", "parameterization.reset_optimizer_on_growth", "parameterization.use_ema_on_growth", "is_di4c", ], ) def test_every_constructor_and_sampler_field_must_match_embedded_config( field: str, ) -> None: runtime = load_runtime_module() authoritative = complete_di4c_sampling_config() authoritative["model"]["causal"] = False embedded = json.loads(json.dumps(authoritative)) cursor = embedded components = field.split(".") for component in components[:-1]: cursor = cursor[component] value = cursor[components[-1]] if isinstance(value, bool): replacement = not value elif isinstance(value, (int, float)): replacement = value + 1 elif value is None: replacement = "different" else: replacement = f"{value}-different" cursor[components[-1]] = replacement with pytest.raises(ValueError, match=re.escape(field)): runtime.validate_embedded_config(authoritative, embedded) def test_tokenizer_binding_cross_checks_data_and_download_manifests(tmp_path: Path) -> None: runtime = load_runtime_module() data_config, downloads, snapshot = prepare_tokenizer_binding(tmp_path, "owt") binding = runtime.load_tokenizer_binding( data_config, downloads, "owt", snapshot ) assert binding.tokenizer_id == "gpt2" assert binding.snapshot == snapshot document = json.loads(downloads.read_text(encoding="utf-8")) document["models"]["gpt2"]["revision"] = "0" * 40 downloads.write_text(json.dumps(document), encoding="utf-8") with pytest.raises(ValueError, match="differs"): runtime.load_tokenizer_binding(data_config, downloads, "owt", snapshot) def test_project_sampling_config_bytes_are_bound_to_checkpoint_provenance( tmp_path: Path, ) -> None: for relative in ( Path("artifacts/checkpoints.yaml"), Path("configs/experiments.yaml"), Path("configs/sampling/di4c_mdlm_owt.yaml"), ): target = tmp_path / relative target.parent.mkdir(parents=True, exist_ok=True) target.write_bytes((ROOT / relative).read_bytes()) manifest = tmp_path / "artifacts/checkpoints.yaml" manifest_sha256 = hashlib.sha256(manifest.read_bytes()).hexdigest() checkpoint_path = "checkpoints/official/mdlm_di4c/owt/sdtt7-di4c2.ckpt" (tmp_path / "artifacts/checkpoint_lock.json").write_text( json.dumps( { "manifest_sha256": manifest_sha256, "resources": { "di4c_mdlm_owt_zenodo": { "status": "downloaded", "files": [ { "path": checkpoint_path, "size_bytes": 3, "sha256": "a" * 64, } ], } }, } ), encoding="utf-8", ) item = request("mdlm_di4c", "owt", samples=1) provenance = _resolve_checkpoint_provenance(tmp_path, item, None) assert provenance.selection["sampling_config_sha256"] == ( "cf627f89b8dede9b25ac8e10407c3b7de203a76c4ba84100723287e571dcec34" ) assert provenance.selection["sampling_config_source_commit"] == ( "ac61ff9fe8e85120f9e1d2a8c5a332f8b8353dd3" ) config = tmp_path / "configs/sampling/di4c_mdlm_owt.yaml" config.write_text(config.read_text(encoding="utf-8") + "\n# mutation\n", encoding="utf-8") with pytest.raises(ValueError, match="project sampling config differs"): _resolve_checkpoint_provenance(tmp_path, item, None) @pytest.mark.parametrize( ("adapter", "model", "dataset"), [ (SDTTAdapter(), "mdlm_sdtt", "lm1b"), (SDTTAdapter(), "mdlm_sdtt", "owt"), (Di4CAdapter("mdlm"), "mdlm_di4c", "lm1b"), (Di4CAdapter("mdlm"), "mdlm_di4c", "owt"), (Di4CAdapter("duo"), "duo_di4c", "lm1b"), (Di4CAdapter("duo"), "duo_di4c", "owt"), ], ) def test_real_commands_reject_missing_or_unverified_checkpoint_bytes( adapter, model: str, dataset: str ) -> None: item = request(model, dataset, samples=1) with pytest.raises(AdapterError, match="checkpoint"): adapter.build_command(item, run_dir(item)) def prepare_official_sdtt_conversion(tmp_path: Path, item: RunRequest) -> RunRequest: for relative in ( Path("artifacts/data.yaml"), Path("artifacts/checkpoints.yaml"), Path("configs/experiments.yaml"), ): target = tmp_path / relative target.parent.mkdir(parents=True, exist_ok=True) target.write_text((ROOT / relative).read_text(encoding="utf-8"), encoding="utf-8") manifest = tmp_path / "artifacts/checkpoints.yaml" manifest_sha256 = hashlib.sha256(manifest.read_bytes()).hexdigest() files = [ { "path": f"checkpoints/official/mdlm_sdtt/owt/{name}", "size_bytes": 3, "sha256": character * 64, } for name, character in (("config.json", "a"), ("model.safetensors", "b")) ] (tmp_path / "artifacts/checkpoint_lock.json").write_text( json.dumps( { "manifest_sha256": manifest_sha256, "resources": {"sdtt_owt_hf": {"status": "downloaded", "files": files}}, } ), encoding="utf-8", ) selector = { "resource": "sdtt_owt_hf", "path": None, "teacher_family": None, "sampling_config": "config.json", "sampling_config_source": "resource", "sampling_config_sha256": None, "sampling_config_source_commit": None, } digest = hashlib.sha256( json.dumps( { "manifest_sha256": manifest_sha256, "selector": selector, "files": sorted(files, key=lambda value: value["path"]), }, sort_keys=True, separators=(",", ":"), ensure_ascii=False, ).encode("utf-8") ).hexdigest() return replace( item, checkpoint_sha256=digest, checkpoint_lock_id=f"sdtt_owt_hf:{manifest_sha256}:all", checkpoint_selection=selector, checkpoint_teacher_family="masked_mdlm", ) def test_conversion_validates_family_before_reading_capture(tmp_path: Path) -> None: item = request("mdlm_sdtt", "owt", samples=1) item = prepare_official_sdtt_conversion(tmp_path, item) item = replace(item, checkpoint_teacher_family="uniform_duo") output = ( tmp_path / "results/samples/owt/mdlm_sdtt" / f"steps_{item.step_count}" ) output.mkdir(parents=True) with pytest.raises(AdapterError, match="teacher family"): list(SDTTAdapter().convert_outputs(item, output))