File size: 11,892 Bytes
784f572 | 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 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 | """Canonical generation matrix construction and TSV serialization.
The matrix is deliberately a small, dependency-light artifact. It contains
only supported registry cells; unsupported cells are written to a separate
inventory so a skipped experiment is distinguishable from a missing row.
"""
from __future__ import annotations
import argparse
import csv
from dataclasses import asdict, dataclass
import os
from pathlib import Path
from typing import Iterable
from dlb.registry import ExperimentRegistry, load_registry, step_grid_for_model
MATRIX_SCHEMA = "dlb-generation-matrix-v1"
MATRIX_COLUMNS = (
"task_id",
"category",
"model",
"dataset",
"steps",
"sample_count",
"seed",
"environment",
"adapter",
"source",
"provenance",
"sample_dir",
"metrics_path",
"timing_path",
)
UNSUPPORTED_COLUMNS = ("status", "model", "dataset", "category", "reason")
SUPPORTED_STATUS = "supported"
@dataclass(frozen=True)
class MatrixTask:
"""One supported model/dataset/step generation request."""
task_id: str
category: str
model: str
dataset: str
steps: int
sample_count: int
seed: int
environment: str
adapter: str
source: str
provenance: str
sample_dir: str
metrics_path: str
timing_path: str
@property
def model_id(self) -> str:
return self.model
@property
def dataset_id(self) -> str:
return self.dataset
@property
def step_count(self) -> int:
return self.steps
def row(self) -> dict[str, str]:
value = asdict(self)
return {key: str(value[key]) for key in MATRIX_COLUMNS}
def _root_path(root: Path | None) -> Path:
return (root or Path(".")).resolve()
def _task_paths(root: Path, model: str, dataset: str, steps: int) -> tuple[str, str, str]:
sample_dir = root / "results" / "samples" / dataset / model / f"steps_{steps}"
metrics = root / "results" / "metrics" / dataset / model / f"steps_{steps}" / "metrics.json"
timing = root / "results" / "timing" / dataset / model / f"steps_{steps}" / "timing.json"
return tuple(path.as_posix() for path in (sample_dir, metrics, timing))
def build_matrix(
registry: ExperimentRegistry,
*,
root: Path | None = None,
sample_count: int = 1024,
seed: int = 42,
) -> list[MatrixTask]:
"""Expand every supported registry cell over its declared step grid."""
if sample_count <= 0:
raise ValueError("sample_count must be positive")
root_path = _root_path(root)
tasks: list[MatrixTask] = []
for model_id, model in registry.models.items():
for dataset_id, support in model.datasets.items():
if support.status != SUPPORTED_STATUS:
continue
for steps in step_grid_for_model(registry, model_id):
task_id = f"{model_id}-{dataset_id}-steps-{steps}"
sample_dir, metrics_path, timing_path = _task_paths(
root_path, model_id, dataset_id, steps
)
tasks.append(
MatrixTask(
task_id=task_id,
category=model.category,
model=model_id,
dataset=dataset_id,
steps=steps,
sample_count=sample_count,
seed=seed,
environment=model.environment,
adapter=model.adapter,
source=model.source,
provenance=support.provenance or "",
sample_dir=sample_dir,
metrics_path=metrics_path,
timing_path=timing_path,
)
)
tasks.sort(key=lambda item: (item.model, item.dataset, item.steps))
_validate_tasks(tasks)
return tasks
def unsupported_inventory(registry: ExperimentRegistry) -> list[dict[str, str]]:
"""Return explicit unsupported cells without expanding them into steps."""
records: list[dict[str, str]] = []
for model_id, model in registry.models.items():
for dataset_id, support in model.datasets.items():
if support.status == SUPPORTED_STATUS:
continue
records.append(
{
"status": "unsupported",
"model": model_id,
"dataset": dataset_id,
"category": model.category,
"reason": support.reason or "no reason recorded",
}
)
records.sort(key=lambda item: (item["model"], item["dataset"]))
return records
def write_unsupported_inventory(
path: Path, records: Iterable[dict[str, str]], *, schema: str = MATRIX_SCHEMA
) -> Path:
"""Write the explicit unsupported model/dataset inventory as TSV."""
path = Path(path).absolute()
values = list(records)
for record in values:
if tuple(record) != UNSUPPORTED_COLUMNS:
raise ValueError("unsupported inventory columns do not match the canonical schema")
if any(
"\t" in value or "\n" in value or "\r" in value
for value in record.values()
):
raise ValueError("unsupported inventory contains a control character")
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:
handle.write(f"# schema={schema}-unsupported\n")
writer = csv.DictWriter(
handle, fieldnames=UNSUPPORTED_COLUMNS, delimiter="\t", lineterminator="\n"
)
writer.writeheader()
writer.writerows(values)
handle.flush()
os.fsync(handle.fileno())
os.replace(temporary, path)
except BaseException:
temporary.unlink(missing_ok=True)
raise
return path
def _validate_tasks(tasks: Iterable[MatrixTask]) -> None:
seen: set[str] = set()
previous: tuple[str, str, int] | None = None
for task in tasks:
if task.task_id in seen:
raise ValueError(f"duplicate matrix task ID: {task.task_id}")
seen.add(task.task_id)
key = (task.model, task.dataset, task.steps)
if previous is not None and key < previous:
raise ValueError("matrix tasks are not in stable order")
previous = key
if task.category not in {"many", "few", "fixed_1024"}:
raise ValueError(f"invalid task category: {task.category}")
if task.steps <= 0 or task.sample_count <= 0:
raise ValueError("steps and sample_count must be positive")
if any("\t" in value or "\n" in value or "\r" in value for value in task.row().values()):
raise ValueError(f"matrix task contains a control character: {task.task_id}")
def write_matrix(path: Path, tasks: Iterable[MatrixTask]) -> Path:
"""Write a deterministic, versioned TSV matrix."""
path = Path(path).absolute()
task_list = list(tasks)
_validate_tasks(task_list)
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:
handle.write(f"# schema={MATRIX_SCHEMA}\n")
writer = csv.DictWriter(
handle, fieldnames=MATRIX_COLUMNS, delimiter="\t", lineterminator="\n"
)
writer.writeheader()
for task in task_list:
writer.writerow(task.row())
handle.flush()
os.fsync(handle.fileno())
os.replace(temporary, path)
except BaseException:
temporary.unlink(missing_ok=True)
raise
return path
def read_matrix(path: Path) -> list[MatrixTask]:
"""Read and validate a matrix TSV without evaluating any field as code."""
path = Path(path)
if path.is_symlink() or not path.is_file():
raise ValueError(f"matrix TSV is missing or unsafe: {path}")
with path.open("r", encoding="utf-8", newline="") as handle:
schema = handle.readline().rstrip("\r\n")
if schema != f"# schema={MATRIX_SCHEMA}":
raise ValueError("matrix TSV has an unsupported schema header")
reader = csv.DictReader(handle, delimiter="\t")
if tuple(reader.fieldnames or ()) != MATRIX_COLUMNS:
raise ValueError("matrix TSV columns do not match the canonical schema")
tasks: list[MatrixTask] = []
for row_number, row in enumerate(reader, start=3):
if None in row or any(value is None for value in row.values()):
raise ValueError(f"matrix TSV row {row_number} has missing fields")
try:
task = MatrixTask(
task_id=row["task_id"],
category=row["category"],
model=row["model"],
dataset=row["dataset"],
steps=int(row["steps"]),
sample_count=int(row["sample_count"]),
seed=int(row["seed"]),
environment=row["environment"],
adapter=row["adapter"],
source=row["source"],
provenance=row["provenance"],
sample_dir=row["sample_dir"],
metrics_path=row["metrics_path"],
timing_path=row["timing_path"],
)
except (KeyError, TypeError, ValueError) as error:
raise ValueError(f"invalid matrix TSV row {row_number}") from error
tasks.append(task)
_validate_tasks(tasks)
if len({task.task_id for task in tasks}) != len(tasks):
raise ValueError("matrix TSV contains duplicate task IDs")
return tasks
def validate_matrix(path: Path, registry: ExperimentRegistry | None = None) -> list[MatrixTask]:
"""Validate a matrix and optionally compare it with the canonical registry."""
tasks = read_matrix(path)
if registry is not None:
expected = build_matrix(
registry,
root=Path("."),
sample_count=tasks[0].sample_count if tasks else 1024,
seed=tasks[0].seed if tasks else 42,
)
actual_keys = [(task.model, task.dataset, task.steps) for task in tasks]
expected_keys = [(task.model, task.dataset, task.steps) for task in expected]
if actual_keys != expected_keys:
raise ValueError("matrix TSV coverage differs from the canonical registry")
return tasks
def _main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--root", type=Path, default=Path.cwd())
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--sample-count", type=int, default=1024)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--unsupported-output", type=Path)
parser.add_argument("--validate", action="store_true")
args = parser.parse_args(argv)
root = args.root.resolve()
registry = load_registry(root / "configs" / "experiments.yaml")
if args.validate:
tasks = validate_matrix(args.output, registry)
else:
tasks = build_matrix(
registry, root=root, sample_count=args.sample_count, seed=args.seed
)
write_matrix(args.output, tasks)
write_unsupported_inventory(
args.unsupported_output
or args.output.with_name("unsupported.tsv"),
unsupported_inventory(registry),
)
print(f"matrix={args.output} tasks={len(tasks)}")
return 0
if __name__ == "__main__":
raise SystemExit(_main())
|