Viewing files at commit 2735d7f5558eCommit detailsView latest

Browse files
scripts/oakbench/baseline_book.py 367 lines · 13.2 KB · python
1
"""Git Baseline Book and target loading.
2
3
Baseline Books are compact committed summaries (ADR-0006). Targets are the
4
red/green contract for scorecards (ADR-0008). Both are JSON here so the
5
harness stays stdlib-only.
6
"""
7
8
from __future__ import annotations
9
10
import json
11
import hashlib
12
from dataclasses import dataclass
13
from datetime import datetime, timezone
14
from pathlib import Path
15
from typing import Any
16
17
from oakbench.rows import number_or_none, output_bytes, row_returncode, token_total, tool_calls
18
from oakbench.stats import bootstrap_ci, named_statistic, validate_interval_settings
19
from oakbench.populations import comparison_identity, independent_run_count, population_conflicts
20
21
22
@dataclass(frozen=True)
23
class BaselineCell:
24
    lane: str
25
    profile: str | None
26
    scenario: str
27
    operation: str
28
    subject: str
29
    metric: str
30
    statistic: str
31
    value: float
32
    ci_low: float | None = None
33
    ci_high: float | None = None
34
    n: int | None = None
35
    noise_floor_pct: float | None = None
36
    comparison_identity: dict[str, Any] | None = None
37
    independent_n: int | None = None
38
    attempted_n: int | None = None
39
    failure_n: int | None = None
40
    skipped_n: int | None = None
41
    confidence: float | None = None
42
43
    @property
44
    def key(self) -> tuple[str, str | None, str, str, str, str, str]:
45
        return (self.lane, self.profile, self.scenario, self.operation, self.subject, self.metric, self.statistic)
46
47
48
@dataclass(frozen=True)
49
class Target:
50
    target_id: str
51
    lane: str
52
    profile: str | None
53
    scenario: str
54
    operation: str
55
    subject: str
56
    comparator: str
57
    metric: str
58
    statistic: str
59
    goal_type: str
60
    goal_value: float
61
    guards: dict[str, Any]
62
63
    @property
64
    def baseline_key(self) -> tuple[str, str | None, str, str, str, str, str]:
65
        return (self.lane, self.profile, self.scenario, self.operation, self.comparator, self.metric, self.statistic)
66
67
68
def load_baseline_book(path: Path) -> dict[tuple[str, str | None, str, str, str, str, str], BaselineCell]:
69
    raw = json.loads(path.read_text())
70
    if raw.get("schema_version") not in (1, 2):
71
        raise ValueError(f"unsupported baseline book schema_version: {raw.get('schema_version')!r}")
72
    cells = raw.get("cells")
73
    if not isinstance(cells, list):
74
        raise ValueError("baseline book requires list-valued 'cells'")
75
    parsed = [_parse_cell(cell) for cell in cells]
76
    result = {}
77
    for cell in parsed:
78
        if cell.key in result:
79
            raise ValueError(f"duplicate baseline cell: {cell.key}")
80
        result[cell.key] = cell
81
    return result
82
83
84
def _parse_cell(cell: dict[str, Any]) -> BaselineCell:
85
    required = ("lane", "scenario", "operation", "subject", "metric", "statistic", "value")
86
    missing = [key for key in required if key not in cell]
87
    if missing:
88
        raise ValueError(f"baseline cell missing keys: {', '.join(missing)}")
89
    identity = cell.get("comparison_identity")
90
    if identity is not None and (not isinstance(identity, dict) or identity.get("schema_version") != 1):
91
        raise ValueError("unsupported baseline comparison_identity")
92
    return BaselineCell(
93
        lane=str(cell["lane"]),
94
        profile=None if cell.get("profile") is None else str(cell.get("profile")),
95
        scenario=str(cell["scenario"]),
96
        operation=str(cell["operation"]),
97
        subject=str(cell["subject"]),
98
        metric=str(cell["metric"]),
99
        statistic=str(cell["statistic"]),
100
        value=float(cell["value"]),
101
        ci_low=_optional_float(cell.get("ci_low")),
102
        ci_high=_optional_float(cell.get("ci_high")),
103
        n=_optional_int(cell.get("n")),
104
        noise_floor_pct=_optional_float(cell.get("noise_floor_pct")),
105
        comparison_identity=identity,
106
        independent_n=_optional_int(cell.get("independent_n")),
107
        attempted_n=_optional_int(cell.get("attempted_n")),
108
        failure_n=_optional_int(cell.get("failure_n")),
109
        skipped_n=_optional_int(cell.get("skipped_n")),
110
        confidence=_optional_float(cell.get("confidence")),
111
    )
112
113
114
def load_targets(path: Path) -> list[Target]:
115
    raw = json.loads(path.read_text())
116
    if raw.get("schema_version") != 1:
117
        raise ValueError(f"unsupported targets schema_version: {raw.get('schema_version')!r}")
118
    targets = raw.get("targets")
119
    if not isinstance(targets, list):
120
        raise ValueError("targets config requires list-valued 'targets'")
121
    return [_parse_target(target) for target in targets]
122
123
124
def _parse_target(target: dict[str, Any]) -> Target:
125
    required = ("id", "lane", "scenario", "operation", "subject", "comparator", "metric", "goal")
126
    missing = [key for key in required if key not in target]
127
    if missing:
128
        raise ValueError(f"target missing keys: {', '.join(missing)}")
129
    goal = target["goal"]
130
    if not isinstance(goal, dict):
131
        raise ValueError(f"target {target.get('id')!r} goal must be an object")
132
    if "type" not in goal:
133
        raise ValueError(f"target {target.get('id')!r} goal requires a 'type' field")
134
    goal_type = str(goal["type"])
135
    if goal_type != "ratio_at_most":
136
        raise ValueError(f"target {target.get('id')!r} unsupported goal type {goal_type!r}")
137
    if number_or_none(goal.get("value")) is None:
138
        raise ValueError(f"target {target.get('id')!r} goal requires a numeric 'value' field")
139
    guards = target.get("guards", {})
140
    if not isinstance(guards, dict):
141
        raise ValueError(f"target {target.get('id')!r} guards must be an object")
142
    return Target(
143
        target_id=str(target["id"]),
144
        lane=str(target["lane"]),
145
        profile=None if target.get("profile") is None else str(target.get("profile")),
146
        scenario=str(target["scenario"]),
147
        operation=str(target["operation"]),
148
        subject=str(target["subject"]),
149
        comparator=str(target["comparator"]),
150
        metric=str(target["metric"]),
151
        statistic=str(target.get("statistic", "median")),
152
        goal_type=goal_type,
153
        goal_value=float(goal["value"]),
154
        guards=guards,
155
    )
156
157
158
def _optional_float(value: Any) -> float | None:
159
    if value is None:
160
        return None
161
    return float(value)
162
163
164
def _optional_int(value: Any) -> int | None:
165
    if value is None:
166
        return None
167
    if isinstance(value, bool):
168
        raise ValueError("boolean is not an integer")
169
    return int(value)
170
171
172
DEFAULT_BASELINE_METRICS = ("elapsed_ms", "tokens", "tool_calls", "output_bytes")
173
DEFAULT_BASELINE_STATISTICS = ("median", "mean")
174
SUBOP_SUFFIXES = (".add", ".commit")
175
176
177
def iter_jsonl(paths: list[Path]) -> list[dict[str, Any]]:
178
    rows: list[dict[str, Any]] = []
179
    for path in paths:
180
        with path.open() as fh:
181
            for raw_line in fh:
182
                line = raw_line.strip()
183
                if line:
184
                    rows.append(json.loads(line))
185
    return rows
186
187
188
def metric_value(row: dict[str, Any], metric: str) -> float | None:
189
    if metric == "elapsed_ms":
190
        return number_or_none(row.get("elapsed_ms"))
191
    if metric == "tokens":
192
        return token_total(row)
193
    if metric == "tool_calls":
194
        return tool_calls(row)
195
    if metric == "output_bytes":
196
        return output_bytes(row)
197
    return number_or_none(row.get(metric))
198
199
200
def cells_from_rows(
201
    rows: list[dict[str, Any]],
202
    *,
203
    lane: str,
204
    metrics: tuple[str, ...] = DEFAULT_BASELINE_METRICS,
205
    statistics_names: tuple[str, ...] = DEFAULT_BASELINE_STATISTICS,
206
    subjects: set[str] | None = None,
207
    include_subops: bool = False,
208
    bootstrap_samples: int = 2000,
209
    confidence: float = 0.95,
210
    seed: int = 20260609,
211
    noise_floor_pct: float | None = None,
212
) -> list[BaselineCell]:
213
    validate_interval_settings(bootstrap_samples, confidence)
214
    grouped: dict[tuple[str | None, str, str, str, str], list[float]] = {}
215
    observations = {}
216
    outcomes = {}
217
    for row in rows:
218
        if str(row.get("subject_kind") or "") != "git":
219
            continue
220
        subject = str(row.get("subject") or "")
221
        if subjects is not None and subject not in subjects:
222
            continue
223
        operation = str(row.get("operation") or "")
224
        if not include_subops and operation.endswith(SUBOP_SUFFIXES):
225
            continue
226
        profile = None if row.get("profile") is None else str(row.get("profile"))
227
        scenario = str(row.get("scenario") or "")
228
        outcomes.setdefault((profile, scenario, operation, subject), []).append(row)
229
        if row_returncode(row) != 0:
230
            continue
231
        for metric in metrics:
232
            value = metric_value(row, metric)
233
            if value is None:
234
                continue
235
            key = (profile, scenario, operation, subject, metric)
236
            grouped.setdefault(key, []).append(float(value))
237
            observations.setdefault(key, []).append(row)
238
239
    cells: list[BaselineCell] = []
240
    for (profile, scenario, operation, subject, metric), values in sorted(grouped.items()):
241
        observed = observations[(profile, scenario, operation, subject, metric)]
242
        all_rows = outcomes[(profile, scenario, operation, subject)]
243
        conflicts = population_conflicts(all_rows)
244
        if conflicts:
245
            raise ValueError(f"incompatible baseline population {scenario}/{operation}: {', '.join(conflicts)}")
246
        for statistic_name in statistics_names:
247
            stat = named_statistic(statistic_name)
248
            ci_low, ci_high = bootstrap_ci(
249
                values,
250
                statistic=statistic_name,
251
                samples=bootstrap_samples,
252
                confidence=confidence,
253
                seed=seed,
254
            )
255
            cells.append(
256
                BaselineCell(
257
                    lane=lane,
258
                    profile=profile,
259
                    scenario=scenario,
260
                    operation=operation,
261
                    subject=subject,
262
                    metric=metric,
263
                    statistic=statistic_name,
264
                    value=float(stat(values)),
265
                    ci_low=ci_low,
266
                    ci_high=ci_high,
267
                    n=len(values),
268
                    noise_floor_pct=noise_floor_pct,
269
                    comparison_identity=comparison_identity(observed[0]),
270
                    independent_n=independent_run_count(observed),
271
                    attempted_n=len(all_rows),
272
                    failure_n=sum(row_returncode(row) not in (0, 77) for row in all_rows),
273
                    skipped_n=sum(row_returncode(row) == 77 for row in all_rows),
274
                    confidence=confidence,
275
                )
276
            )
277
    return cells
278
279
280
def book_json(cells: list[BaselineCell]) -> dict[str, Any]:
281
    return {
282
        "schema_version": 2,
283
        "cells": [
284
            {
285
                "lane": cell.lane,
286
                "profile": cell.profile,
287
                "scenario": cell.scenario,
288
                "operation": cell.operation,
289
                "subject": cell.subject,
290
                "metric": cell.metric,
291
                "statistic": cell.statistic,
292
                "value": cell.value,
293
                "ci_low": cell.ci_low,
294
                "ci_high": cell.ci_high,
295
                "n": cell.n,
296
                "noise_floor_pct": cell.noise_floor_pct,
297
                "comparison_identity": cell.comparison_identity,
298
                "independent_n": cell.independent_n,
299
                "attempted_n": cell.attempted_n,
300
                "failure_n": cell.failure_n,
301
                "skipped_n": cell.skipped_n,
302
                "confidence": cell.confidence,
303
            }
304
            for cell in cells
305
        ],
306
    }
307
308
309
def manifest_json(
310
    *,
311
    book_id: str,
312
    row_paths: list[Path],
313
    runner_class: str | None,
314
    runner_id: str | None = None,
315
    runner_profile: dict[str, Any] | None = None,
316
    git_version: str | None = None,
317
    git_binary_sha256: str | None = None,
318
    raw_storage_keys: list[str] | None = None,
319
    suite_hash: str | None = None,
320
) -> dict[str, Any]:
321
    return {
322
        "schema_version": 1,
323
        "book_id": book_id,
324
        "generated_at_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
325
        "suite_hash": suite_hash,
326
        "runner_id": runner_id,
327
        "runner_class": runner_class,
328
        "runner_profile": runner_profile,
329
        "git_version": git_version,
330
        "git_binary_sha256": git_binary_sha256,
331
        "raw_storage_keys": raw_storage_keys or [],
332
        "raw_inputs": [
333
            {
334
                "path": str(path),
335
                "sha256": sha256_file(path),
336
                "bytes": path.stat().st_size,
337
            }
338
            for path in row_paths
339
        ],
340
    }
341
342
343
def book_markdown(book_id: str, cells: list[BaselineCell]) -> str:
344
    lines = [
345
        f"# Git Baseline Book: `{book_id}`",
346
        "",
347
        "| Subject | Scenario / operation | Metric | Statistic | n | Value | CI |",
348
        "| --- | --- | --- | --- | ---: | ---: | --- |",
349
    ]
350
    for cell in cells:
351
        ci = ""
352
        if cell.ci_low is not None and cell.ci_high is not None:
353
            ci = f"[{cell.ci_low:.3f}, {cell.ci_high:.3f}]"
354
        lines.append(
355
            f"| `{cell.subject}` | `{cell.scenario}/{cell.operation}` | `{cell.metric}` | "
356
            f"`{cell.statistic}` | {cell.n or 0} | {cell.value:.3f} | {ci} |"
357
        )
358
    lines.append("")
359
    return "\n".join(lines)
360
361
362
def sha256_file(path: Path) -> str:
363
    digest = hashlib.sha256()
364
    with path.open("rb") as fh:
365
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
366
            digest.update(chunk)
367
    return digest.hexdigest()