Fix runner resume behaviour, --clean flag, and add --seed

README.md

@@ -141,6 +141,7 @@ or Gemini model and route your request appropriately.
 | Samplers             | `--samplers <names>`        | One or more sampler names (default: All except You.com Research).  |
 | Datasets             | `--datasets <names>`        | One or more datasets (default: `simpleqa`, `frames`).              |
 | Limit                | `--limit <n>`               | Run on at most `n` problems (optional).                            |
+| Seed                 | `--seed <n>`                | Random seed for `--limit` sampling, so two limited runs use the same subset. |
 | Batch size           | `--batch-size 50`           | Number of problems per batch before writing results (default: 50). |
 | Max concurrent tasks | `--max-concurrent-tasks 10` | Concurrency limit (default: 10).                                   |
 | Clean                | `--clean`                   | Remove existing results and run from scratch. (default False)      |

src/evals/eval_runner.py

@@ -52,18 +52,37 @@ def clean_results_folder(results_dir: Path = None):
 
 
 def get_remaining_problems(
-    dataset: datasets.Dataset, sampler: BaseSampler, results_dir: Path = None
+    dataset: datasets.Dataset,
+    sampler: BaseSampler,
+    results_dir: Path = None,
+    problems: pd.DataFrame = None,
 ):
-    """In case of failure, only run problems from the dataset that have not been run yet"""
+    """In case of failure, only run problems from the dataset that have not been run yet.
+
+    Rows previously written as FAILED are treated as not yet run, so a transient
+    error (timeout, rate limit, provider outage) is retried on the next run
+    instead of being permanently excluded from the sampler's results.
+
+    Args:
+        dataset: The dataset being evaluated, used to locate the results file.
+        sampler: The sampler being evaluated.
+        results_dir: Directory holding existing results. Defaults to src/evals/results.
+        problems: Problems to filter. Defaults to dataset.df. Pass explicitly so
+            callers can filter without mutating the shared dataset.
+    """
     if results_dir is None:
         results_dir = get_default_results_dir()
+    if problems is None:
+        problems = dataset.df
     sampler_results_filepath = get_sampler_filepath(sampler, dataset, results_dir)
     if os.path.isdir(results_dir) and os.path.isfile(sampler_results_filepath):
         sampler_results = pd.read_csv(sampler_results_filepath)
-        return dataset.df[
-            ~dataset.df["problem"].isin(sampler_results["query"].tolist())
-        ]
-    return dataset.df
+        completed = sampler_results[
+            (sampler_results["evaluation_result"] != "FAILED")
+            & (sampler_results["generated_answer"] != "FAILED")
+        ]["query"].tolist()
+        return problems[~problems["problem"].isin(completed)]
+    return problems
 
 
 async def process_query_with_semaphore(
@@ -107,12 +126,21 @@ async def run_evals(
         if not dataset:
             raise ValueError(f"Dataset {dataset_name} not found")
         if args.limit:
-            dataset.df = dataset.df.sample(n=args.limit)
+            dataset.df = dataset.df.sample(
+                n=args.limit, random_state=getattr(args, "seed", None)
+            )
+        # Snapshot the problem set once per dataset. Per-sampler filtering below
+        # must not narrow this, or each sampler would inherit the previous
+        # sampler's already-completed set and silently run fewer problems.
+        dataset_problems = dataset.df
         for sampler_name in args.samplers:
             sampler = evals_utils.get_sampler(sampler_name)
             # Only run on problems that are not already in results folder
             remaining_problems = get_remaining_problems(
-                dataset=dataset, sampler=sampler, results_dir=results_dir
+                dataset=dataset,
+                sampler=sampler,
+                results_dir=results_dir,
+                problems=dataset_problems,
             )
             if len(remaining_problems) == 0:
                 logging.info(
@@ -126,18 +154,17 @@ async def run_evals(
             logging.info(
                 f"Running sampler {sampler.sampler_name} on dataset {dataset_name} on {len(remaining_problems)} problems"
             )
-            dataset.df = remaining_problems
 
             with tqdm(
-                total=len(dataset.df),
+                total=len(remaining_problems),
                 desc=f"Running sampler: {sampler.sampler_name} for dataset {dataset.dataset_name}",
                 unit="queries",
             ) as pbar:
                 max_tasks = args.max_concurrent_tasks or sampler.max_concurrency
                 semaphore = asyncio.Semaphore(max_tasks)
                 tasks = []
                 # Create tasks all at once
-                for _, row in dataset.df.iterrows():
+                for _, row in remaining_problems.iterrows():
                     query = row["problem"]
                     ground_truth = row["answer"]
                     task = asyncio.create_task(
@@ -232,6 +259,13 @@ async def main():
         type=int,
         help="Determines the amount of problems to evaluate against",
     )
+    parser.add_argument(
+        "--seed",
+        default=None,
+        type=int,
+        help="Random seed used when --limit samples a subset of problems. Set this to "
+        "make two limited runs comparable; without it each run samples a different subset.",
+    )
     parser.add_argument(
         "--batch-size",
         default=50,
@@ -246,10 +280,9 @@ async def main():
     )
     parser.add_argument(
         "--clean",
-        default=False,
-        type=str,
-        help="If set to True, wipes results folder if it exists to set up a fresh run on all samplers and all problems. If set to False, "
-        "the evaluation is only run for samplers and problems that do not already exist in results folder",
+        action="store_true",
+        help="Wipe the results folder before running, for a fresh run on all samplers and all problems. "
+        "Omit the flag to only run samplers and problems that do not already exist in the results folder.",
     )
 
     args = parser.parse_args()

tests/test_eval_runner.py

@@ -0,0 +1,167 @@
+"""Tests for eval_runner resume behaviour and per-sampler problem isolation.
+
+These tests use a stub sampler and never hit the network, so they run without
+API keys.
+"""
+
+import argparse
+from pathlib import Path
+
+import pandas as pd
+import pytest
+
+from evals.configs import datasets
+from evals.eval_runner import get_remaining_problems, run_evals, get_sampler_filepath
+
+
+def _dataset(tmp_path: Path, problems: list[str]) -> datasets.Dataset:
+    df = pd.DataFrame({"problem": problems, "answer": [f"a-{p}" for p in problems]})
+    csv_path = tmp_path / "fake.csv"
+    df.to_csv(csv_path, index=False)
+    dataset = datasets.Dataset(
+        dataset_name="fake",
+        csv_path=str(csv_path),
+        grader=None,
+        df=df,
+    )
+    return dataset
+
+
+class _StubSampler:
+    """Stands in for a real sampler; records what it was asked to run."""
+
+    def __init__(self, name: str, evaluation_result: str = "is_correct"):
+        self.sampler_name = name
+        self.max_concurrency = 5
+        self.evaluation_result = evaluation_result
+        self.seen: list[str] = []
+
+    async def __call__(self, query, dataset, ground_truth="", overwrite=False):
+        self.seen.append(query)
+        return {
+            "query": query,
+            "internal_response_time_ms": 1.0,
+            "request_response_time_ms": 2.0,
+            "evaluation_result": self.evaluation_result,
+            "generated_answer": f"answer for {query}",
+            "ground_truth": ground_truth,
+        }
+
+
+def _write_results(results_dir: Path, dataset, sampler, rows: list[dict]):
+    results_dir.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame(rows).to_csv(
+        get_sampler_filepath(sampler, dataset, results_dir), index=False
+    )
+
+
+def test_completed_problems_are_skipped(tmp_path):
+    dataset = _dataset(tmp_path, ["q1", "q2", "q3"])
+    sampler = _StubSampler("stub")
+    results_dir = tmp_path / "results"
+    _write_results(
+        results_dir,
+        dataset,
+        sampler,
+        [
+            {
+                "query": "q1",
+                "evaluation_result": "is_correct",
+                "generated_answer": "a",
+            }
+        ],
+    )
+
+    remaining = get_remaining_problems(dataset, sampler, results_dir)
+
+    assert remaining["problem"].tolist() == ["q2", "q3"]
+
+
+def test_failed_problems_are_retried(tmp_path):
+    """A FAILED row must not count as completed, or transient errors become permanent."""
+    dataset = _dataset(tmp_path, ["q1", "q2"])
+    sampler = _StubSampler("stub")
+    results_dir = tmp_path / "results"
+    _write_results(
+        results_dir,
+        dataset,
+        sampler,
+        [
+            {
+                "query": "q1",
+                "evaluation_result": "FAILED",
+                "generated_answer": "FAILED",
+            },
+            {
+                "query": "q2",
+                "evaluation_result": "is_correct",
+                "generated_answer": "a",
+            },
+        ],
+    )
+
+    remaining = get_remaining_problems(dataset, sampler, results_dir)
+
+    assert remaining["problem"].tolist() == ["q1"]
+
+
+def test_explicit_problems_argument_is_not_the_shared_dataframe(tmp_path):
+    """Filtering must read from the passed-in frame, leaving dataset.df untouched."""
+    dataset = _dataset(tmp_path, ["q1", "q2", "q3"])
+    sampler = _StubSampler("stub")
+    results_dir = tmp_path / "results"
+    _write_results(
+        results_dir,
+        dataset,
+        sampler,
+        [{"query": "q1", "evaluation_result": "is_correct", "generated_answer": "a"}],
+    )
+    snapshot = dataset.df
+
+    remaining = get_remaining_problems(
+        dataset, sampler, results_dir, problems=dataset.df
+    )
+
+    assert remaining["problem"].tolist() == ["q2", "q3"]
+    assert dataset.df is snapshot
+    assert dataset.df["problem"].tolist() == ["q1", "q2", "q3"]
+
+
+@pytest.mark.asyncio
+async def test_second_sampler_is_not_narrowed_by_the_first(tmp_path, monkeypatch):
+    """Regression: a resumed run must not shrink the problem set for later samplers.
+
+    sampler_a has already completed q1, so it only needs q2 and q3. sampler_b has
+    no prior results and must still be given all three.
+    """
+    dataset = _dataset(tmp_path, ["q1", "q2", "q3"])
+    sampler_a = _StubSampler("sampler_a")
+    sampler_b = _StubSampler("sampler_b")
+    results_dir = tmp_path / "results"
+
+    _write_results(
+        results_dir,
+        dataset,
+        sampler_a,
+        [{"query": "q1", "evaluation_result": "is_correct", "generated_answer": "a"}],
+    )
+
+    monkeypatch.setattr("evals.eval_runner.evals_utils.get_dataset", lambda _: dataset)
+    monkeypatch.setattr(
+        "evals.eval_runner.evals_utils.get_sampler",
+        lambda name: {"sampler_a": sampler_a, "sampler_b": sampler_b}[name],
+    )
+
+    args = argparse.Namespace(
+        samplers=["sampler_a", "sampler_b"],
+        datasets=["fake"],
+        limit=None,
+        seed=None,
+        batch_size=10,
+        max_concurrent_tasks=5,
+        clean=False,
+    )
+    await run_evals(args, results_dir=results_dir)
+
+    assert sorted(sampler_a.seen) == ["q2", "q3"]
+    assert sorted(sampler_b.seen) == ["q1", "q2", "q3"]