code/local_train/policy_eval.py

"""One-turn eval on live-fault prefixes: next bash must be sed -i, never a marker."""

from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4

from albedo_eval_service.remote.dataset import EvalSample, format_messages
from albedo_eval_service.shared.observation_format import first_bash_block
from local_eval.constants import DEFAULT_RUNS_DIR, TOKENIZER_DIR
from local_eval.live_protocol import generate_retrying_bad_turns, is_live_submit
from local_eval.rollout import build_generator

from .chain_eval import pick_free_gpus
from .constants import DEFAULT_DPO_EXPORT_DIR
from .dpo_pack import iter_prefixes
from .pack import is_edit_command


def run_policy(
    *,
    challenger: Path = DEFAULT_DPO_EXPORT_DIR,
    runs_dir: Path = DEFAULT_RUNS_DIR,
    gpu_ids: list[str] | None = None,
) -> dict:
    from local_eval.cuda_env import apply as apply_cuda

    apply_cuda()
    prefixes = iter_prefixes()
    samples = []
    for prefix in prefixes:
        prompt = format_messages(
            prefix.messages, tokenizer_path=str(TOKENIZER_DIR), enable_thinking=True
        )
        samples.append(
            EvalSample(
                sample_id=prefix.sample_id,
                prompt=prompt,
                messages=list(prefix.messages),
                submit_command=prefix.submit_command,
                submit_marker=prefix.submit_marker,
            )
        )
    gpu_ids = (gpu_ids or pick_free_gpus())[:4]
    generator = build_generator(str(challenger), gpu_ids)
    run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid4().hex[:8]
    out = Path(runs_dir) / f"{run_id}-policy"
    out.mkdir(parents=True, exist_ok=True)
    print(f"policy run={run_id} n={len(samples)} chal={challenger}", flush=True)
    try:
        results = generate_retrying_bad_turns(generator, samples)
    finally:
        generator.close()
    by_id = {prefix.sample_id: prefix for prefix in prefixes}
    rows = []
    for result, sample in zip(results, samples, strict=True):
        prefix = by_id[sample.sample_id]
        bash = first_bash_block(result.text or "") or ""
        submitted = is_live_submit(
            result.text or "", command=prefix.submit_command, marker=prefix.submit_marker
        )
        edited = is_edit_command(bash)
        passed = edited and not submitted and "sed -i" in bash
        rows.append(
            {
                "sample_id": sample.sample_id,
                "kind": prefix.kind,
                "passed": passed,
                "edited": edited,
                "submitted": submitted,
                "bash": bash[:200],
                "error": result.error,
            }
        )
    passed = all(r["passed"] for r in rows) and bool(rows)
    report = {
        "run_id": run_id,
        "go": passed,
        "challenger": str(challenger),
        "n": len(rows),
        "pass_rate": round(sum(1 for r in rows if r["passed"]) / max(len(rows), 1), 4),
        "by_kind": {
            kind: round(
                sum(1 for r in rows if r["kind"] == kind and r["passed"])
                / max(sum(1 for r in rows if r["kind"] == kind), 1),
                4,
            )
            for kind in sorted({r["kind"] for r in rows})
        },
        "samples": rows,
    }
    (out / "policy-report.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2), flush=True)
    print(f"artifacts: {out}", flush=True)
    return report