harness/scripts/select_harness.py

#!/usr/bin/env python3
"""Apply the precommitted sampling-temperature and review selection rules."""

from __future__ import annotations

import argparse
import json
import time
import tomllib
from pathlib import Path

from compare_evals import exact_two_sided_sign_p, read_traces, solved, task_key


def outcomes(path: Path) -> dict[str, bool]:
    run_dir = path if path.is_dir() else path.parent
    config = tomllib.loads((run_dir / "config.toml").read_text())
    if config.get("num_tasks") != 32 or config.get("num_rollouts", 1) != 1:
        raise ValueError(f"selection run is not the precommitted 32x1 protocol: {path}")
    traces = read_traces(path)
    result = {task_key(trace): solved(trace) for trace in traces}
    if len(result) != len(traces):
        raise ValueError(f"selection run contains duplicate task keys: {path}")
    if len(result) != 32:
        raise ValueError(
            f"selection run is incomplete: expected 32 distinct traces, "
            f"observed {len(result)}: {path}"
        )
    return result


def compare(path_a: Path, path_b: Path, label_a: str, label_b: str) -> dict:
    a = outcomes(path_a)
    b = outcomes(path_b)
    if a.keys() != b.keys():
        raise ValueError(
            f"paired selection runs contain different task keys: {path_a} vs {path_b}"
        )
    keys = a.keys() | b.keys()
    a_only = sum(a.get(key, False) and not b.get(key, False) for key in keys)
    b_only = sum(b.get(key, False) and not a.get(key, False) for key in keys)
    return {
        f"{label_a}_successes": sum(a.values()),
        f"{label_b}_successes": sum(b.values()),
        f"{label_a}_only": a_only,
        f"{label_b}_only": b_only,
        "observed_task_union": len(keys),
        f"missing_from_{label_a}": len(b.keys() - a.keys()),
        f"missing_from_{label_b}": len(a.keys() - b.keys()),
    }


def challenger_gate(
    tb2: dict, swe: dict, default_label: str, challenger_label: str
) -> dict:
    default_only = sum(
        result[f"{default_label}_only"] for result in (tb2, swe)
    )
    challenger_only = sum(
        result[f"{challenger_label}_only"] for result in (tb2, swe)
    )
    sign_p = exact_two_sided_sign_p(default_only, challenger_only)
    no_suite_regression = all(
        result[f"{challenger_label}_successes"]
        >= result[f"{default_label}_successes"]
        for result in (tb2, swe)
    )
    positive_paired_advantage = challenger_only > default_only
    significant = sign_p is not None and sign_p <= 0.10
    return {
        "challenger_qualifies": (
            no_suite_regression and positive_paired_advantage and significant
        ),
        "rule": {
            "no_suite_regression": no_suite_regression,
            "positive_combined_paired_advantage": positive_paired_advantage,
            "combined_exact_two_sided_sign_test_p_le_0_10": significant,
        },
        "combined": {
            f"{default_label}_only": default_only,
            f"{challenger_label}_only": challenger_only,
            "discordant_tasks": default_only + challenger_only,
            "exact_two_sided_sign_test_p": sign_p,
        },
        "terminal_bench_2": tb2,
        "swe_bench_verified": swe,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("no_review_t02_tb2", type=Path)
    parser.add_argument("review_t02_tb2", type=Path)
    parser.add_argument("no_review_t06_tb2", type=Path)
    parser.add_argument("review_t06_tb2", type=Path)
    parser.add_argument("no_review_t02_swe", type=Path)
    parser.add_argument("review_t02_swe", type=Path)
    parser.add_argument("no_review_t06_swe", type=Path)
    parser.add_argument("review_t06_swe", type=Path)
    parser.add_argument("stock_t02_tb2", type=Path)
    parser.add_argument("stock_t06_tb2", type=Path)
    parser.add_argument("stock_t02_swe", type=Path)
    parser.add_argument("stock_t06_swe", type=Path)
    args = parser.parse_args()

    # The fixed development wave can finish while the shared sandbox pool is
    # saturated.  An independent watcher resumes only terminal infrastructure
    # failures on the same task indices.  Do not let those transient failures
    # influence harness selection while that repair is still in progress.
    paths = [Path(value) for value in vars(args).values()]
    if any("step150-" in str(path) for path in paths):
        marker = Path(__file__).resolve().parent.parent / "state" / "step150-infra-repair-complete"
        deadline = time.monotonic() + 8 * 60 * 60
        while not marker.exists():
            if time.monotonic() >= deadline:
                raise TimeoutError(f"timed out waiting for infrastructure repair: {marker}")
            time.sleep(10)

    temperature_gate = challenger_gate(
        compare(
            args.no_review_t06_tb2,
            args.no_review_t02_tb2,
            "official_t06",
            "low_t02",
        ),
        compare(
            args.no_review_t06_swe,
            args.no_review_t02_swe,
            "official_t06",
            "low_t02",
        ),
        "official_t06",
        "low_t02",
    )
    temperature = "0.2" if temperature_gate["challenger_qualifies"] else "0.6"

    review_gates = {
        "0.2": challenger_gate(
            compare(
                args.no_review_t02_tb2,
                args.review_t02_tb2,
                "no_review",
                "review",
            ),
            compare(
                args.no_review_t02_swe,
                args.review_t02_swe,
                "no_review",
                "review",
            ),
            "no_review",
            "review",
        ),
        "0.6": challenger_gate(
            compare(
                args.no_review_t06_tb2,
                args.review_t06_tb2,
                "no_review",
                "review",
            ),
            compare(
                args.no_review_t06_swe,
                args.review_t06_swe,
                "no_review",
                "review",
            ),
            "no_review",
            "review",
        ),
    }
    custom_choice = (
        "review" if review_gates[temperature]["challenger_qualifies"] else "noreview"
    )
    custom_paths = {
        ("0.2", "noreview", "tb2"): args.no_review_t02_tb2,
        ("0.2", "review", "tb2"): args.review_t02_tb2,
        ("0.6", "noreview", "tb2"): args.no_review_t06_tb2,
        ("0.6", "review", "tb2"): args.review_t06_tb2,
        ("0.2", "noreview", "swe"): args.no_review_t02_swe,
        ("0.2", "review", "swe"): args.review_t02_swe,
        ("0.6", "noreview", "swe"): args.no_review_t06_swe,
        ("0.6", "review", "swe"): args.review_t06_swe,
    }
    stock_paths = {
        ("0.2", "tb2"): args.stock_t02_tb2,
        ("0.6", "tb2"): args.stock_t06_tb2,
        ("0.2", "swe"): args.stock_t02_swe,
        ("0.6", "swe"): args.stock_t06_swe,
    }
    # The aligned scaffold is the evidence-backed default because the SFT
    # corpus used its exact prompt/schema. Stock may replace it only on strong,
    # paired evidence of non-inferiority on each suite and a combined advantage.
    scaffold_gate = challenger_gate(
        compare(
            custom_paths[(temperature, custom_choice, "tb2")],
            stock_paths[(temperature, "tb2")],
            "selected_custom",
            "stock",
        ),
        compare(
            custom_paths[(temperature, custom_choice, "swe")],
            stock_paths[(temperature, "swe")],
            "selected_custom",
            "stock",
        ),
        "selected_custom",
        "stock",
    )
    choice = "stock" if scaffold_gate["challenger_qualifies"] else custom_choice
    print(
        json.dumps(
            {
                "temperature": temperature,
                "choice": choice,
                "temperature_decision": temperature_gate,
                "review_decisions": review_gates,
                "scaffold_decision": scaffold_gate,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()