Skip to content

feat: adapt orchestrator to index-addressed env server#2856

Closed
mikasenghaas wants to merge 1 commit into
feat/nano-as-v1from
feat/dynamic-tasksets
Closed

feat: adapt orchestrator to index-addressed env server#2856
mikasenghaas wants to merge 1 commit into
feat/nano-as-v1from
feat/dynamic-tasksets

Conversation

@mikasenghaas

Copy link
Copy Markdown
Member

Summary

Companion to verifiers PR #1829, which makes the v1 env-server index-addressed with no required task count so it can serve lazy/unbounded tasksets — InfoResponse.num_tasks is now int | None (None = served lazily; the caller drives task_idx). This adapts the orchestrator to that API.

  • Env.num_tasks is now int | None (None when the taskset is served lazily).
  • EvalEnv.start requires num_examples to bound a lazily-served eval — an unknown/infinite count can't be enumerated (raises a clear error if num_examples < 0).
  • TrainSource:
    • Finite path unchanged — enumerate the env's index range, shuffle, reshuffle each epoch, weight by dataset size.
    • Unbounded path (num_tasks is None) — next_example streams a monotonically increasing task_idx (no epoch/reshuffle; the taskset's own generator, e.g. RNG-seeded pairs, supplies the variety).
    • An unbounded env has no dataset size to weight by, so it requires an explicit ratio (clear error otherwise).

Backward-compatible: against a server that still reports a count (every current taskset that returns a list, and the v0 bridge), every env is bounded and behavior is identical. The change is safe to land before the verifiers pin bump.

Verification

Unit-checked TrainSource against stub envs:

  • finite env → epochs cover all indices, reshuffled (e.g. [2,0,1,3, 1,2,0,3, ...]);
  • unbounded env (num_tasks=None, with a ratio) → streams [0,1,2,3,4,5];
  • unbounded without a ratio alongside a finite env → ValueError (no size to weight by);
  • mixed finite + unbounded with ratios → both envs sampled.

ruff check / ruff format clean (pre-commit).

Follow-up

The verifiers pin bump (to the merged #1829 commit) lands separately once that PR merges; only then does an actual unbounded taskset report num_tasks=None end-to-end.

…ask count)

The env server no longer reports a task count for lazily-served (generator,
possibly unbounded) tasksets — `info().num_tasks` is now `int | None`. Adapt the
orchestrator to drive purely by `task_idx`:

- `Env.num_tasks` is `int | None`; `EvalEnv` requires `num_examples` to bound a
  lazily-served eval (can't enumerate an unknown/infinite count).
- `TrainSource` keeps the finite path unchanged (enumerate the index range,
  shuffle, reshuffle per epoch, weight by size) and adds an unbounded path:
  `num_tasks is None` → stream a monotonically increasing `task_idx` (the
  taskset's own generator supplies the variety). Unbounded envs have no size to
  weight by, so they require an explicit `ratio`.

Backward-compatible: against a server that still reports a count, every env is
bounded and behavior is unchanged. Pairs with verifiers PR #1829 (the env-server
API change); the verifiers pin bump lands when that merges.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@mikasenghaas

Copy link
Copy Markdown
Member Author

Superseded by #2882, which reimplements the same orchestrator change off main (this branch was based on feat/nano-as-v1).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant