Viewing files at commit 25881039a2bbCommit detailsView latest

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