feat: dynamic tasksets (lazy/infinite load_tasks, pull-based env server)

[mikasenghaas]

Summary

• Overload Taskset.load_tasks from -> list[TaskT] to -> Iterable[TaskT] so a taskset can yield its tasks from a generator — lazily built, or genuinely infinite — instead of only returning a fixed list. A list is still a valid Iterable, so existing tasksets are unchanged.

• Add verifiers.v1.taskset.select_tasks(taskset, num_tasks, shuffle), which materializes only the tasks a run needs. Without --shuffle it itertools.islices the first --num-tasks straight off load_tasks, so an infinite generator stops after the subset and the rest is never built. (--shuffle has to see the whole set to sample from it, so it materializes first.) The shuffle seed is a fixed module constant (_SHUFFLE_SEED = 0), not a config field — matching the eval path, which never exposed a seed knob either.

• Add a Taskset.INFINITE capability flag (a ClassVar[bool], mirrors NEEDS_CONTAINER) for a load_tasks that may never terminate. A bare list()/--shuffle on such a taskset would hang — and --num-tasks N --shuffle hangs too, since the list() runs before the slice. With the flag set, select_tasks requires --num-tasks (raises if missing — it would hang otherwise) and ignores --shuffle with a warning rather than hard-failing the run. (--shuffle is moot anyway for a generated taskset: the first -n items are already an arbitrary sample.)

• Make the env server pull-based — it owns task scheduling. Three RPCs replace the old index-addressed run_rollout(task_idx) / run_group(task_idx, n):

  • sample() -> Task advances the server's cursor and returns the next task (by value, as a WireTask).
  • run_rollout(task) -> Trace runs one rollout of a sampled task.
  • run_group(task, n) -> list[Trace] runs n rollouts of one sampled task (for group-scored envs).

  The caller samples a task, then runs it — it never chooses which task by index. Shuffle + epoch live on the server (EnvConfig.shuffle, default False — deterministic, like in-process eval; prime-rl's train envs override to True; seed fixed, not configurable): a finite taskset is materialized once and served as a reshuffled permutation that loops over epochs (_SHUFFLE_SEED + epoch); an INFINITE one is pulled lazily from its generator with a one-task look-ahead (shuffle a no-op, warned at startup). InfoResponse.num_tasks is the count (None ⟺ INFINITE) and only bounds a finite eval. The eval --server path and the v0 legacy bridge both pull via these RPCs.

• The pool broker owns a global task cursor. EnvServerPool stamps a monotonic index onto each sample request; finite workers share the shuffle seed, so their permutations agree and any worker can serve any index. So sample load-balances like a rollout — no worker pinning — while coverage stays exact (each task once per epoch across the pool). A lone server (max_workers <= 1, no broker) falls back to its own cursor. run_rollout / run_group carry the task and route least-busy as before.

• Shuffle policy moves to TasksetConfig.shuffle: ShuffleConfig | None (off EnvConfig). None (default) serves in load order (deterministic — what eval wants); a ShuffleConfig (a seed) shuffles a finite taskset (reshuffled each epoch) and seeds an INFINITE taskset's generation. An INFINITE load_tasks reads the seed via self.config.shuffle.seed (may be None); since infinite streams can't be index-addressed by the broker, the env-server offsets that seed by the worker index so a multi-worker pool draws divergent streams. The eval/validate -s/--shuffle flag is removed; configure via --taskset.shuffle. select_tasks is now a Taskset method.

• Wire it into the eval + validate entrypoints. The in-process eval runner (run_eval, the default vf eval path) and the validate entrypoint select through select_tasks (lazy islice), replacing the old "load all → shuffle → [:num_tasks]" slice; the --server path pulls task-by-task, bounded by the reported count / -n.

• Convert the two built-in generative tasksets to the generator form:

  • textarena — yield one seeded episode per index; for seed-specific games (WordLadder, WordSearch) each episode is a deepcopy(template) + reset(), now built only for the episodes a run touches.
  • harbor — keep the dataset download + dir listing up front (needed to enumerate idxs), but parse each task.toml / instruction.md lazily.

  Note: wikispeedia_v1 was the original worked INFINITE example here, but main has since moved it out of this repo to research-environments, so it no longer ships in this PR. Re-applying the infinite-generator conversion there is a follow-up; the INFINITE machinery (flag + select_tasks guards + GUIDE.md) lands here regardless and is covered by the synthetic itertools.count() test below.

• Dataset-backed tasksets (gsm8k, reverse_text, alphabet_sort, …) are intentionally left as lists: their dominant setup cost is the eager load_dataset, so laziness wouldn't help.

• Update the v1 GUIDE.md "Loading tasks" section with the lazy/infinite pattern, the two INFINITE rules (needs --num-tasks; can't shuffle), and capability-flag tables.

Verification

Setup time — load_tasks for textarena at num_tasks=1000, selecting a subset of 5 (eager "build all then slice" vs. lazy select_tasks), nltk corpora pre-warmed:

game	per-task work	build-all-then-slice	lazy (-n 5)	speedup	subset identical
WordLadder-v0	deepcopy + reset + BFS	91,945 ms	1,086 ms	84.7×	✓
WordSearch-v0	deepcopy + reset + grid	199 ms	2.7 ms	73.3×	✓
Wordle-v0	none (prompt built once)	809 ms	812 ms	1.0×	✓

The lazily-drawn subset is byte-identical (model_dump()) to the first 5 of the full eager build in every case; seed-invariant games (Wordle) correctly show no change, since the prompt is built once regardless.

Still runs correctly — a live gpt-4.1-mini Wordle eval through the real run_eval (num_tasks=1000, -n 2 -r 1) built exactly 2 tasks from the generator (not 1000) and scored both rollouts normally:

tasks built by generator: 2 (config num_tasks=1000, eval -n 2)
  idx=0 turns=7 reward=0.2
  idx=1 turns=6 reward=0.3

An INFINITE taskset over itertools.count() confirms the guards: select_tasks(num_tasks=5) returns 5 (no hang); --shuffle is ignored with a warning (still returns the first 5); a missing --num-tasks raises a clear ValueError; finite tasksets (all / cap / shuffle) are unaffected.

Pull protocol — a live --server eval (i3-logic-v1, deepseek-v4-flash, -n 4) drove rollouts entirely through the pull RPCs: the server reported its count, sample() handed out tasks 0,1,2,3 in order (eval shuffle=False, deterministic), and each Trace carried the served idx — 0 errors. A scheduler unit test covers finite shuffle + per-epoch reshuffle (reproducible), finite in-order looping, eval-subset uniqueness (pull ≤ count never wraps), and infinite monotonic pulls.

Multi-worker cursor (no pin) — a 2-worker static pool, no shuffle, covered tasks 0–5 exactly once (broker-cursor coverage without pinning); the same pool driving a reverse_text training run (4 steps, group_size=16) produced 8 distinct tasks × 16 rollouts per step, shuffled, with no duplication across steps and 0 errors (one coherent global cursor across workers).

ruff check / ruff format clean; tests/v1/test_configs.py and tests/v1/test_trace.py pass.

Breaking

• Taskset.load_tasks return annotation is now Iterable[TaskT] (was list[TaskT]). Returning a list still satisfies it, so existing tasksets need no change; only code that consumed load_tasks() and relied on it being a concrete list (indexing, len) must wrap it — the two in-tree consumers (select_tasks, the env-server) already do.
• InfoResponse.num_tasks is now int | None (was int); None ⟺ the taskset is INFINITE. Consumers that assumed an int must handle None — run_eval_server does; prime-rl's orchestrator gets the matching change in its companion PR.
• The rollout RPC is now pull-based. The old index-addressed RunRolloutRequest / RunGroupRequest (and EnvClient.run_rollout(task_idx) / run_group(task_idx, n)) are replaced by sample() -> Task plus run_rollout(task) / run_group(task, n) (task carried by value). A caller that addressed tasks by index must instead sample() then run; the served task's idx is on the returned Trace. (Breaks the index-addressed orchestrator — coordinated with the prime-rl companion PR.)
• Shuffle config moved + reshaped: shuffling lives on TasksetConfig.shuffle: ShuffleConfig | None (default None = deterministic), not EnvConfig.shuffle: bool. The eval/validate -s/--shuffle flag is removed — use --taskset.shuffle (e.g. --taskset.shuffle.seed 0). Consumers that want shuffled training set a ShuffleConfig. The select_tasks free function is now the Taskset.select_tasks method, and the SHUFFLE_SEED constant / Taskset.seed attribute are gone (read the seed off config.shuffle).

[macroscopeapp]

🟠 High serve/server.py:190

When self._count is None (infinite taskset), _next_task ignores the broker-supplied index and instead returns next(self._iter) from the worker's private generator. In a multi-worker pool, the broker stamps each sample request with a global cursor so any worker can serve the same logical position, but for infinite tasksets the returned task depends on how many prior requests happened to land on that worker — not on the cursor. Two sample requests with the same index served by different workers return different tasks, and the global ordering the cursor is meant to enforce is lost. If streaming without indexing is intentional for infinite tasksets, consider documenting this limitation at the _next_task level so the broker/pool layer can avoid stamping a meaningless index, or adjust the pool to route infinite tasksets without a global cursor.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/serve/server.py around line 190:

When `self._count is None` (infinite taskset), `_next_task` ignores the broker-supplied `index` and instead returns `next(self._iter)` from the worker's private generator. In a multi-worker pool, the broker stamps each `sample` request with a global cursor so any worker can serve the same logical position, but for infinite tasksets the returned task depends on how many prior requests happened to land on that worker — not on the cursor. Two `sample` requests with the same `index` served by different workers return different tasks, and the global ordering the cursor is meant to enforce is lost. If streaming without indexing is intentional for infinite tasksets, consider documenting this limitation at the `_next_task` level so the broker/pool layer can avoid stamping a meaningless index, or adjust the pool to route infinite tasksets without a global cursor.

[macroscopeapp]

🟡 Medium serve/server.py:186

_index_at seeds the first epoch's shuffle with random.Random(f"{self._seed}-0") instead of random.Random(self._seed), so a finite taskset's first-epoch order in server mode differs from the in-process path (e.g. Taskset.select_tasks), which seeds with random.Random(seed). The same shuffle.seed therefore produces different task orderings depending on execution mode, breaking reproducibility across eval paths. Consider seeding epoch 0 with self._seed directly and only appending the epoch for subsequent epochs.

-            if self._shuffle:
-                random.Random(f"{self._seed}-{epoch}").shuffle(self._order)
+            if self._shuffle:
+                seed = self._seed if epoch == 0 else f"{self._seed}-{epoch}"
+                random.Random(seed).shuffle(self._order)

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/serve/server.py around lines 186-187:

`_index_at` seeds the first epoch's shuffle with `random.Random(f"{self._seed}-0")` instead of `random.Random(self._seed)`, so a finite taskset's first-epoch order in server mode differs from the in-process path (e.g. `Taskset.select_tasks`), which seeds with `random.Random(seed)`. The same `shuffle.seed` therefore produces different task orderings depending on execution mode, breaking reproducibility across eval paths. Consider seeding epoch 0 with `self._seed` directly and only appending the epoch for subsequent epochs.

[macroscopeapp]

🟡 Medium serve/types.py:65

RunRolloutRequest.task accepts a full WireTask from the caller, and the server validates and executes it directly (self._task_type.model_validate(req.task.model_dump()) → env.episode(task, ...)). This lets a caller modify prompt, task-specific answer fields, or image/workdir before calling run_rollout, bypassing the server-owned taskset and producing tampered evaluation results. The request should carry only an index so the server resolves the canonical task from its own taskset, not the caller-supplied object.

🚀 Reply "fix it for me" or copy this AI Prompt for your agent:

In file @verifiers/v1/serve/types.py around line 65:

`RunRolloutRequest.task` accepts a full `WireTask` from the caller, and the server validates and executes it directly (`self._task_type.model_validate(req.task.model_dump())` → `env.episode(task, ...)`). This lets a caller modify `prompt`, task-specific answer fields, or `image`/`workdir` before calling `run_rollout`, bypassing the server-owned taskset and producing tampered evaluation results. The request should carry only an index so the server resolves the canonical task from its own taskset, not the caller-supplied object.