From 9a7c742bb30721a71642af8bad647003cfe8c449 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 18:47:58 +0000 Subject: [PATCH 1/7] Add SampleSource for dynamic sample generation within a task MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the TaskSource pattern at the sample level: pass a SampleSource as a Task's `dataset` to drive the task from code — a synchronous seed (`initial_samples`, may be empty), per-result follow-ups returned from `sample_complete`, and a blocking `next_samples()` pull that ends the task when it returns None. Useful for RL loops and adaptive evals. - Live dispatcher in task_run (mirrors run_multiple's feed loop): injected samples start immediately on free capacity; next_samples() is only awaited when fully idle (no lost wakeups) - enqueue_sample() imperative primitive (ContextVar-scoped per task), mirroring enqueue_task; requires a SampleSource-driven task - Auto-ids for injected samples continue the seed's numbering (str-form uniqueness, matching ensure_unique_ids); duplicate ids are an error - Planned totals grow as samples are added: display denominator, fractional fail_on_error threshold, and control-channel counters (new record_samples_added in eval_state) - Injected samples honor the task's epochs, like the seed - Wake moved from _eval/run.py to _util/_async.py for reuse Fixes meridianlabs-ai/inspect_ai#36 Co-authored-by: Ransom Richardson Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 1 + design/sample-source.md | 130 +++++++++ docs/_quarto.yml | 1 + docs/reference/inspect_ai.qmd | 2 + docs/sample-source.qmd | 125 ++++++++ src/inspect_ai/__init__.py | 4 + src/inspect_ai/_control/eval_state.py | 27 +- src/inspect_ai/_display/core/progress.py | 4 +- src/inspect_ai/_eval/run.py | 24 +- src/inspect_ai/_eval/task/__init__.py | 2 + src/inspect_ai/_eval/task/run.py | 192 ++++++++++++- src/inspect_ai/_eval/task/sample_source.py | 221 ++++++++++++++ src/inspect_ai/_eval/task/task.py | 52 +++- src/inspect_ai/_util/_async.py | 19 ++ tests/test_sample_source.py | 320 +++++++++++++++++++++ 15 files changed, 1081 insertions(+), 43 deletions(-) create mode 100644 design/sample-source.md create mode 100644 docs/sample-source.qmd create mode 100644 src/inspect_ai/_eval/task/sample_source.py create mode 100644 tests/test_sample_source.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d1dbd6d4b1..d33a047722 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## Unreleased +- Sample Sources: Generate a task's samples dynamically while it runs by passing a `SampleSource` as the task's `dataset` (with `enqueue_sample()` for imperative additions) — the sample-level mirror of `TaskSource`, for RL loops and adaptive evals. - Control Channel: Added `inspect ctl limits` to view or retune a running eval's `max_samples` / `max_sandboxes` / `max_connections` concurrency limits mid-flight (with `--dry-run`). - vLLM: Keep the connection-pool/adaptive-concurrency scope stable across lazy server startup instead of splitting it on the first generate. - Bugfix: Sample concurrency now honors model-level `max_connections` / `adaptive_connections` settings instead of classifying the adaptive-vs-static path from task-level config alone. diff --git a/design/sample-source.md b/design/sample-source.md new file mode 100644 index 0000000000..2b477c26d7 --- /dev/null +++ b/design/sample-source.md @@ -0,0 +1,130 @@ +# SampleSource: driving a running task from code + +This document describes `SampleSource` — the sample-level mirror of +[`TaskSource`](task-source.md). Where a `TaskSource` produces *tasks* for a +running eval, a `SampleSource` produces *samples* for a running task: a seed +plus result-driven follow-ups, all within one task / eval log. Motivating use +cases are the same (RL loops, adaptive evals), scoped inside a single task. + +## The contract + +`src/inspect_ai/_eval/task/sample_source.py` defines a subclassable base class +with no-op defaults: + +```python +class SampleSource: + def initial_samples(self) -> list[Sample]: ... # sync seed (immediate; may be empty) + async def next_samples(self) -> list[Sample] | None: ... # async; None ends the task + async def sample_complete(self, sample: EvalSample) -> list[Sample] | None: ... # observe + add +``` + +The asymmetry mirrors `TaskSource`: + +- **`initial_samples()` is synchronous** — it is called when the `Task` is + constructed (see below) and becomes the task's up-front dataset, driving + validation / sandbox startup / display totals. Unlike a plain dataset it + **may be empty**: the task then starts by calling `next_samples()`. +- **`next_samples()` is async and may block indefinitely.** It is called only + when the task is fully idle (nothing in flight, nothing buffered), so no + completion can enqueue while it blocks (no lost wakeup). Returning `None` + ends the task. +- **`sample_complete()`** fires per sample (right after the sample is logged / + `emit_sample_end`, alongside the `TaskSource.sample_complete` firing site in + `task_run_sample`) and may **return follow-up samples**, which are routed + onto the task's `SampleEnqueuer` — the same buffer `enqueue_sample` feeds. + +`SampleSource.from_samples(initial_samples, *, next_samples=None, +sample_complete=None)` builds a source from a seed + callbacks without +subclassing (mirrors `TaskSource.from_tasks`). + +There is no `@sample_source` decorator: a `SampleSource` lives *inside* a +`Task`, and tasks are already registerable / loadable by name via `@task` — the +source needs no registry identity of its own. + +## Where it plugs in + +A `SampleSource` is passed as the **`dataset` argument to `Task`** (just as a +`TaskSource` is passed as the `tasks` argument to `eval()`): + +- **Task resolution** — `resolve_dataset_or_source` + ([task.py](../src/inspect_ai/_eval/task/task.py)) returns a `ResolvedDataset` + NamedTuple: `task.dataset` is a `MemoryDataset` of the seed (empty allowed — + the plain-dataset empty check is bypassed) and `task.sample_source` carries + the source. Everything downstream that reads `task.dataset` (log spec, + sandbox startup, auto-id assignment of the seed in `prepare_options`) is + unchanged. +- **No eval-level changes** — the source travels on the `Task` itself; `eval()` + / `eval_run` / `TaskRunOptions` needed no new parameter. (`TaskRunOptions` + already has a `sample_source` field — that is `EvalSampleSource`, the + prior-attempt lookup used by retries. Inside `task_run` the SampleSource is + therefore held in a local named `sample_feed` to avoid confusion.) +- **The dispatcher** — in [task/run.py](../src/inspect_ai/_eval/task/run.py), + `task_run` replaces the fixed `tg_collect` fan-out with a live loop + (`run_samples_dynamic`) when `task.sample_source` is set. It mirrors + `run_multiple`'s feed loop: spawn the seed, then per cycle drain the + enqueuer (spawning injected samples immediately — live, not batched), wait + on a re-armable `Wake` (moved to `_util/_async.py`, shared with the + eval-level dispatchers) while anything is in flight, and only when fully + idle await the blocking `next_samples()`. Sample concurrency is already + bounded by the sample semaphore inside `run_sample`, so the dispatcher + spawns eagerly with no cap of its own. Plain tasks keep the unchanged + `tg_collect` path. +- **Injected sample storage** — injected samples are appended to an in-memory + list indexed after the (possibly disk-paged) seed store; `get_sample(index)` + dispatches between the two. Each injected sample runs for the task's + configured `epochs`, like the seed. +- **Ids** — injected samples without ids get auto-ids continuing the seed's + 1-based numbering, skipping ids already in use; a duplicate explicit id is a + hard error. Ids are compared by their `str()` form, matching + `ensure_unique_ids` (log member names and score grouping key on it). +- **Growing totals** — `total_samples` grows as samples are added, updating the + display denominator (`td.sample_complete`), the fractional `fail_on_error` + threshold (`SampleErrorHandler.total_samples`), the end-of-run + `eval_results` / `_should_eval_fail` counts, and the control-channel state + via `record_samples_added` ([eval_state.py](../src/inspect_ai/_control/eval_state.py)) + — which also un-stamps a provisional `completed_at` (every planned sample + can be terminal while the source is still producing). `register_eval` now + copies its `sample_ids` argument so the state's planned-ids list grows only + via `record_samples_added` (task_run mutates its own list for + carry-forward). + +## `enqueue_sample` + +The imperative primitive (mirrors `enqueue_task`): `enqueue_sample(samples)` +adds samples to the running task from *arbitrary* code (a solver, a scorer, a +tool). Backed by a `ContextVar`-scoped `SampleEnqueuer` registered by +`task_run` for the task's lifetime — per-task, so concurrent tasks in one +process can't cross-enqueue. Unlike `enqueue_task` (which works on plain runs +as a follow-up batch), `enqueue_sample` **requires** the running task to be +SampleSource-driven: a plain task's fixed `tg_collect` path has no loop to run +additions, so the call raises a clear `RuntimeError` instead of silently +dropping samples. + +## Interactions / known limits + +- **`--limit` / `--sample-id`** slice the seed only; samples produced while the + task runs are not subject to them (documented). +- **Task retries** (`task_retry_attempts` / eval-set): the retry attempt + re-drives the source; completed samples are reused via the normal + `EvalSampleSource` lookup only where regenerated ids match — i.e. resume + pays off for deterministic sources, the same determinism contract as + `TaskSource` + eval_set (see task-source.md). +- **Early stopping** managers receive only the seed at `start_task`; + `schedule_sample` / `complete_sample` still fire for injected samples. +- **Sandboxes**: the task-level sandbox startup pass runs once, up front, over + the seed. Injected samples get per-sample sandboxes via `sandboxenv_context` + as usual, but a *sample-level* sandbox spec appearing only on injected + samples won't have had its task-level `task_init` startup pass. +- **Progress bar steps** (`profile.steps`) are fixed at seed size; the + completed/total counter grows correctly (total passed on each update), and a + zero-step seed no longer divides by zero (`RichProgress.update` guards it). + +## Tests + +`tests/test_sample_source.py` — generations under one task/log with auto-ids; +seed-only when `next_samples()` is `None`; `sample_complete` returning +follow-ups chains generations; `from_samples` (callbacks and seed-only); empty +seed; epochs applied to injected samples; explicit + auto id assignment and +duplicate-id error; live injection discriminated from batch-at-a-time (blocker +parks until an injected sample releases it, `fail_after` bounds a regression); +`enqueue_sample` rejected on plain tasks and outside a task. diff --git a/docs/_quarto.yml b/docs/_quarto.yml index 9f4b889b7b..4c675a54f3 100644 --- a/docs/_quarto.yml +++ b/docs/_quarto.yml @@ -120,6 +120,7 @@ inspect-docs: - href: control-channel.qmd - href: early-stopping.qmd - href: task-source.qmd + - href: sample-source.qmd - href: tracing.qmd - text: Analysis href: analysis.qmd diff --git a/docs/reference/inspect_ai.qmd b/docs/reference/inspect_ai.qmd index e8b5b8344f..b0c4a7c15b 100644 --- a/docs/reference/inspect_ai.qmd +++ b/docs/reference/inspect_ai.qmd @@ -18,7 +18,9 @@ description: Tasks, evaluation, and scoring. ### TaskInfo ### Tasks ### TaskSource +### SampleSource ### enqueue_task +### enqueue_sample ## Scanning diff --git a/docs/sample-source.qmd b/docs/sample-source.qmd new file mode 100644 index 0000000000..5a389f42f6 --- /dev/null +++ b/docs/sample-source.qmd @@ -0,0 +1,125 @@ +--- +title: Sample Sources +llms-description: Generate a task's samples dynamically from code — a seed plus result-driven follow-ups — while the task runs. +--- + +::: {.callout-note appearance="simple"} +Sample sources require the development version of Inspect, which you can install from GitHub: + +```bash +pip install git+https://github.com/UKGovernmentBEIS/inspect_ai +``` +::: + +## Overview + +The `dataset` argument to `Task` is normally static: you pass a `Dataset` (or list of samples) and the task runs exactly those. A `SampleSource` generates samples dynamically instead — a seed plus follow-ups that depend on results — all within one task and one log. + +Use it when the next samples to run depend on the results of the previous ones: + +- Reinforcement-learning loops that generate the next samples from scores. +- Adaptive evaluation that branches on model performance (e.g. escalate difficulty until failure). +- Open-ended generation that runs until some external condition stops it. + +A `SampleSource` is just a value the `dataset` parameter accepts, so there is no separate argument. It is the sample-level mirror of [Task Sources](task-source.qmd): use a `TaskSource` to generate whole tasks across a run, a `SampleSource` to generate samples within a task. + +## Defining a source + +Subclass `SampleSource` and override the methods you need (the defaults are no-ops): + +```python +from inspect_ai import SampleSource +from inspect_ai.dataset import Sample +from inspect_ai.log import EvalSample + + +class MySource(SampleSource): + def initial_samples(self) -> list[Sample]: + """Seed samples to run first (synchronous; may be empty).""" + ... + + async def next_samples(self) -> list[Sample] | None: + """More samples, or None when the task is complete.""" + ... + + async def sample_complete(self, sample: EvalSample) -> list[Sample] | None: + """Observe a finished sample; optionally return follow-up samples.""" + ... +``` + +Pass an instance as the task's `dataset`: + +```python +from inspect_ai import Task, eval + +eval(Task(dataset=MySource(), solver=my_solver(), scorer=my_scorer()), + model="openai/gpt-4o") +``` + +`initial_samples()` is synchronous and returns the seed, so it must return immediately rather than `await`. It may be empty, in which case the task starts by calling `next_samples()`. `next_samples()` is async, called whenever no samples remain in flight, and may block (for example, awaiting external input); return `None` to end the task. + +All generated samples run within the task — one log file, one set of results — and each runs for the task's configured number of [epochs](options.qmd#epochs). Samples without an `id` are assigned one automatically, continuing the seed's numbering. + +## Returning follow-up samples + +`sample_complete` fires as each sample completes. Besides observing the result, it can return samples to add to the task, which start as soon as there is free capacity: + +```python +class Adaptive(SampleSource): + def initial_samples(self) -> list[Sample]: + return [make_sample(difficulty=1)] + + async def sample_complete(self, sample: EvalSample) -> list[Sample] | None: + # escalate difficulty while the model keeps passing + difficulty = sample.metadata["difficulty"] + if passed(sample) and difficulty < 10: + return [make_sample(difficulty=difficulty + 1)] + return None +``` + +A source that returns follow-ups from this callback needs no `next_samples()`: the task ends when completions return nothing and `next_samples()` returns `None`. Use `next_samples()` for the blocking case a per-result callback can't express. + +## Sources from callbacks + +`SampleSource.from_samples()` builds a source from a seed and optional callbacks, without subclassing: + +```python +from inspect_ai import SampleSource + +scores: list[float] = [] + +async def on_sample(sample): + scores.append(score_value(sample)) + return [harder_sample()] if sum(scores) / len(scores) >= 0.8 else None + +source = SampleSource.from_samples([easy_sample()], sample_complete=on_sample) +``` + +`from_samples(initial_samples, *, next_samples=None, sample_complete=None)` delegates to the callables. Omitting `next_samples` and returning nothing from the callback stops after the seed. + +## Adding samples imperatively + +`enqueue_sample()` adds samples to the running task from any code — a solver, scorer, or tool — not only the `SampleSource` itself: + +```python +from inspect_ai import enqueue_sample +from inspect_ai.dataset import Sample +from inspect_ai.solver import Generate, TaskState, solver + + +@solver +def spawn_followup(): + async def solve(state: TaskState, generate: Generate) -> TaskState: + enqueue_sample(Sample(input=followup_prompt(state))) + return state + + return solve +``` + +Enqueued samples share a buffer with samples returned from `sample_complete`. `enqueue_sample()` is only available inside a task driven by a `SampleSource` (a plain task's sample set is fixed) and raises otherwise. + +## Concurrency + +A `SampleSource` task is live: a sample added mid-run starts as soon as there is free capacity, rather than waiting for the current samples to finish. Concurrency is bounded by `max_samples` as usual (see [Parallelism](parallelism.qmd)). + +The `--limit` and `--sample-id` options apply to the seed only — samples the source produces while the task runs are not subject to them. diff --git a/src/inspect_ai/__init__.py b/src/inspect_ai/__init__.py index 83106a432c..42dd8a4279 100644 --- a/src/inspect_ai/__init__.py +++ b/src/inspect_ai/__init__.py @@ -9,12 +9,14 @@ from inspect_ai._eval.score import score, score_async from inspect_ai._eval.task import ( Epochs, + SampleSource, Task, TaskInfo, TaskSource, task_with, ) from inspect_ai._eval.task.enqueue import enqueue_task +from inspect_ai._eval.task.sample_source import enqueue_sample from inspect_ai._eval.task.scan import ScannerConfig, Scanners from inspect_ai._eval.task.tasks import Tasks from inspect_ai._util.constants import PKG_NAME @@ -42,6 +44,7 @@ "Epochs", "Scanners", "ScannerConfig", + "SampleSource", "Task", "Tasks", "TaskInfo", @@ -49,6 +52,7 @@ "task", "task_source", "task_with", + "enqueue_sample", "enqueue_task", "view", ] diff --git a/src/inspect_ai/_control/eval_state.py b/src/inspect_ai/_control/eval_state.py index 6b987155bd..57d43a61db 100644 --- a/src/inspect_ai/_control/eval_state.py +++ b/src/inspect_ai/_control/eval_state.py @@ -341,7 +341,10 @@ def register_eval( model=model, log_location=log_location, live=live, - sample_ids=sample_ids or [], + # copy: a SampleSource-driven task appends to its planned-ids list + # as samples are injected; the state's list must grow only via + # record_samples_added, not by aliasing the caller's list + sample_ids=list(sample_ids or []), epochs=epochs, run_id=run_id, will_retry=will_retry, @@ -515,6 +518,28 @@ def record_sample_cancelled( _maybe_mark_finished(state) +def record_samples_added( + eval_id: str, total: int, *, sample_ids: list[str | int] | None = None +) -> None: + """Grow a running eval's planned totals when samples are added dynamically. + + Called by ``task_run`` when a ``SampleSource`` injects samples mid-run: + ``total`` is the number of additional planned runs (samples × epochs) and + ``sample_ids`` the injected ids (so the per-sample listing can surface them + as pending). If the eval had already been provisionally marked finished + (every previously-planned sample terminal while the source was still + producing), the finish stamp is cleared. No-ops if unregistered. + """ + with _lock: + state = _eval_states.get(eval_id) + if state is not None: + state.total += total + if sample_ids: + state.sample_ids.extend(sample_ids) + if not state.is_finished: + state.completed_at = None + + def task_registered(task_id: str) -> bool: """True if any attempt of ``task_id`` is tracked in this process. diff --git a/src/inspect_ai/_display/core/progress.py b/src/inspect_ai/_display/core/progress.py index f4d59fcb8e..4a2f60fead 100644 --- a/src/inspect_ai/_display/core/progress.py +++ b/src/inspect_ai/_display/core/progress.py @@ -52,7 +52,9 @@ def __init__( @override def update(self, n: int = 1) -> None: - advance = (float(n) / float(self.total)) * 100 + # a SampleSource-driven task can start with zero planned steps + # (empty seed) and still make progress as samples are injected + advance = (float(n) / float(self.total)) * 100 if self.total > 0 else 0.0 self.progress.update( task_id=self.task_id, advance=advance, refresh=True, status=self.status() ) diff --git a/src/inspect_ai/_eval/run.py b/src/inspect_ai/_eval/run.py index e3faf1408b..08a53a8609 100644 --- a/src/inspect_ai/_eval/run.py +++ b/src/inspect_ai/_eval/run.py @@ -5,6 +5,7 @@ from typing import Any, Awaitable, Callable, NamedTuple, Set, cast from inspect_ai._eval.task.constants import TASK_ALL_PARAMS_ATTR +from inspect_ai._util._async import Wake from inspect_ai._util.environ import environ_vars from inspect_ai._util.file import cleanup_s3_sessions from inspect_ai._util.task import task_display_name @@ -416,25 +417,6 @@ async def feed_next() -> list[TaskRunOptions] | None: log.warning(f"Error cleaning up S3 sessions: {exception_message(ex)}") -class _Wake: - """One-shot wake signal that can be re-armed (set on completion / injection). - - Safe under cooperative scheduling: the only await is on ``wait()``; the - re-arm assignment afterwards runs without a yield point, so a concurrent - ``set()`` can't be lost between waking and re-arming. - """ - - def __init__(self) -> None: - self._event = anyio.Event() - - def set(self) -> None: - self._event.set() - - async def wait(self) -> None: - await self._event.wait() - self._event = anyio.Event() - - def _empty_feed() -> PreparedFeed: """A feed that supplies nothing — turns the dispatcher into a fixed-set run.""" @@ -556,7 +538,7 @@ def note_models(options: list[TaskRunOptions]) -> None: source_done = False # woken on each task completion and on each injection (enqueue) - wake = _Wake() + wake = Wake() feed.set_wake(wake.set) def add(options: list[TaskRunOptions]) -> None: @@ -700,7 +682,7 @@ def note_models(options: list[TaskRunOptions]) -> None: source_done = False # woken on each task completion and on each injection (enqueue) - wake = _Wake() + wake = Wake() feed.set_wake(wake.set) def add(options: list[TaskRunOptions]) -> None: diff --git a/src/inspect_ai/_eval/task/__init__.py b/src/inspect_ai/_eval/task/__init__.py index dfdfcf52f1..7c96e61de1 100644 --- a/src/inspect_ai/_eval/task/__init__.py +++ b/src/inspect_ai/_eval/task/__init__.py @@ -1,5 +1,6 @@ from .task import Task, TaskInfo, PreviousTask, task_with # noqa: I001, F401 from .epochs import Epochs +from .sample_source import SampleSource from .task_source import TaskSource __all__ = [ @@ -8,5 +9,6 @@ "TaskInfo", "PreviousTask", "task_with", + "SampleSource", "TaskSource", ] diff --git a/src/inspect_ai/_eval/task/run.py b/src/inspect_ai/_eval/task/run.py index d60632586d..dcae8cb666 100644 --- a/src/inspect_ai/_eval/task/run.py +++ b/src/inspect_ai/_eval/task/run.py @@ -2,6 +2,7 @@ import functools import sys import time +from contextvars import Token from copy import copy, deepcopy from dataclasses import dataclass, field from datetime import datetime, timezone @@ -18,6 +19,7 @@ record_sample_cancelled, record_sample_completed, record_sample_errored, + record_samples_added, register_eval, ) from inspect_ai._display import ( @@ -29,7 +31,7 @@ ) from inspect_ai._display.core.display import TaskCancel, TaskDisplayMetric from inspect_ai._eval.task.scan import Scanners -from inspect_ai._util._async import aexit_shielded_when, tg_collect +from inspect_ai._util._async import Wake, aexit_shielded_when, tg_collect from inspect_ai._util.async_zip import AsyncZipReader from inspect_ai._util.asyncfiles import get_async_filesystem from inspect_ai._util.constants import ( @@ -175,6 +177,13 @@ ) from .log import TaskLogger, collect_eval_data, log_start from .results import eval_results +from .sample_source import ( + SampleEnqueuer, + SampleSource, + clear_sample_enqueuer, + get_sample_enqueuer, + register_sample_enqueuer, +) from .sandbox import sandboxenv_context from .scan import ( resume_scan_previous_sample, @@ -321,6 +330,19 @@ def _enqueue_source_tasks(tasks: list[Task] | None) -> None: enqueuer.enqueue(tasks) +def _enqueue_source_samples(samples: list[Sample] | None) -> None: + """Add samples a SampleSource callback returned to the running task's queue. + + Routes through the task's enqueuer (the same buffer ``enqueue_sample`` + feeds), so the task's dispatch loop starts them as capacity frees up. A + no-op if the callback returned nothing or there is no active enqueuer. + """ + if samples: + enqueuer = get_sample_enqueuer() + if enqueuer is not None: + enqueuer.enqueue(samples) + + async def task_run(options: TaskRunOptions, task_cancel: TaskCancel | None) -> EvalLog: from inspect_ai.hooks._hooks import ( emit_task_end, @@ -346,6 +368,10 @@ async def task_run(options: TaskRunOptions, task_cancel: TaskCancel | None) -> E sample_source = options.sample_source kwargs = options.kwargs + # a SampleSource-driven task generates samples while it runs (`sample_feed` + # to distinguish it from `sample_source`, the prior-attempt lookup above) + sample_feed: SampleSource | None = task.sample_source + # resolve default generate_config for task generate_config = task.config.merge(GenerateConfigArgs(**kwargs)) @@ -432,6 +458,26 @@ async def finish_task_log( if sample_store is not dataset: del dataset + # samples a SampleSource injects while the task runs, indexed after the + # seed store (kept in memory — they arrive incrementally, not up front) + store_len = len(sample_store) + injected_samples: list[Sample] = [] + + def get_sample(sample_index: int) -> Sample: + if sample_index < store_len: + return sample_store[sample_index] + return injected_samples[sample_index - store_len] + + # register the sample enqueuer that buffers additions to a + # SampleSource-driven task (callback-returned samples / enqueue_sample); + # ContextVar-scoped so it propagates to this task's samples but not to + # other tasks running in the process. Reset in the finally below. + sample_enqueuer: SampleEnqueuer | None = None + sample_enqueuer_token: Token[SampleEnqueuer | None] | None = None + if sample_feed is not None: + sample_enqueuer = SampleEnqueuer() + sample_enqueuer_token = register_sample_enqueuer(sample_enqueuer) + # resolve the plan (unroll chains) solver = solver or task.solver plan = resolve_plan(task, solver) @@ -619,7 +665,7 @@ async def run_sample( ) -> dict[str, SampleScore] | EarlyStop | None: # check for cached result from previous eval (before # materialization to avoid unnecessary deepcopy + image I/O) - sample_id = sample_store[sample_index].id + sample_id = get_sample(sample_index).id resume_checkpoint: ResumeCheckpoint | None = None # prior task-attempt errors to seed this re-run's # error_retries (empty unless the sample source reports a @@ -685,7 +731,7 @@ async def run_sample( async def create_sample_state( sample_uuid: str | None = None, ) -> tuple[Sample, TaskState]: - sample = deepcopy(sample_store[sample_index]) + sample = deepcopy(get_sample(sample_index)) if log_images: sample = await sample_with_base64_content(sample) state = deepcopy( @@ -732,6 +778,7 @@ async def create_sample_state( sample_complete=sample_complete, early_stopping=options.task.early_stopping, task_source=options.task_source, + sample_feed=sample_feed, fails_on_error=( config.fail_on_error is not False and config.continue_on_fail is not True @@ -751,13 +798,120 @@ async def create_sample_state( scan_id=options.scan_id, ) - sample_results = await tg_collect( - [ - functools.partial(run_sample, sample_index, epoch) - for epoch in range(1, epochs + 1) - for sample_index in range(len(sample_store)) - ] - ) + async def run_samples_dynamic( + feed: SampleSource, enqueuer: SampleEnqueuer + ) -> list[dict[str, SampleScore] | EarlyStop | None]: + """Run the seed samples plus every sample the source adds. + + Mirrors the eval-level live dispatcher (``run_multiple``): + injected samples start immediately (concurrency is bounded + by the sample semaphore inside ``run_sample``), and the + blocking ``next_samples()`` is only awaited when nothing is + in flight or buffered — so no completion can enqueue while + it blocks (no lost wakeup). + """ + nonlocal total_samples + + results: list[dict[str, SampleScore] | EarlyStop | None] = [] + in_flight = 0 + wake = Wake() + enqueuer.on_enqueue = wake.set + + # injected samples without ids continue the seed's 1-based + # numbering, skipping ids already in use; ids are compared + # by their str() form (matching ensure_unique_ids, since + # log member names and score grouping key on it) + seen_ids = {str(get_sample(i).id) for i in range(store_len)} + auto_id = store_len + + def add_samples(samples: list[Sample]) -> list[int]: + nonlocal total_samples, auto_id + indexes: list[int] = [] + for sample in samples: + if sample.id is None: + auto_id += 1 + while str(auto_id) in seen_ids: + auto_id += 1 + sample.id = auto_id + if str(sample.id) in seen_ids: + raise ValueError( + f"SampleSource added a sample with duplicate " + f"id '{sample.id}'. Please ensure each sample " + "has a unique id." + ) + seen_ids.add(str(sample.id)) + sample_ids.append(sample.id) + injected_samples.append(sample) + indexes.append(store_len + len(injected_samples) - 1) + # grow the planned totals (display denominator, + # fail_on_error threshold, control-channel counters) + total_samples += len(samples) * epochs + sample_error_handler.total_samples = total_samples + record_samples_added( + logger.eval.eval_id, + len(samples) * epochs, + sample_ids=[s.id for s in samples if s.id is not None], + ) + td.sample_complete( + complete=len(progress_results), total=total_samples + ) + return indexes + + async with anyio.create_task_group() as tg: + + async def run_one(sample_index: int, epoch: int) -> None: + nonlocal in_flight + try: + results.append(await run_sample(sample_index, epoch)) + finally: + in_flight -= 1 + wake.set() + + def spawn(indexes: list[int]) -> None: + nonlocal in_flight + for sample_index in indexes: + for epoch in range(1, epochs + 1): + in_flight += 1 + tg.start_soon(run_one, sample_index, epoch) + + spawn(list(range(store_len))) + while True: + # samples buffered since the last cycle (callbacks + # returning samples / enqueue_sample) start first + buffered = enqueuer.drain() + if buffered: + spawn(add_samples(buffered)) + continue + if in_flight > 0: + await wake.wait() + continue + # fully idle: ask the source for more (may block) + # and finish when it is exhausted + more = await feed.next_samples() + if more is None: + break + if more: + spawn(add_samples(more)) + + return results + + if sample_feed is not None and sample_enqueuer is not None: + try: + sample_results = await run_samples_dynamic( + sample_feed, sample_enqueuer + ) + except Exception as ex: + # match tg_collect: surface the first child exception + # rather than an ExceptionGroup + raise inner_exception(ex) + else: + sample_results = await tg_collect( + [ + functools.partial(run_sample, sample_index, epoch) + for epoch in range(1, epochs + 1) + for sample_index in range(len(sample_store)) + ] + ) # compute and record metrics if we have scores completed_scores = [ @@ -784,7 +938,7 @@ async def create_sample_state( if len(completed_scores) > 0: results, reductions = eval_results( - samples=profile.samples, + samples=total_samples, scores=completed_scores, reducers=task.epochs_reducer, scorers=scorers, @@ -803,7 +957,7 @@ async def create_sample_state( # `result is None` would miss them. sample_error_count = sample_error_handler.error_count mark_log_as_error = _should_eval_fail( - sample_error_count, profile.samples, config.fail_on_error + sample_error_count, total_samples, config.fail_on_error ) # finish @@ -833,7 +987,7 @@ async def create_sample_state( # compute partial results from samples that completed if len(progress_results) > 0: results, reductions = eval_results( - samples=profile.samples, + samples=total_samples, scores=progress_results, reducers=task.epochs_reducer, scorers=scorers, @@ -901,6 +1055,9 @@ async def create_sample_state( td.complete(TaskError(logger.samples_completed, type, value, traceback)) finally: + if sample_enqueuer_token is not None: + clear_sample_enqueuer(sample_enqueuer_token) + # every sample task has exited by here (the try encloses the task # group), so any still-unaccounted samples can no longer record finalize_eval(logger.eval.eval_id) @@ -1065,6 +1222,7 @@ async def task_run_sample( fails_on_error: bool, early_stopping: EarlyStopping | None, task_source: TaskSource | None, + sample_feed: SampleSource | None, retry_on_error: int, score_on_error: bool, error_retries: list[EvalRetryError], @@ -1705,6 +1863,13 @@ def make_eval_sample(include_events: bool = True) -> EvalSample: await emit_sample_end( eval_set_id, run_id, task_id, state.uuid, eval_sample ) + # notify the task's SampleSource (if it has one) as each + # sample completes, so it can react in real time (and add + # samples to the running task) + if sample_feed is not None: + _enqueue_source_samples( + await sample_feed.sample_complete(eval_sample) + ) # notify a TaskSource (if the run has one) as each sample # completes, so it can react in real time (and add tasks) if task_source is not None: @@ -1755,6 +1920,7 @@ def make_eval_sample(include_events: bool = True) -> EvalSample: sample_complete=sample_complete, early_stopping=early_stopping, task_source=task_source, + sample_feed=sample_feed, fails_on_error=fails_on_error, # tick retry count down retry_on_error=retry_on_error - 1, diff --git a/src/inspect_ai/_eval/task/sample_source.py b/src/inspect_ai/_eval/task/sample_source.py new file mode 100644 index 0000000000..369e5117ab --- /dev/null +++ b/src/inspect_ai/_eval/task/sample_source.py @@ -0,0 +1,221 @@ +"""A dynamic source of samples for a running task. + +Pass a :class:`SampleSource` instance as the ``dataset`` argument to a ``Task`` +to generate the task's samples from code rather than a fixed dataset — the +sample-level mirror of passing a :class:`~inspect_ai.TaskSource` as the +``tasks`` argument to ``eval()``. The task starts with ``initial_samples()`` +and, whenever no samples remain in flight, calls ``next_samples()`` for more — +until it returns ``None``. + +The simplest sources spawn follow-up work straight from results: override +``sample_complete`` and **return** the samples to run next (e.g. an RL loop +that spawns samples from scores, or an adaptive eval that branches on model +performance). A returned list is added to the task exactly as +``enqueue_sample`` would — it starts as soon as there is free capacity — so a +source can be driven entirely by this callback, with ``next_samples()`` +reserved for the blocking / explicit-pull case. + +``next_samples()`` is async and may block — awaiting external input or more +results — and returns ``None`` to end the task, so an open-ended "keep +producing until told to stop" loop is expressible. ``initial_samples()`` is +synchronous and returns immediately (it's the seed the task sets up and starts +from, not subject to that blocking); it may be empty, in which case the task +starts by calling ``next_samples()``. + +For the common case where subclassing is more than you need, +:meth:`SampleSource.from_samples` builds a ``SampleSource`` from a seed plus +optional callbacks. +""" + +from __future__ import annotations + +from contextvars import ContextVar, Token +from dataclasses import dataclass, field +from threading import Lock +from typing import TYPE_CHECKING, Awaitable, Callable + +if TYPE_CHECKING: + from inspect_ai.dataset import Sample + from inspect_ai.log._log import EvalSample + + +class SampleSource: + """Drives a running task from code: a seed plus result-driven follow-ups. + + Subclass and override the methods you need. The default implementations are + no-ops / empty, so a bare ``SampleSource`` runs nothing — override at least + ``initial_samples`` and ``next_samples``. + """ + + def initial_samples(self) -> list["Sample"]: + """Samples to run first (the seed). + + Called once, synchronously, when the ``Task`` is created — so it must + return immediately (no awaiting / blocking). The returned samples drive + the task's up-front setup (validation, sandbox startup) and are the + first batch. May be empty, in which case the task starts by calling + ``next_samples()``. + """ + return [] + + async def next_samples(self) -> list["Sample"] | None: + """More samples to run, or ``None`` when the task is complete. + + Called whenever no samples remain in flight or buffered (after those + samples' ``sample_complete`` notifications). May ``await`` — for more + results or external input — and may block indefinitely; return ``None`` + to end the task. + """ + return None + + async def sample_complete(self, sample: "EvalSample") -> list["Sample"] | None: + """A sample finished — observe it and optionally return follow-up samples. + + Return a list of samples to add to the task (equivalent to calling + ``enqueue_sample`` with them): they start as soon as there is free + capacity. Return ``None`` (the default) to add nothing. + """ + return None + + @classmethod + def from_samples( + cls, + initial_samples: list["Sample"], + *, + next_samples: Callable[[], Awaitable[list["Sample"] | None]] | None = None, + sample_complete: Callable[["EvalSample"], Awaitable[list["Sample"] | None]] + | None = None, + ) -> "SampleSource": + """Create a :class:`SampleSource` from a seed plus optional callbacks. + + A convenience for when subclassing is more than you need: provide the + initial samples directly and, optionally, callbacks that react to + results. The ``sample_complete`` callback may **return** a list of + follow-up samples to add to the task (see that method); + ``next_samples`` is the blocking / explicit-pull alternative. Callbacks + typically close over shared state (e.g. accumulated scores) to decide + what to run next. + + Args: + initial_samples: The seed samples to run first (see + :meth:`initial_samples`). + next_samples: Optional async callback returning more samples, or + ``None`` to end the task (see :meth:`next_samples`). If omitted + (and no callback returns samples), the task stops after the + seed — equivalent to passing ``initial_samples`` directly as + the dataset. + sample_complete: Optional async callback invoked as each sample + finishes; may return follow-up samples to add to the task. + + Returns: + A ``SampleSource`` that delegates to the provided seed and callbacks. + """ + return _CallableSampleSource(initial_samples, next_samples, sample_complete) + + +class _CallableSampleSource(SampleSource): + """A :class:`SampleSource` backed by a fixed seed and optional callables.""" + + def __init__( + self, + initial_samples: list["Sample"], + next_samples: Callable[[], Awaitable[list["Sample"] | None]] | None, + sample_complete: Callable[["EvalSample"], Awaitable[list["Sample"] | None]] + | None, + ) -> None: + self._initial_samples = list(initial_samples) + self._next_samples = next_samples + self._sample_complete = sample_complete + + def initial_samples(self) -> list["Sample"]: + return self._initial_samples + + async def next_samples(self) -> list["Sample"] | None: + if self._next_samples is not None: + return await self._next_samples() + return None + + async def sample_complete(self, sample: "EvalSample") -> list["Sample"] | None: + if self._sample_complete is not None: + return await self._sample_complete(sample) + return None + + +@dataclass +class SampleEnqueuer: + """Buffers samples added to one running task and hands them to its loop. + + Owned by ``task_run`` for the task's lifetime (only when the task has a + :class:`SampleSource`). ``enqueue`` buffers; ``drain`` is the non-blocking + pull the task's dispatch loop does between cycles. + """ + + _pending: list["Sample"] = field(default_factory=list) + _lock: Lock = field(default_factory=Lock) + on_enqueue: Callable[[], None] | None = None + """Fired after samples are buffered — used to wake the task's dispatcher.""" + + def enqueue(self, samples: list["Sample"]) -> None: + """Queue ``samples`` to run in the task.""" + with self._lock: + self._pending.extend(samples) + if self.on_enqueue is not None: + self.on_enqueue() + + def drain(self) -> list["Sample"]: + """Remove and return all currently-buffered samples (empty if none).""" + with self._lock: + batch, self._pending = self._pending, [] + return batch + + +# The running task's enqueuer, scoped to the task's async context (and +# propagated to its samples — solvers / scorers / tools). ``None`` when the +# current task is not driven by a SampleSource. +_sample_enqueuer: ContextVar[SampleEnqueuer | None] = ContextVar( + "sample_enqueuer", default=None +) + + +def register_sample_enqueuer(enqueuer: SampleEnqueuer) -> Token[SampleEnqueuer | None]: + """Install ``enqueuer`` as the active one; returns a token to restore with.""" + return _sample_enqueuer.set(enqueuer) + + +def get_sample_enqueuer() -> SampleEnqueuer | None: + return _sample_enqueuer.get() + + +def clear_sample_enqueuer(token: Token[SampleEnqueuer | None]) -> None: + """Restore the enqueuer in scope before the matching ``register`` call.""" + _sample_enqueuer.reset(token) + + +def enqueue_sample(samples: "Sample | list[Sample]") -> None: + """Add one or more samples to the running task. + + The samples run in the current task as soon as there is free capacity + (bounded by ``max_samples``), each for the task's configured number of + epochs. Samples without an ``id`` are assigned one automatically. + + Only available inside a task driven by a :class:`SampleSource` (i.e. a + ``Task`` whose ``dataset`` is a ``SampleSource``) — a plain task's sample + set is fixed, so there is no loop to run additions. Callable from any code + running within such a task: a solver, a scorer, a tool. + + Args: + samples: A ``Sample`` (or list of samples) to add to the running task. + + Raises: + RuntimeError: If the current task is not driven by a ``SampleSource`` + (or no task is running in this context). + """ + from inspect_ai.dataset import Sample + + enqueuer = get_sample_enqueuer() + if enqueuer is None: + raise RuntimeError( + "enqueue_sample() can only be called from within a running task " + "whose dataset is a SampleSource." + ) + enqueuer.enqueue([samples] if isinstance(samples, Sample) else list(samples)) diff --git a/src/inspect_ai/_eval/task/task.py b/src/inspect_ai/_eval/task/task.py index 1cd7ea9e43..09e9d9c0f2 100644 --- a/src/inspect_ai/_eval/task/task.py +++ b/src/inspect_ai/_eval/task/task.py @@ -2,7 +2,16 @@ from dataclasses import dataclass from logging import getLogger -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Sequence, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + NamedTuple, + Sequence, + cast, + overload, +) from inspect_ai.util._early_stopping import EarlyStopping @@ -47,6 +56,7 @@ from inspect_ai.viewer import ViewerConfig from .epochs import Epochs +from .sample_source import SampleSource logger = getLogger(__name__) @@ -66,7 +76,7 @@ class Task: def __init__( self, - dataset: Dataset | Sequence[Sample] | None = None, + dataset: Dataset | Sequence[Sample] | SampleSource | None = None, setup: Solver | list[Solver] | None = None, solver: Solver | Agent | list[Solver] = generate(), cleanup: Callable[[TaskState], Awaitable[None]] | None = None, @@ -102,7 +112,8 @@ def __init__( """Create a task. Args: - dataset: Dataset to evaluate + dataset: Dataset to evaluate, or a `SampleSource` that generates + samples dynamically while the task runs. setup: Setup step (always run even when the main `solver` is replaced). solver: Solver or list of solvers. Defaults to generate(), a normal call to the model. cleanup: Optional cleanup function for task. Called after @@ -176,7 +187,9 @@ def __init__( f"DEPRECATED: the '{arg}' parameter is deprecated (please use the '{newarg}' parameter instead)", ) - self.dataset = resolve_dataset(dataset) + resolved_dataset = resolve_dataset_or_source(dataset) + self.dataset = resolved_dataset.dataset + self.sample_source = resolved_dataset.sample_source self.setup = setup self.solver = resolve_solver(solver) self.cleanup = cleanup @@ -246,7 +259,7 @@ def attribs(self) -> dict[str, Any]: def task_with( task: Task, *, - dataset: Dataset | Sequence[Sample] | None | NotGiven = NOT_GIVEN, + dataset: Dataset | Sequence[Sample] | SampleSource | None | NotGiven = NOT_GIVEN, setup: Solver | list[Solver] | None | NotGiven = NOT_GIVEN, solver: Solver | Agent | list[Solver] | NotGiven = NOT_GIVEN, cleanup: Callable[[TaskState], Awaitable[None]] | None | NotGiven = NOT_GIVEN, @@ -290,7 +303,8 @@ def task_with( Args: task: Task to adapt - dataset: Dataset to evaluate + dataset: Dataset to evaluate, or a `SampleSource` that generates + samples dynamically while the task runs. setup: Setup step (always run even when the main `solver` is replaced). solver: Solver or list of solvers. Defaults to generate(), a normal call to the model. cleanup: Optional cleanup function for task. Called after @@ -344,7 +358,9 @@ def task_with( Task: Passed `task` with modifications. """ if not isinstance(dataset, NotGiven): - task.dataset = resolve_dataset(dataset) + resolved_dataset = resolve_dataset_or_source(dataset) + task.dataset = resolved_dataset.dataset + task.sample_source = resolved_dataset.sample_source if not isinstance(setup, NotGiven): task.setup = setup if not isinstance(solver, NotGiven): @@ -458,6 +474,28 @@ def resolve_epochs(epochs: int | Epochs | None) -> Epochs | None: return epochs +class ResolvedDataset(NamedTuple): + """A task's dataset plus the `SampleSource` that seeded it (if any).""" + + dataset: Dataset + sample_source: SampleSource | None + + +def resolve_dataset_or_source( + dataset: Dataset | Sequence[Sample] | SampleSource | None, +) -> ResolvedDataset: + """Resolve a task's `dataset` argument, which may be a `SampleSource`. + + A `SampleSource` contributes its (possibly empty) `initial_samples()` as + the task's up-front dataset — the empty-dataset check doesn't apply since + the source can produce samples while the task runs. + """ + if isinstance(dataset, SampleSource): + return ResolvedDataset(MemoryDataset(list(dataset.initial_samples())), dataset) + else: + return ResolvedDataset(resolve_dataset(dataset), None) + + def resolve_dataset(dataset: Dataset | Sequence[Sample] | None) -> Dataset: # this is a convenience for tests that don't want to define a dummy sample if dataset is None: diff --git a/src/inspect_ai/_util/_async.py b/src/inspect_ai/_util/_async.py index 994ab091b7..ad1c854a6f 100644 --- a/src/inspect_ai/_util/_async.py +++ b/src/inspect_ai/_util/_async.py @@ -77,6 +77,25 @@ async def run_task(func: Callable[[], Awaitable[T]], index: int) -> None: raise ex.exceptions[0] from None +class Wake: + """One-shot wake signal that can be re-armed (set on completion / injection). + + Safe under cooperative scheduling: the only await is on ``wait()``; the + re-arm assignment afterwards runs without a yield point, so a concurrent + ``set()`` can't be lost between waking and re-arming. + """ + + def __init__(self) -> None: + self._event = anyio.Event() + + def set(self) -> None: + self._event.set() + + async def wait(self) -> None: + await self._event.wait() + self._event = anyio.Event() + + class aexit_shielded_when(contextlib.AbstractAsyncContextManager[Any]): """Wrap an async context manager so its `__aexit__` runs shielded when `shield()` is True. diff --git a/tests/test_sample_source.py b/tests/test_sample_source.py new file mode 100644 index 0000000000..ccba952e26 --- /dev/null +++ b/tests/test_sample_source.py @@ -0,0 +1,320 @@ +"""Tests for driving a running task from a `SampleSource`. + +The source seeds the task with `initial_samples()`, observes each result via +`sample_complete` (which may return follow-up samples), and produces more from +`next_samples()` until it returns `None` — all within one task / log. Added +samples are live: they start as soon as there is free capacity. +""" + +import anyio +import pytest + +from inspect_ai import SampleSource, Task, enqueue_sample, eval +from inspect_ai.dataset import Sample +from inspect_ai.log import EvalLog, EvalSample +from inspect_ai.solver import Generate, Solver, TaskState, generate, solver + + +def _sample_inputs(log: EvalLog) -> list[str]: + return sorted(str(sample.input) for sample in (log.samples or [])) + + +class _Generations(SampleSource): + """Produces `count` samples one at a time; each produced after the prior.""" + + def __init__(self, count: int) -> None: + self.count = count + self.samples_seen: list[str] = [] + self._produced = 1 # s0 is the initial seed + + def initial_samples(self) -> list[Sample]: + return [Sample(input="s0", target="ok")] + + async def sample_complete(self, sample: EvalSample) -> None: + self.samples_seen.append(str(sample.input)) + + async def next_samples(self) -> list[Sample] | None: + if self._produced < self.count: + name = f"s{self._produced}" + self._produced += 1 + return [Sample(input=name, target="ok")] + return None + + +def test_sample_source_runs_generations_in_one_task() -> None: + source = _Generations(3) + logs = eval( + Task(dataset=source, solver=[generate()]), + model="mockllm/model", + display="none", + ) + + # all generations ran within a single task/log, observed by the source + assert len(logs) == 1 + log = logs[0] + assert log.status == "success" + assert _sample_inputs(log) == ["s0", "s1", "s2"] + assert sorted(source.samples_seen) == ["s0", "s1", "s2"] + # injected samples continue the seed's auto-id numbering + assert sorted(sample.id for sample in (log.samples or [])) == [1, 2, 3] + + +def test_sample_source_seed_only_when_next_is_none() -> None: + # next_samples() returns None immediately -> only the seed runs + logs = eval( + Task(dataset=_Generations(1), solver=[generate()]), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["s0"] + + +def test_sample_complete_returning_samples_chains_generations() -> None: + # the simplest source: return follow-up samples straight from + # sample_complete, no next_samples(). The task chains until a completion + # returns nothing. + class _Chain(SampleSource): + def __init__(self, count: int) -> None: + self.count = count + self._produced = 1 # s0 is the seed + + def initial_samples(self) -> list[Sample]: + return [Sample(input="s0", target="ok")] + + async def sample_complete(self, sample: EvalSample) -> list[Sample] | None: + if self._produced < self.count: + name = f"s{self._produced}" + self._produced += 1 + return [Sample(input=name, target="ok")] + return None + + logs = eval( + Task(dataset=_Chain(3), solver=[generate()]), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["s0", "s1", "s2"] + + +def test_sample_source_factory_seed_and_callbacks() -> None: + # from_samples() builds a SampleSource from a seed + callbacks (no + # subclass); next_samples closes over shared state for two more samples + produced = [1] + samples_seen: list[str] = [] + + async def next_samples() -> list[Sample] | None: + if produced[0] < 3: + name = f"s{produced[0]}" + produced[0] += 1 + return [Sample(input=name, target="ok")] + return None + + async def on_sample(sample: EvalSample) -> None: + samples_seen.append(str(sample.input)) + + source = SampleSource.from_samples( + [Sample(input="s0", target="ok")], + next_samples=next_samples, + sample_complete=on_sample, + ) + logs = eval( + Task(dataset=source, solver=[generate()]), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["s0", "s1", "s2"] + assert sorted(samples_seen) == ["s0", "s1", "s2"] + + +def test_sample_source_factory_seed_only_runs_once() -> None: + # omitting next_samples stops after the seed (equivalent to a plain dataset) + logs = eval( + Task( + dataset=SampleSource.from_samples([Sample(input="only", target="ok")]), + solver=[generate()], + ), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["only"] + + +def test_sample_source_empty_seed_pulls_from_next_samples() -> None: + # an empty seed is allowed: the task starts by calling next_samples() + produced = [False] + + class _EmptySeed(SampleSource): + async def next_samples(self) -> list[Sample] | None: + if not produced[0]: + produced[0] = True + return [Sample(input="only", target="ok")] + return None + + logs = eval( + Task(dataset=_EmptySeed(), solver=[generate()]), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["only"] + + +def test_sample_source_epochs_apply_to_injected_samples() -> None: + # injected samples run for the task's configured epochs, like the seed + class _Two(SampleSource): + def __init__(self) -> None: + self._done = False + + def initial_samples(self) -> list[Sample]: + return [Sample(input="seed", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + if not self._done: + self._done = True + return [Sample(input="added", target="ok")] + return None + + logs = eval( + Task(dataset=_Two(), solver=[generate()], epochs=2), + model="mockllm/model", + display="none", + ) + log = logs[0] + assert log.status == "success" + runs = sorted((str(sample.input), sample.epoch) for sample in (log.samples or [])) + assert runs == [("added", 1), ("added", 2), ("seed", 1), ("seed", 2)] + + +def test_sample_source_explicit_and_auto_ids() -> None: + # explicit injected ids are honored; auto-ids skip ids already in use + class _Ids(SampleSource): + def __init__(self) -> None: + self._produced = False + + def initial_samples(self) -> list[Sample]: + return [Sample(input="seed", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + return [ + Sample(id=2, input="explicit", target="ok"), + Sample(input="auto", target="ok"), + ] + return None + + logs = eval( + Task(dataset=_Ids(), solver=[generate()]), + model="mockllm/model", + display="none", + ) + log = logs[0] + assert log.status == "success" + ids = {str(sample.input): sample.id for sample in (log.samples or [])} + assert ids["seed"] == 1 + assert ids["explicit"] == 2 + assert ids["auto"] == 3 # auto-id skipped the explicit 2 + + +def test_sample_source_duplicate_id_errors() -> None: + # an injected sample whose id collides with an existing one fails the task + class _Dup(SampleSource): + def __init__(self) -> None: + self._produced = False + + def initial_samples(self) -> list[Sample]: + return [Sample(id=1, input="seed", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + return [Sample(id=1, input="dup", target="ok")] + return None + + logs = eval( + Task(dataset=_Dup(), solver=[generate()]), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "error" + assert "duplicate" in (logs[0].error.message if logs[0].error else "") + + +def test_live_injection_runs_concurrently_with_in_flight_sample() -> None: + # Discriminates live injection from batch-at-a-time: a sample injected + # mid-run must start while another sample is still in flight. "blocker" + # parks until "injected" releases it, and "injected" is only enqueued (by + # "injector") once the task is already underway — so only live injection + # lets this finish (the fail_after turns a regression into a fast failure + # instead of a hang). + released: anyio.Event | None = None + order: list[str] = [] + + @solver + def router() -> Solver: + async def solve(state: TaskState, generate: Generate) -> TaskState: + assert released is not None + if state.input_text == "blocker": + order.append("blocker-start") + with anyio.fail_after(10): + await released.wait() + order.append("blocker-end") + elif state.input_text == "injector": + enqueue_sample(Sample(input="injected", target="ok")) + elif state.input_text == "injected": + order.append("injected") + released.set() + return state + + return solve + + class _Src(SampleSource): + def initial_samples(self) -> list[Sample]: + return [ + Sample(input="blocker", target="ok"), + Sample(input="injector", target="ok"), + ] + + async def next_samples(self) -> list[Sample] | None: + return None + + released = anyio.Event() + logs = eval( + Task(dataset=_Src(), solver=[router()]), + model="mockllm/model", + display="none", + ) + log = logs[0] + assert log.status == "success" + # blocker only reaches its end if the injected sample ran while it blocked + assert "blocker-end" in order + assert _sample_inputs(log) == ["blocker", "injected", "injector"] + + +def test_enqueue_sample_rejected_outside_sample_source_task() -> None: + # enqueue_sample() requires a running SampleSource-driven task: a plain + # task has a fixed sample set (no loop to run additions) + @solver + def bad() -> Solver: + async def solve(state: TaskState, generate: Generate) -> TaskState: + enqueue_sample(Sample(input="x")) + return state + + return solve + + logs = eval( + Task(dataset=[Sample(input="plain")], solver=[bad()]), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "error" + assert "SampleSource" in (logs[0].error.message if logs[0].error else "") + + +def test_enqueue_sample_rejected_outside_task() -> None: + with pytest.raises(RuntimeError, match="SampleSource"): + enqueue_sample(Sample(input="x")) From 502dc3fce7a5fbae1ab2008ac4a68c358a7a3210 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:51:46 +0000 Subject: [PATCH 2/7] Support --limit / --sample-id for SampleSource tasks and address review feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --limit now caps a SampleSource task's total samples (seed + produced): the seed is sliced as usual and the remainder is the budget for samples added while the task runs — additions beyond it are ignored with a warning, and once exhausted the dispatcher finishes without consulting next_samples() again. --sample-id filters produced samples with the same normalise+fnmatch predicate as the seed, and ids missing from the seed no longer warn or error for SampleSource tasks (the source may produce them at runtime; seed datasets are marked via DATASET_SAMPLE_SOURCE_ATTR). Review fixes: samples enqueued during a terminal next_samples() are drained and run instead of silently dropped; an anyio checkpoint keeps the dispatch loop cancellable if a source never blocks; the dispatcher's inner-exception unwrap now raises `from None` matching tg_collect; docs note that sample_complete fires per epoch and follow-ups multiply. Co-authored-by: Ransom Richardson Co-Authored-By: Claude Fable 5 --- design/sample-source.md | 29 +++- docs/sample-source.qmd | 8 +- src/inspect_ai/_eval/task/constants.py | 1 + src/inspect_ai/_eval/task/run.py | 97 ++++++++++--- src/inspect_ai/_eval/task/sample_source.py | 7 +- src/inspect_ai/_eval/task/task.py | 10 +- src/inspect_ai/_eval/task/util.py | 69 ++++++--- tests/test_sample_source.py | 155 +++++++++++++++++++++ 8 files changed, 330 insertions(+), 46 deletions(-) diff --git a/design/sample-source.md b/design/sample-source.md index 2b477c26d7..91177001e7 100644 --- a/design/sample-source.md +++ b/design/sample-source.md @@ -68,7 +68,11 @@ A `SampleSource` is passed as the **`dataset` argument to `Task`** (just as a idle await the blocking `next_samples()`. Sample concurrency is already bounded by the sample semaphore inside `run_sample`, so the dispatcher spawns eagerly with no cap of its own. Plain tasks keep the unchanged - `tg_collect` path. + `tg_collect` path. When `next_samples()` returns `None`, the enqueuer is + drained once more so samples enqueued during the terminal call still run + (nothing is silently dropped; the source may then be consulted again). Each + loop iteration starts with an `anyio.lowlevel.checkpoint()` so a source + that never blocks (e.g. returns `[]` forever) stays cancellable. - **Injected sample storage** — injected samples are appended to an in-memory list indexed after the (possibly disk-paged) seed store; `get_sample(index)` dispatches between the two. Each injected sample runs for the task's @@ -102,8 +106,21 @@ dropping samples. ## Interactions / known limits -- **`--limit` / `--sample-id`** slice the seed only; samples produced while the - task runs are not subject to them (documented). +- **`--limit`** caps the *total* number of samples (seed + produced): the seed + is sliced as usual, and the dispatcher spends the remainder + (`sample_limit_count(limit) - seed`) as its budget for added samples — + additions beyond it are ignored with a warning, and once the budget is + exhausted the loop finishes without consulting `next_samples()` again. A + tuple limit's budget is the size of its slice (`stop - start`). +- **`--sample-id`** filters produced samples with the same + normalise+`fnmatch` predicate that filters the seed (`sample_id_filter` in + task/util.py). The seed of a SampleSource task is marked with + `DATASET_SAMPLE_SOURCE_ATTR` so `slice_dataset` (log spec, sandbox startup, + task_run) neither warns nor raises for requested ids missing from the seed — + the source may produce them at runtime. Filtered-out samples still register + their ids (duplicate ids stay a hard error) but don't run and don't grow the + planned totals. As in `slice_dataset`, `--sample-id` and `--limit` are + mutually exclusive (the filter wins). - **Task retries** (`task_retry_attempts` / eval-set): the retry attempt re-drives the source; completed samples are reused via the normal `EvalSampleSource` lookup only where regenerated ids match — i.e. resume @@ -127,4 +144,8 @@ follow-ups chains generations; `from_samples` (callbacks and seed-only); empty seed; epochs applied to injected samples; explicit + auto id assignment and duplicate-id error; live injection discriminated from batch-at-a-time (blocker parks until an injected sample releases it, `fail_after` bounds a regression); -`enqueue_sample` rejected on plain tasks and outside a task. +`enqueue_sample` rejected on plain tasks and outside a task; `--limit` caps +totals (budget spent, seed-consumed limit never consults the source, batch +truncation, samples-not-runs with epochs); `--sample-id` filters produced +samples and tolerates ids missing from the seed; samples enqueued during a +terminal `next_samples()` still run. diff --git a/docs/sample-source.qmd b/docs/sample-source.qmd index 5a389f42f6..b45fa9449a 100644 --- a/docs/sample-source.qmd +++ b/docs/sample-source.qmd @@ -79,6 +79,8 @@ class Adaptive(SampleSource): A source that returns follow-ups from this callback needs no `next_samples()`: the task ends when completions return nothing and `next_samples()` returns `None`. Use `next_samples()` for the blocking case a per-result callback can't express. +Note that with `epochs > 1`, `sample_complete` fires once per *epoch*, and each returned follow-up itself runs all epochs — so a per-completion follow-up pattern multiplies (the example above with `epochs=2` would spawn two follow-ups per passed difficulty level, each running twice). Use `sample.epoch` (or dedupe on the source's own state) to react once per logical sample. + ## Sources from callbacks `SampleSource.from_samples()` builds a source from a seed and optional callbacks, without subclassing: @@ -122,4 +124,8 @@ Enqueued samples share a buffer with samples returned from `sample_complete`. `e A `SampleSource` task is live: a sample added mid-run starts as soon as there is free capacity, rather than waiting for the current samples to finish. Concurrency is bounded by `max_samples` as usual (see [Parallelism](parallelism.qmd)). -The `--limit` and `--sample-id` options apply to the seed only — samples the source produces while the task runs are not subject to them. +## Limits and sample filtering + +The `--limit` option caps the total number of samples in the task — the seed plus everything the source produces. For example, `--limit 10` on a source with 5 seed samples allows 5 more generated samples; once the cap is reached, further additions are ignored (with a warning) and the task ends when the in-flight samples finish (`next_samples()` is not consulted again). The cap counts samples, not runs: each sample within the limit still runs for the task's configured number of epochs. + +The `--sample-id` option filters generated samples the same way it filters the seed: only samples whose ids match run. Since the source may produce a requested id while the task runs, it is not an error for a `--sample-id` value to be missing from the seed. diff --git a/src/inspect_ai/_eval/task/constants.py b/src/inspect_ai/_eval/task/constants.py index f395b7b1b9..5be30b28ac 100644 --- a/src/inspect_ai/_eval/task/constants.py +++ b/src/inspect_ai/_eval/task/constants.py @@ -1,3 +1,4 @@ TASK_FILE_ATTR = "__task_file__" TASK_RUN_DIR_ATTR = "__task_run_dir__" TASK_ALL_PARAMS_ATTR = "__task_all_params__" +DATASET_SAMPLE_SOURCE_ATTR = "__dataset_sample_source__" diff --git a/src/inspect_ai/_eval/task/run.py b/src/inspect_ai/_eval/task/run.py index dcae8cb666..090b23a85a 100644 --- a/src/inspect_ai/_eval/task/run.py +++ b/src/inspect_ai/_eval/task/run.py @@ -192,7 +192,7 @@ ) from .store import DiskSampleStore, maybe_page_to_disk from .task_source import TaskSource -from .util import sample_messages, slice_dataset +from .util import sample_id_filter, sample_limit_count, sample_messages, slice_dataset py_logger = getLogger(__name__) @@ -809,6 +809,13 @@ async def run_samples_dynamic( blocking ``next_samples()`` is only awaited when nothing is in flight or buffered — so no completion can enqueue while it blocks (no lost wakeup). + + ``--limit`` caps the total samples (seed + added): once the + cap is reached further additions are ignored (with a + warning) and the loop finishes without consulting the + source again. ``--sample-id`` filters added samples the + same way it filters the seed (the filter and the cap are + mutually exclusive, matching ``slice_dataset``). """ nonlocal total_samples @@ -817,6 +824,26 @@ async def run_samples_dynamic( wake = Wake() enqueuer.on_enqueue = wake.set + # --sample-id: added samples must also match the filter + include_id = ( + sample_id_filter(config.sample_id) + if config.sample_id is not None + else None + ) + + # --limit: cap on total samples (seed + added); the seed + # was already sliced above, so the cap's remainder is the + # budget for added samples. None when --sample-id is set + # (slice_dataset likewise ignores limit then). + limit_count = ( + sample_limit_count(config.limit) if include_id is None else None + ) + remaining = ( + max(0, limit_count - store_len) + if limit_count is not None + else None + ) + # injected samples without ids continue the seed's 1-based # numbering, skipping ids already in use; ids are compared # by their str() form (matching ensure_unique_ids, since @@ -825,8 +852,10 @@ async def run_samples_dynamic( auto_id = store_len def add_samples(samples: list[Sample]) -> list[int]: - nonlocal total_samples, auto_id + nonlocal total_samples, auto_id, remaining indexes: list[int] = [] + added_ids: list[int | str] = [] + over_limit = 0 for sample in samples: if sample.id is None: auto_id += 1 @@ -840,21 +869,35 @@ def add_samples(samples: list[Sample]) -> list[int]: "has a unique id." ) seen_ids.add(str(sample.id)) + if include_id is not None and not include_id(sample.id): + continue + if remaining is not None: + if remaining <= 0: + over_limit += 1 + continue + remaining -= 1 sample_ids.append(sample.id) + added_ids.append(sample.id) injected_samples.append(sample) indexes.append(store_len + len(injected_samples) - 1) - # grow the planned totals (display denominator, - # fail_on_error threshold, control-channel counters) - total_samples += len(samples) * epochs - sample_error_handler.total_samples = total_samples - record_samples_added( - logger.eval.eval_id, - len(samples) * epochs, - sample_ids=[s.id for s in samples if s.id is not None], - ) - td.sample_complete( - complete=len(progress_results), total=total_samples - ) + if over_limit: + py_logger.warning( + f"Sample limit ({limit_count}) reached: ignoring " + f"{over_limit} sample(s) added to the task." + ) + if indexes: + # grow the planned totals (display denominator, + # fail_on_error threshold, control-channel counters) + total_samples += len(indexes) * epochs + sample_error_handler.total_samples = total_samples + record_samples_added( + logger.eval.eval_id, + len(indexes) * epochs, + sample_ids=added_ids, + ) + td.sample_complete( + complete=len(progress_results), total=total_samples + ) return indexes async with anyio.create_task_group() as tg: @@ -876,6 +919,10 @@ def spawn(indexes: list[int]) -> None: spawn(list(range(store_len))) while True: + # checkpoint so a misbehaving source that never + # blocks (e.g. next_samples() returning []) keeps + # the loop cancellable instead of spinning hot + await anyio.lowlevel.checkpoint() # samples buffered since the last cycle (callbacks # returning samples / enqueue_sample) start first buffered = enqueuer.drain() @@ -885,12 +932,24 @@ def spawn(indexes: list[int]) -> None: if in_flight > 0: await wake.wait() continue - # fully idle: ask the source for more (may block) - # and finish when it is exhausted + # fully idle: the task is complete once the sample + # limit is exhausted (don't consult the source for + # samples that could never run) + if remaining is not None and remaining <= 0: + break + # ask the source for more (may block) and finish + # when it is exhausted more = await feed.next_samples() if more is None: - break - if more: + # run samples enqueued while next_samples() + # was finishing rather than dropping them + # (enqueue_sample never drops silently); if + # any run, the source may be consulted again + final = add_samples(enqueuer.drain()) + if not final: + break + spawn(final) + elif more: spawn(add_samples(more)) return results @@ -903,7 +962,7 @@ def spawn(indexes: list[int]) -> None: except Exception as ex: # match tg_collect: surface the first child exception # rather than an ExceptionGroup - raise inner_exception(ex) + raise inner_exception(ex) from None else: sample_results = await tg_collect( [ diff --git a/src/inspect_ai/_eval/task/sample_source.py b/src/inspect_ai/_eval/task/sample_source.py index 369e5117ab..eb84864375 100644 --- a/src/inspect_ai/_eval/task/sample_source.py +++ b/src/inspect_ai/_eval/task/sample_source.py @@ -64,7 +64,8 @@ async def next_samples(self) -> list["Sample"] | None: Called whenever no samples remain in flight or buffered (after those samples' ``sample_complete`` notifications). May ``await`` — for more results or external input — and may block indefinitely; return ``None`` - to end the task. + to end the task. (If samples were enqueued while a ``None`` return was + in progress they still run, and this method may then be called again.) """ return None @@ -203,6 +204,10 @@ def enqueue_sample(samples: "Sample | list[Sample]") -> None: set is fixed, so there is no loop to run additions. Callable from any code running within such a task: a solver, a scorer, a tool. + When the eval was run with ``--limit``, samples beyond the limit are + ignored (with a warning); with ``--sample-id``, only samples matching the + filter run. + Args: samples: A ``Sample`` (or list of samples) to add to the running task. diff --git a/src/inspect_ai/_eval/task/task.py b/src/inspect_ai/_eval/task/task.py index 09e9d9c0f2..58c68fc890 100644 --- a/src/inspect_ai/_eval/task/task.py +++ b/src/inspect_ai/_eval/task/task.py @@ -55,6 +55,7 @@ ) from inspect_ai.viewer import ViewerConfig +from .constants import DATASET_SAMPLE_SOURCE_ATTR from .epochs import Epochs from .sample_source import SampleSource @@ -488,10 +489,15 @@ def resolve_dataset_or_source( A `SampleSource` contributes its (possibly empty) `initial_samples()` as the task's up-front dataset — the empty-dataset check doesn't apply since - the source can produce samples while the task runs. + the source can produce samples while the task runs. The seed is marked + with `DATASET_SAMPLE_SOURCE_ATTR` so consumers that treat a dataset as + the complete sample set (e.g. `slice_dataset` validating `--sample-id`) + know more samples may be produced at runtime. """ if isinstance(dataset, SampleSource): - return ResolvedDataset(MemoryDataset(list(dataset.initial_samples())), dataset) + seed = MemoryDataset(list(dataset.initial_samples())) + setattr(seed, DATASET_SAMPLE_SOURCE_ATTR, True) + return ResolvedDataset(seed, dataset) else: return ResolvedDataset(resolve_dataset(dataset), None) diff --git a/src/inspect_ai/_eval/task/util.py b/src/inspect_ai/_eval/task/util.py index 16b1995cdf..cd243942ad 100644 --- a/src/inspect_ai/_eval/task/util.py +++ b/src/inspect_ai/_eval/task/util.py @@ -3,7 +3,7 @@ from copy import deepcopy from fnmatch import fnmatch from logging import getLogger -from typing import cast +from typing import Callable, cast from inspect_ai._util.error import PrerequisiteError from inspect_ai._util.logger import warn_once @@ -14,7 +14,11 @@ from inspect_ai.model import ChatMessage, ChatMessageUser from ..task import Task -from .constants import TASK_FILE_ATTR, TASK_RUN_DIR_ATTR +from .constants import ( + DATASET_SAMPLE_SOURCE_ATTR, + TASK_FILE_ATTR, + TASK_RUN_DIR_ATTR, +) logger = getLogger(__name__) @@ -44,36 +48,63 @@ def task_file(task: Task, relative: bool = False) -> str | None: return None +def sample_id_filter( + sample_id: str | int | list[str] | list[int] | list[str | int], +) -> Callable[[str | int | None], bool]: + """Predicate matching a sample id against `sample_id` patterns. + + Ids are normalised (`normalise_sample_id`) and matched with `fnmatch`, + the same semantics `slice_dataset` applies to a dataset. + """ + sample_ids = sample_id if isinstance(sample_id, list) else [sample_id] + patterns = [normalise_sample_id(id) for id in sample_ids] + + def matches(id: str | int | None) -> bool: + return any(fnmatch(normalise_sample_id(id), pat) for pat in patterns) + + return matches + + +def sample_limit_count(limit: int | tuple[int, int] | None) -> int | None: + """Number of samples a `limit` option selects (the size of its slice).""" + if limit is None: + return None + return max(0, limit[1] - limit[0]) if isinstance(limit, tuple) else limit + + def slice_dataset( dataset: Dataset, limit: int | tuple[int, int] | None, sample_id: str | int | list[str] | list[int] | list[str | int] | None, ) -> Dataset: + # a SampleSource seed is not the complete sample set — samples matching + # a requested id may be produced while the task runs, so unmatched + # sample_id entries neither warn nor error (the filter still applies to + # the seed; the task's dispatcher applies it to produced samples) + dynamic = getattr(dataset, DATASET_SAMPLE_SOURCE_ATTR, False) + if sample_id is not None: - # reduce to list of normalized sample ids + include_id = sample_id_filter(sample_id) sample_ids = sample_id if isinstance(sample_id, list) else [sample_id] - sample_id = [normalise_sample_id(id) for id in sample_ids] + normalised = [normalise_sample_id(id) for id in sample_ids] # validate all the sample ids and warn if they aren't in the dataset - all_sample_ids_raw = [sample.id for sample in dataset] - all_sample_ids = [normalise_sample_id(id) for id in all_sample_ids_raw] - for id in sample_id: - if id not in all_sample_ids: - warn_once( - logger, f"sample id '{id}' not found in dataset '{dataset.name}'." - ) - - # helper to check for a matching sample id - def include_sample(sample: Sample) -> bool: - id = normalise_sample_id(sample.id) - return any(fnmatch(id, pat) for pat in sample_id) + if not dynamic: + all_sample_ids = [normalise_sample_id(sample.id) for sample in dataset] + for id in normalised: + if id not in all_sample_ids: + warn_once( + logger, + f"sample id '{id}' not found in dataset '{dataset.name}'.", + ) # filter the dataset - filtered = dataset.filter(include_sample) + filtered = dataset.filter(lambda sample: include_id(sample.id)) # raise error if we got no hits - if len(filtered) == 0: - filter = ",".join([str(id) for id in sample_id]) + if len(filtered) == 0 and not dynamic: + filter = ",".join([str(id) for id in normalised]) + all_sample_ids_raw = [sample.id for sample in dataset] r = reprlib.Repr() r.maxlist = 8 raise PrerequisiteError( diff --git a/tests/test_sample_source.py b/tests/test_sample_source.py index ccba952e26..a8c497e5f0 100644 --- a/tests/test_sample_source.py +++ b/tests/test_sample_source.py @@ -318,3 +318,158 @@ async def solve(state: TaskState, generate: Generate) -> TaskState: def test_enqueue_sample_rejected_outside_task() -> None: with pytest.raises(RuntimeError, match="SampleSource"): enqueue_sample(Sample(input="x")) + + +class _TrackedGenerations(_Generations): + """`_Generations` that counts `next_samples()` calls.""" + + def __init__(self, count: int) -> None: + super().__init__(count) + self.next_calls = 0 + + async def next_samples(self) -> list[Sample] | None: + self.next_calls += 1 + return await super().next_samples() + + +def test_sample_source_limit_caps_total_samples() -> None: + # --limit caps the total samples (seed + produced): limit 3 with a + # 1-sample seed leaves a budget of 2, so a source that would produce 10 + # runs only s0..s2 (and the task still succeeds) + source = _TrackedGenerations(10) + logs = eval( + Task(dataset=source, solver=[generate()]), + model="mockllm/model", + display="none", + limit=3, + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["s0", "s1", "s2"] + + +def test_sample_source_limit_consumed_by_seed_skips_source() -> None: + # a limit fully consumed by the seed ends the task without consulting + # the source at all + source = _TrackedGenerations(10) + logs = eval( + Task(dataset=source, solver=[generate()]), + model="mockllm/model", + display="none", + limit=1, + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["s0"] + assert source.next_calls == 0 + + +def test_sample_source_limit_truncates_batch() -> None: + # a produced batch larger than the remaining budget is truncated + class _Batch(SampleSource): + def __init__(self) -> None: + self._produced = False + + def initial_samples(self) -> list[Sample]: + return [Sample(input="s0", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + return [Sample(input=f"b{i}", target="ok") for i in range(3)] + return None + + logs = eval( + Task(dataset=_Batch(), solver=[generate()]), + model="mockllm/model", + display="none", + limit=2, + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["b0", "s0"] + + +def test_sample_source_limit_counts_samples_not_epoch_runs() -> None: + # the limit counts samples: each sample within it still runs all epochs + logs = eval( + Task(dataset=_TrackedGenerations(10), solver=[generate()], epochs=2), + model="mockllm/model", + display="none", + limit=2, + ) + log = logs[0] + assert log.status == "success" + runs = sorted((str(sample.input), sample.epoch) for sample in (log.samples or [])) + assert runs == [("s0", 1), ("s0", 2), ("s1", 1), ("s1", 2)] + + +class _IdsSource(SampleSource): + """Seed ids 1/2; produces ids 3/4 once.""" + + def __init__(self) -> None: + self._produced = False + + def initial_samples(self) -> list[Sample]: + return [ + Sample(id=1, input="one", target="ok"), + Sample(id=2, input="two", target="ok"), + ] + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + return [ + Sample(id=3, input="three", target="ok"), + Sample(id=4, input="four", target="ok"), + ] + return None + + +def test_sample_source_sample_id_filters_produced_samples() -> None: + # --sample-id filters samples the source produces, not just the seed + logs = eval( + Task(dataset=_IdsSource(), solver=[generate()]), + model="mockllm/model", + display="none", + sample_id=[1, 3], + ) + log = logs[0] + assert log.status == "success" + assert sorted(sample.id for sample in (log.samples or [])) == [1, 3] + + +def test_sample_source_sample_id_missing_from_seed_ok() -> None: + # a requested id absent from the seed is not an error — the source may + # produce it while the task runs + logs = eval( + Task(dataset=_IdsSource(), solver=[generate()]), + model="mockllm/model", + display="none", + sample_id=4, + ) + log = logs[0] + assert log.status == "success" + assert [sample.id for sample in (log.samples or [])] == [4] + + +def test_sample_source_enqueue_during_terminal_next_samples() -> None: + # samples enqueued while next_samples() is returning None still run + # (the dispatcher drains once more before finishing) + class _Late(SampleSource): + def __init__(self) -> None: + self.calls = 0 + + def initial_samples(self) -> list[Sample]: + return [Sample(input="seed", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + self.calls += 1 + if self.calls == 1: + enqueue_sample(Sample(input="late", target="ok")) + return None + + logs = eval( + Task(dataset=_Late(), solver=[generate()]), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sample_inputs(logs[0]) == ["late", "seed"] From ca5f78a6fb96ff3e036ee87d6ec08245e89efcb5 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:10:16 +0000 Subject: [PATCH 3/7] Drop SampleEnqueuer lock (document loop-only contract) and document seed-sized log dataset spec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the two review nits from the latest review round: - SampleEnqueuer paired a threading.Lock with a wake callback (anyio.Event.set) that is not thread-safe, signalling two conflicting intents. Resolve by picking the loop-only intent: drop the lock (the buffer ops have no yield points under the event loop) and document in SampleEnqueuer and enqueue_sample() that calls must come from the task's event loop, not a worker thread. - Document the seed-sized log dataset spec as a known limit: the finished log's eval.dataset stays seed-sized (intentionally — it is load-bearing for retry reuse via eval_log_sample_source), so spec consumers (analysis evals dataframe, crash recovery) under-count for dynamic tasks. Noted in design/sample-source.md known limits and as a user-facing "Logs" note in docs/sample-source.qmd. Co-authored-by: Ransom Richardson <1209015+ransomr@users.noreply.github.com> Co-Authored-By: Claude Fable 5 --- design/sample-source.md | 12 ++++++++++++ docs/sample-source.qmd | 4 ++++ src/inspect_ai/_eval/task/sample_source.py | 17 ++++++++++------- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/design/sample-source.md b/design/sample-source.md index 91177001e7..62d40f6aa0 100644 --- a/design/sample-source.md +++ b/design/sample-source.md @@ -128,6 +128,18 @@ dropping samples. `TaskSource` + eval_set (see task-source.md). - **Early stopping** managers receive only the seed at `start_task`; `schedule_sample` / `complete_sample` still fire for injected samples. +- **The log's dataset spec stays seed-sized**: the finished log's + `eval.dataset.samples` / `sample_ids` describe the seed, not the grown set + (`log.samples` and `results.total_samples` do reflect everything that ran). + Consumers that read the spec as the planned set therefore under-count for + dynamic tasks — the analysis evals dataframe (`dataset_samples` / + `dataset_sample_ids` columns and the completion denominator in + `analysis/_dataframe/evals/table.py`) and crash recovery + (`log/_recover/_api.py` computes `total = dataset.samples * epochs`). This + is intentional, not an oversight: the seed-sized spec is load-bearing for + retry reuse — `eval_log_sample_source` rejects a prior log when + `eval.dataset.samples != len(dataset)`, and a retry's fresh seed is + seed-sized — so the spec must not be rewritten to the grown size at finish. - **Sandboxes**: the task-level sandbox startup pass runs once, up front, over the seed. Injected samples get per-sample sandboxes via `sandboxenv_context` as usual, but a *sample-level* sandbox spec appearing only on injected diff --git a/docs/sample-source.qmd b/docs/sample-source.qmd index b45fa9449a..0cd1f7d75c 100644 --- a/docs/sample-source.qmd +++ b/docs/sample-source.qmd @@ -129,3 +129,7 @@ A `SampleSource` task is live: a sample added mid-run starts as soon as there is The `--limit` option caps the total number of samples in the task — the seed plus everything the source produces. For example, `--limit 10` on a source with 5 seed samples allows 5 more generated samples; once the cap is reached, further additions are ignored (with a warning) and the task ends when the in-flight samples finish (`next_samples()` is not consulted again). The cap counts samples, not runs: each sample within the limit still runs for the task's configured number of epochs. The `--sample-id` option filters generated samples the same way it filters the seed: only samples whose ids match run. Since the source may produce a requested id while the task runs, it is not an error for a `--sample-id` value to be missing from the seed. + +## Logs + +All samples — seed and generated — are written to one log file, and `results.total_samples` reflects everything that ran. Note however that the log's *dataset* field (`eval.dataset.samples` and `eval.dataset.sample_ids`) describes only the seed: tools that read it as the planned sample count (for example the `dataset_samples` column in analysis dataframes) will under-count for a dynamic task, so use the results or the samples themselves for the true total. diff --git a/src/inspect_ai/_eval/task/sample_source.py b/src/inspect_ai/_eval/task/sample_source.py index eb84864375..8af6122760 100644 --- a/src/inspect_ai/_eval/task/sample_source.py +++ b/src/inspect_ai/_eval/task/sample_source.py @@ -31,7 +31,6 @@ from contextvars import ContextVar, Token from dataclasses import dataclass, field -from threading import Lock from typing import TYPE_CHECKING, Awaitable, Callable if TYPE_CHECKING: @@ -149,24 +148,26 @@ class SampleEnqueuer: Owned by ``task_run`` for the task's lifetime (only when the task has a :class:`SampleSource`). ``enqueue`` buffers; ``drain`` is the non-blocking pull the task's dispatch loop does between cycles. + + Not thread-safe: ``enqueue`` must be called from the task's event loop + (the ``on_enqueue`` wake callback — ``anyio.Event.set`` — is not safe to + fire from other threads). The buffer needs no lock under the event loop + since neither method has a yield point. """ _pending: list["Sample"] = field(default_factory=list) - _lock: Lock = field(default_factory=Lock) on_enqueue: Callable[[], None] | None = None """Fired after samples are buffered — used to wake the task's dispatcher.""" def enqueue(self, samples: list["Sample"]) -> None: """Queue ``samples`` to run in the task.""" - with self._lock: - self._pending.extend(samples) + self._pending.extend(samples) if self.on_enqueue is not None: self.on_enqueue() def drain(self) -> list["Sample"]: """Remove and return all currently-buffered samples (empty if none).""" - with self._lock: - batch, self._pending = self._pending, [] + batch, self._pending = self._pending, [] return batch @@ -202,7 +203,9 @@ def enqueue_sample(samples: "Sample | list[Sample]") -> None: Only available inside a task driven by a :class:`SampleSource` (i.e. a ``Task`` whose ``dataset`` is a ``SampleSource``) — a plain task's sample set is fixed, so there is no loop to run additions. Callable from any code - running within such a task: a solver, a scorer, a tool. + running within such a task — a solver, a scorer, a tool — but it must be + called from the task's event loop (where those all run), not from a + worker thread. When the eval was run with ``--limit``, samples beyond the limit are ignored (with a warning); with ``--sample-id``, only samples matching the From d2d69bd9c484dd091073a3d05407d05c0468b445 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:18:20 +0000 Subject: [PATCH 4/7] Re-trigger CI after transient 504 downloading Python in test jobs CI run 29032396298 failed only because uv got a 504 Gateway Timeout fetching the python-build-standalone 3.10 archive from GitHub releases; no code changes needed. Co-Authored-By: Claude Fable 5 From 5c812f2a1353c0892594be4f8a592b66bf6b8b43 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:33:01 +0000 Subject: [PATCH 5/7] Re-trigger CI again after repeated 504 on python-build-standalone download Co-Authored-By: Claude Fable 5 From b3d5e9483114739b2e36b9c3c5eb10922fe99639 Mon Sep 17 00:00:00 2001 From: Ransom Richardson Date: Mon, 13 Jul 2026 18:10:42 -0400 Subject: [PATCH 6/7] Run sandbox startup for samples a SampleSource adds mid-run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously only the seed dataset drove the run-level sandbox startup pass, so a sandbox config first seen in an added sample (or a task-level config with an empty seed) never got task_init — no up-front image build/pull, no fail-fast validation, and no registered task_cleanup. Route added samples through the same incremental SandboxManager pass used for injected tasks: the dynamic dispatcher now awaits sandbox startup for each added batch before spawning it, with already-started configs a cheap no-op and initialization serialized against concurrent live tasks. Co-Authored-By: Claude Fable 5 --- docs/sample-source.qmd | 4 + src/inspect_ai/_eval/run.py | 122 +++++++++---- src/inspect_ai/_eval/task/run.py | 27 ++- src/inspect_ai/_eval/task/sample_source.py | 5 +- tests/test_sample_source.py | 202 +++++++++++++++++++++ 5 files changed, 320 insertions(+), 40 deletions(-) diff --git a/docs/sample-source.qmd b/docs/sample-source.qmd index 0cd1f7d75c..82e90194f2 100644 --- a/docs/sample-source.qmd +++ b/docs/sample-source.qmd @@ -124,6 +124,10 @@ Enqueued samples share a buffer with samples returned from `sample_complete`. `e A `SampleSource` task is live: a sample added mid-run starts as soon as there is free capacity, rather than waiting for the current samples to finish. Concurrency is bounded by `max_samples` as usual (see [Parallelism](parallelism.qmd)). +## Sandboxes + +Samples added mid-run get the same [sandbox](sandboxing.qmd) startup as the seed: a sandbox configuration first seen in an added sample (including the task's own configuration when the seed is empty) is initialized — images built/pulled, configuration validated, cleanup registered — before that sample runs, and configurations the seed already initialized are not re-initialized. The seed is therefore not required for sandboxes: a source driven entirely by `next_samples()` works with sandboxed samples too. + ## Limits and sample filtering The `--limit` option caps the total number of samples in the task — the seed plus everything the source produces. For example, `--limit 10` on a source with 5 seed samples allows 5 more generated samples; once the cap is reached, further additions are ignored (with a warning) and the task ends when the in-flight samples finish (`next_samples()` is not consulted again). The cap counts samples, not runs: each sample within the limit still runs for the task's configured number of epochs. diff --git a/src/inspect_ai/_eval/run.py b/src/inspect_ai/_eval/run.py index b3703849c6..09af608628 100644 --- a/src/inspect_ai/_eval/run.py +++ b/src/inspect_ai/_eval/run.py @@ -1,3 +1,4 @@ +import functools import logging import os import sys @@ -29,7 +30,7 @@ from inspect_ai._eval.task.scan import Scanners from inspect_ai._util.error import PrerequisiteError, exception_message from inspect_ai._util.path import chdir -from inspect_ai.dataset._dataset import Dataset +from inspect_ai.dataset._dataset import Dataset, Sample from inspect_ai.log import EvalConfig, EvalLog from inspect_ai.log._file import EvalLogInfo from inspect_ai.log._log import eval_error @@ -322,6 +323,16 @@ async def prepare_options( ) await logger.init() + # samples a SampleSource adds mid-run go through the same + # incremental sandbox startup pass as the seed dataset above + startup_sandboxes: Callable[[list[Sample]], Awaitable[None]] | None = ( + None + ) + if run_samples and task.sample_source is not None: + startup_sandboxes = functools.partial( + sandbox_manager.start_for_samples, resolved_task + ) + # append task task_run_options.append( TaskRunOptions( @@ -346,6 +357,7 @@ async def prepare_options( initial_model_usage=resolved_task.initial_model_usage, initial_role_usage=resolved_task.initial_role_usage, task_source=task_source, + startup_sandboxes=startup_sandboxes, ) ) # register the prepared task so a failed run can clean it up @@ -902,9 +914,10 @@ def collect_for_task(sample: str | int) -> str | int | None: class SandboxManager: """Starts sandbox environments incrementally and tears them all down. - Tasks injected into a live run arrive in batches, so startup must be - callable repeatedly — :meth:`start` initializes only sandboxenvs not already - started and accumulates their cleanups; :meth:`shutdown` runs every + Tasks injected into a live run — and samples a `SampleSource` adds to a + live task — arrive in batches, so startup must be callable repeatedly: + :meth:`start` / :meth:`start_for_samples` initialize only sandboxenvs not + already started and accumulate their cleanups; :meth:`shutdown` runs every accumulated cleanup once, at the end of the run. """ @@ -919,9 +932,10 @@ def __init__( self._cleanups: list[ tuple[TaskCleanup, SandboxEnvironmentConfigType | None, str] ] = [] + self._init_lock = anyio.Lock() async def start(self, tasks: list[ResolvedTask]) -> None: - # find unique sandboxenvs not already started + # find unique sandboxenvs to start sandboxenvs: Set[TaskSandboxEnvironment] = set() for task in tasks: # resolve each sample and add to sandboxenvs @@ -935,45 +949,81 @@ async def start(self, tasks: list[ResolvedTask]) -> None: sandbox = await resolve_sandbox_for_task_and_sample( task.sandbox, task.task, sample ) - if ( - sandbox is not None - and sandbox not in self._started - and sandbox not in sandboxenvs - ): + if sandbox is not None and sandbox not in sandboxenvs: sandboxenvs.add(sandbox) + await self._start_sandboxenvs(sandboxenvs) + + async def start_for_samples( + self, task: ResolvedTask, samples: list[Sample] + ) -> None: + """Initialize the sandboxenvs of samples added to a live task. + + The counterpart of :meth:`start` for samples a `SampleSource` adds + while its task runs: a sandbox config first seen in an added sample + (including the task-level config when the seed dataset was empty) + gets the same `task_init` startup as seed configs — up-front image + build/pull, fail-fast validation, and a registered `task_cleanup` — + before the sample runs. Configs already started are a no-op. + """ + sandboxenvs: Set[TaskSandboxEnvironment] = set() + for sample in samples: + sandbox = await resolve_sandbox_for_task_and_sample( + task.sandbox, task.task, sample + ) + if sandbox is not None and sandbox not in sandboxenvs: + sandboxenvs.add(sandbox) + + await self._start_sandboxenvs(sandboxenvs) + + async def _start_sandboxenvs( + self, sandboxenvs: Set[TaskSandboxEnvironment] + ) -> None: + """Run `task_init` for sandboxenvs not already started. + + Serialized on a lock with the started-set re-checked under it, so + concurrent callers (e.g. two live tasks adding samples at once) + can't double-init the same sandboxenv. + """ if not sandboxenvs: return - # initialiase sandboxenvs (track cleanups) - with display().suspend_task_app(): - for sandboxenv in sandboxenvs: - # find type - sandboxenv_type = registry_find_sandboxenv(sandboxenv.sandbox.type) - - # pre-register the type's resizable concurrency limiter before - # task_init (image pulls/builds can take minutes) so a `ctl - # limits --max-sandboxes` issued during startup isn't dropped - await ensure_sandbox_limiter( - sandboxenv_type, sandboxenv.sandbox.type, self._config.max_sandboxes - ) + async with self._init_lock: + sandboxenvs = {env for env in sandboxenvs if env not in self._started} + if not sandboxenvs: + return + + # initialiase sandboxenvs (track cleanups) + with display().suspend_task_app(): + for sandboxenv in sandboxenvs: + # find type + sandboxenv_type = registry_find_sandboxenv(sandboxenv.sandbox.type) + + # pre-register the type's resizable concurrency limiter before + # task_init (image pulls/builds can take minutes) so a `ctl + # limits --max-sandboxes` issued during startup isn't dropped + await ensure_sandbox_limiter( + sandboxenv_type, + sandboxenv.sandbox.type, + self._config.max_sandboxes, + ) - # run startup - task_init = cast(TaskInit, getattr(sandboxenv_type, "task_init")) - with chdir(sandboxenv.run_dir), environ_vars(dict(sandboxenv.env)): - await task_init("startup", sandboxenv.sandbox.config) + # run startup + task_init = cast(TaskInit, getattr(sandboxenv_type, "task_init")) + with chdir(sandboxenv.run_dir), environ_vars(dict(sandboxenv.env)): + await task_init("startup", sandboxenv.sandbox.config) - # track as started and append cleanup method - self._started.add(sandboxenv) - task_cleanup = cast( - TaskCleanup, getattr(sandboxenv_type, "task_cleanup") - ) - self._cleanups.append( - (task_cleanup, sandboxenv.sandbox.config, sandboxenv.run_dir) - ) + # track as started and append cleanup method + self._started.add(sandboxenv) + task_cleanup = cast( + TaskCleanup, getattr(sandboxenv_type, "task_cleanup") + ) + self._cleanups.append( + (task_cleanup, sandboxenv.sandbox.config, sandboxenv.run_dir) + ) - # provide some space above task display - print("") + # provide some space above task display + print("") async def shutdown(self) -> None: with anyio.CancelScope(shield=True): diff --git a/src/inspect_ai/_eval/task/run.py b/src/inspect_ai/_eval/task/run.py index 73a4782f86..73efde17cf 100644 --- a/src/inspect_ai/_eval/task/run.py +++ b/src/inspect_ai/_eval/task/run.py @@ -284,6 +284,10 @@ class TaskRunOptions: initial_role_usage: dict[str, ModelUsage] | None = field(default=None) task_source: "TaskSource | None" = field(default=None) """Run-level task source notified as this task's samples/task complete.""" + startup_sandboxes: Callable[[list[Sample]], Awaitable[None]] | None = field( + default=None + ) + """Run-level incremental sandbox startup for samples a SampleSource adds.""" def resolve_plan(task: Task, solver: Solver | None) -> Plan: @@ -914,6 +918,23 @@ def add_samples(samples: list[Sample]) -> list[int]: ) return indexes + async def add_and_start(samples: list[Sample]) -> list[int]: + """Add samples, starting any sandboxenvs they need first. + + Added samples get the same run-level sandbox startup as + the seed (``task_init`` for configs not seen before: + image build/pull, validation, cleanup registration) + before they spawn; already-started configs are a cheap + no-op. A startup failure propagates and fails the task, + matching a seed config failing startup. + """ + indexes = add_samples(samples) + if indexes and options.startup_sandboxes is not None: + await options.startup_sandboxes( + [get_sample(index) for index in indexes] + ) + return indexes + async with anyio.create_task_group() as tg: async def run_one(sample_index: int, epoch: int) -> None: @@ -941,7 +962,7 @@ def spawn(indexes: list[int]) -> None: # returning samples / enqueue_sample) start first buffered = enqueuer.drain() if buffered: - spawn(add_samples(buffered)) + spawn(await add_and_start(buffered)) continue if in_flight > 0: await wake.wait() @@ -959,12 +980,12 @@ def spawn(indexes: list[int]) -> None: # was finishing rather than dropping them # (enqueue_sample never drops silently); if # any run, the source may be consulted again - final = add_samples(enqueuer.drain()) + final = await add_and_start(enqueuer.drain()) if not final: break spawn(final) elif more: - spawn(add_samples(more)) + spawn(await add_and_start(more)) return results diff --git a/src/inspect_ai/_eval/task/sample_source.py b/src/inspect_ai/_eval/task/sample_source.py index 8af6122760..1944a39cfa 100644 --- a/src/inspect_ai/_eval/task/sample_source.py +++ b/src/inspect_ai/_eval/task/sample_source.py @@ -53,7 +53,10 @@ def initial_samples(self) -> list["Sample"]: return immediately (no awaiting / blocking). The returned samples drive the task's up-front setup (validation, sandbox startup) and are the first batch. May be empty, in which case the task starts by calling - ``next_samples()``. + ``next_samples()``. The seed isn't required for sandboxes: a sandbox + config first seen in a later-added sample gets the same startup + (image build/pull, validation, registered cleanup) before that sample + runs. """ return [] diff --git a/tests/test_sample_source.py b/tests/test_sample_source.py index a8c497e5f0..d1f91c3376 100644 --- a/tests/test_sample_source.py +++ b/tests/test_sample_source.py @@ -6,13 +6,24 @@ samples are live: they start as soon as there is free capacity. """ +from typing import Any, Literal, overload + import anyio import pytest +from pydantic import BaseModel +from typing_extensions import override from inspect_ai import SampleSource, Task, enqueue_sample, eval from inspect_ai.dataset import Sample from inspect_ai.log import EvalLog, EvalSample from inspect_ai.solver import Generate, Solver, TaskState, generate, solver +from inspect_ai.util import ( + ExecResult, + SandboxEnvironment, + SandboxEnvironmentConfigType, + SandboxEnvironmentSpec, + sandboxenv, +) def _sample_inputs(log: EvalLog) -> list[str]: @@ -473,3 +484,194 @@ async def next_samples(self) -> list[Sample] | None: ) assert logs[0].status == "success" assert _sample_inputs(logs[0]) == ["late", "seed"] + + +# frozen so specs embedding it stay hashable (sandbox startup sets them) +class _SandboxConfig(BaseModel, frozen=True): + name: str + + +_sandbox_events: list[str] = [] + + +def _config_name(config: SandboxEnvironmentConfigType | None) -> str: + assert isinstance(config, _SandboxConfig) + return config.name + + +@sandboxenv(name="source_events") +class _EventsSandbox(SandboxEnvironment): + """Sandboxenv that records task/sample lifecycle events by config name.""" + + @override + @classmethod + def config_deserialize(cls, config: dict[str, Any]) -> BaseModel: + return _SandboxConfig(**config) + + @override + @classmethod + async def task_init( + cls, task_name: str, config: SandboxEnvironmentConfigType | None + ) -> None: + _sandbox_events.append(f"task_init:{_config_name(config)}") + + @override + @classmethod + async def sample_init( + cls, + task_name: str, + config: SandboxEnvironmentConfigType | None, + metadata: dict[str, str], + ) -> dict[str, SandboxEnvironment]: + _sandbox_events.append(f"sample_init:{_config_name(config)}") + return {"default": _EventsSandbox()} + + @override + @classmethod + async def sample_cleanup( + cls, + task_name: str, + config: SandboxEnvironmentConfigType | None, + environments: dict[str, SandboxEnvironment], + interrupted: bool, + ) -> None: + pass + + @override + @classmethod + async def task_cleanup( + cls, task_name: str, config: SandboxEnvironmentConfigType | None, cleanup: bool + ) -> None: + _sandbox_events.append(f"task_cleanup:{_config_name(config)}") + + @override + async def exec( + self, + cmd: list[str], + input: str | bytes | None = None, + cwd: str | None = None, + env: dict[str, str] | None = None, + user: str | None = None, + timeout: int | None = None, + timeout_retry: bool = True, + concurrency: bool = True, + ) -> ExecResult[str]: + raise NotImplementedError + + @override + async def write_file(self, file: str, contents: str | bytes) -> None: + raise NotImplementedError + + @overload + async def read_file(self, file: str, text: Literal[True] = True) -> str: ... + + @overload + async def read_file(self, file: str, text: Literal[False]) -> bytes: ... + + @override + async def read_file(self, file: str, text: bool = True) -> str | bytes: + raise NotImplementedError + + +def test_sample_source_sandbox_startup_without_seed() -> None: + # sandboxes don't require a seed: with an empty seed and a task-level + # sandbox, the config is initialized (task_init) before the first added + # sample runs, and its task_cleanup runs at the end of the run + _sandbox_events.clear() + + class _Src(SampleSource): + def __init__(self) -> None: + self._produced = False + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + return [Sample(input="added", target="ok")] + return None + + logs = eval( + Task( + dataset=_Src(), + solver=[generate()], + sandbox=SandboxEnvironmentSpec("source_events", _SandboxConfig(name="A")), + ), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sandbox_events.index("task_init:A") < _sandbox_events.index("sample_init:A") + assert "task_cleanup:A" in _sandbox_events + + +def test_sample_source_sandbox_startup_for_added_config() -> None: + # a sandbox config first seen in added samples gets task_init once (not + # per sample) before any of them runs; the seed's config isn't re-inited + _sandbox_events.clear() + + class _Src(SampleSource): + def __init__(self) -> None: + self._produced = False + + def initial_samples(self) -> list[Sample]: + return [Sample(input="seed", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + sandbox = SandboxEnvironmentSpec( + "source_events", _SandboxConfig(name="B") + ) + return [ + Sample(input="added1", target="ok", sandbox=sandbox), + Sample(input="added2", target="ok", sandbox=sandbox), + ] + return None + + logs = eval( + Task( + dataset=_Src(), + solver=[generate()], + sandbox=SandboxEnvironmentSpec("source_events", _SandboxConfig(name="A")), + ), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sandbox_events.count("task_init:A") == 1 + assert _sandbox_events.count("task_init:B") == 1 + assert _sandbox_events.index("task_init:B") < _sandbox_events.index("sample_init:B") + assert _sandbox_events.count("sample_init:B") == 2 + assert "task_cleanup:A" in _sandbox_events + assert "task_cleanup:B" in _sandbox_events + + +def test_sample_source_sandbox_same_config_not_reinited() -> None: + # added samples reusing the seed's config don't trigger another task_init + _sandbox_events.clear() + + class _Src(SampleSource): + def __init__(self) -> None: + self._produced = False + + def initial_samples(self) -> list[Sample]: + return [Sample(input="seed", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + return [Sample(input="added", target="ok")] + return None + + logs = eval( + Task( + dataset=_Src(), + solver=[generate()], + sandbox=SandboxEnvironmentSpec("source_events", _SandboxConfig(name="A")), + ), + model="mockllm/model", + display="none", + ) + assert logs[0].status == "success" + assert _sandbox_events.count("task_init:A") == 1 + assert _sandbox_events.count("sample_init:A") == 2 + assert _sandbox_events.count("task_cleanup:A") == 1 From df9c1f5afb1316ac00d0fcdeec7272980af12596 Mon Sep 17 00:00:00 2001 From: Ransom Richardson Date: Tue, 14 Jul 2026 10:55:35 -0400 Subject: [PATCH 7/7] Review fixes for SampleSource: retries, early-stop/range-limit, control channel Correctness fixes from a full review pass of the SampleSource feature: - Task retries: the prior-attempt reuse path now re-fires sample_complete for reused samples, so completion-driven sources regenerate their follow-ups instead of silently dropping them from the retry log; the same-instance retry contract is documented. - Reject early_stopping for SampleSource tasks (managers register the seed only, and halted samples complete without notifying the source, hanging sources blocked on completions) and range --limit (positional for the seed but a bare count budget for additions). Both raise PrerequisiteError in prepare_options, before the task logger exists. - Defer fractional fail_on_error thresholds to the end-of-run check for SampleSource tasks: the planned total grows while the task runs, so a mid-run check measured early errors against a transiently small denominator. - Skip the source notification for cancelled samples: it was awaited under a cancellation shield, so a never-resolving user callback made teardown uncancellable. - EvalState gains a dynamic flag: a SampleSource eval's counters reaching total no longer stamps completed_at (an idle eval blocked in next_samples() read as finished, making ctl task cancel a no-op); finalize_eval stamps at the task's true finish point. Plus efficiency and cleanup fixes: seen_ids no longer re-reads the disk-paged seed store; injected sample slots are released after their last epoch (bounded memory for open-ended sources); sandbox resolution is cached for metadata-less samples and the init lock has a pre-lock fast path; slice_dataset takes an explicit dynamic parameter instead of a setattr marker on the dataset; sample_id_filter returns a NamedTuple so slice_dataset stops duplicating its normalization; SandboxManager's resolve loop is shared between start and start_for_samples. Co-Authored-By: Claude Fable 5 --- design/sample-source.md | 73 ++++++---- docs/sample-source.qmd | 8 +- src/inspect_ai/_control/eval_state.py | 38 +++-- src/inspect_ai/_eval/run.py | 93 +++++++++--- src/inspect_ai/_eval/task/constants.py | 1 - src/inspect_ai/_eval/task/error.py | 36 ++++- src/inspect_ai/_eval/task/log.py | 6 +- src/inspect_ai/_eval/task/run.py | 142 ++++++++++++------ src/inspect_ai/_eval/task/sample_source.py | 5 + src/inspect_ai/_eval/task/task.py | 13 +- src/inspect_ai/_eval/task/util.py | 47 +++--- tests/_control/test_eval_state.py | 19 +++ tests/test_sample_source.py | 159 ++++++++++++++++++++- 13 files changed, 512 insertions(+), 128 deletions(-) diff --git a/design/sample-source.md b/design/sample-source.md index 62d40f6aa0..c10568c4b9 100644 --- a/design/sample-source.md +++ b/design/sample-source.md @@ -85,12 +85,16 @@ A `SampleSource` is passed as the **`dataset` argument to `Task`** (just as a display denominator (`td.sample_complete`), the fractional `fail_on_error` threshold (`SampleErrorHandler.total_samples`), the end-of-run `eval_results` / `_should_eval_fail` counts, and the control-channel state - via `record_samples_added` ([eval_state.py](../src/inspect_ai/_control/eval_state.py)) - — which also un-stamps a provisional `completed_at` (every planned sample - can be terminal while the source is still producing). `register_eval` now - copies its `sample_ids` argument so the state's planned-ids list grows only - via `record_samples_added` (task_run mutates its own list for - carry-forward). + via `record_samples_added` ([eval_state.py](../src/inspect_ai/_control/eval_state.py)). + The eval registers as `dynamic`, which suppresses the provisional + `completed_at` stamp entirely — every planned sample can be terminal while + the source is still producing (or the seed can be empty, `total == 0`), and + a stamped `completed_at` would make `ctl task cancel` no-op and listings + read "completed" while the task idles in `next_samples()`; only + `finalize_eval` (the task's true finish point) clears the flag and stamps. + `register_eval` also copies its `sample_ids` argument so the state's + planned-ids list grows only via `record_samples_added` (task_run mutates + its own list for carry-forward). ## `enqueue_sample` @@ -111,23 +115,40 @@ dropping samples. (`sample_limit_count(limit) - seed`) as its budget for added samples — additions beyond it are ignored with a warning, and once the budget is exhausted the loop finishes without consulting `next_samples()` again. A - tuple limit's budget is the size of its slice (`stop - start`). + *range* limit (`start,end`) is rejected with a `PrerequisiteError`: it + selects seed samples by position, and added samples have no position, so + there is no coherent way to apply it (it would select nothing from a short + seed while still admitting `stop - start` additions). +- **`fail_on_error`**: a fractional threshold is *deferred* to the end-of-run + check for SampleSource tasks (`SampleErrorHandler(defer_fractional=True)`) — + the planned total grows while the task runs, so a mid-run check would + measure early errors against a transiently small denominator (a 1-sample + seed erroring first would trip `1 >= 0.5*1` before the source produced the + rest). Absolute-count and any-error thresholds still abort mid-run. - **`--sample-id`** filters produced samples with the same normalise+`fnmatch` predicate that filters the seed (`sample_id_filter` in - task/util.py). The seed of a SampleSource task is marked with - `DATASET_SAMPLE_SOURCE_ATTR` so `slice_dataset` (log spec, sandbox startup, - task_run) neither warns nor raises for requested ids missing from the seed — - the source may produce them at runtime. Filtered-out samples still register - their ids (duplicate ids stay a hard error) but don't run and don't grow the - planned totals. As in `slice_dataset`, `--sample-id` and `--limit` are - mutually exclusive (the filter wins). + task/util.py). Callers that hold the task pass `dynamic=True` to + `slice_dataset` (log spec, sandbox startup, task_run) so it neither warns + nor raises for requested ids missing from the seed — the source may produce + them at runtime. Filtered-out samples still register their ids (duplicate + ids stay a hard error) but don't run and don't grow the planned totals. As + in `slice_dataset`, `--sample-id` and `--limit` are mutually exclusive (the + filter wins). - **Task retries** (`task_retry_attempts` / eval-set): the retry attempt - re-drives the source; completed samples are reused via the normal - `EvalSampleSource` lookup only where regenerated ids match — i.e. resume - pays off for deterministic sources, the same determinism contract as - `TaskSource` + eval_set (see task-source.md). -- **Early stopping** managers receive only the seed at `start_task`; - `schedule_sample` / `complete_sample` still fire for injected samples. + re-drives the *same source instance* (the retry rebuilds `TaskRunOptions` + around the same `Task`); completed samples are reused via the normal + `EvalSampleSource` lookup, and the reuse path re-fires + `sample_feed.sample_complete` for each reused sample so a completion-driven + source regenerates its follow-ups (which are then themselves reused where + regenerated ids match). A source that instead holds internal + `next_samples()` state resumes mid-state on retry — it must be resumable + (or derive its follow-ups from `sample_complete`) for retries to + reconstruct the run; this is the same determinism contract as `TaskSource` + + eval_set (see task-source.md). +- **Early stopping** is rejected (`PrerequisiteError`): managers register a + fixed sample set at `start_task` (added samples would never be registered), + and samples a manager halts complete without notifying the source, which + would hang a source blocked awaiting completions. - **The log's dataset spec stays seed-sized**: the finished log's `eval.dataset.samples` / `sample_ids` describe the seed, not the grown set (`log.samples` and `results.total_samples` do reflect everything that ran). @@ -140,10 +161,14 @@ dropping samples. retry reuse — `eval_log_sample_source` rejects a prior log when `eval.dataset.samples != len(dataset)`, and a retry's fresh seed is seed-sized — so the spec must not be rewritten to the grown size at finish. -- **Sandboxes**: the task-level sandbox startup pass runs once, up front, over - the seed. Injected samples get per-sample sandboxes via `sandboxenv_context` - as usual, but a *sample-level* sandbox spec appearing only on injected - samples won't have had its task-level `task_init` startup pass. +- **Sandboxes**: the task-level sandbox startup pass runs up front over the + seed, and *incrementally* for samples the source adds: each added batch goes + through `SandboxManager.start_for_samples` (threaded into the dispatcher as + `TaskRunOptions.startup_sandboxes`) before spawning, so a sandbox config + first seen in an added sample — including the task-level config when the + seed is empty — still gets `task_init` (image build/pull, validation) and a + registered `task_cleanup`. Already-started configs are a set-membership + no-op. Per-sample sandboxes then run via `sandboxenv_context` as usual. - **Progress bar steps** (`profile.steps`) are fixed at seed size; the completed/total counter grows correctly (total passed on each update), and a zero-step seed no longer divides by zero (`RichProgress.update` guards it). diff --git a/docs/sample-source.qmd b/docs/sample-source.qmd index 82e90194f2..d9e18ad609 100644 --- a/docs/sample-source.qmd +++ b/docs/sample-source.qmd @@ -130,10 +130,16 @@ Samples added mid-run get the same [sandbox](sandboxing.qmd) startup as the seed ## Limits and sample filtering -The `--limit` option caps the total number of samples in the task — the seed plus everything the source produces. For example, `--limit 10` on a source with 5 seed samples allows 5 more generated samples; once the cap is reached, further additions are ignored (with a warning) and the task ends when the in-flight samples finish (`next_samples()` is not consulted again). The cap counts samples, not runs: each sample within the limit still runs for the task's configured number of epochs. +The `--limit` option caps the total number of samples in the task — the seed plus everything the source produces. For example, `--limit 10` on a source with 5 seed samples allows 5 more generated samples; once the cap is reached, further additions are ignored (with a warning) and the task ends when the in-flight samples finish (`next_samples()` is not consulted again). The cap counts samples, not runs: each sample within the limit still runs for the task's configured number of epochs. A *range* limit (`--limit 10,20`) is not supported for `SampleSource` tasks — it selects samples by dataset position, which generated samples don't have. The `--sample-id` option filters generated samples the same way it filters the seed: only samples whose ids match run. Since the source may produce a requested id while the task runs, it is not an error for a `--sample-id` value to be missing from the seed. +A *fractional* `fail_on_error` threshold (e.g. `0.5`) is evaluated at the end of the run rather than mid-run for `SampleSource` tasks: the planned total grows while the task runs, so a mid-run check would measure early errors against a transiently small denominator. Absolute-count and any-error thresholds behave as usual. The `early_stopping` task option is not supported with a `SampleSource` (managers require a fixed sample set). + +## Retries + +On a task retry (`eval_set` / task retry attempts), completed samples from the prior attempt are reused, and the source's `sample_complete` is called for each reused sample — so a source that derives its follow-ups from `sample_complete` (the patterns above) regenerates them naturally, and regenerated samples whose ids match the prior attempt are themselves reused rather than re-run. Note that the retry re-drives the *same source instance*: a source that instead holds internal `next_samples()` state resumes where it left off, so such sources should be written to be resumable if used with retries. + ## Logs All samples — seed and generated — are written to one log file, and `results.total_samples` reflects everything that ran. Note however that the log's *dataset* field (`eval.dataset.samples` and `eval.dataset.sample_ids`) describes only the seed: tools that read it as the planned sample count (for example the `dataset_samples` column in analysis dataframes) will under-count for a dynamic task, so use the results or the samples themselves for the true total. diff --git a/src/inspect_ai/_control/eval_state.py b/src/inspect_ai/_control/eval_state.py index 373bfb2fd8..20b35a12ef 100644 --- a/src/inspect_ai/_control/eval_state.py +++ b/src/inspect_ai/_control/eval_state.py @@ -246,7 +246,19 @@ class EvalState: """Unix timestamp when this eval's last sample finished (i.e. when ``completed + errored`` first reached ``total``). ``None`` while the eval is still running. Used by the control endpoint to surface - completion to agents without forcing them to derive it from counters.""" + completion to agents without forcing them to derive it from counters. + For a :attr:`dynamic` eval this is only stamped by :func:`finalize_eval` + (the task's finish point) — counters reaching ``total`` doesn't mean + done when the source can still add samples.""" + + dynamic: bool = False + """Whether this eval's planned sample set can grow while it runs (a + ``SampleSource``-driven task). While set, ``terminal >= total`` is not + proof of completion — the task may be idle awaiting its source (or have + an empty seed, ``total == 0``, at registration) — so the provisional + ``completed_at`` stamp is suppressed and consumers (task cancel, status + listings) correctly see the eval as running. Cleared by + :func:`finalize_eval`, the task's single true finish point.""" started_at: float | None = None """Earliest observed sample-start time, tracked as a running minimum. @@ -351,6 +363,7 @@ def register_eval( run_id: str | None = None, will_retry: bool = False, task_cancel: "TaskCancel | None" = None, + dynamic: bool = False, ) -> EvalState: """Initialize tracking for a new eval. @@ -379,12 +392,15 @@ def register_eval( run_id=run_id, will_retry=will_retry, task_cancel=task_cancel, + dynamic=dynamic, ) _eval_states[eval_id] = state # A zero-sample eval (``total == 0``, eg. a limit past the dataset) is # already finished — no sample will ever run to fire a terminal counter # and stamp ``completed_at`` via record_sample_*, so do it now. A no-op - # for the normal ``total > 0`` case (not yet finished at registration). + # for the normal ``total > 0`` case (not yet finished at registration) + # and for a dynamic eval (an empty seed just means the source hasn't + # produced yet — finalize_eval stamps it when the task truly ends). _maybe_mark_finished(state) return state @@ -559,9 +575,9 @@ def record_samples_added( Called by ``task_run`` when a ``SampleSource`` injects samples mid-run: ``total`` is the number of additional planned runs (samples × epochs) and ``sample_ids`` the injected ids (so the per-sample listing can surface them - as pending). If the eval had already been provisionally marked finished - (every previously-planned sample terminal while the source was still - producing), the finish stamp is cleared. No-ops if unregistered. + as pending). The eval is :attr:`EvalState.dynamic`, so no provisional + finish stamp needs clearing here — ``completed_at`` stays ``None`` until + :func:`finalize_eval`. No-ops if unregistered. """ with _lock: state = _eval_states.get(eval_id) @@ -569,8 +585,6 @@ def record_samples_added( state.total += total if sample_ids: state.sample_ids.extend(sample_ids) - if not state.is_finished: - state.completed_at = None def latest_eval_for_task(task_id: str) -> "EvalState | None": @@ -658,6 +672,9 @@ def finalize_eval(eval_id: str) -> None: shortfall = state.total - state.terminal if shortfall > 0: state.cancelled += shortfall + # the task has truly finished — no source can add samples now, so + # the dynamic suppression of the finish stamp no longer applies + state.dynamic = False _maybe_mark_finished(state) @@ -667,12 +684,15 @@ def _maybe_mark_finished(state: EvalState) -> None: Fires the first time the terminal sum (``completed + errored + cancelled``) reaches ``total``; later updates are no-ops so a late counter update from a teardown race doesn't overwrite the original - finish time. Also drops + finish time. Suppressed for a :attr:`EvalState.dynamic` eval — its + counters reaching ``total`` doesn't mean done (the source may add more + samples); ``finalize_eval`` clears the flag at the task's true finish + point. Also drops ``sample_ids`` — a finished eval has no pending samples, so the planned-id list is dead weight (it's retained on the state until the run boundary clears it). Caller must hold the registry lock. """ - if state.completed_at is None and state.is_finished: + if state.completed_at is None and not state.dynamic and state.is_finished: state.completed_at = time.time() state.sample_ids = [] diff --git a/src/inspect_ai/_eval/run.py b/src/inspect_ai/_eval/run.py index 20542fe44e..ec3fba31a8 100644 --- a/src/inspect_ai/_eval/run.py +++ b/src/inspect_ai/_eval/run.py @@ -3,7 +3,7 @@ import os import sys from dataclasses import dataclass, replace -from typing import Any, Awaitable, Callable, NamedTuple, Set, cast +from typing import Any, Awaitable, Callable, Iterable, NamedTuple, Set, cast from inspect_ai._eval.task.constants import TASK_ALL_PARAMS_ATTR from inspect_ai._util._async import Wake @@ -49,6 +49,7 @@ from inspect_ai.util._display import display_type from inspect_ai.util._sandbox.environment import ( SandboxEnvironmentConfigType, + SandboxEnvironmentSpec, TaskCleanup, TaskInit, ) @@ -70,8 +71,10 @@ from .task.sandbox import ( TaskSandboxEnvironment, ensure_sandbox_limiter, + resolve_sandbox, resolve_sandbox_for_task_and_sample, ) +from .task.task import Task from .task.task_source import TaskSource from .task.util import slice_dataset, task_run_dir @@ -201,6 +204,31 @@ async def prepare_options( resolved_task.task.name, task_eval_config.sample_id ) + # reject options that assume a fixed sample set for a + # SampleSource-driven task — here, before the task's logger + # is created, so the config error surfaces pre-run rather + # than as a task error log + if task.sample_source is not None: + # early stopping managers register the seed only (samples + # the source adds are never registered) and samples they + # halt complete without notifying the source — a source + # blocked awaiting completions would hang the task + if task.early_stopping is not None: + raise PrerequisiteError( + "early_stopping is not currently supported for " + "tasks whose dataset is a SampleSource." + ) + # a range limit selects seed samples by position, but + # samples the source adds have no position — there is no + # coherent way to apply it + if isinstance(task_eval_config.limit, tuple): + raise PrerequisiteError( + "A range --limit (start,end) is not supported for " + "tasks whose dataset is a SampleSource (added " + "samples have no dataset position); use a plain " + "numeric --limit to cap total samples." + ) + # resolve the task scorers eval_scorer_specs = ( [as_scorer_spec(scorer) for scorer in task.scorer] @@ -322,6 +350,7 @@ async def prepare_options( viewer=task.viewer, recorder=recorder, header_only=header_only, + dynamic_dataset=task.sample_source is not None, ) await logger.init() @@ -942,24 +971,26 @@ def __init__( tuple[TaskCleanup, SandboxEnvironmentConfigType | None, str] ] = [] self._init_lock = anyio.Lock() + # keyed by (task, spec): resolution folds in per-task state (run_dir), + # so a spec-only key would leak one task's resolution to another + self._resolved_no_metadata: dict[ + tuple[Task, SandboxEnvironmentSpec], TaskSandboxEnvironment + ] = {} async def start(self, tasks: list[ResolvedTask]) -> None: # find unique sandboxenvs to start sandboxenvs: Set[TaskSandboxEnvironment] = set() for task in tasks: - # resolve each sample and add to sandboxenvs resolved_task_sample_ids = resolve_task_sample_ids( task.task.name, self._config.sample_id ) dataset = slice_dataset( - task.task.dataset, self._config.limit, resolved_task_sample_ids + task.task.dataset, + self._config.limit, + resolved_task_sample_ids, + dynamic=task.task.sample_source is not None, ) - for sample in dataset: - sandbox = await resolve_sandbox_for_task_and_sample( - task.sandbox, task.task, sample - ) - if sandbox is not None and sandbox not in sandboxenvs: - sandboxenvs.add(sandbox) + sandboxenvs |= await self._resolve_sandboxenvs(task, dataset) await self._start_sandboxenvs(sandboxenvs) @@ -975,15 +1006,39 @@ async def start_for_samples( build/pull, fail-fast validation, and a registered `task_cleanup` — before the sample runs. Configs already started are a no-op. """ + await self._start_sandboxenvs(await self._resolve_sandboxenvs(task, samples)) + + async def _resolve_sandboxenvs( + self, task: ResolvedTask, samples: Iterable[Sample] + ) -> Set[TaskSandboxEnvironment]: + """Resolve each sample's sandboxenv (deduped; sandbox-less samples skipped). + + Full resolution reads the sandbox config (and, for docker configs with + metadata interpolation, spawns a subprocess) per sample — meaningful + when a source adds samples batch after batch. Metadata-less samples + resolve deterministically per spec, so those resolutions are cached; + samples with metadata must always fully resolve (their config may + interpolate it into a per-sample init environment). + """ sandboxenvs: Set[TaskSandboxEnvironment] = set() for sample in samples: - sandbox = await resolve_sandbox_for_task_and_sample( - task.sandbox, task.task, sample - ) - if sandbox is not None and sandbox not in sandboxenvs: - sandboxenvs.add(sandbox) - - await self._start_sandboxenvs(sandboxenvs) + sandboxenv: TaskSandboxEnvironment | None = None + cache_key: tuple[Task, SandboxEnvironmentSpec] | None = None + if not sample.metadata: + spec = await resolve_sandbox(task.sandbox, sample, task.task.name) + if spec is None: + continue + cache_key = (task.task, spec) + sandboxenv = self._resolved_no_metadata.get(cache_key) + if sandboxenv is None: + sandboxenv = await resolve_sandbox_for_task_and_sample( + task.sandbox, task.task, sample + ) + if sandboxenv is not None and cache_key is not None: + self._resolved_no_metadata[cache_key] = sandboxenv + if sandboxenv is not None: + sandboxenvs.add(sandboxenv) + return sandboxenvs async def _start_sandboxenvs( self, sandboxenvs: Set[TaskSandboxEnvironment] @@ -992,8 +1047,12 @@ async def _start_sandboxenvs( Serialized on a lock with the started-set re-checked under it, so concurrent callers (e.g. two live tasks adding samples at once) - can't double-init the same sandboxenv. + can't double-init the same sandboxenv. The pre-lock filter lets a + caller whose configs are all already started return without waiting + out another caller's (possibly minutes-long) image pull — safe + because the started set only ever grows. """ + sandboxenvs = {env for env in sandboxenvs if env not in self._started} if not sandboxenvs: return diff --git a/src/inspect_ai/_eval/task/constants.py b/src/inspect_ai/_eval/task/constants.py index 5be30b28ac..f395b7b1b9 100644 --- a/src/inspect_ai/_eval/task/constants.py +++ b/src/inspect_ai/_eval/task/constants.py @@ -1,4 +1,3 @@ TASK_FILE_ATTR = "__task_file__" TASK_RUN_DIR_ATTR = "__task_run_dir__" TASK_ALL_PARAMS_ATTR = "__task_all_params__" -DATASET_SAMPLE_SOURCE_ATTR = "__dataset_sample_source__" diff --git a/src/inspect_ai/_eval/task/error.py b/src/inspect_ai/_eval/task/error.py index 092d2dc7f5..f13ab6ccc9 100644 --- a/src/inspect_ai/_eval/task/error.py +++ b/src/inspect_ai/_eval/task/error.py @@ -21,10 +21,32 @@ def _should_eval_fail( class SampleErrorHandler: - def __init__(self, fail_on_error: bool | float | None, total_samples: int) -> None: + def __init__( + self, + fail_on_error: bool | float | None, + total_samples: int, + defer_fractional: bool = False, + ) -> None: + """Handle sample errors, aborting the eval when `fail_on_error` says to. + + Args: + fail_on_error: `True` to fail on any error, a fraction (< 1) of + `total_samples`, or an absolute count (>= 1). + total_samples: Planned sample runs (denominator for fractional + thresholds). May be grown while running (`SampleSource`). + defer_fractional: Don't abort mid-run on a fractional threshold — + leave it to the end-of-run check. Used by `SampleSource`-driven + tasks, whose planned total grows while they run: an early error + would otherwise be measured against a transiently small + denominator (e.g. a 1-sample seed erroring trips `1 >= 0.5*1` + before the source has produced the rest). Absolute-count and + any-error thresholds don't depend on the denominator and still + abort mid-run. + """ self.error_count = 0 self.fail_on_error = True if fail_on_error is None else fail_on_error self.total_samples = total_samples + self.defer_fractional = defer_fractional def __call__(self, ex: BaseException) -> tuple[EvalError, BaseException | None]: # increment error count @@ -39,7 +61,13 @@ def sample_error( ), ex if raise_error else None # check against limits - raise_error = _should_eval_fail( - self.error_count, self.total_samples, self.fail_on_error - ) + if self.defer_fractional and self._is_fractional(): + raise_error = False + else: + raise_error = _should_eval_fail( + self.error_count, self.total_samples, self.fail_on_error + ) return sample_error(raise_error=raise_error) + + def _is_fractional(self) -> bool: + return not isinstance(self.fail_on_error, bool) and 0 < self.fail_on_error < 1 diff --git a/src/inspect_ai/_eval/task/log.py b/src/inspect_ai/_eval/task/log.py index 3a541ca5cf..9711fd67d2 100644 --- a/src/inspect_ai/_eval/task/log.py +++ b/src/inspect_ai/_eval/task/log.py @@ -171,6 +171,7 @@ def __init__( viewer: ViewerConfig | None, recorder: Recorder, header_only: bool, + dynamic_dataset: bool = False, ) -> None: packages = { PKG_NAME: importlib_metadata.version(PKG_NAME), @@ -195,7 +196,10 @@ def __init__( [ sample.id for sample in slice_dataset( - dataset, eval_config.limit, eval_config.sample_id + dataset, + eval_config.limit, + eval_config.sample_id, + dynamic=dynamic_dataset, ) ], ) diff --git a/src/inspect_ai/_eval/task/run.py b/src/inspect_ai/_eval/task/run.py index b1ffc95e07..1893eacde3 100644 --- a/src/inspect_ai/_eval/task/run.py +++ b/src/inspect_ai/_eval/task/run.py @@ -40,7 +40,10 @@ DEFAULT_MAX_CONNECTIONS_BATCH, ) from inspect_ai._util.dateutil import iso_now -from inspect_ai._util.error import exception_message, is_cancellation_message +from inspect_ai._util.error import ( + exception_message, + is_cancellation_message, +) from inspect_ai._util.exception import TerminateSampleError, TerminateTaskError from inspect_ai._util.json import to_json_str_safe from inspect_ai._util.notgiven import NOT_GIVEN @@ -429,7 +432,9 @@ async def task_run(options: TaskRunOptions, task_cancel: TaskCancel | None) -> E # slice dataset (but don't materialize all sample+state pairs upfront -- # they are created lazily inside run_sample to keep memory at # O(concurrent_samples) instead of O(total_samples * epochs)) - dataset = slice_dataset(task.dataset, config.limit, config.sample_id) + dataset = slice_dataset( + task.dataset, config.limit, config.sample_id, dynamic=sample_feed is not None + ) total_samples = len(dataset) * epochs # capture sample ids now, before `dataset` may be paged to disk and @@ -464,10 +469,14 @@ async def finish_task_log( # handle sample errors (raise as required). use total_samples (sliced # dataset * epochs) as the denominator for fractional fail_on_error so - # the mid-run abort threshold matches the end-of-run check below. + # the mid-run abort threshold matches the end-of-run check below. for a + # SampleSource-driven task that denominator grows while the task runs, + # so fractional thresholds are deferred to the end-of-run check (see + # SampleErrorHandler). sample_error_handler = SampleErrorHandler( config.fail_on_error if config.continue_on_fail is not True else False, total_samples, + defer_fractional=sample_feed is not None, ) # optionally page dataset to disk if it exceeds the memory budget @@ -478,14 +487,18 @@ async def finish_task_log( del dataset # samples a SampleSource injects while the task runs, indexed after the - # seed store (kept in memory — they arrive incrementally, not up front) + # seed store (kept in memory — they arrive incrementally, not up front). + # a slot is released (set to None) once all its epochs have run, so an + # open-ended source doesn't accumulate every sample it ever produced store_len = len(sample_store) - injected_samples: list[Sample] = [] + injected_samples: list[Sample | None] = [] def get_sample(sample_index: int) -> Sample: if sample_index < store_len: return sample_store[sample_index] - return injected_samples[sample_index - store_len] + sample = injected_samples[sample_index - store_len] + assert sample is not None, "sample accessed after all its epochs completed" + return sample # register the sample enqueuer that buffers additions to a # SampleSource-driven task (callback-returned samples / enqueue_sample); @@ -593,6 +606,10 @@ def get_sample(sample_index: int) -> Sample: # the cancel handle the control channel's task-cancel # directive fires (with "abort" — the display's user-cancel) task_cancel=task_cancel, + # a SampleSource-driven eval's totals grow while it runs, so + # counters reaching total must not read as "finished" (e.g. + # while blocked in next_samples() with an empty seed) + dynamic=sample_feed is not None, ) # call early stopping if we have it @@ -729,6 +746,15 @@ async def run_sample( eval_spec=logger.eval, ) await sample_complete(sample_id, epoch, sample_scores) + # notify the task's SampleSource of the reused + # sample: a completion-driven source regenerates + # its follow-ups on retry from these notifications + # (the follow-ups are then themselves reused via + # this same prior-attempt lookup) + if sample_feed is not None: + _enqueue_source_samples( + await sample_feed.sample_complete(previous_sample) + ) # reused sample: accumulate its own logged usage record_sample_completed( logger.eval.eval_id, @@ -850,7 +876,7 @@ async def run_samples_dynamic( # --sample-id: added samples must also match the filter include_id = ( - sample_id_filter(config.sample_id) + sample_id_filter(config.sample_id).matches if config.sample_id is not None else None ) @@ -871,14 +897,19 @@ async def run_samples_dynamic( # injected samples without ids continue the seed's 1-based # numbering, skipping ids already in use; ids are compared # by their str() form (matching ensure_unique_ids, since - # log member names and score grouping key on it) - seen_ids = {str(get_sample(i).id) for i in range(store_len)} + # log member names and score grouping key on it). the seed + # ids come from sample_ids (captured before the store may + # have been paged to disk) rather than re-reading the store + seen_ids = {str(id) for id in sample_ids} auto_id = store_len - def add_samples(samples: list[Sample]) -> list[int]: + class AddedSamples(NamedTuple): + indexes: list[int] + samples: list[Sample] + + def add_samples(samples: list[Sample]) -> AddedSamples: nonlocal total_samples, auto_id, remaining - indexes: list[int] = [] - added_ids: list[int | str] = [] + added = AddedSamples([], []) over_limit = 0 for sample in samples: if sample.id is None: @@ -901,47 +932,39 @@ def add_samples(samples: list[Sample]) -> list[int]: continue remaining -= 1 sample_ids.append(sample.id) - added_ids.append(sample.id) injected_samples.append(sample) - indexes.append(store_len + len(injected_samples) - 1) + added.samples.append(sample) + added.indexes.append(store_len + len(injected_samples) - 1) if over_limit: py_logger.warning( f"Sample limit ({limit_count}) reached: ignoring " f"{over_limit} sample(s) added to the task." ) - if indexes: + if added.indexes: # grow the planned totals (display denominator, # fail_on_error threshold, control-channel counters) - total_samples += len(indexes) * epochs + total_samples += len(added.indexes) * epochs sample_error_handler.total_samples = total_samples record_samples_added( logger.eval.eval_id, - len(indexes) * epochs, - sample_ids=added_ids, + len(added.indexes) * epochs, + sample_ids=[ + sample.id + for sample in added.samples + if sample.id is not None + ], ) td.sample_complete( complete=len(progress_results), total=total_samples ) - return indexes - - async def add_and_start(samples: list[Sample]) -> list[int]: - """Add samples, starting any sandboxenvs they need first. - - Added samples get the same run-level sandbox startup as - the seed (``task_init`` for configs not seen before: - image build/pull, validation, cleanup registration) - before they spawn; already-started configs are a cheap - no-op. A startup failure propagates and fails the task, - matching a seed config failing startup. - """ - indexes = add_samples(samples) - if indexes and options.startup_sandboxes is not None: - await options.startup_sandboxes( - [get_sample(index) for index in indexes] - ) - return indexes + return added async with anyio.create_task_group() as tg: + # runs outstanding per injected index: when an index's + # last epoch completes its slot is released (the seed + # store pages to disk under memory pressure; injected + # samples would otherwise stay resident forever) + injected_runs_left: dict[int, int] = {} async def run_one(sample_index: int, epoch: int) -> None: nonlocal in_flight @@ -949,15 +972,43 @@ async def run_one(sample_index: int, epoch: int) -> None: results.append(await run_sample(sample_index, epoch)) finally: in_flight -= 1 + left = injected_runs_left.get(sample_index) + if left is not None: + if left == 1: + del injected_runs_left[sample_index] + injected_samples[sample_index - store_len] = ( + None + ) + else: + injected_runs_left[sample_index] = left - 1 wake.set() def spawn(indexes: list[int]) -> None: nonlocal in_flight for sample_index in indexes: + if sample_index >= store_len: + injected_runs_left[sample_index] = epochs for epoch in range(1, epochs + 1): in_flight += 1 tg.start_soon(run_one, sample_index, epoch) + async def add_and_start(samples: list[Sample]) -> bool: + """Add samples and start them; True if any started. + + Added samples get the same run-level sandbox startup + as the seed (``task_init`` for configs not seen + before: image build/pull, validation, cleanup + registration) before they spawn; already-started + configs are a cheap no-op. A startup failure + propagates and fails the task, matching a seed + config failing startup. + """ + added = add_samples(samples) + if added.samples and options.startup_sandboxes is not None: + await options.startup_sandboxes(added.samples) + spawn(added.indexes) + return bool(added.indexes) + spawn(list(range(store_len))) while True: # checkpoint so a misbehaving source that never @@ -968,7 +1019,7 @@ def spawn(indexes: list[int]) -> None: # returning samples / enqueue_sample) start first buffered = enqueuer.drain() if buffered: - spawn(await add_and_start(buffered)) + await add_and_start(buffered) continue if in_flight > 0: await wake.wait() @@ -986,16 +1037,16 @@ def spawn(indexes: list[int]) -> None: # was finishing rather than dropping them # (enqueue_sample never drops silently); if # any run, the source may be consulted again - final = await add_and_start(enqueuer.drain()) - if not final: + if not await add_and_start(enqueuer.drain()): break - spawn(final) elif more: - spawn(await add_and_start(more)) + await add_and_start(more) return results - if sample_feed is not None and sample_enqueuer is not None: + if sample_feed is not None: + # created together with sample_feed's registration above + assert sample_enqueuer is not None try: sample_results = await run_samples_dynamic( sample_feed, sample_enqueuer @@ -1978,8 +2029,11 @@ def make_eval_sample(include_events: bool = True) -> EvalSample: ) # notify the task's SampleSource (if it has one) as each # sample completes, so it can react in real time (and add - # samples to the running task) - if sample_feed is not None: + # samples to the running task). skipped for a cancelled + # sample: the task is unwinding (any follow-ups could + # never run) and this scope is shielded, so awaiting user + # callback code here would be uncancellable + if sample_feed is not None and cancelled_error is None: _enqueue_source_samples( await sample_feed.sample_complete(eval_sample) ) diff --git a/src/inspect_ai/_eval/task/sample_source.py b/src/inspect_ai/_eval/task/sample_source.py index 1944a39cfa..b42b414076 100644 --- a/src/inspect_ai/_eval/task/sample_source.py +++ b/src/inspect_ai/_eval/task/sample_source.py @@ -77,6 +77,11 @@ async def sample_complete(self, sample: "EvalSample") -> list["Sample"] | None: Return a list of samples to add to the task (equivalent to calling ``enqueue_sample`` with them): they start as soon as there is free capacity. Return ``None`` (the default) to add nothing. + + On a task retry this is also called for samples reused from the prior + attempt, so a completion-driven source regenerates its follow-ups + (returned samples whose ids match the prior attempt are themselves + reused rather than re-run). """ return None diff --git a/src/inspect_ai/_eval/task/task.py b/src/inspect_ai/_eval/task/task.py index f4236bedad..9f930333bc 100644 --- a/src/inspect_ai/_eval/task/task.py +++ b/src/inspect_ai/_eval/task/task.py @@ -60,7 +60,6 @@ ) from inspect_ai.viewer import ViewerConfig -from .constants import DATASET_SAMPLE_SOURCE_ATTR from .epochs import Epochs from .sample_source import SampleSource @@ -542,15 +541,13 @@ def resolve_dataset_or_source( A `SampleSource` contributes its (possibly empty) `initial_samples()` as the task's up-front dataset — the empty-dataset check doesn't apply since - the source can produce samples while the task runs. The seed is marked - with `DATASET_SAMPLE_SOURCE_ATTR` so consumers that treat a dataset as - the complete sample set (e.g. `slice_dataset` validating `--sample-id`) - know more samples may be produced at runtime. + the source can produce samples while the task runs. Consumers that treat + a dataset as the complete sample set (e.g. `slice_dataset` validating + `--sample-id`) learn that more samples may be produced at runtime via + their `dynamic` flag, passed by callers that hold the task. """ if isinstance(dataset, SampleSource): - seed = MemoryDataset(list(dataset.initial_samples())) - setattr(seed, DATASET_SAMPLE_SOURCE_ATTR, True) - return ResolvedDataset(seed, dataset) + return ResolvedDataset(MemoryDataset(list(dataset.initial_samples())), dataset) else: return ResolvedDataset(resolve_dataset(dataset), None) diff --git a/src/inspect_ai/_eval/task/util.py b/src/inspect_ai/_eval/task/util.py index cd243942ad..4460fd97c2 100644 --- a/src/inspect_ai/_eval/task/util.py +++ b/src/inspect_ai/_eval/task/util.py @@ -3,7 +3,7 @@ from copy import deepcopy from fnmatch import fnmatch from logging import getLogger -from typing import Callable, cast +from typing import Callable, NamedTuple, cast from inspect_ai._util.error import PrerequisiteError from inspect_ai._util.logger import warn_once @@ -15,7 +15,6 @@ from ..task import Task from .constants import ( - DATASET_SAMPLE_SOURCE_ATTR, TASK_FILE_ATTR, TASK_RUN_DIR_ATTR, ) @@ -48,13 +47,22 @@ def task_file(task: Task, relative: bool = False) -> str | None: return None +class SampleIdMatcher(NamedTuple): + """A `--sample-id` filter: the predicate plus its normalised patterns.""" + + matches: Callable[[str | int | None], bool] + patterns: list[str] + + def sample_id_filter( sample_id: str | int | list[str] | list[int] | list[str | int], -) -> Callable[[str | int | None], bool]: +) -> SampleIdMatcher: """Predicate matching a sample id against `sample_id` patterns. Ids are normalised (`normalise_sample_id`) and matched with `fnmatch`, - the same semantics `slice_dataset` applies to a dataset. + the same semantics `slice_dataset` applies to a dataset. The normalised + patterns ride along so callers reporting on the filter (warnings/errors) + use exactly what the predicate matches. """ sample_ids = sample_id if isinstance(sample_id, list) else [sample_id] patterns = [normalise_sample_id(id) for id in sample_ids] @@ -62,7 +70,7 @@ def sample_id_filter( def matches(id: str | int | None) -> bool: return any(fnmatch(normalise_sample_id(id), pat) for pat in patterns) - return matches + return SampleIdMatcher(matches, patterns) def sample_limit_count(limit: int | tuple[int, int] | None) -> int | None: @@ -76,22 +84,27 @@ def slice_dataset( dataset: Dataset, limit: int | tuple[int, int] | None, sample_id: str | int | list[str] | list[int] | list[str | int] | None, + dynamic: bool = False, ) -> Dataset: - # a SampleSource seed is not the complete sample set — samples matching - # a requested id may be produced while the task runs, so unmatched - # sample_id entries neither warn nor error (the filter still applies to - # the seed; the task's dispatcher applies it to produced samples) - dynamic = getattr(dataset, DATASET_SAMPLE_SOURCE_ATTR, False) - + """Apply `--limit` / `--sample-id` to a dataset. + + Args: + dataset: Dataset to slice. + limit: Count or `(start, stop)` positional range of samples to select. + sample_id: Sample id pattern(s) to filter by (exclusive with `limit`). + dynamic: The dataset is a `SampleSource` seed, not the complete sample + set — samples matching a requested id may be produced while the + task runs, so unmatched `sample_id` entries neither warn nor error + (the filter still applies to the seed; the task's dispatcher + applies it to produced samples). + """ if sample_id is not None: - include_id = sample_id_filter(sample_id) - sample_ids = sample_id if isinstance(sample_id, list) else [sample_id] - normalised = [normalise_sample_id(id) for id in sample_ids] + matcher = sample_id_filter(sample_id) # validate all the sample ids and warn if they aren't in the dataset if not dynamic: all_sample_ids = [normalise_sample_id(sample.id) for sample in dataset] - for id in normalised: + for id in matcher.patterns: if id not in all_sample_ids: warn_once( logger, @@ -99,11 +112,11 @@ def slice_dataset( ) # filter the dataset - filtered = dataset.filter(lambda sample: include_id(sample.id)) + filtered = dataset.filter(lambda sample: matcher.matches(sample.id)) # raise error if we got no hits if len(filtered) == 0 and not dynamic: - filter = ",".join([str(id) for id in normalised]) + filter = ",".join([str(id) for id in matcher.patterns]) all_sample_ids_raw = [sample.id for sample in dataset] r = reprlib.Repr() r.maxlist = 8 diff --git a/tests/_control/test_eval_state.py b/tests/_control/test_eval_state.py index da2f520f33..7dfe3c6936 100644 --- a/tests/_control/test_eval_state.py +++ b/tests/_control/test_eval_state.py @@ -17,6 +17,7 @@ get_eval_state, record_sample_completed, record_sample_errored, + record_samples_added, register_eval, ) @@ -44,6 +45,24 @@ def test_finalize_folds_unaccounted_samples_into_cancelled() -> None: assert state.completed_at is not None +def test_dynamic_eval_not_provisionally_finished() -> None: + # a SampleSource-driven eval's totals grow while it runs: counters + # reaching total (or an empty seed's 0 >= 0 at registration) must not + # read as finished — the task may be idle awaiting its source — so the + # stamp waits for finalize_eval, keeping task cancel and status honest + state = register_eval("e1", 0, dynamic=True) + assert state.completed_at is None + + record_samples_added("e1", 2) + record_sample_completed("e1") + record_sample_completed("e1") + assert state.terminal == state.total + assert state.completed_at is None + + finalize_eval("e1") + assert state.completed_at is not None + + def test_finalize_is_a_noop_when_counters_already_complete() -> None: register_eval("e1", 2) record_sample_completed("e1") diff --git a/tests/test_sample_source.py b/tests/test_sample_source.py index d1f91c3376..3466bfe623 100644 --- a/tests/test_sample_source.py +++ b/tests/test_sample_source.py @@ -6,6 +6,8 @@ samples are live: they start as soon as there is free capacity. """ +import tempfile +from pathlib import Path from typing import Any, Literal, overload import anyio @@ -13,9 +15,11 @@ from pydantic import BaseModel from typing_extensions import override -from inspect_ai import SampleSource, Task, enqueue_sample, eval +from inspect_ai import SampleSource, Task, enqueue_sample, eval, eval_set, task +from inspect_ai._util.error import PrerequisiteError from inspect_ai.dataset import Sample -from inspect_ai.log import EvalLog, EvalSample +from inspect_ai.log import EvalLog, EvalSample, read_eval_log +from inspect_ai.scorer import SampleScore from inspect_ai.solver import Generate, Solver, TaskState, generate, solver from inspect_ai.util import ( ExecResult, @@ -675,3 +679,154 @@ async def next_samples(self) -> list[Sample] | None: assert _sandbox_events.count("task_init:A") == 1 assert _sandbox_events.count("sample_init:A") == 2 assert _sandbox_events.count("task_cleanup:A") == 1 + + +class _EarlyStoppingStub: + """Minimal EarlyStopping implementation (never stops anything).""" + + async def start_task(self, task: Any, samples: Any, epochs: int) -> str: + return "stub" + + async def schedule_sample(self, id: str | int, epoch: int) -> None: + return None + + async def complete_sample( + self, id: str | int, epoch: int, scores: dict[str, SampleScore] + ) -> None: + return None + + async def complete_task(self) -> dict[str, Any]: + return {} + + +def test_sample_source_early_stopping_rejected() -> None: + # early stopping assumes a fixed sample set (the manager registers the + # seed only, and halted samples complete without notifying the source), + # so the combination is rejected pre-run with a clear error + with pytest.raises(PrerequisiteError, match="early_stopping"): + eval( + Task( + dataset=SampleSource.from_samples([Sample(input="x", target="y")]), + solver=[generate()], + early_stopping=_EarlyStoppingStub(), + ), + model="mockllm/model", + display="none", + ) + + +def test_sample_source_range_limit_rejected() -> None: + # a range limit selects seed samples by position; added samples have no + # position, so the combination is rejected pre-run with a clear error + with pytest.raises(PrerequisiteError, match="range"): + eval( + Task( + dataset=SampleSource.from_samples([Sample(input="x", target="y")]), + solver=[generate()], + ), + model="mockllm/model", + display="none", + limit=(2, 5), + ) + + +def test_sample_source_fractional_fail_on_error_uses_final_total() -> None: + # a fractional fail_on_error is deferred to the end-of-run check: the + # planned total grows while the task runs, so an early error must not be + # measured against a transiently small denominator (here the 1-sample + # seed errors first — 1/1 mid-run — but the final ratio is 1/4) + @solver + def fail_for_fail_input() -> Solver: + async def solve(state: TaskState, generate: Generate) -> TaskState: + if state.input_text == "fail": + raise RuntimeError("boom") + return state + + return solve + + class _Src(SampleSource): + def __init__(self) -> None: + self._produced = False + + def initial_samples(self) -> list[Sample]: + return [Sample(input="fail", target="ok")] + + async def next_samples(self) -> list[Sample] | None: + if not self._produced: + self._produced = True + return [Sample(input=f"ok{i}", target="ok") for i in range(3)] + return None + + logs = eval( + Task(dataset=_Src(), solver=[fail_for_fail_input()]), + model="mockllm/model", + display="none", + fail_on_error=0.5, + retry_on_error=0, + ) + log = logs[0] + assert log.status == "success" + errors = [str(sample.input) for sample in (log.samples or []) if sample.error] + assert errors == ["fail"] + assert len(log.samples or []) == 4 + + +def test_sample_source_task_retry_regenerates_followups() -> None: + # on a task retry, reused samples re-fire the source's sample_complete, + # so a completion-driven source regenerates its follow-ups: the reused + # seed produces the follow-up again, which (having errored last attempt) + # re-runs and succeeds — nothing is silently dropped from the retry log + attempts = {"n": 0} + + @solver + def fail_followup_once() -> Solver: + async def solve(state: TaskState, generate: Generate) -> TaskState: + if state.input_text == "followup": + attempts["n"] += 1 + if attempts["n"] == 1: + raise RuntimeError("transient failure") + return state + + return solve + + class _Src(SampleSource): + # stateless on purpose: the same instance is re-driven on retry, so + # follow-ups derive from the completion, not internal state + async def sample_complete(self, sample: EvalSample) -> list[Sample] | None: + if sample.id == 1: + return [Sample(id=2, input="followup", target="ok")] + return None + + def initial_samples(self) -> list[Sample]: + return [Sample(id=1, input="seed", target="ok")] + + @task + def retry_source_task() -> Task: + return Task( + dataset=_Src(), + solver=[fail_followup_once()], + name="retry_source_task", + ) + + with tempfile.TemporaryDirectory() as d: + log_dir = str(Path(d) / "logs") + Path(log_dir).mkdir() + ok, logs = eval_set( + tasks=[retry_source_task()], + log_dir=log_dir, + model="mockllm/model", + retry_attempts=2, + retry_on_error=0, # no sample-level retry -> task-level retry + ) + assert ok, "eval-set did not succeed after task retry" + log = read_eval_log(logs[0].location) + + assert sorted(str(sample.input) for sample in (log.samples or [])) == [ + "followup", + "seed", + ] + followup = next(sample for sample in (log.samples or []) if sample.id == 2) + assert followup.error is None + # the re-run follow-up carries the prior attempt's error history + assert followup.error_retries + assert "transient failure" in followup.error_retries[0].message