runner.py

"""Run a model over the suite and score it.

pass@k uses the unbiased estimator from Chen et al. (2021), "Evaluating Large
Language Models Trained on Code" -- not the naive "did any of k samples pass",
which is biased upward and not comparable across different n.
"""

from __future__ import annotations

import json
import time
from collections import defaultdict
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from pyspark.sql import SparkSession

from .harness import ExecResult, evaluate_candidate
from .models import DummyModel, Model, ModelError
from .prompting import build_prompt, extract_code
from .schema import Task


def pass_at_k(n: int, c: int, k: int) -> float:
    """Probability that at least one of k samples drawn from n passes.

    n = samples generated, c = samples that passed.
    """
    if n < k:
        raise ValueError(f"cannot estimate pass@{k} from only {n} samples")
    if n - c < k:
        return 1.0
    # Product form avoids overflow in the binomial coefficients.
    prob = 1.0
    for i in range(k):
        prob *= (n - c - i) / (n - i)
    return 1.0 - prob


@dataclass
class SampleRecord:
    task_id: str
    category: str
    difficulty: str
    sample_index: int
    passed: bool
    status: str
    detail: str
    raw_response: str
    extracted_code: str
    latency_s: float


@dataclass
class TaskRecord:
    task_id: str
    category: str
    difficulty: str
    probes: str
    n: int
    c: int
    statuses: dict[str, int] = field(default_factory=dict)


@dataclass
class RunReport:
    model: str
    n_samples: int
    temperature: float
    started_at: str
    duration_s: float
    n_tasks: int
    pass_at_1: float
    pass_at_k: dict[str, float]
    by_category: dict[str, dict]
    by_difficulty: dict[str, dict]
    failure_modes: dict[str, int]
    tasks: list[TaskRecord]
    samples: list[SampleRecord]

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=2)


def get_spark(app_name: str = "spark-eval") -> SparkSession:
    """A small, deterministic, local Spark session.

    Single shuffle partition is deliberate: it makes float aggregation order
    reproducible and cuts per-task overhead by more than half. Tasks are tiny;
    parallelism buys nothing here.
    """
    return (
        SparkSession.builder.appName(app_name)
        .master("local[2]")
        .config("spark.sql.shuffle.partitions", "1")
        .config("spark.default.parallelism", "2")
        .config("spark.sql.adaptive.enabled", "false")
        .config("spark.ui.enabled", "false")
        .config("spark.sql.session.timeZone", "UTC")
        .config("spark.driver.memory", "2g")
        .getOrCreate()
    )


def run_suite(
    tasks: list[Task],
    model: Model,
    *,
    n_samples: int = 1,
    ks: tuple[int, ...] = (1,),
    temperature: float = 0.2,
    max_tokens: int = 1024,
    timeout: int = 60,
    spark: SparkSession | None = None,
    keep_responses: bool = True,
    progress: bool = True,
) -> RunReport:
    owns_spark = spark is None
    spark = spark or get_spark()
    spark.sparkContext.setLogLevel("ERROR")

    started = datetime.now(timezone.utc)
    t0 = time.time()

    samples: list[SampleRecord] = []
    task_records: list[TaskRecord] = []
    failure_modes: dict[str, int] = defaultdict(int)

    for idx, task in enumerate(tasks, 1):
        prompt = build_prompt(task)
        # The dummy backend needs the gold answer keyed by the exact prompt.
        if isinstance(model, DummyModel):
            model.register(prompt, task.solution)

        statuses: dict[str, int] = defaultdict(int)
        passed_count = 0

        for s in range(n_samples):
            s_t0 = time.time()
            try:
                raw = model.generate(prompt, temperature, max_tokens)
            except ModelError as exc:
                raw = ""
                result = ExecResult(False, "model_error", str(exc))
                code = ""
            else:
                code = extract_code(raw)
                if not code.strip():
                    result = ExecResult(False, "empty_response", "no code in response")
                else:
                    result = evaluate_candidate(spark, task, code, timeout=timeout)

            latency = time.time() - s_t0
            statuses[result.status] += 1
            if result.ok:
                passed_count += 1
            else:
                failure_modes[result.status] += 1

            samples.append(
                SampleRecord(
                    task_id=task.id,
                    category=task.category,
                    difficulty=task.difficulty,
                    sample_index=s,
                    passed=result.ok,
                    status=result.status,
                    detail=result.detail,
                    raw_response=raw if keep_responses else "",
                    extracted_code=code if keep_responses else "",
                    latency_s=round(latency, 3),
                )
            )

            # A broken reference solution means the benchmark is lying. Stop.
            if result.status == "reference_broken":
                raise RuntimeError(result.detail)

        task_records.append(
            TaskRecord(
                task_id=task.id,
                category=task.category,
                difficulty=task.difficulty,
                probes=task.probes,
                n=n_samples,
                c=passed_count,
                statuses=dict(statuses),
            )
        )

        if progress:
            mark = "PASS" if passed_count == n_samples else (
                "FAIL" if passed_count == 0 else f"{passed_count}/{n_samples}"
            )
            print(
                f"[{idx:>3}/{len(tasks)}] {task.id:<40} {mark}",
                flush=True,
            )

    def _agg(records: list[TaskRecord]) -> dict:
        if not records:
            return {"n_tasks": 0, "pass_at_1": 0.0}
        out = {
            "n_tasks": len(records),
            "pass_at_1": round(
                sum(pass_at_k(r.n, r.c, 1) for r in records) / len(records), 4
            ),
        }
        for k in ks:
            if k > 1 and n_samples >= k:
                out[f"pass_at_{k}"] = round(
                    sum(pass_at_k(r.n, r.c, k) for r in records) / len(records), 4
                )
        return out

    by_category: dict[str, dict] = {}
    for cat in sorted({r.category for r in task_records}):
        by_category[cat] = _agg([r for r in task_records if r.category == cat])

    by_difficulty: dict[str, dict] = {}
    for diff in ("easy", "medium", "hard"):
        subset = [r for r in task_records if r.difficulty == diff]
        if subset:
            by_difficulty[diff] = _agg(subset)

    overall = _agg(task_records)
    report = RunReport(
        model=model.name,
        n_samples=n_samples,
        temperature=temperature,
        started_at=started.isoformat(),
        duration_s=round(time.time() - t0, 2),
        n_tasks=len(task_records),
        pass_at_1=overall["pass_at_1"],
        pass_at_k={
            f"pass_at_{k}": overall[f"pass_at_{k}"]
            for k in ks
            if k > 1 and f"pass_at_{k}" in overall
        },
        by_category=by_category,
        by_difficulty=by_difficulty,
        failure_modes=dict(sorted(failure_modes.items(), key=lambda kv: -kv[1])),
        tasks=task_records,
        samples=samples,
    )

    if owns_spark:
        spark.stop()

    return report


def format_report(report: RunReport) -> str:
    """Human-readable summary. The JSON is the machine-readable artifact."""
    lines = [
        "",
        "=" * 68,
        f"  spark-eval  |  {report.model}",
        "=" * 68,
        f"  tasks       {report.n_tasks}",
        f"  samples     {report.n_samples} per task @ temperature {report.temperature}",
        f"  duration    {report.duration_s}s",
        "",
        f"  pass@1      {report.pass_at_1:.1%}",
    ]
    for key, v in report.pass_at_k.items():
        label = "pass@" + key.rsplit("_", 1)[-1]
        lines.append(f"  {label:<11} {v:.1%}")

    def _plural(n: int) -> str:
        return f"{n} task" if n == 1 else f"{n} tasks"

    # Worst category first: the point of the breakdown is finding the weakness.
    lines += ["", "  By category", "  " + "-" * 46]
    for cat, stats in sorted(
        report.by_category.items(), key=lambda kv: kv[1]["pass_at_1"]
    ):
        lines.append(
            f"    {cat:<22} {stats['pass_at_1']:>6.1%}  ({_plural(stats['n_tasks'])})"
        )

    if report.by_difficulty:
        lines += ["", "  By difficulty", "  " + "-" * 46]
        for diff, stats in report.by_difficulty.items():
            lines.append(
                f"    {diff:<22} {stats['pass_at_1']:>6.1%}  ({_plural(stats['n_tasks'])})"
            )

    if report.failure_modes:
        lines += ["", "  Failure modes", "  " + "-" * 46]
        total = sum(report.failure_modes.values())
        for mode, count in report.failure_modes.items():
            lines.append(f"    {mode:<22} {count:>4}  ({count / total:.0%})")

    lines += ["", "=" * 68, ""]
    return "\n".join(lines)


def write_report(report: RunReport, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(report.to_json())