tenant.openagents/omega

1"""Re-grade a finished run with a different judge, without redoing agent work.
2
3A rejudge takes the agent output already stored in a parent run (the expensive
4part: patches/answers the agent produced) and re-runs only the *judge* step with
5a new judge model, producing a brand-new derived run. The original run is read
6only and never modified.
7
8This is deliberately not a re-implementation of the grader. It runs the exact
9SWE-Atlas verifier script shipped in each task package (`evaluate_rubrics.py` for
10RF, `evaluate_answer.py` for QnA) through the same judge proxy shim the live
11in-sandbox path uses (`judge_proxy.JUDGE_PROXY_SCRIPT`), so a rejudged verdict is
12produced identically to a first-pass verdict — only the judge model differs.
13
14The verifier scripts read their inputs from fixed absolute paths (`/tests`,
15`/logs/...`) and pick the judge up from `EVAL_MODEL`/`EVAL_BASE_URL`/
16`EVAL_API_KEY`, so this module is meant to run inside a throwaway container (the
17Modal `rejudge_controller`), where staging those absolute paths is safe. The pure
18helpers (resolution, reward recombination, result patching) have no such
19requirement and are unit-tested directly.
20"""
21
22from __future__ import annotations
23
24import json
25import os
26import shutil
27import socket
28import subprocess
29import sys
30import time
31from pathlib import Path
32from typing import Any
33
34from . import config
35from .judge_proxy import JUDGE_PROXY_SCRIPT
36
37# Each task package ships exactly one of these verifier scripts; its presence is
38# how we tell an RF task from a QnA one without trusting the experiment config.
39RF_VERIFIER = "evaluate_rubrics.py"
40QNA_VERIFIER = "evaluate_answer.py"
41
42PART_RF = "rf"
43PART_QNA = "qna"
44
45# Absolute paths the cached verifier scripts hardcode. Staged per trial.
46TESTS_DIR = Path("/tests")
47LOGS_AGENT_DIR = Path("/logs/agent")
48LOGS_VERIFIER_DIR = Path("/logs/verifier")
49RF_PATCH_PATH = LOGS_VERIFIER_DIR / "agent.patch"
50RF_RESULTS_PATH = LOGS_VERIFIER_DIR / "rubrics_results.json"
51QNA_ANSWER_PATH = LOGS_AGENT_DIR / "answer.txt"
52QNA_RESULTS_PATH = LOGS_VERIFIER_DIR / "evaluation_results.json"
53
54
55def load_json(path: Path) -> Any:
56    try:
57        return json.loads(Path(path).read_text())
58    except (OSError, json.JSONDecodeError):
59        return None
60
61
62def task_tests_dir(task_name: str, task_ref: str, tasks_root: Path) -> Path:
63    """Resolve a trial's cached task package `tests/` dir from its name + ref.
64
65    Mirrors Harbor's content-addressed layout:
66    `<tasks_root>/<org>/<task-id>/<ref-hash>/tests`.
67    """
68    org, task_id = task_name.split("/", 1)
69    ref_hash = task_ref.split(":", 1)[1] if ":" in task_ref else task_ref
70    return tasks_root / org / task_id / ref_hash / "tests"
71
72
73def tests_dir_for_trial(trial_dir: Path, tasks_root: Path) -> Path | None:
74    config_data = load_json(trial_dir / "config.json")
75    if not isinstance(config_data, dict):
76        return None
77    task = config_data.get("task") or {}
78    name, ref = task.get("name"), task.get("ref")
79    if not (name and ref):
80        return None
81    tests_dir = task_tests_dir(name, ref, tasks_root)
82    return tests_dir if tests_dir.is_dir() else None
83
84
85def detect_part(tests_dir: Path) -> str | None:
86    if (tests_dir / RF_VERIFIER).exists():
87        return PART_RF
88    if (tests_dir / QNA_VERIFIER).exists():
89        return PART_QNA
90    return None
91
92
93def recompute_rf_rewards(
94    rubrics_results: dict[str, Any], prior_rewards: dict[str, Any]
95) -> dict[str, float]:
96    """Combine the freshly re-judged rubric verdict with the *unchanged* saved
97    deterministic test reward. The agent patch didn't change, so `tests_reward`
98    is invariant under rejudge; only the rubric (`must_have_pass`/agg) moves.
99
100    The overall RF reward is `must_have_pass AND tests_reward == 1`, verified to
101    hold across every cached RF trial.
102    """
103    must_have_pass = bool(rubrics_results.get("must_have_pass"))
104    agg_score = rubrics_results.get("agg_score")
105    tests_reward = float(prior_rewards.get("tests_reward") or 0.0)
106    overall_pass = 1.0 if (must_have_pass and tests_reward >= 1.0) else 0.0
107    return {
108        "reward": overall_pass,
109        "overall_pass": overall_pass,
110        "tests_reward": tests_reward,
111        "must_have_pass": 1.0 if must_have_pass else 0.0,
112        "rubrics_agg_score": float(agg_score) if agg_score is not None else 0.0,
113    }
114
115
116def recompute_qna_rewards(evaluation_results: dict[str, Any]) -> dict[str, float]:
117    """QnA has no deterministic test component; the reward is just the judged
118    pass verdict, matching the single-key `{"reward": ...}` production shape."""
119    return {"reward": 1.0 if bool(evaluation_results.get("pass")) else 0.0}
120
121
122def proxy_environment(judge: config.JudgeConfig) -> dict[str, str]:
123    """Env the judge proxy subprocess reads. The auth key it forwards
124    (`judge.auth_env`, e.g. BASETEN_API_KEY) must already be present in the
125    process environment via the mounted LLM-providers secret."""
126    env = {
127        "ZED_JUDGE_UPSTREAM": judge.upstream,
128        "ZED_JUDGE_AUTH_ENV": judge.auth_env,
129    }
130    if judge.max_tokens is not None:
131        env["ZED_JUDGE_MAX_TOKENS"] = str(judge.max_tokens)
132    return env
133
134
135def verifier_environment(judge_model: str, port: int) -> dict[str, str]:
136    """Env the cached verifier reads to find the judge.
137
138    This mirrors the verifier env injected by the benchmark harness command,
139    pointed at the local rejudge proxy.
140    """
141    return {
142        "EVAL_MODEL": judge_model,
143        "EVAL_BASE_URL": f"http://127.0.0.1:{port}/v1",
144        "EVAL_API_KEY": "unused-by-agent-evals-proxy",
145    }
146
147
148class JudgeProxy:
149    """Runs `judge_proxy.JUDGE_PROXY_SCRIPT` as a local subprocess for the life
150    of a rejudge job. One proxy serves every trial (they all target the same
151    judge), mirroring the single in-sandbox daemon."""
152
153    def __init__(self, judge: config.JudgeConfig, port: int = 8089) -> None:
154        self._judge = judge
155        self._port = port
156        self._process: subprocess.Popen[bytes] | None = None
157        self._script_path = Path("/tmp/zed_judge_proxy.py")
158
159    @property
160    def port(self) -> int:
161        return self._port
162
163    def __enter__(self) -> JudgeProxy:
164        self._script_path.write_text(JUDGE_PROXY_SCRIPT)
165        env = os.environ.copy()
166        env.update(proxy_environment(self._judge))
167        env["ZED_JUDGE_PROXY_PORT"] = str(self._port)
168        self._process = subprocess.Popen(
169            [sys.executable, str(self._script_path)],
170            env=env,
171            stdout=subprocess.DEVNULL,
172            stderr=subprocess.DEVNULL,
173        )
174        self._wait_for_port()
175        return self
176
177    def __exit__(self, *_exc: object) -> None:
178        if self._process is not None:
179            self._process.terminate()
180            try:
181                self._process.wait(timeout=10)
182            except subprocess.TimeoutExpired:
183                self._process.kill()
184
185    def _wait_for_port(self, timeout_secs: float = 30.0) -> None:
186        deadline = time.time() + timeout_secs
187        while time.time() < deadline:
188            if self._process is not None and self._process.poll() is not None:
189                raise RuntimeError("judge proxy exited before becoming ready")
190            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
191                probe.settimeout(0.5)
192                if probe.connect_ex(("127.0.0.1", self._port)) == 0:
193                    return
194            time.sleep(0.25)
195        raise TimeoutError(f"judge proxy did not open port {self._port} in time")
196
197
198def _stage_tests_dir(tests_dir: Path) -> None:
199    """Point the absolute `/tests` path at this trial's package. Reset per trial
200    so each verifier reads its own rubrics/prompts."""
201    if TESTS_DIR.is_symlink() or TESTS_DIR.exists():
202        if TESTS_DIR.is_dir() and not TESTS_DIR.is_symlink():
203            shutil.rmtree(TESTS_DIR)
204        else:
205            TESTS_DIR.unlink()
206    TESTS_DIR.symlink_to(tests_dir)
207
208
209def _run_verifier(
210    part: str, tests_dir: Path, judge_model: str, port: int
211) -> dict[str, Any]:
212    """Stage the verifier's expected inputs and run the real cached script."""
213    LOGS_AGENT_DIR.mkdir(parents=True, exist_ok=True)
214    LOGS_VERIFIER_DIR.mkdir(parents=True, exist_ok=True)
215    _stage_tests_dir(tests_dir)
216
217    env = os.environ.copy()
218    env.update(verifier_environment(judge_model, port))
219    script = tests_dir / (RF_VERIFIER if part == PART_RF else QNA_VERIFIER)
220    results_path = RF_RESULTS_PATH if part == PART_RF else QNA_RESULTS_PATH
221    if results_path.exists():
222        results_path.unlink()
223
224    completed = subprocess.run(
225        [sys.executable, str(script)],
226        env=env,
227        capture_output=True,
228        text=True,
229    )
230    results = load_json(results_path)
231    if not isinstance(results, dict):
232        raise RuntimeError(
233            f"verifier produced no results (exit {completed.returncode}): "
234            f"{completed.stderr.strip()[:500]}"
235        )
236    return results
237
238
239def rejudge_trial(
240    trial_dir: Path,
241    tasks_root: Path,
242    judge: config.JudgeConfig,
243    judge_model: str,
244    port: int,
245) -> dict[str, Any]:
246    """Re-grade one trial in place (the trial dir is a copy in the derived run).
247
248    Returns a per-trial record; on any recoverable problem it returns
249    `{"ok": False, "error": ...}` so one bad trial can't sink the whole job.
250    """
251    task_name = trial_dir.name
252    tests_dir = tests_dir_for_trial(trial_dir, tasks_root)
253    if tests_dir is None:
254        return {"task": task_name, "ok": False, "error": "task package not found"}
255    part = detect_part(tests_dir)
256    if part is None:
257        return {"task": task_name, "ok": False, "error": "no verifier in task package"}
258
259    # The agent output is already on disk from the parent run; copy it to the
260    # path the verifier expects. RF reads the patch, QnA the delivered answer.
261    if part == PART_RF:
262        source = trial_dir / "verifier" / "agent.patch"
263        if not source.exists():
264            source = trial_dir / "verifier" / "agent_source_only.patch"
265        if not source.exists():
266            return {"task": task_name, "ok": False, "error": "no agent.patch"}
267        RF_PATCH_PATH.parent.mkdir(parents=True, exist_ok=True)
268        shutil.copyfile(source, RF_PATCH_PATH)
269    else:
270        source = trial_dir / "agent" / "answer.txt"
271        if not source.exists():
272            return {"task": task_name, "ok": False, "error": "no answer.txt"}
273        QNA_ANSWER_PATH.parent.mkdir(parents=True, exist_ok=True)
274        shutil.copyfile(source, QNA_ANSWER_PATH)
275
276    try:
277        results = _run_verifier(part, tests_dir, judge_model, port)
278    except (RuntimeError, OSError) as error:
279        return {"task": task_name, "ok": False, "error": str(error)}
280
281    prior_rewards = (
282        (load_json(trial_dir / "result.json") or {}).get("verifier_result") or {}
283    ).get("rewards") or {}
284    if part == PART_RF:
285        new_rewards = recompute_rf_rewards(results, prior_rewards)
286        verifier_results_name = "rubrics_results.json"
287    else:
288        new_rewards = recompute_qna_rewards(results)
289        verifier_results_name = "evaluation_results.json"
290
291    # Persist the fresh verifier output and verdict into the derived trial dir so
292    # report.py reads the new judge's result with no special-casing.
293    (trial_dir / "verifier").mkdir(parents=True, exist_ok=True)
294    (trial_dir / "verifier" / verifier_results_name).write_text(
295        json.dumps(results, indent=2)
296    )
297    patch_result_rewards(trial_dir / "result.json", new_rewards)
298    return {"task": task_name, "ok": True, "part": part, "rewards": new_rewards}
299
300
301def patch_result_rewards(result_path: Path, new_rewards: dict[str, float]) -> None:
302    """Overwrite only `verifier_result.rewards` in a trial's `result.json`,
303    leaving the agent's resource metrics (tokens/steps/tool calls) untouched —
304    the agent didn't re-run, only the judge did."""
305    result = load_json(result_path)
306    if not isinstance(result, dict):
307        result = {}
308    verifier_result = result.get("verifier_result")
309    if not isinstance(verifier_result, dict):
310        verifier_result = {}
311    verifier_result["rewards"] = new_rewards
312    result["verifier_result"] = verifier_result
313    Path(result_path).write_text(json.dumps(result, indent=2))
314
315
316def rejudge_job(
317    parent_job_dir: Path,
318    out_job_dir: Path,
319    tasks_root: Path,
320    judge: config.JudgeConfig,
321    judge_model: str,
322    port: int = 8089,
323    log: Any = print,
324) -> dict[str, Any]:
325    """Copy a parent run's Harbor job dir into the derived run and re-judge every
326    trial in place. Returns a summary suitable for the run's `summary.json`."""
327    if out_job_dir.exists():
328        shutil.rmtree(out_job_dir)
329    shutil.copytree(parent_job_dir, out_job_dir)
330
331    trial_dirs = sorted(
332        path
333        for path in out_job_dir.iterdir()
334        if path.is_dir() and (path / "result.json").exists()
335    )
336    records: list[dict[str, Any]] = []
337    with JudgeProxy(judge, port=port) as proxy:
338        for index, trial_dir in enumerate(trial_dirs, start=1):
339            record = rejudge_trial(
340                trial_dir, tasks_root, judge, judge_model, proxy.port
341            )
342            records.append(record)
343            status = "ok" if record["ok"] else f"FAILED: {record.get('error')}"
344            log(f"[{index}/{len(trial_dirs)}] {trial_dir.name}: {status}")
345
346    rejudged = sum(1 for record in records if record["ok"])
347    passed = sum(
348        1
349        for record in records
350        if record["ok"] and record["rewards"].get("reward", 0.0) >= 1.0
351    )
352    return {
353        "trial_count": len(trial_dirs),
354        "rejudged_count": rejudged,
355        "failed_count": len(trial_dirs) - rejudged,
356        "passed_count": passed,
357        "judge_model": judge_model,
358        "trials": records,
359    }
360