import json import copy from pathlib import Path import subprocess import sys import pytest import yaml from dlb.data import ( build_data_manifest, build_processing_contract, disk_preflight_evidence, processing_contract_sha256, publication_is_reusable, publish_staged_output, recover_incomplete_publication, validate_download_manifest, validate_manifest_contract, verify_processed_dataset, ) from dlb.io import atomic_json_write class TinySplit: def __init__(self, rows): self._rows = rows def __len__(self): return len(self._rows) def __getitem__(self, index): return self._rows[index] class TinyTokenizer: def decode(self, token_ids, skip_special_tokens=False): assert skip_special_tokens is False return " ".join(str(token_id) for token_id in token_ids) def _test_contract( *, dataset, dataset_id, source_revision, tokenizer_id, tokenizer_revision, sequence_length, split_expression, document_counts, vocab_size, materialization_revision=None, ): return { "schema_version": 1, "dataset": dataset, "dataset_id": dataset_id, "source_revision": source_revision, "source_materialization_revision": materialization_revision, "tokenizer_id": tokenizer_id, "tokenizer_revision": tokenizer_revision, "sequence_length": sequence_length, "split_expression": split_expression, "expected_document_counts": document_counts, "vocab_size": vocab_size, "packing": { "algorithm": "document_eos_continuous_bos_eos", "version": 1, }, } def test_manifest_records_reproducibility_fields_and_file_digests(tmp_path): """Catch manifests that cannot identify and integrity-check their data cache.""" output_dir = tmp_path / "processed" output_dir.mkdir() (output_dir / "train.arrow").write_bytes(b"train rows") (output_dir / "validation.arrow").write_bytes(b"validation rows") contract = _test_contract( dataset="owt", dataset_id="openwebtext", source_revision="a" * 40, tokenizer_id="gpt2", tokenizer_revision="b" * 40, sequence_length=1024, split_expression={ "train": "train[:-100000]", "validation": "train[-100000:]", }, document_counts={"train": 7_913_769, "validation": 100_000}, vocab_size=50_257, ) manifest = build_data_manifest( dataset="owt", dataset_id="openwebtext", source_revision="a" * 40, tokenizer_id="gpt2", tokenizer_revision="b" * 40, sequence_length=1024, split_expression={ "train": "train[:-100000]", "validation": "train[-100000:]", }, document_counts={"train": 7_913_769, "validation": 100_000}, packed_sequence_counts={"train": 2, "validation": 1}, vocab_size=50_257, min_token_id=1, max_token_id=50_256, output_dir=output_dir, root=tmp_path, processing_contract=contract, ) assert manifest["source_revision"] == "a" * 40 assert manifest["tokenizer_revision"] == "b" * 40 assert manifest["split_expression"]["validation"] == "train[-100000:]" assert manifest["document_counts"] == {"train": 7_913_769, "validation": 100_000} assert manifest["packed_sequence_counts"] == {"train": 2, "validation": 1} assert manifest["vocabulary"] == { "size": 50_257, "min_token_id": 1, "max_token_id": 50_256, } assert manifest["created_at"].endswith("Z") assert manifest["processed_path"] == "processed" assert [entry["path"] for entry in manifest["files"]] == [ "train.arrow", "validation.arrow", ] assert all(len(entry["sha256"]) == 64 for entry in manifest["files"]) assert manifest["verified"] is False assert manifest["processing_contract"] == contract assert manifest["processing_contract_sha256"] == processing_contract_sha256(contract) def test_verifier_checks_lengths_bounds_counts_digests_and_decoding(tmp_path): """Catch a verifier that trusts manifest claims without reading processed rows.""" output_dir = tmp_path / "processed" output_dir.mkdir() (output_dir / "state.json").write_text("complete\n", encoding="utf-8") dataset = { "train": TinySplit([{"input_ids": [101, 10, 11, 102]}]), "validation": TinySplit([{"input_ids": [101, 12, 13, 102]}]), } contract = _test_contract( dataset="lm1b", dataset_id="lm1b", source_revision="c" * 40, tokenizer_id="bert-base-uncased", tokenizer_revision="d" * 40, sequence_length=4, split_expression={"train": "train", "validation": "test"}, document_counts={"train": 2, "validation": 1}, vocab_size=30_522, ) manifest = build_data_manifest( dataset="lm1b", dataset_id="lm1b", source_revision="c" * 40, tokenizer_id="bert-base-uncased", tokenizer_revision="d" * 40, sequence_length=4, split_expression={"train": "train", "validation": "test"}, document_counts={"train": 2, "validation": 1}, packed_sequence_counts={"train": 1, "validation": 1}, vocab_size=30_522, min_token_id=10, max_token_id=102, output_dir=output_dir, root=tmp_path, processing_contract=contract, ) result = verify_processed_dataset( dataset=dataset, manifest=manifest, output_dir=output_dir, tokenizer=TinyTokenizer(), expected_contract=contract, ) assert result["verified"] is True assert result["verification"]["checked_sequences"] == 2 assert result["verification"]["decoded_samples"] == 2 def test_verifier_rejects_out_of_range_tokens(tmp_path): """Catch processed IDs that exceed the pinned tokenizer vocabulary.""" output_dir = tmp_path / "processed" output_dir.mkdir() (output_dir / "state.json").write_text("complete\n", encoding="utf-8") contract = _test_contract( dataset="tiny", dataset_id="tiny", source_revision="e" * 40, tokenizer_id="tiny", tokenizer_revision="f" * 40, sequence_length=4, split_expression={"train": "train", "validation": "validation"}, document_counts={"train": 1, "validation": 1}, vocab_size=20, ) manifest = build_data_manifest( dataset="tiny", dataset_id="tiny", source_revision="e" * 40, tokenizer_id="tiny", tokenizer_revision="f" * 40, sequence_length=4, split_expression={"train": "train", "validation": "validation"}, document_counts={"train": 1, "validation": 1}, packed_sequence_counts={"train": 1, "validation": 1}, vocab_size=20, min_token_id=1, max_token_id=19, output_dir=output_dir, root=tmp_path, processing_contract=contract, ) dataset = { "train": TinySplit([{"input_ids": [1, 2, 20, 3]}]), "validation": TinySplit([{"input_ids": [1, 2, 3, 4]}]), } with pytest.raises(ValueError, match="outside vocabulary"): verify_processed_dataset( dataset=dataset, manifest=manifest, output_dir=output_dir, tokenizer=TinyTokenizer(), expected_contract=contract, ) def test_verified_manifest_is_atomically_persistable(tmp_path): """Catch verification results that cannot be stored as valid JSON metadata.""" path = tmp_path / "manifest.json" value = {"dataset": "tiny", "verified": True, "verification": {"ok": True}} atomic_json_write(path, value) assert json.loads(path.read_text(encoding="utf-8")) == value def test_data_configuration_pins_exact_revisions_and_semantics(): """Catch moving Hub references or accidental changes to split/tokenizer semantics.""" root = Path(__file__).parents[1] configuration = yaml.safe_load((root / "artifacts" / "data.yaml").read_text()) assert configuration["datasets"]["lm1b"] == { "repo_id": "billion-word-benchmark/lm1b", "revision": "35161838ea9e05371a25a8db037f94fcae4c2064", "parquet_revision": "8d52bfd3cc2819fc1166dbb8144c328e2690de3e", "config": "plain_text", "splits": {"train": "train", "validation": "test"}, "sequence_length": 128, "tokenizer": "bert-base-uncased", } assert configuration["datasets"]["owt"] == { "repo_id": "Skylion007/openwebtext", "revision": "79d93d786212f7344586290adb811d4ae6a1762c", "config": "plain_text", "documents": 8_013_769, "splits": { "train": "train[:-100000]", "validation": "train[-100000:]", }, "sequence_length": 1024, "tokenizer": "gpt2", } assert configuration["models"] == { "bert-base-uncased": "86b5e0934494bd15c9632b12f734a8a67f723594", "gpt2": "607a30d783dfa663caf39e06633721c8d4cfcd7e", "gpt2-large": "32b71b12589c2f8d625668d2335a01cac3249519", } assert configuration["lm1b_archive"] == { "url": "http://www.statmt.org/lm-benchmark/1-billion-word-language-modeling-benchmark-r13output.tar.gz", "size_bytes": 1_792_209_805, "sha256": "01ba60381110baf7f189dfd2b8374de371e8c9a340835793f190bdae9e90a34e", "train_documents": 30_301_028, "test_documents": 306_688, } def test_fetch_dry_run_lists_every_target_without_writing_data(tmp_path): """Catch dry-run mode that omits a required snapshot or creates cache files.""" root = Path(__file__).parents[1] data_path = tmp_path / "data" result = subprocess.run( [ sys.executable, str(root / "scripts" / "fetch_data.py"), "--root", str(tmp_path), "--config", str(root / "artifacts" / "data.yaml"), "--dry-run", ], check=False, capture_output=True, text=True, ) assert result.returncode == 0, result.stderr assert result.stdout.splitlines() == [ "DATASET lm1b billion-word-benchmark/lm1b@35161838ea9e05371a25a8db037f94fcae4c2064", "DATASET owt Skylion007/openwebtext@79d93d786212f7344586290adb811d4ae6a1762c", "MODEL bert-base-uncased@86b5e0934494bd15c9632b12f734a8a67f723594", "MODEL gpt2@607a30d783dfa663caf39e06633721c8d4cfcd7e", "MODEL gpt2-large@32b71b12589c2f8d625668d2335a01cac3249519", ] assert not data_path.exists() def test_owt_disk_preflight_requires_55_gib_unless_explicitly_overridden(): """Catch starting the large OWT fetch on unsafe disk without explicit consent.""" required = 55 * 1024**3 with pytest.raises(RuntimeError, match="at least 55 GiB"): disk_preflight_evidence(required - 1, allow_low_disk=False) def test_owt_disk_preflight_records_requested_and_actual_override_separately(): """Catch metadata claiming a bypass when the flag was unnecessary.""" required = 55 * 1024**3 sufficient = disk_preflight_evidence(required, allow_low_disk=True) low = disk_preflight_evidence(required - 1, allow_low_disk=True) assert sufficient == { "observed_free_bytes": required, "required_free_bytes": required, "below_threshold": False, "override_requested": True, "override_used": False, } assert low == { "observed_free_bytes": required - 1, "required_free_bytes": required, "below_threshold": True, "override_requested": True, "override_used": True, } def _configuration(): root = Path(__file__).parents[1] return yaml.safe_load((root / "artifacts" / "data.yaml").read_text()) def _valid_downloads(configuration): lm1b = configuration["datasets"]["lm1b"] owt = configuration["datasets"]["owt"] return { "schema_version": 1, "datasets": { "lm1b": { "repo_id": lm1b["repo_id"], "source_revision": lm1b["revision"], "materialization_revision": lm1b["parquet_revision"], "split_rows": { "train": configuration["lm1b_archive"]["train_documents"], "test": configuration["lm1b_archive"]["test_documents"], }, "dataset_snapshot": ( "data/raw/huggingface/hub/datasets--lm1b/snapshots/" + lm1b["parquet_revision"] ), "source_metadata_snapshot": ( "data/raw/huggingface/hub/datasets--lm1b/snapshots/" + lm1b["revision"] ), "cache_files": {"train": ["train.arrow"], "test": ["test.arrow"]}, }, "owt": { "repo_id": owt["repo_id"], "source_revision": owt["revision"], "materialization_revision": None, "split_rows": {"train": owt["documents"]}, "dataset_snapshot": None, "cache_files": {"train": ["owt.arrow"]}, }, }, "models": { name: { "repo_id": name, "revision": revision, "snapshot_path": ( f"data/raw/huggingface/hub/models--{name}/snapshots/{revision}" ), } for name, revision in configuration["models"].items() }, } @pytest.mark.parametrize( ("field", "bad_value"), [ ("repo_id", "wrong/lm1b"), ("source_revision", "0" * 40), ("materialization_revision", "1" * 40), ("split_rows", {"train": 1, "test": 1}), ], ) def test_preprocessing_rejects_stale_lm1b_download_provenance(field, bad_value): """Catch preprocessing data that does not match its claimed immutable source.""" configuration = _configuration() downloads = _valid_downloads(configuration) downloads["datasets"]["lm1b"][field] = bad_value with pytest.raises(ValueError, match=field): validate_download_manifest(configuration, downloads, "lm1b") def test_preprocessing_validates_owt_rows_and_tokenizer_snapshot_revision(): """Catch truncated OWT caches or a tokenizer snapshot from another revision.""" configuration = _configuration() downloads = _valid_downloads(configuration) downloads["datasets"]["owt"]["split_rows"] = {"train": 8_000_000} with pytest.raises(ValueError, match="split_rows"): validate_download_manifest(configuration, downloads, "owt") downloads = _valid_downloads(configuration) downloads["models"]["gpt2"]["revision"] = "2" * 40 with pytest.raises(ValueError, match="tokenizer revision"): validate_download_manifest(configuration, downloads, "owt") def test_preprocessing_rejects_snapshot_paths_for_other_immutable_revisions(): """Catch correct labels pointing preprocessing at stale snapshot content.""" configuration = _configuration() downloads = _valid_downloads(configuration) downloads["datasets"]["lm1b"]["dataset_snapshot"] = ( "data/raw/huggingface/hub/datasets--lm1b/snapshots/" + "4" * 40 ) with pytest.raises(ValueError, match="dataset_snapshot"): validate_download_manifest(configuration, downloads, "lm1b") @pytest.mark.parametrize( ("path", "bad_value"), [ (("sequence_length",), 999), (("split_expression", "validation"), "train[:100000]"), (("tokenizer_id",), "wrong-tokenizer"), (("source_materialization_revision",), "3" * 40), (("packing", "version"), 999), ], ) def test_manifest_contract_rejects_processing_semantic_mutations(path, bad_value): """Catch reuse or verification under stale processing semantics.""" expected = build_processing_contract(_configuration(), "lm1b") manifest_contract = copy.deepcopy(expected) target = manifest_contract for key in path[:-1]: target = target[key] target[path[-1]] = bad_value manifest = { "processing_contract": manifest_contract, "processing_contract_sha256": processing_contract_sha256(manifest_contract), } with pytest.raises(ValueError, match="processing contract"): validate_manifest_contract(manifest, expected) def test_restart_recovers_final_output_missing_manifest_and_publishes_cleanly(tmp_path): """Catch a crash-after-rename state requiring manual deletion before rerun.""" configuration = _configuration() contract = build_processing_contract(configuration, "lm1b") final = tmp_path / "lm1b-bert-128" final.mkdir() (final / "interrupted.arrow").write_bytes(b"incomplete") manifest_path = tmp_path / "lm1b.json" assert recover_incomplete_publication(final, manifest_path, contract) is False assert not final.exists() staged = tmp_path / "lm1b-bert-128.staging" staged.mkdir() (staged / "train.arrow").write_bytes(b"complete") manifest = build_data_manifest( dataset="lm1b", dataset_id=contract["dataset_id"], source_revision=contract["source_revision"], tokenizer_id=contract["tokenizer_id"], tokenizer_revision=contract["tokenizer_revision"], sequence_length=contract["sequence_length"], split_expression=contract["split_expression"], document_counts=contract["expected_document_counts"], packed_sequence_counts={"train": 1, "validation": 1}, vocab_size=contract["vocab_size"], min_token_id=1, max_token_id=102, output_dir=staged, root=tmp_path, processing_contract=contract, published_output_dir=final, ) publish_staged_output(staged, final, manifest_path, manifest) assert publication_is_reusable(final, manifest_path, contract) assert json.loads(manifest_path.read_text())["processed_path"] == "lm1b-bert-128"