feat(v1): lazy and infinite tasksets (load() may yield)

[mikasenghaas]

Summary

Redo of #1829 on top of the Task/TaskData/Taskset split (#1948), scoped to the taskset side: load() may now be an iterator.

• Taskset.load is overloaded from -> list[TaskT] to -> Iterable[TaskT], so a taskset can yield its tasks from a generator — built lazily, or genuinely infinite. A list is still a valid Iterable, so existing dataset-backed tasksets are unchanged.

• Taskset.INFINITE class sentinel (default False, mirroring NEEDS_CONTAINER) marks a load that never terminates. Infinity is inherent to the taskset class, not a config knob — how many tasks a run takes is the selector's decision (-n), so the procedural tasksets carry no num_tasks config field at all.

• Taskset.select(num_tasks, shuffle, seed) is what runs call now (replacing the copy-pasted load → shuffle → slice triplet in eval, validate, and debug): it islices the first num_tasks off load, so a generator only builds what the run takes. shuffle still samples from the whole (finite) taskset — that requires materializing it first. On an infinite taskset, select requires num_tasks (clear ValueError instead of a hang) and ignores shuffle with a warning (the first n generated tasks are already an arbitrary sample).

• The env server supports infinite tasksets without a protocol change: a finite taskset is materialized up front exactly as before; an infinite one is pulled off its generator on demand and cached up to the highest task_idx requested (bounded by the driver's num_tasks, and hard-capped at MAX_LAZY_TASKS = 1_000_000 per worker so a runaway driver idx fails its request instead of generating and caching forever). task_idx is validated ge=0 on the rollout RPCs (a negative idx used to silently index from the end of the materialized list). Generation must be deterministic — pool workers each run their own load() and rely on producing the same sequence, which all ported tasksets satisfy (constant seeds / seed-per-idx). InfoResponse.num_tasks is now int | None, None ⟺ infinite; the --server eval path bounds such runs with -n and errors clearly without it.

• Ported the procedurally generated example tasksets to infinite generators (INFINITE = True):

  • textarena (built-in; wordle_v1 inherits) — one seeded episode per index, yielded lazily; seed-specific games (WordLadder, WordSearch) no longer deepcopy+reset the game for episodes a run never touches.
  • alphabet_sort_v1 — cycles the source name lists forever, with the rng advancing across passes so every pass draws fresh turn splits and the episode stream never repeats; a full pass yielding nothing (degenerate turn config) raises instead of spinning.
  • color_codeword_v1 — streams synthetic episodes (the num_examples knob is gone).

  Dataset-backed tasksets (gsm8k, reverse_text, lean, harbor, …) intentionally stay lists: their dominant cost is the eager load_dataset, so laziness wouldn't buy anything.

• Docs: a "Lazy and infinite tasksets" section in docs/v1/environments.md (yield form, infinite, the two rules, determinism requirement), notes on num_tasks/shuffle in docs/v1/evaluation.md; configs/wordle.toml and configs/textarena.toml drop their now-redundant finite episode pools.

• num_tasks is validated ge=1 on the eval/validate/debug configs — a negative value, with islice, raised an opaque ValueError from the eval runner instead of a config error.

• Tests: tests/v1/test_taskset.py covers select over finite/generator/infinite loads (bounding, the missing-num_tasks error, shuffle no-op + reproducible finite sampling, and that only the selected tasks are ever built).

Relative to #1829, the pull-based sample() RPC / broker cursor / TasksetConfig.shuffle redesign is dropped — idx-addressing works unchanged for infinite tasksets as long as generation is deterministic, which keeps this PR orchestrator-compatible for finite tasksets with no prime-rl companion.

Verification

• In-process eval on the now-infinite wordle-v1 (-n 2 -r 1, gpt-4.1-mini via PI): built exactly tasks 0–1 from the generator, both scored, 0 errors.
• --server eval on the same: EnvServer up: … tasks=infinite, tasks 0–1 pulled on demand, both scored, 0 errors; finite path re-smoked with reverse-text-v1 --server (tasks=1000, 0 errors).
• Guards: eval without -n fails with WordleTaskset is infinite - select a bounded subset with num_tasks (-n on the CLI) (in-process), wordle-v1 is infinite - bound the run with -n/--num-tasks (--server), and same via validate; --shuffle -n 1 warns shuffle is a no-op on an infinite taskset and runs task 0.
• alphabet-sort-v1 in-process eval (-n 2 -r 1): both episodes reward 1.0, 0 errors. Its generator was driven past a full source pass (112,248 episodes > 112,245 entries): idx stays monotonic, second-pass episodes differ, and two independent load() iterations produce identical heads (pool-worker determinism). color_codeword_v1 and lazy WordLadder (3 seed-specific tasks in 0.9s) checked the same way.
• Bounds guards: -n -1 and --taskset.num-tasks -5 now fail config validation (Input should be greater than or equal to 1); EnvServer._task past the lazy cap raises task_idx 1000007 exceeds the lazy-generation cap (1000000) while a materialized taskset larger than the cap still serves fine (the cap only gates generation); out-of-range on a finite taskset still IndexErrors.
• pytest tests/v1 -m "not e2e": 63 passed. ruff check / ruff format clean.

Breaking

• Taskset.load is annotated -> Iterable[TaskT] (was list[TaskT]). Returning a list still satisfies it; only consumers that indexed/len()ed the result must go through select (all in-tree ones now do).
• InfoResponse.num_tasks is int | None (was int); None ⟺ infinite. Finite tasksets still report an int, so existing consumers (e.g. prime-rl's orchestrator) keep working for finite tasksets; supporting infinite tasksets in training lands in the companion feat(orchestrator): support infinite tasksets prime-rl#2993 (verified live against this branch).
• The procedural example tasksets are now unconditionally infinite and their size knobs are removed: textarena/wordle_v1 (num_tasks: int = 1000 gone), color_codeword_v1 (num_examples: int = 1000 gone), and alphabet_sort_v1 (was one pass over its source dataset). Runs bound them with the selector (-n on eval/validate/debug; num_examples on a prime-rl eval env); a config still passing the old taskset field fails validation.

Closes #1829 (superseded — that branch predates the Task/TaskData/Taskset split).

🤖 Generated with Claude Code

Note

Add lazy and infinite taskset support via generator-based load() and Taskset.select()

• Adds INFINITE: ClassVar[bool] to the Taskset base class and broadens load() to return Iterable[TaskT], allowing subclasses to yield tasks indefinitely.
• Adds Taskset.select(num_tasks, shuffle, seed) to centralize task materialization: infinite tasksets require num_tasks, ignore shuffle with a warning, and use itertools.islice for lazy consumption; finite tasksets support seeded shuffle.
• Converts AlphabetSortTaskset, ColorCodewordTaskset, and TextArenaTaskset to infinite generators using itertools.count().
• Updates EnvServer to lazily build and cache tasks per-index for infinite tasksets, reporting num_tasks=None in InfoResponse, with a cap of 1,000,000 tasks per worker.
• CLI commands (eval, debug, validate) now delegate task selection to Taskset.select() instead of manually loading and slicing.
• Behavioral Change: num_tasks is now required when running against infinite tasksets; passing --shuffle with an infinite taskset is silently ignored.

^{Macroscope summarized 264a43f.}

---

Note

Medium Risk
Breaking API/config changes (InfoResponse.num_tasks, removed taskset size fields) affect external orchestrators and old TOML; infinite server paths depend on deterministic generation across pool workers.

Overview
Taskset.load() may now return an Iterable (generator) instead of only a list; INFINITE = True marks procedural tasksets that never end. Taskset.select(num_tasks, shuffle, seed) centralizes bounded materialization (islice for lazy loads, full materialize + shuffle for finite sets) and replaces duplicated load/shuffle/slice logic in eval, validate, and debug.

Env server keeps finite tasksets fully loaded; infinite ones use on-demand _task(idx) generation with caching and MAX_LAZY_TASKS. InfoResponse.num_tasks is int | None (None = infinite). The --server eval path requires -n for infinite tasksets and treats shuffle as a no-op with a warning.

textarena, alphabet_sort_v1, and color_codeword_v1 are converted to infinite generators; per-taskset num_tasks / num_examples config fields are removed. num_tasks on eval/validate/debug configs gains ge=1; rollout RPC task_idx is ge=0.

Docs and sample configs are updated; tests/v1/test_taskset.py covers select behavior.

^{Reviewed by Cursor Bugbot for commit 264a43f. Bugbot is set up for automated code reviews on this repo. Configure here.}

[macroscopeapp]

Approvability

Verdict: Needs human review

This PR introduces a significant new feature (lazy/infinite tasksets) that changes core infrastructure behavior - how tasks are loaded, materialized, and served. The changes span the Taskset base class, EnvServer, CLI commands, and multiple taskset implementations, warranting human review.

^{You can customize Macroscope's approvability policy. Learn more.}