| """Audit generation, evaluation and timing artifacts into reproducible tables.""" |
|
|
| from __future__ import annotations |
|
|
| import csv |
| from dataclasses import dataclass |
| import json |
| import math |
| import os |
| from pathlib import Path |
| from typing import Any, Iterable, Mapping |
|
|
| from dlb.io import validate_samples |
| from dlb.matrix import MatrixTask, build_matrix, unsupported_inventory |
| from dlb.registry import ExperimentRegistry, load_registry |
|
|
|
|
| QUALITY_METRICS = ( |
| "generative_perplexity", |
| "unigram_entropy", |
| "self_bleu", |
| ) |
| ALL_METRICS = QUALITY_METRICS + ("generation_seconds_per_sample",) |
| AGGREGATE_SCHEMA = "dlb-results-summary-v1" |
|
|
|
|
| class IncompleteMatrixError(RuntimeError): |
| """Raised when strict aggregation cannot prove a complete matrix.""" |
|
|
| def __init__(self, failures: Iterable[Mapping[str, Any]]) -> None: |
| self.failures = [dict(item) for item in failures] |
| super().__init__( |
| f"matrix is incomplete: {len(self.failures)} task(s) failed validation" |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class AggregateReport: |
| """In-memory result of an aggregation pass.""" |
|
|
| rows: tuple[dict[str, Any], ...] |
| failures: tuple[dict[str, Any], ...] |
| unsupported: tuple[dict[str, Any], ...] |
| expected_tasks: int |
| complete: bool |
| output_dir: Path | None = None |
|
|
|
|
| def _json(path: Path) -> dict[str, Any]: |
| if path.is_symlink() or not path.is_file(): |
| raise ValueError(f"artifact is missing or unsafe: {path}") |
| try: |
| value = json.loads(path.read_text(encoding="utf-8")) |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: |
| raise ValueError(f"artifact is not valid JSON: {path}") from error |
| if not isinstance(value, dict): |
| raise ValueError(f"artifact must be a JSON object: {path}") |
| return value |
|
|
|
|
| def _finite(value: object, *, positive: bool = False) -> float: |
| if type(value) not in {int, float} or isinstance(value, bool): |
| raise ValueError("metric value must be a finite number") |
| result = float(value) |
| if not math.isfinite(result) or (positive and result <= 0) or (not positive and result < 0): |
| raise ValueError("metric value is outside the allowed finite range") |
| return result |
|
|
|
|
| def _lookup_metric(metrics: Mapping[str, Any], name: str) -> float: |
| aliases = { |
| "generative_perplexity": ("generative_perplexity", "gen_ppl", "ppl"), |
| "unigram_entropy": ("unigram_entropy", "entropy"), |
| "self_bleu": ("self_bleu", "self_bleu_score"), |
| } |
| value: object = None |
| found = False |
| for key in aliases[name]: |
| if key in metrics: |
| value = metrics[key] |
| found = True |
| break |
| if not found: |
| raise ValueError(f"missing metric {name}") |
| if isinstance(value, Mapping): |
| field = { |
| "generative_perplexity": "perplexity", |
| "unigram_entropy": "mean_entropy", |
| "self_bleu": "score", |
| }[name] |
| if field not in value: |
| raise ValueError(f"metric {name} is missing field {field}") |
| value = value[field] |
| result = _finite(value, positive=name == "generative_perplexity") |
| if name == "self_bleu" and result > 1: |
| raise ValueError("self_bleu must be in [0, 1]") |
| return result |
|
|
|
|
| def _provenance( |
| metadata: Mapping[str, Any], |
| task: MatrixTask, |
| *, |
| require_sample_count: bool = True, |
| ) -> dict[str, Any]: |
| identity = metadata.get("identity") |
| identity = identity if isinstance(identity, Mapping) else {} |
| nested = metadata.get("provenance") |
| nested = nested if isinstance(nested, Mapping) else {} |
|
|
| def first(*keys: str) -> object: |
| for source in (metadata, identity, nested): |
| for key in keys: |
| if key in source and source[key] not in (None, ""): |
| return source[key] |
| return None |
|
|
| expected = { |
| "model_id": (task.model, ("model_id", "model")), |
| "dataset_id": (task.dataset, ("dataset_id", "dataset")), |
| "step_count": (task.steps, ("step_count", "steps")), |
| "seed": (task.seed, ("seed",)), |
| } |
| if require_sample_count: |
| expected["sample_count"] = (task.sample_count, ("sample_count",)) |
|
|
| for key, (value, aliases) in expected.items(): |
| observed = first(*aliases) |
| if observed != value: |
| raise ValueError(f"run metadata {key} does not match matrix task") |
|
|
| required = { |
| "source_sha256": ("source_sha256", "source_commit"), |
| "config_sha256": ("config_sha256",), |
| "checkpoint_sha256": ("checkpoint_sha256",), |
| "checkpoint_lock_id": ("checkpoint_lock_id",), |
| "checkpoint_selection": ("checkpoint_selection",), |
| "checkpoint_teacher_family": ("checkpoint_teacher_family",), |
| "adapter_identity": ("adapter_identity",), |
| "environment": ("environment",), |
| } |
|
|
| result: dict[str, Any] = {} |
| for key, aliases in required.items(): |
| value = first(*aliases) |
| if value in (None, "", {}): |
| raise ValueError(f"run metadata is missing provenance field {key}") |
| result[key] = value |
| return result |
|
|
| def _validate_task(task: MatrixTask) -> dict[str, Any]: |
| sample_dir = Path(task.sample_dir) |
| samples = sample_dir / "samples.jsonl" |
| metadata_path = sample_dir / "run_metadata.json" |
| metrics_path = Path(task.metrics_path) |
| timing_path = Path(task.timing_path) |
| validate_samples(samples, expected=task.sample_count) |
| metadata = _json(metadata_path) |
| if metadata.get("status") != "succeeded": |
| raise ValueError("run metadata is not marked succeeded") |
| provenance = _provenance(metadata, task) |
|
|
| metrics_document = _json(metrics_path) |
| if metrics_document.get("sample_count") != task.sample_count: |
| raise ValueError("metrics sample_count does not match matrix task") |
| if metrics_document.get("partial") is True: |
| raise ValueError("partial metrics cannot be used for a complete aggregation") |
| metric_document = metrics_document.get("metrics") |
| if not isinstance(metric_document, Mapping): |
| raise ValueError("metrics document has no metrics object") |
| quality = { |
| name: _lookup_metric(metric_document, name) for name in QUALITY_METRICS |
| } |
|
|
| timing_document = _json(timing_path) |
| if timing_document.get("schema") != "dlb-generation-timing-v1": |
| raise ValueError("timing document has an unsupported schema") |
| timing = timing_document.get("timing") |
| if not isinstance(timing, Mapping): |
| raise ValueError("timing document has no timing object") |
| protocol = { |
| "mode": "primary_latency", |
| "warmups": 5, |
| "repeats": 32, |
| "batch_size": 1, |
| "num_timed_samples": 32, |
| } |
| if any(timing.get(key) != value for key, value in protocol.items()): |
| raise ValueError("timing document does not use the pinned latency protocol") |
| seconds = timing.get("seconds_per_sample") |
| seconds_value = _finite(seconds) |
| timing_metadata = timing_document.get("metadata") |
| if not isinstance(timing_metadata, Mapping): |
| raise ValueError("timing document has no metadata object") |
| _provenance(timing_metadata, task, require_sample_count=False) |
| if timing_metadata.get("attempt_id") in (None, ""): |
| raise ValueError("timing metadata is missing attempt_id") |
|
|
| return { |
| "task_id": task.task_id, |
| "model": task.model, |
| "dataset": task.dataset, |
| "category": task.category, |
| "steps": task.steps, |
| "sample_count": task.sample_count, |
| "seed": task.seed, |
| "provenance": task.provenance, |
| **quality, |
| "generation_seconds_per_sample": seconds_value, |
| "source_sha256": provenance["source_sha256"], |
| "config_sha256": provenance["config_sha256"], |
| "checkpoint_sha256": provenance["checkpoint_sha256"], |
| "checkpoint_lock_id": provenance["checkpoint_lock_id"], |
| "checkpoint_teacher_family": provenance["checkpoint_teacher_family"], |
| "adapter_identity": provenance["adapter_identity"], |
| "environment": ( |
| provenance["environment"] |
| if isinstance(provenance["environment"], str) |
| else json.dumps(provenance["environment"], sort_keys=True) |
| ), |
| } |
|
|
|
|
| def _failure(task: MatrixTask, error: Exception) -> dict[str, Any]: |
| return { |
| "status": "failed", |
| "task_id": task.task_id, |
| "model": task.model, |
| "dataset": task.dataset, |
| "category": task.category, |
| "steps": task.steps, |
| "reason": str(error), |
| } |
|
|
|
|
| def _write_csv(path: Path, rows: list[Mapping[str, Any]], columns: tuple[str, ...]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.partial") |
| try: |
| with temporary.open("w", encoding="utf-8", newline="") as handle: |
| writer = csv.DictWriter( |
| handle, fieldnames=columns, extrasaction="ignore", lineterminator="\n" |
| ) |
| writer.writeheader() |
| for row in rows: |
| writer.writerow({column: row.get(column, "") for column in columns}) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| os.replace(temporary, path) |
| except BaseException: |
| temporary.unlink(missing_ok=True) |
| raise |
|
|
|
|
| def _write_outputs( |
| output_dir: Path, |
| report: AggregateReport, |
| ) -> None: |
| output_dir = output_dir.absolute() |
| output_dir.mkdir(parents=True, exist_ok=True) |
| rows = list(report.rows) |
| long_rows = [ |
| { |
| **{ |
| key: row[key] |
| for key in ( |
| "task_id", |
| "model", |
| "dataset", |
| "category", |
| "steps", |
| "sample_count", |
| "seed", |
| "provenance", |
| ) |
| }, |
| "metric": metric, |
| "value": row[metric], |
| "source_sha256": row["source_sha256"], |
| "config_sha256": row["config_sha256"], |
| "checkpoint_sha256": row["checkpoint_sha256"], |
| "checkpoint_lock_id": row["checkpoint_lock_id"], |
| "adapter_identity": row["adapter_identity"], |
| "environment": row["environment"], |
| } |
| for row in rows |
| for metric in ALL_METRICS |
| ] |
| long_columns = ( |
| "task_id", |
| "model", |
| "dataset", |
| "category", |
| "steps", |
| "sample_count", |
| "seed", |
| "provenance", |
| "metric", |
| "value", |
| "source_sha256", |
| "config_sha256", |
| "checkpoint_sha256", |
| "checkpoint_lock_id", |
| "adapter_identity", |
| "environment", |
| ) |
| wide_columns = ( |
| "task_id", |
| "model", |
| "dataset", |
| "category", |
| "steps", |
| "sample_count", |
| "seed", |
| *ALL_METRICS, |
| "provenance", |
| "source_sha256", |
| "config_sha256", |
| "checkpoint_sha256", |
| "checkpoint_lock_id", |
| "checkpoint_teacher_family", |
| "adapter_identity", |
| "environment", |
| ) |
| _write_csv(output_dir / "results_long.csv", long_rows, long_columns) |
| _write_csv(output_dir / "results_wide.csv", rows, wide_columns) |
| failures = list(report.failures) + [ |
| {**item, "status": "unsupported"} for item in report.unsupported |
| ] |
| failure_columns = ( |
| "status", |
| "task_id", |
| "model", |
| "dataset", |
| "category", |
| "steps", |
| "reason", |
| ) |
| _write_csv(output_dir / "failures.csv", failures, failure_columns) |
| _write_csv( |
| output_dir / "unsupported.csv", |
| list(report.unsupported), |
| ("status", "model", "dataset", "category", "reason"), |
| ) |
| (output_dir / "README.md").write_text( |
| "\n".join( |
| [ |
| f"# Baseline Results ({AGGREGATE_SCHEMA})", |
| "", |
| f"- complete: `{str(report.complete).lower()}`", |
| f"- expected supported tasks: `{report.expected_tasks}`", |
| f"- valid result rows: `{len(report.rows)}`", |
| f"- failures: `{len(report.failures)}`", |
| f"- unsupported cells: `{len(report.unsupported)}`", |
| "", |
| "A strict result set requires every supported task to have 1,024 " |
| "samples, three quality metrics, one independent timing artifact, " |
| "and matching provenance.", |
| "", |
| ] |
| ), |
| encoding="utf-8", |
| ) |
|
|
|
|
| def aggregate( |
| root: Path, |
| *, |
| strict: bool = True, |
| partial: bool = False, |
| output_dir: Path | None = None, |
| registry: ExperimentRegistry | None = None, |
| ) -> AggregateReport: |
| """Validate all supported tasks and optionally publish summary tables.""" |
|
|
| root = Path(root).resolve() |
| registry = registry or load_registry(root / "configs" / "experiments.yaml") |
| tasks = build_matrix(registry, root=root) |
| rows: list[dict[str, Any]] = [] |
| failures: list[dict[str, Any]] = [] |
| for task in tasks: |
| try: |
| rows.append(_validate_task(task)) |
| except (OSError, ValueError, TypeError) as error: |
| failures.append(_failure(task, error)) |
| unsupported = unsupported_inventory(registry) |
| complete = not failures and len(rows) == len(tasks) |
| report = AggregateReport( |
| rows=tuple(rows), |
| failures=tuple(failures), |
| unsupported=tuple(unsupported), |
| expected_tasks=len(tasks), |
| complete=complete, |
| output_dir=Path(output_dir).resolve() if output_dir is not None else None, |
| ) |
| if strict and not partial and not complete: |
| raise IncompleteMatrixError(failures) |
| if output_dir is not None: |
| _write_outputs(Path(output_dir), report) |
| return report |
|
|
|
|
| __all__ = [ |
| "AGGREGATE_SCHEMA", |
| "ALL_METRICS", |
| "AggregateReport", |
| "IncompleteMatrixError", |
| "QUALITY_METRICS", |
| "aggregate", |
| ] |
|
|
