| import json |
| from pathlib import Path |
|
|
| import pytest |
|
|
| import dlb.aggregate as aggregate_module |
| from dlb.aggregate import IncompleteMatrixError, aggregate |
| from dlb.io import write_samples_atomic |
| from dlb.matrix import MatrixTask |
|
|
| PROJECT_ROOT = Path(__file__).parents[1] |
|
|
|
|
| def prepare_registry(root: Path) -> None: |
| (root / "configs").mkdir(parents=True, exist_ok=True) |
| (root / "configs/experiments.yaml").write_text( |
| (PROJECT_ROOT / "configs/experiments.yaml").read_text(encoding="utf-8"), |
| encoding="utf-8", |
| ) |
|
|
|
|
| def fake_task(root: Path, sample_count: int = 2) -> MatrixTask: |
| sample_dir = root / "results/samples/lm1b/flm/steps_2" |
| metrics = root / "results/metrics/lm1b/flm/steps_2/metrics.json" |
| timing = root / "results/timing/lm1b/flm/steps_2/timing.json" |
| return MatrixTask( |
| task_id="flm-lm1b-steps-2", |
| category="many", |
| model="flm", |
| dataset="lm1b", |
| steps=2, |
| sample_count=sample_count, |
| seed=42, |
| environment="dlb-flm", |
| adapter="flm", |
| source="flm", |
| provenance="official", |
| sample_dir=str(sample_dir), |
| metrics_path=str(metrics), |
| timing_path=str(timing), |
| ) |
|
|
|
|
| def write_complete_fake_run(root: Path, task: MatrixTask) -> None: |
| sample_dir = Path(task.sample_dir) |
| write_samples_atomic( |
| sample_dir / "samples.jsonl", |
| [ |
| { |
| "sample_id": index, |
| "text": f"sample {index}", |
| "token_ids": [101, index + 1, 102], |
| "seed": task.seed, |
| "generation_seconds": 0.01, |
| } |
| for index in range(task.sample_count) |
| ], |
| expected=task.sample_count, |
| ) |
| provenance = { |
| "source_sha256": "a" * 40, |
| "config_sha256": "b" * 64, |
| "checkpoint_sha256": "c" * 64, |
| "checkpoint_lock_id": "lock:fake", |
| "checkpoint_selection": {"resource": "fake"}, |
| "checkpoint_teacher_family": "continuous_flm", |
| "adapter_identity": "tests.fake:v1", |
| "environment": task.environment, |
| } |
| identity = { |
| "model_id": task.model, |
| "dataset_id": task.dataset, |
| "step_count": task.steps, |
| "seed": task.seed, |
| "sample_count": task.sample_count, |
| } |
| metadata = { |
| "status": "succeeded", |
| "identity": {**identity, **provenance}, |
| **provenance, |
| } |
| sample_dir.mkdir(parents=True, exist_ok=True) |
| (sample_dir / "run_metadata.json").write_text( |
| json.dumps(metadata), encoding="utf-8" |
| ) |
| metrics_path = Path(task.metrics_path) |
| metrics_path.parent.mkdir(parents=True, exist_ok=True) |
| metrics_path.write_text( |
| json.dumps( |
| { |
| "schema_version": 1, |
| "sample_count": task.sample_count, |
| "partial": False, |
| "metrics": { |
| "generative_perplexity": {"perplexity": 12.0}, |
| "unigram_entropy": {"mean_entropy": 2.0}, |
| "self_bleu": {"score": 0.4}, |
| }, |
| } |
| ), |
| encoding="utf-8", |
| ) |
| timing_path = Path(task.timing_path) |
| timing_path.parent.mkdir(parents=True, exist_ok=True) |
| timing_path.write_text( |
| json.dumps( |
| { |
| "schema": "dlb-generation-timing-v1", |
| "timing": { |
| "mode": "primary_latency", |
| "warmups": 5, |
| "repeats": 32, |
| "batch_size": 1, |
| "num_timed_samples": 32, |
| "seconds_per_sample": 0.25, |
| }, |
| "metadata": { |
| "attempt_id": "d" * 32, |
| "identity": {**identity, **provenance}, |
| **provenance, |
| }, |
| } |
| ), |
| encoding="utf-8", |
| ) |
|
|
|
|
| def patch_single_task(monkeypatch, task: MatrixTask): |
| monkeypatch.setattr( |
| aggregate_module, |
| "build_matrix", |
| lambda registry, **kwargs: [task], |
| ) |
|
|
|
|
| def test_aggregate_requires_all_metrics(monkeypatch, tmp_path): |
| prepare_registry(tmp_path) |
| task = fake_task(tmp_path) |
| write_complete_fake_run(tmp_path, task) |
| metrics = json.loads(Path(task.metrics_path).read_text(encoding="utf-8")) |
| del metrics["metrics"]["self_bleu"] |
| Path(task.metrics_path).write_text(json.dumps(metrics), encoding="utf-8") |
| patch_single_task(monkeypatch, task) |
| with pytest.raises(IncompleteMatrixError): |
| aggregate(tmp_path, strict=True) |
|
|
|
|
| def test_partial_aggregate_publishes_failures_and_provenance(monkeypatch, tmp_path): |
| prepare_registry(tmp_path) |
| task = fake_task(tmp_path) |
| write_complete_fake_run(tmp_path, task) |
| patch_single_task(monkeypatch, task) |
| report = aggregate( |
| tmp_path, |
| strict=False, |
| partial=True, |
| output_dir=tmp_path / "results/summary", |
| ) |
| assert report.complete is True |
| assert len(report.rows) == 1 |
| assert set( |
| row["metric"] |
| for row in __import__("csv").DictReader( |
| (tmp_path / "results/summary/results_long.csv").open() |
| ) |
| ) == { |
| "generative_perplexity", |
| "unigram_entropy", |
| "self_bleu", |
| "generation_seconds_per_sample", |
| } |
| assert "checkpoint_sha256" in ( |
| tmp_path / "results/summary/results_wide.csv" |
| ).read_text(encoding="utf-8").splitlines()[0] |
|
|
