tenant.openagents/omega

1"""Baseline-of-record management.
2
3A "baseline" is a completed run promoted to the canonical reference for its
4`(experiment, model)` pair. It only qualifies if its build came from a *clean*
5commit reachable from `origin/main` (no local patch) — so baselines are both
6obviously discoverable (the `baselines/` tree + `zed-eval baseline list`) and
7provably pinned to a commit on main. The judge preset is recorded but not gated,
8so non-leaderboard baselines are allowed.
9
10Records live on the shared volume at:
11
12    baselines/index.json                                  # roll-up for discovery
13    baselines/<experiment>/<model-slug>/current.json      # canonical pointer
14    baselines/<experiment>/<model-slug>/history/<sha>.json # superseded baselines
15"""
16
17from __future__ import annotations
18
19import argparse
20import json
21from typing import Any
22
23from . import config, source
24from .common import default_namespace, deployed_function, utc_now
25
26
27def _model_slug(model: str) -> str:
28    return source.sanitize_namespace(model)
29
30
31def _record_from_provenance(
32    *,
33    namespace: str,
34    fallback_experiment: str,
35    run_id: str,
36    provenance: dict[str, Any],
37) -> tuple[dict[str, Any], bool]:
38    request = provenance["request"]
39    summary = provenance["summary"]
40    build_info = provenance["build_info"]
41    source_info = build_info.get("source") or {}
42
43    experiment = request.get("experiment_name") or fallback_experiment
44    model = request.get("agent_model") or config.DEFAULT_MODEL
45    base_sha = build_info.get("base_sha")
46
47    record = {
48        "experiment": experiment,
49        "model": model,
50        "model_slug": _model_slug(model),
51        "judge": request.get("judge_preset") or request.get("judge"),
52        "base_sha": base_sha,
53        "base_ref": source_info.get("base_ref"),
54        "on_main": None,
55        "clean": not bool(build_info.get("patch_sha256")),
56        "build_id": build_info.get("build_id") or request.get("build_id"),
57        "run": {
58            "namespace": namespace,
59            "experiment_name": experiment,
60            "run_id": run_id,
61        },
62        "resources": {
63            "override_cpus": request.get("override_cpus"),
64            "override_memory_mb": request.get("override_memory_mb"),
65        },
66        "summary": {
67            "status": summary.get("status"),
68            "trial_count": summary.get("trial_count"),
69        },
70        "recorded_at": utc_now(),
71        "recorded_by": namespace,
72    }
73    return record, bool(build_info)
74
75
76def _baseline_problems(
77    *,
78    record: dict[str, Any],
79    has_build_info: bool,
80    repo_url: str,
81    allow_dirty: bool,
82    allow_off_main: bool,
83) -> tuple[list[str], bool | None]:
84    problems: list[str] = []
85    run_id = (record.get("run") or {}).get("run_id")
86    status = (record.get("summary") or {}).get("status")
87    base_sha = record.get("base_sha")
88
89    if not has_build_info:
90        problems.append(
91            f"no build-info found for run {run_id}; cannot verify provenance"
92        )
93    if status != "completed":
94        problems.append(f"run status is {status!r}, expected 'completed'")
95    if not record.get("clean") and not allow_dirty:
96        problems.append(
97            "build carries a local patch (not a clean commit); pass --allow-dirty to override"
98        )
99
100    on_main: bool | None = None
101    if base_sha:
102        try:
103            on_main = source.base_sha_on_main(base_sha, repo_url)
104        except Exception as error:  # noqa: BLE001 - surface any verification failure
105            if not allow_off_main:
106                problems.append(f"could not verify base_sha is on origin/main: {error}")
107        if on_main is False and not allow_off_main:
108            problems.append(
109                f"base_sha {base_sha[:12]} is not reachable from origin/main; "
110                "pass --allow-off-main to override"
111            )
112    elif not allow_off_main:
113        problems.append("build-info has no base_sha to verify against origin/main")
114
115    return problems, on_main
116
117
118def _print_refusal(run_id: str, problems: list[str]) -> None:
119    print(f"refusing to record baseline for {run_id}:")
120    for problem in problems:
121        print(f"  - {problem}")
122
123
124def _print_recorded(record: dict[str, Any]) -> None:
125    print(
126        f"recorded baseline: {record.get('experiment')} / {record.get('model')}  "
127        f"base_sha={str(record.get('base_sha'))[:12]} "
128        f"on_main={record.get('on_main')} clean={record.get('clean')} "
129        f"judge={record.get('judge')}"
130    )
131
132
133def command_baseline_record(args: argparse.Namespace) -> int:
134    namespace = default_namespace(args)
135    experiment_slug = source.sanitize_namespace(args.experiment_name)
136    repo_url = getattr(args, "repo_url", None) or source.DEFAULT_REPO_URL
137    read_run_provenance = deployed_function(args, "read_run_provenance")
138    record_baseline = deployed_function(args, "record_baseline")
139
140    exit_code = 0
141    for run_id in args.run_id:
142        provenance = read_run_provenance.remote(namespace, experiment_slug, run_id)
143        record, has_build_info = _record_from_provenance(
144            namespace=namespace,
145            fallback_experiment=args.experiment_name,
146            run_id=run_id,
147            provenance=provenance,
148        )
149        problems, on_main = _baseline_problems(
150            record=record,
151            has_build_info=has_build_info,
152            repo_url=repo_url,
153            allow_dirty=args.allow_dirty,
154            allow_off_main=args.allow_off_main,
155        )
156        record["on_main"] = on_main
157
158        if problems:
159            _print_refusal(run_id, problems)
160            exit_code = 1
161            continue
162
163        record_baseline.remote(record)
164        _print_recorded(record)
165    return exit_code
166
167
168def command_baseline_list(args: argparse.Namespace) -> int:
169    index = deployed_function(args, "read_baselines").remote(None, None, False)
170    rows = index.get("baselines") or []
171    if getattr(args, "json", False):
172        print(json.dumps(index, indent=2))
173        return 0
174    if not rows:
175        print("no baselines recorded")
176        return 0
177    header = (
178        f"{'EXPERIMENT':<20} {'MODEL':<32} {'BASE_REF':<14} "
179        f"{'SHA':<12} {'ON_MAIN':<8} {'JUDGE':<14} RUN"
180    )
181    print(header)
182    for row in rows:
183        print(
184            f"{(row.get('experiment') or ''):<20} "
185            f"{(row.get('model') or ''):<32} "
186            f"{(row.get('base_ref') or '') or '-':<14} "
187            f"{str(row.get('base_sha') or '')[:12]:<12} "
188            f"{str(row.get('on_main')):<8} "
189            f"{(row.get('judge') or '') or '-':<14} "
190            f"{row.get('run_id') or ''}"
191        )
192    return 0
193
194
195def command_baseline_show(args: argparse.Namespace) -> int:
196    result = deployed_function(args, "read_baselines").remote(
197        args.experiment_name,
198        _model_slug(args.model),
199        bool(getattr(args, "history", False)),
200    )
201    if not result or result.get("current") is None:
202        print(f"no baseline recorded for {args.experiment_name} / {args.model}")
203        return 1
204    print(json.dumps(result, indent=2))
205    return 0
206