File size: 6,694 Bytes
79a20c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | import hashlib
from pathlib import Path
import pytest
import dlb.conditional_matrix as conditional_matrix
import dlb.runner as runner_module
from dlb.conditional_prompts import PromptManifest, load_protocol
from dlb.matrix import build_matrix
from dlb.registry import load_registry
from dlb.runner import RunRequest
from tests.test_runner import FakeAdapter, prepare_canonical_root
ROOT = Path(__file__).parents[1]
@pytest.fixture
def registry():
return load_registry(ROOT / "configs/experiments.yaml")
def _prompt_manifest(dataset: str) -> PromptManifest:
return PromptManifest(
schema_version=1,
protocol="c64_zs_v1",
dataset=dataset,
source_split="validation",
source_processed_path=f"data/processed/{dataset}/validation",
source_manifest_path=f"data/manifests/{dataset}.json",
source_manifest_sha256="a" * 64,
tokenizer_id="test-tokenizer",
tokenizer_revision="revision",
vocabulary_size=10,
selection_algorithm="test-selection",
selection_seed=42,
source_row_count=1024,
prompt_count=1024,
prefix_length=64,
evaluation_continuation_length=64,
model_length=128 if dataset == "lm1b" else 1024,
prompt_file=f"data/conditional/{dataset}-c64/prompts.jsonl",
prompt_file_sha256="b" * 64,
)
def _install_verified_sidecars(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> list[str]:
calls: list[str] = []
for dataset in ("lm1b", "owt"):
sidecar = tmp_path / "data" / "manifests" / f"conditional-{dataset}-c64.json"
sidecar.parent.mkdir(parents=True, exist_ok=True)
sidecar.write_text(f'{{"dataset":"{dataset}"}}\n', encoding="utf-8")
def verify(root: Path, dataset: str, protocol):
assert root == tmp_path
assert protocol.protocol == "c64_zs_v1"
calls.append(dataset)
return _prompt_manifest(dataset)
monkeypatch.setattr(conditional_matrix, "verify_prompts", verify)
return calls
def test_conditional_matrix_matches_supported_unconditional_tasks(registry, tmp_path, monkeypatch):
"""Catch a conditional matrix that diverges from the supported registry cells."""
_install_verified_sidecars(tmp_path, monkeypatch)
ordinary = build_matrix(registry, root=tmp_path)
conditional = conditional_matrix.build_conditional_matrix(
registry, root=tmp_path, protocol=load_protocol(ROOT / "configs/conditional.yaml")
)
assert len(conditional) == len(ordinary) == 132
assert [(task.model, task.dataset, task.steps) for task in conditional] == [
(task.model, task.dataset, task.steps) for task in ordinary
]
assert all("/results/conditional/" in task.sample_dir for task in conditional)
assert all(task.sample_count == 2048 for task in conditional)
def test_conditional_matrix_verifies_both_prompts_and_binds_sidecar_digest(
registry, tmp_path, monkeypatch
):
"""Catch tasks published from unverified prompts or a prompt-file digest."""
calls = _install_verified_sidecars(tmp_path, monkeypatch)
tasks = conditional_matrix.build_conditional_matrix(
registry, root=tmp_path, protocol=load_protocol(ROOT / "configs/conditional.yaml")
)
assert calls == ["lm1b", "owt"]
for dataset in ("lm1b", "owt"):
sidecar = tmp_path / "data" / "manifests" / f"conditional-{dataset}-c64.json"
dataset_tasks = [task for task in tasks if task.dataset == dataset]
assert {task.conditioning_manifest for task in dataset_tasks} == {str(sidecar)}
assert {task.conditioning_manifest_sha256 for task in dataset_tasks} == {
hashlib.sha256(sidecar.read_bytes()).hexdigest()
}
def test_conditional_matrix_task_resolves_through_runner_contract(
registry, tmp_path, monkeypatch
):
"""Catch matrix task fields that cannot be submitted to the shared runner."""
prepare_canonical_root(tmp_path)
config = tmp_path / "configs" / "conditional.yaml"
config.write_text(
(ROOT / "configs" / "conditional.yaml").read_text(encoding="utf-8"),
encoding="utf-8",
)
_install_verified_sidecars(tmp_path, monkeypatch)
manifests = {dataset: _prompt_manifest(dataset) for dataset in ("lm1b", "owt")}
monkeypatch.setattr(
runner_module, "verify_prompts", lambda root, dataset, protocol: manifests[dataset]
)
tasks = conditional_matrix.build_conditional_matrix(
registry, root=tmp_path, protocol=load_protocol(config)
)
task = next(item for item in tasks if (item.model, item.dataset, item.steps) == ("flm", "lm1b", 2))
request = RunRequest(
run_id=task.task_id,
model_id=task.model,
dataset_id=task.dataset,
step_count=task.steps,
seed=task.seed,
sample_count=task.sample_count,
generation_mode="conditional_prefix",
conditioning_manifest=task.conditioning_manifest,
conditioning_manifest_sha256=task.conditioning_manifest_sha256,
)
resolved, _ = runner_module._resolve_request(
request, tmp_path, FakeAdapter([], [])
)
assert resolved.conditioning_manifest == task.conditioning_manifest
assert resolved.conditioning_manifest_sha256 == task.conditioning_manifest_sha256
@pytest.mark.parametrize("failure", ["missing", "tampered"])
def test_conditional_matrix_rejects_missing_or_tampered_prompt_manifest(
registry, tmp_path, monkeypatch, failure
):
"""Catch matrix construction that emits cells after prompt verification fails."""
if failure == "tampered":
monkeypatch.setattr(
conditional_matrix,
"verify_prompts",
lambda root, dataset, protocol: (_ for _ in ()).throw(ValueError("prompt manifest tampered")),
)
else:
monkeypatch.setattr(conditional_matrix, "verify_prompts", lambda root, dataset, protocol: _prompt_manifest(dataset))
with pytest.raises(ValueError, match="manifest|tampered|No such file"):
conditional_matrix.build_conditional_matrix(
registry, root=tmp_path, protocol=load_protocol(ROOT / "configs/conditional.yaml")
)
def test_conditional_unsupported_inventory_is_rdlm_owt(registry):
"""Catch conditional reporting that omits the explicit unsupported RDLM/OWT cell."""
from dlb.conditional_matrix import conditional_unsupported_inventory
assert conditional_unsupported_inventory(registry) == [
{
"status": "unsupported",
"model": "rdlm",
"dataset": "owt",
"category": "few",
"reason": registry.models["rdlm"].datasets["owt"].reason,
}
]
|