diff --git a/CHANGELOG.md b/CHANGELOG.md index 076ddab858..4cc4cce005 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. - Eval: `EvalSample` and `EvalSampleSummary` now record `turn_count` and the sample's token limit (`token_limit`, `token_limit_type`, and metered `token_limit_usage`). - Analysis: `samples_df` gains default `turn_count` and `token_limit_usage` columns, and `evals_df` configuration columns gain `token_limit_type`. - Control Channel: `inspect ctl sample list` now shows per-sample turn count and, when a token limit is configured, its computed usage and configured ceiling. diff --git a/design/sample-source.md b/design/sample-source.md new file mode 100644 index 0000000000..c10568c4b9 --- /dev/null +++ b/design/sample-source.md @@ -0,0 +1,188 @@ +# 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. 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 + 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)). + 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` + +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`** 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 + *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). 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 *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). + 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 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). + +## 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; `--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/_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..d9e18ad609 --- /dev/null +++ b/docs/sample-source.qmd @@ -0,0 +1,145 @@ +--- +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. + +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: + +```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)). + +## 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. 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/__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 28014f496e..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. @@ -371,17 +384,23 @@ def register_eval( solver=solver or "", 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, 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 @@ -548,6 +567,26 @@ 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). 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) + if state is not None: + state.total += total + if sample_ids: + state.sample_ids.extend(sample_ids) + + def latest_eval_for_task(task_id: str) -> "EvalState | None": """The last-registered attempt of ``task_id``, or ``None`` if untracked. @@ -633,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) @@ -642,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/_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 4c57f6026b..ec3fba31a8 100644 --- a/src/inspect_ai/_eval/run.py +++ b/src/inspect_ai/_eval/run.py @@ -1,10 +1,12 @@ +import functools import logging 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 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 @@ -29,7 +31,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 @@ -47,6 +49,7 @@ from inspect_ai.util._display import display_type from inspect_ai.util._sandbox.environment import ( SandboxEnvironmentConfigType, + SandboxEnvironmentSpec, TaskCleanup, TaskInit, ) @@ -68,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 @@ -199,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] @@ -320,9 +350,20 @@ async def prepare_options( viewer=task.viewer, recorder=recorder, header_only=header_only, + dynamic_dataset=task.sample_source is not None, ) 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( @@ -347,6 +388,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 @@ -421,25 +463,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.""" @@ -579,7 +602,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: @@ -723,7 +746,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: @@ -929,9 +952,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. """ @@ -946,63 +970,130 @@ def __init__( self._cleanups: list[ 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 not already started + # 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( + sandboxenvs |= await self._resolve_sandboxenvs(task, dataset) + + 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. + """ + 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: + 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 ( - sandbox is not None - and sandbox not in self._started - and sandbox not in sandboxenvs - ): - sandboxenvs.add(sandbox) - + 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] + ) -> 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. 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 - # 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 ("none" has no task - # display and must keep stdout machine-readable, e.g. --json) - if display_type() != "none": - print("") + # provide some space above task display ("none" has no task + # display and must keep stdout machine-readable, e.g. --json) + if display_type() != "none": + print("") async def shutdown(self) -> None: with anyio.CancelScope(shield=True): 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/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 ee6e729dcb..1893eacde3 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 ( @@ -38,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 @@ -179,6 +184,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, @@ -187,7 +199,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__) @@ -278,6 +290,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: @@ -336,6 +352,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, @@ -361,6 +390,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)) @@ -399,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 @@ -434,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 @@ -447,6 +486,30 @@ 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). + # 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 | None] = [] + + def get_sample(sample_index: int) -> Sample: + if sample_index < store_len: + return sample_store[sample_index] + 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); + # 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) @@ -543,6 +606,10 @@ async def finish_task_log( # 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 @@ -638,7 +705,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 @@ -679,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, @@ -704,7 +780,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( @@ -752,6 +828,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 @@ -771,13 +848,221 @@ 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). + + ``--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 + + results: list[dict[str, SampleScore] | EarlyStop | None] = [] + in_flight = 0 + wake = Wake() + enqueuer.on_enqueue = wake.set + + # --sample-id: added samples must also match the filter + include_id = ( + sample_id_filter(config.sample_id).matches + 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 + # 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 + + class AddedSamples(NamedTuple): + indexes: list[int] + samples: list[Sample] + + def add_samples(samples: list[Sample]) -> AddedSamples: + nonlocal total_samples, auto_id, remaining + added = AddedSamples([], []) + over_limit = 0 + 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)) + 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) + injected_samples.append(sample) + 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 added.indexes: + # grow the planned totals (display denominator, + # fail_on_error threshold, control-channel counters) + total_samples += len(added.indexes) * epochs + sample_error_handler.total_samples = total_samples + record_samples_added( + logger.eval.eval_id, + 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 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 + try: + 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 + # 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() + if buffered: + await add_and_start(buffered) + continue + if in_flight > 0: + await wake.wait() + continue + # 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: + # 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 + if not await add_and_start(enqueuer.drain()): + break + elif more: + await add_and_start(more) + + return results + + 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 + ) + except Exception as ex: + # match tg_collect: surface the first child exception + # rather than an ExceptionGroup + raise inner_exception(ex) from None + 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 = [ @@ -804,7 +1089,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, @@ -823,7 +1108,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 @@ -853,7 +1138,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, @@ -921,6 +1206,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) @@ -1085,6 +1373,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], @@ -1738,6 +2027,16 @@ 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). 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) + ) # 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: @@ -1788,6 +2087,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..b42b414076 --- /dev/null +++ b/src/inspect_ai/_eval/task/sample_source.py @@ -0,0 +1,237 @@ +"""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 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()``. 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 [] + + 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. (If samples were enqueued while a ``None`` return was + in progress they still run, and this method may then be called again.) + """ + 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. + + 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 + + @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. + + 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) + 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.""" + 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).""" + 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 — 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 + filter run. + + 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 3024d232fa..9f930333bc 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 @@ -52,6 +61,7 @@ from inspect_ai.viewer import ViewerConfig from .epochs import Epochs +from .sample_source import SampleSource logger = getLogger(__name__) @@ -71,7 +81,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, @@ -109,7 +119,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 @@ -201,7 +212,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 @@ -273,7 +286,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, @@ -319,7 +332,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 @@ -393,7 +407,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): @@ -511,6 +527,31 @@ 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. 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): + 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/_eval/task/util.py b/src/inspect_ai/_eval/task/util.py index 16b1995cdf..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 cast +from typing import Callable, NamedTuple, cast from inspect_ai._util.error import PrerequisiteError from inspect_ai._util.logger import warn_once @@ -14,7 +14,10 @@ from inspect_ai.model import ChatMessage, ChatMessageUser from ..task import Task -from .constants import TASK_FILE_ATTR, TASK_RUN_DIR_ATTR +from .constants import ( + TASK_FILE_ATTR, + TASK_RUN_DIR_ATTR, +) logger = getLogger(__name__) @@ -44,36 +47,77 @@ 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], +) -> 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 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] + + def matches(id: str | int | None) -> bool: + return any(fnmatch(normalise_sample_id(id), pat) for pat in patterns) + + return SampleIdMatcher(matches, patterns) + + +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, + dynamic: bool = False, ) -> Dataset: + """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: - # reduce to list of normalized sample ids - sample_ids = sample_id if isinstance(sample_id, list) else [sample_id] - sample_id = [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 - 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 matcher.patterns: + 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: matcher.matches(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 matcher.patterns]) + all_sample_ids_raw = [sample.id for sample in dataset] r = reprlib.Repr() r.maxlist = 8 raise PrerequisiteError( 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/_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 new file mode 100644 index 0000000000..3466bfe623 --- /dev/null +++ b/tests/test_sample_source.py @@ -0,0 +1,832 @@ +"""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 tempfile +from pathlib import Path +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, eval_set, task +from inspect_ai._util.error import PrerequisiteError +from inspect_ai.dataset import Sample +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, + SandboxEnvironment, + SandboxEnvironmentConfigType, + SandboxEnvironmentSpec, + sandboxenv, +) + + +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")) + + +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"] + + +# 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 + + +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