diff --git a/deps/verifiers b/deps/verifiers index 192fd8ef0d..b13ba60da6 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 192fd8ef0d2fdcefc8566b0ec793d73631e7b51a +Subproject commit b13ba60da63cea91389e7575766b7270d0d11fc5 diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 91cc3c2185..998ff6fd53 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -387,6 +387,7 @@ def next_fresh_group(self, kind: RolloutKind, envs) -> GroupState | None: kind=kind, env_name=env_name, task_idx=example["task_idx"], + task=example.get("task"), rollouts_to_schedule=group_size, target_rollouts=group_size, eval_step=eval_step, @@ -438,6 +439,7 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - cache_salt = None if env.requires_group_scoring: + # Legacy-only route (a v1 env never group-scores) — addressed by row. permits = group.rollouts_to_schedule group.rollouts_to_schedule = 0 await self.acquire(permits) @@ -451,15 +453,20 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - ) ) else: + # A v1 env takes the task itself; the legacy bridge its dataset row. + if group.task is not None: + addressing = {"task_data": group.task.data.model_dump(mode="json")} + else: + addressing = {"task_idx": group.task_idx} permits = 1 group.rollouts_to_schedule -= 1 await self.acquire(permits) task = asyncio.create_task( env.run( client=client, - task_idx=group.task_idx, model_name=model_name, cache_salt=cache_salt, + **addressing, ) ) diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 9948d8f07b..2cae9d185f 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -2,15 +2,19 @@ Each ``Env`` owns a v1 ``EnvServer`` (spawned as a child process, or an external one given by ``config.address``) and an ``EnvClient`` to drive it. The -orchestrator never *runs* an environment: it asks the server for ``info`` -(``num_tasks`` + whether group scoring is needed), then runs rollouts purely by -**task index**. The server answers one ``Episode`` per env-rollout, whose traces -we validate into ``Trace[WireTaskData]`` — real ``vf.Trace``\\ s (never loose -dicts) whose task keeps the env's task-specific -fields as extras (``WireTaskData`` allows them). The orchestrator never imports the -env package: the env's *type* and *runtime* both live only in the server, and the orchestrator -drives it purely by task index. (Nothing here reads typed env task fields — only ``task.idx`` -and a full ``task.model_dump``, both of which ``WireTaskData`` preserves.) +orchestrator never *runs* an environment — the agents and their runtimes live only +in the server — but it does own the *taskset*: a v1 env's tasks are loaded here, +once, and each dispatched env-rollout ships its task's data on the request +(``task_data``); the server pydantic-validates it into the taskset's declared +``TaskData`` type and runs it. That keeps the server (and every worker in its +pool) stateless about data — no per-worker dataset loads, no idx-addressed task +cache — and gives the orchestrator real tasks to cycle, shuffle, and filter. Only +the legacy (v0) bridge, whose dataset genuinely lives server-side, is still driven +by ``task_idx`` (its count comes from ``info``). + +The server answers one ``Episode`` per env-rollout, whose traces we validate into +``Trace[WireTaskData]`` — real ``vf.Trace``\\ s (never loose dicts) whose task +keeps the env's task-specific fields as extras (``WireTaskData`` allows them). """ from __future__ import annotations @@ -22,6 +26,7 @@ import queue import sys from collections.abc import Iterator, Sequence +from itertools import islice from multiprocessing.process import BaseProcess from pathlib import Path from typing import Generic, TypeVar @@ -40,9 +45,8 @@ # task.idx + task.model_dump). ROLLOUT_TYPE = Rollout[vf.WireTaskData] -# Max wait for a spawned env server to bind and report its address. The child -# loads the taskset (possibly downloading a dataset) before reporting, so this -# is generous. +# Max wait for a spawned env server to bind and report its address. A legacy +# child loads its dataset before reporting, so this is generous. ENV_SERVER_SPAWN_TIMEOUT = 600.0 @@ -78,14 +82,20 @@ def _run_env_server( class Env: - """Wraps a v1 env server + client. The orchestrator never loads the env.""" + """Wraps a v1 env server + client. The orchestrator owns the taskset (loaded once, + client-side); the server owns agent/harness execution.""" def __init__(self, config: EnvConfig): self.config = config self.sampling_args: dict = {} self.num_tasks: int | None = 0 - """Task count reported by the server; ``None`` means the taskset is infinite.""" + """Task count; ``None`` means the taskset is infinite.""" self.requires_group_scoring: bool = False + self.tasks: Iterator[vf.Task] | None = None + """The env's tasks, client-side; ``None`` for legacy (its dataset lives on the + server). A finite taskset is materialized at ``start()`` (``num_tasks`` is its + count) and iterated from there; an infinite one streams off its generator. + Consumed once — by ``TrainSource`` (train) or ``EvalEnv.start`` (eval).""" self._env_client: EnvClient | None = None self._env_server_process: BaseProcess | None = None @@ -100,25 +110,36 @@ def env_client(self) -> EnvClient: return self._env_client async def start(self, log_dir: Path, log_level: str | None = None, json_logging: bool = False) -> None: - """Spawn the env server (if needed), connect, and cache its ``info``.""" + """Spawn the env server (if needed), connect, and load the taskset client-side + (legacy instead asks the server for ``info`` — its dataset is server-side).""" external = self.config.address is not None address = self.config.address or await self._spawn(log_dir, log_level or "INFO", json_logging) get_logger().debug(f"Connecting {self.name} to env server {address}") self._env_client = EnvClient(address=address) - # A spawned server already reported its address *after* binding + loading, - # so it's up — the untimed ``info`` below is enough. An external server has - # no such handshake, so poll until it answers before we block on ``info``. + # A spawned server already reported its address *after* binding, so it's up. An + # external server has no such handshake, so poll until it answers. if external: await self.env_client.wait_for_server_startup() - info = await self.env_client.info() - self.num_tasks = info.num_tasks - self.requires_group_scoring = info.requires_group_scoring + if self.config.is_legacy: + info = await self.env_client.info() + self.num_tasks = info.num_tasks + self.requires_group_scoring = info.requires_group_scoring + else: + taskset = vf.load_taskset(self.config.env.taskset) + if type(taskset).INFINITE: + self.tasks = iter(taskset.load()) + self.num_tasks = None + else: + # Materialize off the event loop — load() may pull a dataset. + materialized = await asyncio.to_thread(lambda: list(taskset.load())) + self.tasks = iter(materialized) + self.num_tasks = len(materialized) num_tasks = self.num_tasks if self.num_tasks is not None else "infinite" get_logger().info(f"Env {self.name} ready: num_tasks={num_tasks} group_scoring={self.requires_group_scoring}") async def _spawn(self, log_dir: Path, log_level: str, json_logging: bool) -> str: - """Spawn a v1 EnvServer child process (it loads the env; we never do). - The server binds an OS-assigned port (``:0``) and reports the concrete + """Spawn a v1 EnvServer child process (it runs the agents; the tasks come from + us). The server binds an OS-assigned port (``:0``) and reports the concrete address back over a queue — no free-port guess, no TOCTOU race. Its output goes to ``/.log`` (``log_dir`` is already the train/eval-split ``.../logs/envs/{train,eval}`` the orchestrator passes in).""" @@ -170,12 +191,20 @@ def _sampling(self, cache_salt: str | None) -> vf.SamplingConfig: return vf.SamplingConfig(**sampling) async def run( - self, client: vf.ClientConfig, task_idx: int, model_name: str, cache_salt: str | None + self, + client: vf.ClientConfig, + model_name: str, + cache_salt: str | None, + task_data: dict | None = None, + task_idx: int | None = None, ) -> list[Rollout]: - """Run one episode for ``task_idx``; return its typed Traces. A zero-trace - episode raises (the dispatcher synthesizes the error marker); a not-``ok`` - episode marks its clean traces failed so partial episodes never train.""" + """Run one episode; return its typed Traces. A v1 env takes the task itself + (``task_data``); the legacy bridge is addressed by dataset row (``task_idx``). + A zero-trace episode raises (the dispatcher synthesizes the error marker); a + not-``ok`` episode marks its clean traces failed so partial episodes never + train.""" episode = await self.env_client.run( + task_data=task_data, task_idx=task_idx, client=client, model=model_name, @@ -236,13 +265,16 @@ def __init__(self, config: EvalEnvConfig): async def start(self, log_dir: Path, log_level: str | None = None, json_logging: bool = False) -> None: await super().start(log_dir=log_dir, log_level=log_level, json_logging=json_logging) - if self.num_tasks is None: - if self.config.num_examples < 0: - raise ValueError(f"Eval env {self.name} has an infinite taskset — set num_examples to bound it") - n = self.config.num_examples - else: - n = self.num_tasks if self.config.num_examples < 0 else min(self.config.num_examples, self.num_tasks) - self.examples = [{"task_idx": i} for i in range(n)] + n = self.config.num_examples + if self.tasks is None: # legacy: the dataset lives on the server — address it by row + count = self.num_tasks if n < 0 else min(n, self.num_tasks) + self.examples = [{"task_idx": i} for i in range(count)] + return + if self.num_tasks is None and n < 0: + raise ValueError(f"Eval env {self.name} has an infinite taskset — set num_examples to bound it") + # A fixed eval set, pulled off the tasks once and reused every epoch. + tasks = list(self.tasks) if n < 0 else list(islice(self.tasks, n)) + self.examples = [{"task_idx": task.data.idx, "task": task} for task in tasks] EnvT = TypeVar("EnvT", bound=Env) diff --git a/src/prime_rl/orchestrator/train_source.py b/src/prime_rl/orchestrator/train_source.py index d351672023..8c005e1047 100644 --- a/src/prime_rl/orchestrator/train_source.py +++ b/src/prime_rl/orchestrator/train_source.py @@ -1,14 +1,18 @@ """TrainSource: weighted round-robin across train envs, infinite pull. Weights are each env's configured ``ratio`` (default 1, i.e. equal weight -per env). A finite env serves a shuffled task-index table, reshuffled on -cursor exhaustion; an infinite env (``num_tasks is None``) streams a -monotonic ``task_idx`` — the server generates tasks on demand, so every -pull is a fresh task and there are no epochs to shuffle.""" +per env). A v1 env serves the tasks the orchestrator loaded client-side: a +finite one as a shuffled table (reshuffled on cursor exhaustion), an +infinite one (``num_tasks is None``) straight off its generator — every +pull is a fresh task and there are no epochs to shuffle. A legacy env's +dataset lives on its server, so it serves shuffled task *indices*.""" from __future__ import annotations import random +from collections.abc import Iterator + +import verifiers.v1 as vf from prime_rl.orchestrator.envs import TrainEnvs @@ -17,7 +21,8 @@ class TrainSource: """``next_example(available_permits)`` picks a weighted-RR env and returns its next example (or ``None`` when the env's per-call permit cost doesn't fit — the dispatch loop retries when permits free up). - Returned dicts carry ``env_name`` + ``task_idx``.""" + Returned dicts carry ``env_name`` + ``task_idx`` (+ ``task`` for v1 envs, + whose data is shipped to the env server at dispatch).""" def __init__(self, train_envs: TrainEnvs, *, seed: int | None) -> None: self.rng = random.Random(seed) @@ -25,20 +30,24 @@ def __init__(self, train_envs: TrainEnvs, *, seed: int | None) -> None: if not self.envs: raise ValueError("TrainSource needs at least one train env") - # A finite env's shuffled index table; ``None`` for an infinite env, - # whose cursor alone is the (monotonic) task_idx stream. + # A finite env's shuffled example table; ``None`` for an infinite env, + # whose generator (``self.iters``) is pulled per example. self.examples: dict[str, list[dict] | None] = {} + self.iters: dict[str, Iterator[vf.Task]] = {} self.cursors: dict[str, int] = {} # Group-scoring envs reserve ``group_size`` permits up front; # per-rollout envs need 1 self.env_costs: dict[str, int] = {} for env in self.envs: - # The orchestrator never loads the env: sample over the task-index - # range the server reported via info() (num_tasks; None = infinite). - if env.num_tasks is None: + if env.tasks is None: # legacy: sample over the index range from info() + rows: list[dict] = [{"task_idx": i, "env_name": env.name} for i in range(env.num_tasks)] + self.rng.shuffle(rows) + self.examples[env.name] = rows + elif env.num_tasks is None: # infinite: pull the generator per example self.examples[env.name] = None + self.iters[env.name] = env.tasks else: - rows: list[dict] = [{"task_idx": i, "env_name": env.name} for i in range(env.num_tasks)] + rows = [{"task_idx": task.data.idx, "task": task, "env_name": env.name} for task in env.tasks] self.rng.shuffle(rows) self.examples[env.name] = rows self.cursors[env.name] = 0 @@ -53,9 +62,9 @@ def next_example(self, available_permits: int) -> dict | None: return None rows = self.examples[env_name] cursor = self.cursors[env_name] - if rows is None: # infinite env: the cursor is the task_idx - self.cursors[env_name] = cursor + 1 - return {"task_idx": cursor, "env_name": env_name} + if rows is None: # infinite env: pull the next generated task + task = next(self.iters[env_name]) + return {"task_idx": task.data.idx, "task": task, "env_name": env_name} if cursor >= len(rows): self.rng.shuffle(rows) cursor = 0 diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 2b3e3989db..1efc349224 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -63,6 +63,9 @@ class GroupState: task_idx: int rollouts_to_schedule: int target_rollouts: int + task: vf.Task | None = None + """The group's task (v1 envs — its data is shipped on every dispatch). ``None`` for + legacy envs, which are addressed by ``task_idx`` alone.""" emitted: int = 0 eval_step: int | None = None pinned_client: vf.ClientConfig | None = None