code/local_eval/run.py

from __future__ import annotations

import json
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path
from statistics import mean
from uuid import uuid4

from .constants import (
    DEFAULT_GPU_MEMORY_UTILIZATION,
    DEFAULT_LOCAL_TURNS,
    DEFAULT_MODEL_DIR,
    DEFAULT_RUNS_DIR,
    WIN_MARGIN,
)
from .gates import evaluate_side
from .rollout import build_generator, generate_side
from .samples import load_samples


def run_duel(
    *,
    challenger: Path,
    king: Path = DEFAULT_MODEL_DIR,
    dataset_root: Path,
    sample_count: int = 8,
    seed: str = "local-eval",
    sample_ids: list[str] | None = None,
    max_turns: int = DEFAULT_LOCAL_TURNS,
    king_gpus: list[str] | None = None,
    chal_gpus: list[str] | None = None,
    max_model_len: int = 65536,
    runs_dir: Path = DEFAULT_RUNS_DIR,
    skip_king: bool = False,
    enforce_eager: bool = True,
    gpu_memory_utilization: float = DEFAULT_GPU_MEMORY_UTILIZATION,
) -> dict:
    samples = load_samples(
        dataset_root, sample_ids=sample_ids, sample_count=sample_count, seed=seed
    )
    king_gpus = king_gpus or ["0", "1", "2", "3"]
    chal_gpus = chal_gpus or ["4", "5", "6", "7"]
    run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid4().hex[:8]
    out = Path(runs_dir) / run_id
    out.mkdir(parents=True, exist_ok=True)

    def _roll(model: Path, gpus: list[str]):
        generator = build_generator(
            str(model),
            gpus,
            max_model_len=max_model_len,
            enforce_eager=enforce_eager,
            gpu_memory_utilization=gpu_memory_utilization,
        )
        return generate_side(
            generator=generator,
            samples=samples,
            dataset_root=dataset_root,
            max_turns=max_turns,
        )

    print(
        f"local duel run={run_id} samples={len(samples)} turns<={max_turns} "
        f"king={king} chal={challenger}",
        flush=True,
    )
    if skip_king:
        king_results = []
        chal_results = _roll(challenger, chal_gpus)
    else:
        with ThreadPoolExecutor(max_workers=2) as pool:
            king_future = pool.submit(_roll, king, king_gpus)
            chal_future = pool.submit(_roll, challenger, chal_gpus)
            king_results = king_future.result()
            chal_results = chal_future.result()

    king_by_id = {r.sample_id: r for r in king_results}
    chal_by_id = {r.sample_id: r for r in chal_results}
    rows = []
    for sample in samples:
        king_r = king_by_id.get(sample.sample_id)
        chal_r = chal_by_id.get(sample.sample_id)
        king_gate = (
            evaluate_side(sample, king_r.text if king_r else "", king_r.turns if king_r else None, truncated=bool(king_r and king_r.truncated))
            if king_r
            else None
        )
        chal_gate = evaluate_side(
            sample,
            chal_r.text if chal_r else "",
            chal_r.turns if chal_r else None,
            truncated=bool(chal_r and chal_r.truncated),
        )
        rows.append(
            {
                "sample_id": sample.sample_id,
                "submit_command": sample.submit_command,
                "submit_marker": sample.submit_marker,
                "rewrite_mode": sample.rewrite_mode,
                "king": king_gate.as_dict() if king_gate else None,
                "challenger": chal_gate.as_dict(),
                "king_error": king_r.error if king_r else "skipped",
                "chal_error": chal_r.error if chal_r else "missing",
                "previous_king_output": king_r.text if king_r else "",
                "challenger_output": chal_r.text if chal_r else "",
            }
        )

    chal_scores = [r["challenger"]["proxy_score"] for r in rows]
    king_scores = [r["king"]["proxy_score"] for r in rows if r["king"]]
    chal_pass = sum(1 for r in rows if r["challenger"]["passed"])
    verdict = {
        "run_id": run_id,
        "mode": "local_gates_proxy",
        "note": (
            "proxy_score is a pre-eval / behaviour stand-in, not the official GLM checklist. "
            "A model that fails these gates will not reach the live duel."
        ),
        "sample_count": len(rows),
        "max_turns": max_turns,
        "seed": seed,
        "king_model": str(king),
        "challenger_model": str(challenger),
        "score_challenger": round(mean(chal_scores), 6) if chal_scores else None,
        "score_king": round(mean(king_scores), 6) if king_scores else None,
        "challenger_gate_pass_rate": round(chal_pass / len(rows), 4) if rows else 0.0,
        "required_win_margin": WIN_MARGIN,
        "sample_ids": [s.sample_id for s in samples],
    }
    if verdict["score_challenger"] is not None and verdict["score_king"] is not None:
        delta = verdict["score_challenger"] - verdict["score_king"]
        verdict["win_margin"] = round(delta, 6)
        verdict["challenger_won_proxy"] = delta >= WIN_MARGIN
    (out / "verdict.json").write_text(json.dumps(verdict, indent=2) + "\n")
    with (out / "gates.jsonl").open("w") as handle:
        for row in rows:
            handle.write(
                json.dumps({k: row[k] for k in row if k not in {"previous_king_output", "challenger_output"}})
                + "\n"
            )
    with (out / "generated-samples.jsonl").open("w") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(json.dumps({k: verdict[k] for k in verdict if k != "sample_ids"}, indent=2), flush=True)
    print(f"artifacts: {out}", flush=True)
    return verdict