"""Re-grade a finished run with a different judge, without redoing agent work. A rejudge takes the agent output already stored in a parent run (the expensivepart: patches/answers the agent produced) and re-runs only the *judge* step witha new judge model, producing a brand-new derived run. The original run is readonly and never modified. This is deliberately not a re-implementation of the grader. It runs the exactSWE-Atlas verifier script shipped in each task package (`evaluate_rubrics.py` forRF, `evaluate_answer.py` for QnA) through the same judge proxy shim the livein-sandbox path uses (`judge_proxy.JUDGE_PROXY_SCRIPT`), so a rejudged verdict isproduced identically to a first-pass verdict — only the judge model differs. The verifier scripts read their inputs from fixed absolute paths (`/tests`,`/logs/...`) and pick the judge up from `EVAL_MODEL`/`EVAL_BASE_URL`/`EVAL_API_KEY`, so this module is meant to run inside a throwaway container (theModal `rejudge_controller`), where staging those absolute paths is safe. The purehelpers (resolution, reward recombination, result patching) have no suchrequirement and are unit-tested directly.""" from __future__ import annotations import jsonimport osimport shutilimport socketimport subprocessimport sysimport timefrom pathlib import Pathfrom typing import Any from . import configfrom .judge_proxy import JUDGE_PROXY_SCRIPT # Each task package ships exactly one of these verifier scripts; its presence is# how we tell an RF task from a QnA one without trusting the experiment config.RF_VERIFIER = "evaluate_rubrics.py"QNA_VERIFIER = "evaluate_answer.py" PART_RF = "rf"PART_QNA = "qna" # Absolute paths the cached verifier scripts hardcode. Staged per trial.TESTS_DIR = Path("/tests")LOGS_AGENT_DIR = Path("/logs/agent")LOGS_VERIFIER_DIR = Path("/logs/verifier")RF_PATCH_PATH = LOGS_VERIFIER_DIR / "agent.patch"RF_RESULTS_PATH = LOGS_VERIFIER_DIR / "rubrics_results.json"QNA_ANSWER_PATH = LOGS_AGENT_DIR / "answer.txt"QNA_RESULTS_PATH = LOGS_VERIFIER_DIR / "evaluation_results.json"  def load_json(path: Path) -> Any:    try:        return json.loads(Path(path).read_text())    except (OSError, json.JSONDecodeError):        return None  def task_tests_dir(task_name: str, task_ref: str, tasks_root: Path) -> Path:    """Resolve a trial's cached task package `tests/` dir from its name + ref.     Mirrors Harbor's content-addressed layout:    `<tasks_root>/<org>/<task-id>/<ref-hash>/tests`.    """    org, task_id = task_name.split("/", 1)    ref_hash = task_ref.split(":", 1)[1] if ":" in task_ref else task_ref    return tasks_root / org / task_id / ref_hash / "tests"  def tests_dir_for_trial(trial_dir: Path, tasks_root: Path) -> Path | None:    config_data = load_json(trial_dir / "config.json")    if not isinstance(config_data, dict):        return None    task = config_data.get("task") or {}    name, ref = task.get("name"), task.get("ref")    if not (name and ref):        return None    tests_dir = task_tests_dir(name, ref, tasks_root)    return tests_dir if tests_dir.is_dir() else None  def detect_part(tests_dir: Path) -> str | None:    if (tests_dir / RF_VERIFIER).exists():        return PART_RF    if (tests_dir / QNA_VERIFIER).exists():        return PART_QNA    return None  def recompute_rf_rewards(    rubrics_results: dict[str, Any], prior_rewards: dict[str, Any]) -> dict[str, float]:    """Combine the freshly re-judged rubric verdict with the *unchanged* saved    deterministic test reward. The agent patch didn't change, so `tests_reward`    is invariant under rejudge; only the rubric (`must_have_pass`/agg) moves.     The overall RF reward is `must_have_pass AND tests_reward == 1`, verified to    hold across every cached RF trial.    """    must_have_pass = bool(rubrics_results.get("must_have_pass"))    agg_score = rubrics_results.get("agg_score")    tests_reward = float(prior_rewards.get("tests_reward") or 0.0)    overall_pass = 1.0 if (must_have_pass and tests_reward >= 1.0) else 0.0    return {        "reward": overall_pass,        "overall_pass": overall_pass,        "tests_reward": tests_reward,        "must_have_pass": 1.0 if must_have_pass else 0.0,        "rubrics_agg_score": float(agg_score) if agg_score is not None else 0.0,    }  def recompute_qna_rewards(evaluation_results: dict[str, Any]) -> dict[str, float]:    """QnA has no deterministic test component; the reward is just the judged    pass verdict, matching the single-key `{"reward": ...}` production shape."""    return {"reward": 1.0 if bool(evaluation_results.get("pass")) else 0.0}  def proxy_environment(judge: config.JudgeConfig) -> dict[str, str]:    """Env the judge proxy subprocess reads. The auth key it forwards    (`judge.auth_env`, e.g. BASETEN_API_KEY) must already be present in the    process environment via the mounted LLM-providers secret."""    env = {        "ZED_JUDGE_UPSTREAM": judge.upstream,        "ZED_JUDGE_AUTH_ENV": judge.auth_env,    }    if judge.max_tokens is not None:        env["ZED_JUDGE_MAX_TOKENS"] = str(judge.max_tokens)    return env  def verifier_environment(judge_model: str, port: int) -> dict[str, str]:    """Env the cached verifier reads to find the judge.     This mirrors the verifier env injected by the benchmark harness command,    pointed at the local rejudge proxy.    """    return {        "EVAL_MODEL": judge_model,        "EVAL_BASE_URL": f"http://127.0.0.1:{port}/v1",        "EVAL_API_KEY": "unused-by-agent-evals-proxy",    }  class JudgeProxy:    """Runs `judge_proxy.JUDGE_PROXY_SCRIPT` as a local subprocess for the life    of a rejudge job. One proxy serves every trial (they all target the same    judge), mirroring the single in-sandbox daemon."""     def __init__(self, judge: config.JudgeConfig, port: int = 8089) -> None:        self._judge = judge        self._port = port        self._process: subprocess.Popen[bytes] | None = None        self._script_path = Path("/tmp/zed_judge_proxy.py")     @property    def port(self) -> int:        return self._port     def __enter__(self) -> JudgeProxy:        self._script_path.write_text(JUDGE_PROXY_SCRIPT)        env = os.environ.copy()        env.update(proxy_environment(self._judge))        env["ZED_JUDGE_PROXY_PORT"] = str(self._port)        self._process = subprocess.Popen(            [sys.executable, str(self._script_path)],            env=env,            stdout=subprocess.DEVNULL,            stderr=subprocess.DEVNULL,        )        self._wait_for_port()        return self     def __exit__(self, *_exc: object) -> None:        if self._process is not None:            self._process.terminate()            try:                self._process.wait(timeout=10)            except subprocess.TimeoutExpired:                self._process.kill()     def _wait_for_port(self, timeout_secs: float = 30.0) -> None:        deadline = time.time() + timeout_secs        while time.time() < deadline:            if self._process is not None and self._process.poll() is not None:                raise RuntimeError("judge proxy exited before becoming ready")            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:                probe.settimeout(0.5)                if probe.connect_ex(("127.0.0.1", self._port)) == 0:                    return            time.sleep(0.25)        raise TimeoutError(f"judge proxy did not open port {self._port} in time")  def _stage_tests_dir(tests_dir: Path) -> None:    """Point the absolute `/tests` path at this trial's package. Reset per trial    so each verifier reads its own rubrics/prompts."""    if TESTS_DIR.is_symlink() or TESTS_DIR.exists():        if TESTS_DIR.is_dir() and not TESTS_DIR.is_symlink():            shutil.rmtree(TESTS_DIR)        else:            TESTS_DIR.unlink()    TESTS_DIR.symlink_to(tests_dir)  def _run_verifier(    part: str, tests_dir: Path, judge_model: str, port: int) -> dict[str, Any]:    """Stage the verifier's expected inputs and run the real cached script."""    LOGS_AGENT_DIR.mkdir(parents=True, exist_ok=True)    LOGS_VERIFIER_DIR.mkdir(parents=True, exist_ok=True)    _stage_tests_dir(tests_dir)     env = os.environ.copy()    env.update(verifier_environment(judge_model, port))    script = tests_dir / (RF_VERIFIER if part == PART_RF else QNA_VERIFIER)    results_path = RF_RESULTS_PATH if part == PART_RF else QNA_RESULTS_PATH    if results_path.exists():        results_path.unlink()     completed = subprocess.run(        [sys.executable, str(script)],        env=env,        capture_output=True,        text=True,    )    results = load_json(results_path)    if not isinstance(results, dict):        raise RuntimeError(            f"verifier produced no results (exit {completed.returncode}): "            f"{completed.stderr.strip()[:500]}"        )    return results  def rejudge_trial(    trial_dir: Path,    tasks_root: Path,    judge: config.JudgeConfig,    judge_model: str,    port: int,) -> dict[str, Any]:    """Re-grade one trial in place (the trial dir is a copy in the derived run).     Returns a per-trial record; on any recoverable problem it returns    `{"ok": False, "error": ...}` so one bad trial can't sink the whole job.    """    task_name = trial_dir.name    tests_dir = tests_dir_for_trial(trial_dir, tasks_root)    if tests_dir is None:        return {"task": task_name, "ok": False, "error": "task package not found"}    part = detect_part(tests_dir)    if part is None:        return {"task": task_name, "ok": False, "error": "no verifier in task package"}     # The agent output is already on disk from the parent run; copy it to the    # path the verifier expects. RF reads the patch, QnA the delivered answer.    if part == PART_RF:        source = trial_dir / "verifier" / "agent.patch"        if not source.exists():            source = trial_dir / "verifier" / "agent_source_only.patch"        if not source.exists():            return {"task": task_name, "ok": False, "error": "no agent.patch"}        RF_PATCH_PATH.parent.mkdir(parents=True, exist_ok=True)        shutil.copyfile(source, RF_PATCH_PATH)    else:        source = trial_dir / "agent" / "answer.txt"        if not source.exists():            return {"task": task_name, "ok": False, "error": "no answer.txt"}        QNA_ANSWER_PATH.parent.mkdir(parents=True, exist_ok=True)        shutil.copyfile(source, QNA_ANSWER_PATH)     try:        results = _run_verifier(part, tests_dir, judge_model, port)    except (RuntimeError, OSError) as error:        return {"task": task_name, "ok": False, "error": str(error)}     prior_rewards = (        (load_json(trial_dir / "result.json") or {}).get("verifier_result") or {}    ).get("rewards") or {}    if part == PART_RF:        new_rewards = recompute_rf_rewards(results, prior_rewards)        verifier_results_name = "rubrics_results.json"    else:        new_rewards = recompute_qna_rewards(results)        verifier_results_name = "evaluation_results.json"     # Persist the fresh verifier output and verdict into the derived trial dir so    # report.py reads the new judge's result with no special-casing.    (trial_dir / "verifier").mkdir(parents=True, exist_ok=True)    (trial_dir / "verifier" / verifier_results_name).write_text(        json.dumps(results, indent=2)    )    patch_result_rewards(trial_dir / "result.json", new_rewards)    return {"task": task_name, "ok": True, "part": part, "rewards": new_rewards}  def patch_result_rewards(result_path: Path, new_rewards: dict[str, float]) -> None:    """Overwrite only `verifier_result.rewards` in a trial's `result.json`,    leaving the agent's resource metrics (tokens/steps/tool calls) untouched —    the agent didn't re-run, only the judge did."""    result = load_json(result_path)    if not isinstance(result, dict):        result = {}    verifier_result = result.get("verifier_result")    if not isinstance(verifier_result, dict):        verifier_result = {}    verifier_result["rewards"] = new_rewards    result["verifier_result"] = verifier_result    Path(result_path).write_text(json.dumps(result, indent=2))  def rejudge_job(    parent_job_dir: Path,    out_job_dir: Path,    tasks_root: Path,    judge: config.JudgeConfig,    judge_model: str,    port: int = 8089,    log: Any = print,) -> dict[str, Any]:    """Copy a parent run's Harbor job dir into the derived run and re-judge every    trial in place. Returns a summary suitable for the run's `summary.json`."""    if out_job_dir.exists():        shutil.rmtree(out_job_dir)    shutil.copytree(parent_job_dir, out_job_dir)     trial_dirs = sorted(        path        for path in out_job_dir.iterdir()        if path.is_dir() and (path / "result.json").exists()    )    records: list[dict[str, Any]] = []    with JudgeProxy(judge, port=port) as proxy:        for index, trial_dir in enumerate(trial_dirs, start=1):            record = rejudge_trial(                trial_dir, tasks_root, judge, judge_model, proxy.port            )            records.append(record)            status = "ok" if record["ok"] else f"FAILED: {record.get('error')}"            log(f"[{index}/{len(trial_dirs)}] {trial_dir.name}: {status}")     rejudged = sum(1 for record in records if record["ok"])    passed = sum(        1        for record in records        if record["ok"] and record["rewards"].get("reward", 0.0) >= 1.0    )    return {        "trial_count": len(trial_dirs),        "rejudged_count": rejudged,        "failed_count": len(trial_dirs) - rejudged,        "passed_count": passed,        "judge_model": judge_model,        "trials": records,    }