"""Contracts for the pinned LangFlow and RDLM continuous teachers.""" from dataclasses import replace import hashlib import json import os 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.langflow import LangFlowAdapter from dlb.adapters.rdlm import RDLMAdapter from dlb.command import main as command_main 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=42, 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 option(command: list[str], key: str) -> str: index = command.index(key) return command[index + 1] 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 prepare_conversion_root(tmp_path: Path) -> Path: (tmp_path / "artifacts").mkdir() (tmp_path / "configs").mkdir() for relative in ( Path("artifacts/data.yaml"), Path("artifacts/checkpoints.yaml"), Path("configs/experiments.yaml"), ): target = tmp_path / relative target.write_text((ROOT / relative).read_text(encoding="utf-8"), encoding="utf-8") return tmp_path def prepare_command_root(tmp_path: Path) -> Path: root = prepare_conversion_root(tmp_path) for relative in ( Path("adapters/sample_langflow.py"), Path("upstreams/rdlm/main.py"), ): target = root / relative target.parent.mkdir(parents=True, exist_ok=True) target.write_text("# fixture\n", encoding="utf-8") return root def canonical_conversion_request(root: Path, item: RunRequest) -> RunRequest: langflow_resource = ( "langflow_lm1b_hf" if item.dataset_id == "lm1b" else "langflow_owt_hf" ) resources = { "langflow": ( langflow_resource, "continuous_langflow", [ "config.json", "config.py", "model.py", "model.safetensors", ], ), "rdlm": ( "rdlm_lm1b_drive", "continuous_rdlm", [ "LM1B/checkpoint.pth", "LM1B/config.yaml", "LM1B/sde.pkl", "Text8/checkpoint.pth", "Text8/config.yaml", "Text8/sde.pkl", ], ), } resource, family, relative_files = resources[item.model_id] manifest = root / "artifacts" / "checkpoints.yaml" manifest_sha256 = hashlib.sha256(manifest.read_bytes()).hexdigest() files = [ { "path": f"checkpoints/fixture/{resource}/{relative}", "size_bytes": index + 1, "sha256": f"{index + 1:064x}", } for index, relative in enumerate(relative_files) ] (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, None) 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_capture(path: Path, count: int, length: int) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text( json.dumps( { "schema": "dlb-upstream-token-capture-v1", "samples": [ { "sample_id": index, "text": f"sample {index}", "token_ids": [index + 1] * length, } for index in range(count) ], } ), encoding="utf-8", ) def prepare_tokenizer_binding( root: Path, dataset: str ) -> tuple[Path, Path, Path]: """Create the minimal immutable tokenizer lock consumed by capture 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 test_registry_exposes_langflow_and_rdlm_supported_continuous_cells() -> None: """Catch supported continuous cells being silently substituted or rejected.""" registry = load_registry(ROOT / "configs" / "experiments.yaml") assert registry.models["langflow"].datasets["owt"].status == "supported" assert registry.models["langflow"].datasets["lm1b"].status == "supported" assert registry.models["rdlm"].datasets["lm1b"].status == "supported" assert registry.models["rdlm"].datasets["owt"].status == "unsupported" @pytest.mark.parametrize( ("dataset", "steps", "tokenizer_name", "tokenizer_revision", "sequence_length"), [ ("lm1b", 32, "bert-base-uncased", "86b5e0934494bd15c9632b12f734a8a67f723594", 128), ("owt", 1024, "gpt2", "607a30d783dfa663caf39e06633721c8d4cfcd7e", 1024), ], ) def test_langflow_renders_the_pinned_inference_argparse_contract( tmp_path: Path, dataset: str, steps: int, tokenizer_name: str, tokenizer_revision: str, sequence_length: int, ) -> None: """Catch flags that are not accepted by the pinned inference argparse parser.""" root = prepare_conversion_root(tmp_path) wrapper = root / "adapters" / "sample_langflow.py" wrapper.parent.mkdir(parents=True) wrapper.write_text("# fixture\n", encoding="utf-8") entrypoint = root / "upstreams" / "langflow" / "inference.py" entrypoint.parent.mkdir(parents=True) entrypoint.write_text("# fixture\n", encoding="utf-8") item = request("langflow", dataset, steps=steps, samples=17) command = LangFlowAdapter().render_command(item, run_dir(root, item), dry_run=True) assert command[1:3] == ["-B", "-u"] assert Path(override(command, "--upstream-entrypoint")) == wrapper assert override(command, "--capture-kind") == "langflow" assert Path(override(command, "--data-config-path")) == root / "artifacts/data.yaml" assert Path(override(command, "--downloads-manifest-path")) == ( root / "data/manifests/downloads.json" ) assert override(command, "--dataset-id") == dataset assert Path(override(command, "--tokenizer-snapshot")) == ( root / "data/raw/huggingface/hub" / f"models--{tokenizer_name}" / "snapshots" / tokenizer_revision ) checkpoint = Path(option(command, "--checkpoint")) assert checkpoint == root / f"checkpoints/official/langflow/{dataset}/model.safetensors" assert option(command, "--num_samples") == "17" assert option(command, "--batch_size") == "1" assert option(command, "--num_steps") == str(steps) assert option(command, "--seq_length") == str(sequence_length) assert option(command, "--seed") == "42" assert Path(option(command, "--output")) == run_dir(root, item) / "upstream_samples.txt" assert option(command, "--tokenizer") == tokenizer_name assert not any( argument in { "--num-samples", "--num-steps", "--seq-length", } for argument in command ) def test_langflow_accepts_many_step_grid_values( tmp_path: Path, ) -> None: """Catch regressing LangFlow back to the fixed 1024-step-only cell.""" root = prepare_conversion_root(tmp_path) wrapper = root / "adapters" / "sample_langflow.py" wrapper.parent.mkdir(parents=True) wrapper.write_text("# fixture\n", encoding="utf-8") item = request("langflow", "owt", steps=32, samples=17) command = LangFlowAdapter().render_command(item, run_dir(root, item), dry_run=True) assert option(command, "--num_steps") == "32" def test_rdlm_renders_the_official_sde_sampler_and_saved_asset_trio( tmp_path: Path, ) -> None: """Catch RDLM bypassing its official Hydra sampler or omitting saved release assets.""" root = prepare_command_root(tmp_path) item = request("rdlm", "lm1b", steps=1000, samples=17) command = RDLMAdapter().render_command(item, run_dir(root, item), dry_run=True) asset_root = root / "checkpoints/official/rdlm/lm1b/LM1B" assert command[1:3] == ["-B", "-u"] assert Path(override(command, "--upstream-entrypoint")) == root / "upstreams/rdlm/main.py" assert override(command, "--capture-kind") == "rdlm" assert Path(override(command, "--data-config-path")) == root / "artifacts/data.yaml" assert Path(override(command, "--downloads-manifest-path")) == ( root / "data/manifests/downloads.json" ) assert override(command, "--dataset-id") == "lm1b" assert Path(override(command, "--tokenizer-snapshot")) == ( root / "data/raw/huggingface/hub/models--bert-base-uncased/snapshots" / "86b5e0934494bd15c9632b12f734a8a67f723594" ) assert Path(override(command, "--saved-config-path")) == asset_root / "config.yaml" assert Path(override(command, "--saved-sde-path")) == asset_root / "sde.pkl" assert override(command, "--expected-samples") == "17" assert override(command, "run_mode") == "sample" assert override(command, "server") == "sample" assert override(command, "exp") == "sample_lm1b" assert Path(override(command, "model_path")) == asset_root / "checkpoint.pth" assert override(command, "sampling.predictor") == "grw" assert override(command, "sampling.steps") == "1000" assert override(command, "sampling.batch_per_gpu") == "8" assert override(command, "seed") == "42" assert override(command, "ngpus") == "1" assert override(command, "eval.entropy") == "False" assert override(command, "eval.nll") == "False" assert override(command, "eval.gen_ppl") == "False" @pytest.mark.parametrize( ("adapter", "model", "dataset"), [ (LangFlowAdapter(), "langflow", "owt"), (RDLMAdapter(), "rdlm", "lm1b"), ], ) def test_real_render_requires_canonical_checkpoint_bytes(adapter, model: str, dataset: str) -> None: """Catch real mode accepting dry-run-only expected paths.""" item = request(model, dataset, steps=1000 if model == "rdlm" else 32) with pytest.raises(AdapterError, match="checkpoint"): adapter.build_command(item, run_dir(ROOT, item)) @pytest.mark.parametrize( ("adapter", "model", "dataset", "length", "family"), [ (LangFlowAdapter(), "langflow", "owt", 1024, "continuous_langflow"), (RDLMAdapter(), "rdlm", "lm1b", 128, "continuous_rdlm"), ], ) def test_capture_conversion_preserves_exact_upstream_ids_and_latency_sentinel( tmp_path: Path, adapter, model: str, dataset: str, length: int, family: str ) -> None: """Catch lossy parsing, wrong lengths/counts, or 0.0 being presented as measured latency.""" root = prepare_conversion_root(tmp_path) item = canonical_conversion_request(root, request(model, dataset, samples=17)) output_dir = run_dir(root, item) write_capture(output_dir / "upstream_token_ids.json", 17, length) records = list(adapter.convert_outputs(item, output_dir)) metadata = json.loads((output_dir / "conversion_metadata.json").read_text()) assert len(records) == 17 assert records[0].token_ids == [1] * length assert metadata["format"] == "dlb-upstream-token-capture-v1" assert metadata["token_ids_source"] == "upstream" assert metadata["generated_samples"] == 17 assert metadata["trimmed_samples"] == 0 assert metadata["teacher_family"] == family assert metadata["generation_seconds_source"] == "unavailable_excluded_sentinel" assert metadata["exclude_from_latency"] is True assert all(record.generation_seconds == 0.0 for record in records) @pytest.mark.parametrize("changed_name", ["LM1B/config.yaml", "LM1B/sde.pkl"]) def test_rdlm_conversion_binds_saved_config_and_sde_checksums_before_output_reads( tmp_path: Path, changed_name: str ) -> None: """Catch provenance that hashes only checkpoint.pth and ignores the saved trio.""" root = prepare_conversion_root(tmp_path) item = canonical_conversion_request(root, request("rdlm", "lm1b", samples=1)) lock_path = root / "artifacts/checkpoint_lock.json" lock = json.loads(lock_path.read_text()) files = lock["resources"]["rdlm_lm1b_drive"]["files"] target = next(file for file in files if file["path"].endswith("/" + changed_name)) target["sha256"] = "f" * 64 lock_path.write_text(json.dumps(lock), encoding="utf-8") output_dir = run_dir(root, item) output_dir.mkdir(parents=True) with pytest.raises(AdapterError, match="checkpoint SHA"): list(RDLMAdapter().convert_outputs(item, output_dir)) def test_rdlm_retokenizes_variable_upstream_rows_offline_and_records_why( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """Catch variable post-BOS RDLM rows being padded without explicit tokenizer provenance.""" class FakeTokenizer: eos_token = "[SEP]" pad_token = "[PAD]" pad_token_id = 0 def __call__(self, texts, **settings): assert settings["max_length"] == 128 return {"input_ids": [[11] + [0] * 127 for _ in texts]} class FakeAutoTokenizer: @staticmethod def from_pretrained(name, **settings): assert name == "bert-base-uncased" assert settings == { "revision": "86b5e0934494bd15c9632b12f734a8a67f723594", "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("rdlm", "lm1b", samples=1)) output_dir = run_dir(root, item) write_capture(output_dir / "upstream_token_ids.json", 1, 127) records = list(RDLMAdapter().convert_outputs(item, output_dir)) metadata = json.loads((output_dir / "conversion_metadata.json").read_text()) assert records[0].token_ids == [11] + [0] * 127 assert metadata["token_ids_source"] == "retokenized" assert metadata["token_ids_transformation"]["reason"] == "upstream_ids_not_canonical_length" def test_rdlm_retokenizes_out_of_vocab_upstream_rows_offline_and_records_why( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """Catch RDLM base-vocabulary IDs being validated as canonical BERT IDs.""" class FakeTokenizer: eos_token = "[SEP]" pad_token = "[PAD]" pad_token_id = 0 def __call__(self, texts, **settings): assert settings["max_length"] == 128 return {"input_ids": [[101, 102] + [0] * 126 for _ in texts]} class FakeAutoTokenizer: @staticmethod def from_pretrained(name, **settings): assert name == "bert-base-uncased" assert settings == { "revision": "86b5e0934494bd15c9632b12f734a8a67f723594", "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("rdlm", "lm1b", samples=1)) output_dir = run_dir(root, item) write_capture(output_dir / "upstream_token_ids.json", 1, 128) capture = json.loads((output_dir / "upstream_token_ids.json").read_text()) capture["samples"][0]["token_ids"] = [35008] + [1] * 127 (output_dir / "upstream_token_ids.json").write_text(json.dumps(capture)) records = list(RDLMAdapter().convert_outputs(item, output_dir)) metadata = json.loads((output_dir / "conversion_metadata.json").read_text()) assert records[0].token_ids == [101, 102] + [0] * 126 assert metadata["token_ids_source"] == "retokenized" assert metadata["token_ids_transformation"]["reason"] == "upstream_ids_outside_canonical_vocab" def test_rdlm_conditional_retokenizes_suffix_out_of_vocab_and_preserves_prefix( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """Catch RDLM conditional capture rejecting base-vocabulary suffix IDs.""" class FakeTokenizer: eos_token = "[SEP]" pad_token = "[PAD]" pad_token_id = 0 def __call__(self, texts, **settings): assert texts == ["generated continuation"] assert settings["max_length"] == 64 return {"input_ids": [[22] * 64]} class FakeAutoTokenizer: @staticmethod def from_pretrained(name, **settings): assert name == "bert-base-uncased" assert settings == { "revision": "86b5e0934494bd15c9632b12f734a8a67f723594", "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("rdlm", "lm1b", steps=1000, samples=1)) item = replace( item, generation_mode="conditional_prefix", prefix_length=64, evaluation_continuation_length=64, prompt_count=1024, diversity_prompt_count=256, completions_per_diversity_prompt=5, completion_schedule="quality_then_diversity_v1", ) output_dir = run_dir(root, item) prefix = [1] * 64 suffix = [35008] + [2] * 63 record = { "sample_id": 0, "text": "fixed prefix generated continuation", "token_ids": prefix + suffix, "prompt_id": 0, "completion_id": 0, "source_index": 100, "prefix_token_ids": prefix, "reference_token_ids": [3] * 64, "full_token_ids": prefix + suffix, "prefix_text": "fixed prefix", "continuation_text": "generated continuation", "reference_text": "source continuation", "full_text": "fixed prefix generated continuation", "prefix_exact_match": True, } write_capture(output_dir / "upstream_token_ids.json", 1, 128) (output_dir / "upstream_token_ids.json").write_text( json.dumps({"schema": "dlb-upstream-token-capture-v1", "samples": [record]}), encoding="utf-8", ) records = list(RDLMAdapter().convert_outputs(item, output_dir)) metadata = json.loads((output_dir / "conversion_metadata.json").read_text()) assert records[0].prefix_token_ids == prefix assert records[0].continuation_token_ids == [22] * 64 assert records[0].full_token_ids == prefix + [22] * 64 assert records[0].prefix_exact_match is True assert metadata["token_ids_source"] == "retokenized_conditional_suffix" assert metadata["token_ids_transformation"]["reason"] == ( "upstream_ids_outside_canonical_vocab" ) def test_shared_capture_runs_langflow_with_locked_offline_tokenizer_and_records_ids( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """Catch LangFlow resolving its movable GPT-2 alias or parsing its text file.""" data_config, downloads, tokenizer_snapshot = prepare_tokenizer_binding(tmp_path, "owt") monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) entrypoint = tmp_path / "inference.py" (entrypoint.parent / "configs").mkdir() source = """ import argparse import os class Rows: def __init__(self, rows): self.rows = rows def detach(self): return self def cpu(self): return self def tolist(self): return self.rows class LangFlow: def generate_samples(self, num_samples, **kwargs): return Rows([[index + 1, index + 2] for index in range(num_samples)]) class Tokenizer: def batch_decode(self, rows, **kwargs): return [f\"text {index}\" for index, _ in enumerate(rows.rows)] class AutoTokenizer: @staticmethod def from_pretrained(name, **settings): assert name == __SNAPSHOT__ assert settings == {'local_files_only': True} assert os.environ['HF_HUB_OFFLINE'] == '1' assert os.environ['TRANSFORMERS_OFFLINE'] == '1' return Tokenizer() def main(): parser = argparse.ArgumentParser() parser.add_argument('--num_samples', type=int, required=True) parser.add_argument('--output', required=True) args = parser.parse_args() rows = LangFlow().generate_samples(args.num_samples) texts = AutoTokenizer.from_pretrained('fixture').batch_decode(rows) with open(args.output, 'w', encoding='utf-8') as handle: handle.write('\\n'.join(texts)) """ entrypoint.write_text( source.replace("__SNAPSHOT__", repr(str(tokenizer_snapshot))), encoding="utf-8" ) capture_path = tmp_path / "capture.json" output_path = tmp_path / "samples.txt" result = capture_module.main( [ f"--upstream-entrypoint={entrypoint}", f"--capture-path={capture_path}", "--capture-kind=langflow", "--expected-samples=2", f"--data-config-path={data_config}", f"--downloads-manifest-path={downloads}", "--dataset-id=owt", f"--tokenizer-snapshot={tokenizer_snapshot}", "--", "--num_samples", "2", "--output", str(output_path), ] ) capture = json.loads(capture_path.read_text()) assert result == 0 assert "HF_HUB_OFFLINE" not in os.environ assert "TRANSFORMERS_OFFLINE" not in os.environ assert output_path.read_text() == "text 0\ntext 1" assert capture == { "schema": "dlb-upstream-token-capture-v1", "samples": [ {"sample_id": 0, "text": "text 0", "token_ids": [1, 2]}, {"sample_id": 1, "text": "text 1", "token_ids": [2, 3]}, ], } def test_rdlm_capture_routes_saved_assets_and_uses_bounded_exact_batches( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """Catch movable tokenizers, fixed batches, or fallback assets replacing locked inputs.""" data_config, downloads, tokenizer_snapshot = prepare_tokenizer_binding(tmp_path, "lm1b") monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) saved_config = tmp_path / "config.yaml" saved_config.write_text("model:\n length: 128\ndata:\n train: lm1b\n", encoding="utf-8") saved_sde = tmp_path / "sde.pkl" saved_sde.write_bytes(b"saved-sde") checkpoint = tmp_path / "checkpoint.pth" checkpoint.write_bytes(b"checkpoint") capture_path = tmp_path / "capture.json" class Rows: def __init__(self, rows): self.rows = rows def detach(self): return self def cpu(self): return self def tolist(self): return self.rows class FakeTorch: @staticmethod def load(path, **kwargs): assert Path(path) == checkpoint return {"config": "embedded-config"} class FakeSeqUtils: @staticmethod def find_bos_and_shift_fn(*args, **kwargs): def shift(rows): return ["one", "two", "excess"], [ [1] * 128, [2] * 128, [3] * 128, ] return shift class FakeBertTokenizer: @staticmethod def from_pretrained(name, **settings): assert name == str(tokenizer_snapshot) assert settings == {"local_files_only": True} assert os.environ["HF_HUB_OFFLINE"] == "1" assert os.environ["TRANSFORMERS_OFFLINE"] == "1" return object() fake_data = SimpleNamespace( transformers=SimpleNamespace(BertTokenizer=FakeBertTokenizer), AutoTokenizer=FakeBertTokenizer, ) def get_tokenizer(dataset): assert dataset == "lm1b" return fake_data.transformers.BertTokenizer.from_pretrained( "bert-base-uncased" ) fake_data.get_tokenizer = get_tokenizer fake_run_sample = SimpleNamespace( torch=FakeTorch(), sutils=FakeSeqUtils(), data=fake_data, tqdm=lambda iterable, **kwargs: iterable, instantiate=lambda config, **kwargs: (config, kwargs), ) monkeypatch.setitem(sys.modules, "run_sample", fake_run_sample) class FakeOmegaConf: @staticmethod def load(path): assert Path(path) == saved_config return SimpleNamespace( model=SimpleNamespace(length=128), data=SimpleNamespace(train="lm1b"), ) @staticmethod def select(config, key): section, field = key.split(".") return getattr(getattr(config, section), field) monkeypatch.setitem( sys.modules, "omegaconf", SimpleNamespace(OmegaConf=FakeOmegaConf) ) class FakeMP: @staticmethod def spawn(*args, **kwargs): raise AssertionError("capture must replace process spawning") fake_entrypoint = SimpleNamespace(mp=FakeMP()) def fake_main(): import run_sample def worker(rank, marker): assert rank == 0 and marker == "worker" run_sample.data.get_tokenizer("lm1b") state = run_sample.torch.load(checkpoint) assert state["config"].model.length == 128 with run_sample.open(tmp_path / "wrong" / "sde.pkl", "rb") as handle: assert handle.read() == b"saved-sde" with pytest.raises(ValueError, match="saved SDE"): run_sample.instantiate( "sde", manifold=object(), scheduler=object(), prior_dist=object(), device="cpu", ) run_sample.instantiate( "sde", manifold=object(), scheduler=object(), prior_dist=object(), preprocessed=(object(), object()), ) assert list(run_sample.tqdm(range(8), leave=False)) == [0, 1] shift = run_sample.sutils.find_bos_and_shift_fn(3, 10, object()) shift(Rows([[0], [0]])) fake_entrypoint.mp.spawn(worker, args=("worker",), nprocs=1, join=True) fake_entrypoint.main = fake_main invocation = capture_module.CaptureInvocation( entrypoint=tmp_path / "main.py", capture_path=capture_path, kind="rdlm", expected_samples=2, saved_config_path=saved_config, saved_sde_path=saved_sde, data_config_path=data_config, downloads_manifest_path=downloads, dataset_id="lm1b", tokenizer_snapshot=tokenizer_snapshot, ) capture_module._capture_rdlm( fake_entrypoint, invocation, [f"model_path={checkpoint}", "sampling.batch_per_gpu=1"], ) capture = json.loads(capture_path.read_text()) assert [sample["text"] for sample in capture["samples"]] == ["one", "two"] assert capture["samples"][1]["token_ids"] == [2] * 128 assert "HF_HUB_OFFLINE" not in os.environ assert "TRANSFORMERS_OFFLINE" not in os.environ def test_continuous_cli_dry_run_emits_valid_commands_and_step_override_rejection( tmp_path: Path, capsys: pytest.CaptureFixture[str], ) -> None: """Catch unsupported cells being substituted or dry-run trying to write artifacts.""" root = prepare_command_root(tmp_path) result = command_main( [ "--root", str(root), "--models", "langflow,rdlm", "--datasets", "lm1b,owt", "--dry-run", ] ) records = [json.loads(line) for line in capsys.readouterr().out.splitlines()] assert result == 1 assert len(records) == 4 supported = [record for record in records if record["status"] == "supported"] unsupported = [record for record in records if record["status"] == "unsupported"] errors = [record for record in records if record["status"] == "error"] assert {(record["model"], record["dataset"]) for record in supported} == { ("langflow", "lm1b"), ("langflow", "owt"), } assert {(record["model"], record["dataset"]) for record in unsupported} == { ("rdlm", "owt"), } assert [(record["model"], record["dataset"]) for record in errors] == [ ("rdlm", "lm1b") ] assert "allowed: 1000,1024" in errors[0]["reason"] assert all(record["command"] for record in supported) assert all(record["reason"] for record in unsupported) assert all( override(record["command"], "seed") == "42" if record["model"] == "rdlm" else option(record["command"], "--seed") == "42" for record in supported ) rdlm_result = command_main( [ "--root", str(root), "--models", "rdlm", "--datasets", "lm1b", "--steps", "1000", "--dry-run", ] ) assert rdlm_result == 0 rdlm_record = json.loads(capsys.readouterr().out.strip()) assert rdlm_record["status"] == "supported" assert override(rdlm_record["command"], "sampling.batch_per_gpu") == "8"