From 15965923e24ab1bae61a2856699f07d9aad4fec1 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sat, 3 Oct 2026 21:24:13 +0000 Subject: [PATCH] fix(orchestrator): one rollout-age definition, staleness enforced per episode --- docs/training.md | 2 +- .../src/prime_rl/configs/orchestrator.py | 2 +- src/prime_rl/orchestrator/dispatcher.py | 152 +++++++++--------- src/prime_rl/orchestrator/train_sink.py | 90 +++++------ src/prime_rl/orchestrator/types.py | 4 +- src/prime_rl/orchestrator/utils.py | 22 +-- 6 files changed, 127 insertions(+), 145 deletions(-) diff --git a/docs/training.md b/docs/training.md index bb3a2db2a73..d40ffad0e59 100644 --- a/docs/training.md +++ b/docs/training.md @@ -60,7 +60,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli | `orchestrator.batch_size` | Tasks per trainer step. | | `orchestrator.constant_trainer_batch_size` | Keep trainer batches constant when samples have no training signal, such as zero advantage on all tokens. Enabled by default. Disable it for faster collection with variable trainer batch sizes. | | `orchestrator.train.group_size` | Rollouts generated per task. | -| `orchestrator.max_off_policy_steps` | Maximum staleness of a trained rollout (default 8): the version a batch trains on minus the oldest version that generated the rollout, queue time included. Episodes past the bound are dropped; a group shares one dispatch version, so its episodes age out together. The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `off_policy/*` and `mismatch_kl/all/mean` when tuning. | +| `orchestrator.max_off_policy_steps` | Maximum staleness of a trained rollout (default 8): the version a batch trains on minus the version the rollout was dispatched at, queue time included. Each episode past the bound is dropped alone; its group is scored from the members that remain. The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `off_policy/*` and `mismatch_kl/all/mean` when tuning. | | `[orchestrator.train.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `rae`, `hierarchical_grpo`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). | | `[[orchestrator.train.source]]` | Training sources. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Training sources](configuration.md#training-sources-orchestratortrainsource). | | `[[orchestrator.eval.source]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). | diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index da6f6826daa..490786247a1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -592,7 +592,7 @@ class OrchestratorConfig(BaseConfig): """Maximum training steps. If None, runs indefinitely.""" max_off_policy_steps: int = Field(8, ge=0) - """Maximum staleness of a trained rollout: the version a batch trains on (v{step-1}) minus the oldest version that generated the rollout (a rollout can span several weight updates), queue time included. Episodes past the bound are dropped, in-flight and queued; a group shares one dispatch version, so its episodes age out together. Higher values yield better throughput at the cost of off-policy noise.""" + """Maximum staleness of a trained rollout: the version a batch trains on (v{step-1}) minus the version the rollout was dispatched at (a rollout can span several weight updates), queue time included. Each episode is checked on its own and dropped alone, in flight or queued; its group is scored from the members that remain. Higher values yield better throughput at the cost of off-policy noise.""" heartbeat: HeartbeatConfig | None = None """BetterStack heartbeat configuration for monitoring training progress.""" diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index e3cfcede58f..493036d8915 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -15,7 +15,8 @@ source's emptiness), so in-flight episodes of the opposite kind drain naturally on either side of an eval boundary. - ``on_version_pending`` (called by the watcher before the engines pause for - the weight update) drops train groups already past ``max_off_policy_steps`` — a + the weight update) cancels train episodes already past + ``max_off_policy_steps``, leaving their group's other members running — a compute-saving early cancel; the sink's queue sweep is what guarantees the bound. Eval episodes are measurements for the policy version they started with. Online evals may explicitly cancel them when a newer checkpoint is ready. Train @@ -55,7 +56,7 @@ Progress, WorkKind, ) -from prime_rl.orchestrator.utils import min_fresh_version +from prime_rl.orchestrator.utils import rollout_age from prime_rl.utils.async_utils import safe_cancel, safe_cancel_all from prime_rl.utils.logger import get_logger @@ -308,17 +309,23 @@ def disable_train_scheduling(self) -> None: self.train_scheduling_disabled = True def inflight_staleness(self) -> list[int]: - """Current staleness of each in-flight live-sourced train episode: the - version the batch being collected trains on (v{step-1}) minus the - episode's dispatch version.""" + """``rollout_age`` of each in-flight live-sourced train episode for the + batch being collected.""" if self.train_envs is None or self.progress is None: return [] return [ - (self.progress.step - 1) - meta.policy_version + rollout_age(meta.policy_version, self.progress.step) for meta in self.inflight.values() if meta.kind == "train" and self.train_envs.get(meta.env_name).generation_source.uses_live_policy ] + def _is_stale(self, meta: InflightEpisode) -> bool: + return ( + meta.kind == "train" + and self.train_envs.get(meta.env_name).generation_source.uses_live_policy + and rollout_age(meta.policy_version, self.progress.step) > self.max_off_policy_steps + ) + # ── lifecycle ────────────────────────────────────────────────────────── async def start(self) -> None: @@ -397,14 +404,14 @@ async def publish_live(self) -> None: await self.flush_live() async def on_version_pending(self, step: int) -> None: - """Drop train groups past ``max_off_policy_steps``: a group dispatched at - v{k} ships at earliest in the batch currently collecting, at staleness - ``(progress.step - 1) - k`` — beyond the bound it can never train, so - cut it before more inference sinks in. This is a compute saver; the - sink's queue sweep is what guarantees no stale episode ships. Eval - groups never go stale (they are tied to their start-time policy - version), nor do frozen-sourced train groups (their generation source - doesn't change with policy updates). + """Cancel in-flight train episodes past ``max_off_policy_steps``: an + episode ships at earliest in the batch currently collecting, so once + its ``rollout_age`` exceeds the bound it can never train — cut it + before more inference sinks in. The rest of its group keeps running. + This is a compute saver; the sink's queue sweep is what guarantees no stale + episode ships. Eval episodes never go stale (they are tied to their + start-time policy version), nor do frozen-sourced train episodes + (their generation source doesn't change with policy updates). Runs *before* the inference engines are paused for the weight update so the resulting aborts are processed while the engine is still stepping — @@ -418,17 +425,9 @@ async def on_version_pending(self, step: int) -> None: if self.train_envs is None or self.progress is None: return - min_version = min_fresh_version(self.progress.step, self.max_off_policy_steps) - stale_groups = [ - gid - for gid, group in self.groups.items() - if group.kind == "train" - and self.train_envs.get(group.env_name).generation_source.uses_live_policy - and group.policy_version_at_start < min_version - ] cancelled = 0 - for gid in stale_groups: - cancelled += await self.drop_group(gid, reason="stale") + for gid in {meta.group_id for meta in self.inflight.values() if self._is_stale(meta)}: + cancelled += await self.drop_group(gid, reason="stale", stale_only=True) if cancelled: get_logger().warning( @@ -556,9 +555,10 @@ async def schedule_group_episode(self, group_id: uuid.UUID, group: GroupState) - if env_collection is None: return False env = env_collection.get(group.env_name) - # Frozen-sourced train rollouts hit a frozen pool; salting per policy - # version would invalidate its prefix cache every weight update for - # no reason. + # Members share the group's salt so they reuse the prompt's prefix + # cache. Frozen-sourced train rollouts hit a frozen pool; salting per + # policy version would invalidate its prefix cache every weight update + # for no reason. if live_sourced: cache_salt = str(group.policy_version_at_start) else: @@ -572,7 +572,7 @@ async def schedule_group_episode(self, group_id: uuid.UUID, group: GroupState) - env_name=group.env_name, group_id=group_id, task=group.task, - policy_version=group.policy_version_at_start, + policy_version=self.policy.version, step=group.step, client_config=client, started_at=time.monotonic(), @@ -700,14 +700,12 @@ async def handle_completed_request(self, task: asyncio.Task) -> None: await self.emit_episode(meta, group, episode) def complete_group_member(self, meta: InflightEpisode, group: GroupState | None) -> int: - """Advance group accounting and return the attempt's pinned policy version.""" - policy_version = meta.policy_version + """Advance group accounting and return the attempt's dispatch version.""" if group is not None: - policy_version = group.policy_version_at_start group.emitted += 1 if group.emitted >= group.target_episodes: self.groups.pop(meta.group_id, None) - return policy_version + return meta.policy_version async def emit_episode( self, @@ -735,54 +733,54 @@ async def emit_episode( episode.record_run(run) await self.out_q.put(episode) - async def drop_group(self, group_id: uuid.UUID, *, reason: CancelReason) -> int: - """Cancel this group's remaining in-flight tasks and emit one - ``GroupCancellation`` covering every episode it still owes the sink (both - in-flight and never-dispatched), so count-to-``group_size`` - finalization still fires. Returns the owed count for metrics.""" - group = self.groups.pop(group_id, None) - # Sync claim phase: pop matching tasks from ``self.inflight`` and - # release their permits in one non-yielding sweep. After this loop - # the dropped tasks are no longer reachable from ``self.inflight``, - # so ``handle_completed_request``'s existing None-guard makes the - # subsequent async emit phase race-free. - claimed: list[tuple[asyncio.Task, InflightEpisode]] = [] - for task, meta in list(self.inflight.items()): - if meta.group_id != group_id: - continue - del self.inflight[task] + async def drop_group(self, group_id: uuid.UUID, *, reason: CancelReason, stale_only: bool = False) -> int: + """Cancel what this group still owes the sink, in flight and never + dispatched, and emit one ``GroupCancellation`` covering it, so + count-to-``group_size`` finalization still fires. With ``stale_only``, + only its in-flight episodes past ``max_off_policy_steps`` go and the + rest of the group keeps running. Returns the cancelled count.""" + group = self.groups.get(group_id) + if group is None: + return 0 + # Sync claim phase: pop the tasks from ``self.inflight`` and release + # their permits in one non-yielding sweep. After this the cancelled + # tasks are unreachable from ``self.inflight``, so + # ``handle_completed_request``'s None-guard makes the async emit phase + # race-free. + claimed = [ + task + for task, meta in self.inflight.items() + if meta.group_id == group_id and (not stale_only or self._is_stale(meta)) + ] + for task in claimed: + self.retire(self.inflight.pop(task)) self.release() - self.retire(meta) - claimed.append((task, meta)) - - inflight_cancelled = len(claimed) - unscheduled_cancelled = group.episodes_to_schedule if group is not None else 0 - cancelled = inflight_cancelled + unscheduled_cancelled - - if cancelled > 0: - # ``group`` can be ``None`` only if every episode was already - # emitted — then nothing is in flight or unscheduled, so kind/env - # always resolve from the group or a claimed meta. - kind = group.kind if group is not None else claimed[-1][1].kind - env_name = group.env_name if group is not None else claimed[-1][1].env_name - self.metrics.record_cancellation(kind=kind, env_name=env_name, n=cancelled) - get_logger().debug( - f"Dropped {kind} group | group={str(group_id)[:8]} env={env_name} reason={reason} | " - f"cancelled={cancelled} (inflight={inflight_cancelled} unscheduled={unscheduled_cancelled})" - ) - await self.out_q.put( - GroupCancellation( - kind=kind, - env_name=env_name, - group_id=str(group_id), - step=group.step if group is not None else claimed[-1][1].step, - count=cancelled, - reason=reason, - ) + unscheduled = 0 if stale_only else group.episodes_to_schedule + group.episodes_to_schedule -= unscheduled + cancelled = len(claimed) + unscheduled + if cancelled == 0: + return 0 + group.emitted += cancelled + if group.emitted >= group.target_episodes: + del self.groups[group_id] + + self.metrics.record_cancellation(kind=group.kind, env_name=group.env_name, n=cancelled) + get_logger().debug( + f"Cancelled {group.kind} episodes | group={str(group_id)[:8]} env={group.env_name} reason={reason} | " + f"cancelled={cancelled} (inflight={len(claimed)} unscheduled={unscheduled})" + ) + await self.out_q.put( + GroupCancellation( + kind=group.kind, + env_name=group.env_name, + group_id=str(group_id), + step=group.step, + count=cancelled, + reason=reason, ) - + ) if claimed: - await safe_cancel_all([task for task, _ in claimed]) + await safe_cancel_all(claimed) return cancelled async def cancel_inflight_episodes(self) -> None: diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 48167e094b6..e09f2205c75 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -1,10 +1,12 @@ """Training-side episode, group, and batch assembly. ``add()`` takes one completed episode, ``fail()`` a request that produced no -episode, and ``cancel()`` a dropped group's ``GroupCancellation``. Before every -readiness check the sink sweeps ``pending_batch`` for traces past -``max_off_policy_steps`` — this sweep, not the dispatcher's in-flight cancel, -is what guarantees nothing stale ships.""" +episode, and ``cancel()`` a ``GroupCancellation`` for some of a group's +episodes. Before every readiness check the sink sweeps ``pending_batch`` for +traces past ``max_off_policy_steps`` — this sweep, not the dispatcher's +in-flight cancel, is what guarantees nothing stale ships. Staleness is per +episode: a stale episode is dropped alone, and its group is scored from the +members that remain.""" from __future__ import annotations @@ -21,7 +23,7 @@ from prime_rl.orchestrator.metrics import TrainEpisodes from prime_rl.orchestrator.trajectories import trace_to_samples from prime_rl.orchestrator.types import DispatchFailure, GroupCancellation, Progress, TrainBatch -from prime_rl.orchestrator.utils import episode_env_name, episode_group_id, min_fresh_version, train_work +from prime_rl.orchestrator.utils import episode_env_name, episode_group_id, rollout_age, train_work from prime_rl.transports.batch import TrainingSample from prime_rl.utils.logger import get_logger @@ -80,9 +82,8 @@ def __init__( self.pending_stale_attempts = 0 self.pending_groups: dict[str, list[vf.Episode]] = defaultdict(list) self.pending_group_failures: dict[str, list[DispatchFailure]] = defaultdict(list) - # A dropped group's terminal marker; its ``count`` fills in for the - # episodes the group will never deliver. - self.pending_group_cancellations: dict[str, GroupCancellation] = {} + # Their ``count``s fill in for the episodes the group will never deliver. + self.pending_group_cancellations: dict[str, list[GroupCancellation]] = defaultdict(list) self.pending_batch: dict[str, list[TrainingSample]] = {} self.episode_by_trace: dict[str, vf.Episode] = {} # Queued traces voided by the staleness sweep since the last ship; @@ -124,12 +125,9 @@ async def add(self, episode: vf.Episode) -> TrainBatch | None: return self._maybe_batch() async def cancel(self, cancellation: GroupCancellation) -> TrainBatch | None: - """Process a dropped group's terminal marker: its ``count`` completes - the group's episode accounting so finalization still fires. A - ``stale`` drop also voids the group's already-arrived episodes in - ``process_group`` — they share the group's dispatch version, so they - are equally stale.""" - self.pending_group_cancellations[cancellation.group_id] = cancellation + """Process a cancellation marker: its ``count`` completes the group's + episode accounting so finalization still fires.""" + self.pending_group_cancellations[cancellation.group_id].append(cancellation) if not self._group_complete(cancellation.group_id, cancellation.env_name): return None await self.process_group(cancellation.group_id) @@ -147,8 +145,7 @@ async def fail(self, failure: DispatchFailure) -> TrainBatch | None: return self._maybe_batch() def _group_complete(self, group_id: str, env_name: str) -> bool: - cancellation = self.pending_group_cancellations.get(group_id) - cancelled = cancellation.count if cancellation is not None else 0 + cancelled = sum(cancellation.count for cancellation in self.pending_group_cancellations[group_id]) failed = len(self.pending_group_failures[group_id]) return len(self.pending_groups[group_id]) + failed + cancelled >= self.group_size_for(env_name) @@ -159,11 +156,9 @@ def _maybe_batch(self) -> TrainBatch | None: return self.process_batch() if len(self.pending_batch) >= self.batch_size else None def _drop_stale(self, trace_ids: Iterable[str] | None = None) -> None: - """Void queued traces past ``max_off_policy_steps``. The batch being - collected is ``progress.step`` and trains policy v{step-1}, so a - queued trace generated from v{k} would ship at staleness - ``(step-1) - k``. This sweep is the hard guarantee on trained - staleness; the dispatcher's in-flight cancel only saves compute. + """Void queued traces past ``max_off_policy_steps``. This sweep is the + hard guarantee on trained staleness; the dispatcher's in-flight cancel + only saves compute. Queued traces only age when ``progress.step`` advances, so the full sweep runs once per step; ``trace_ids`` scopes the check to a freshly inserted group, whose traces may already be stale on arrival. @@ -177,14 +172,10 @@ def _drop_stale(self, trace_ids: Iterable[str] | None = None) -> None: return self._swept_step = self.progress.step trace_ids = list(self.pending_batch) - min_version = min_fresh_version(self.progress.step, self.config.max_off_policy_steps) - if min_version <= 0: - return dropped = 0 for trace_id in trace_ids: episode = self.episode_by_trace[trace_id] - policy = train_work(episode).policy - if policy is None or policy.start >= min_version: + if not self._is_stale(episode): continue del self.pending_batch[trace_id] del self.episode_by_trace[trace_id] @@ -197,6 +188,10 @@ def _drop_stale(self, trace_ids: Iterable[str] | None = None) -> None: "Consider increasing it to avoid this." ) + def _is_stale(self, episode: vf.Episode) -> bool: + policy = train_work(episode).policy + return policy is not None and rollout_age(policy.start, self.progress.step) > self.config.max_off_policy_steps + async def process_episode(self, episode: vf.Episode) -> None: """Run rollout-local algorithm work on one native episode.""" env_name = episode_env_name(episode) @@ -205,12 +200,12 @@ async def process_episode(self, episode: vf.Episode) -> None: async def process_group(self, group_id: str) -> None: group = self.pending_groups.pop(group_id, []) failures = self.pending_group_failures.pop(group_id, []) - cancellation = self.pending_group_cancellations.pop(group_id, None) - if not group and not failures and cancellation is None: + cancellations = self.pending_group_cancellations.pop(group_id, []) + if not group and not failures and not cancellations: return env_name = ( - episode_env_name(group[0]) if group else (failures[0].env_name if failures else cancellation.env_name) + episode_env_name(group[0]) if group else (failures[0].env_name if failures else cancellations[0].env_name) ) env = self.train_envs.get(env_name) traces = [trace for episode in group for trace in episode.traces] @@ -220,26 +215,17 @@ async def process_group(self, group_id: str) -> None: + sum(not episode.ok for episode in group if not episode.traces) + len(failures) ) - n_owed = len(group) + len(failures) + (cancellation.count if cancellation is not None else 0) + cancelled = sum(cancellation.count for cancellation in cancellations) + n_owed = len(group) + len(failures) + cancelled self.pending_failures.extend(failures) - if cancellation is not None: - self.pending_cancelled_attempts += cancellation.count - if cancellation.reason == "stale": - self.pending_stale_attempts += cancellation.count - - # A stale drop voids the whole group: every member shares the dispatch - # version, so the arrived episodes are exactly as stale as the - # cancelled tail. Stale groups bypass the curriculum — a pipeline - # decision is not a task result. - if cancellation is not None and cancellation.reason == "stale": - self.pending_episodes.extend(group, admitted=False, cancelled=True) - self._record_zero_output(group, [], n_owed) - get_logger().debug( - f"Dropped group | env={env_name} task_idx={task_idx} | " - f"episodes={len(group)} traces={len(traces)} (errored={num_errored}) | reason=cancelled (stale)" - ) - return - + self.pending_cancelled_attempts += cancelled + self.pending_stale_attempts += sum(c.count for c in cancellations if c.reason == "stale") + + # Stale members are left out like errored ones, before the algorithm + # and the curriculum see the group. + stale = [episode for episode in group if self._is_stale(episode)] + group = [episode for episode in group if not self._is_stale(episode)] + self.pending_episodes.extend(stale, admitted=False, cancelled=True) survivors = [trace for _, trace in iter_trainable_traces(group)] if survivors: await env.algorithm.finalize_group(group) @@ -286,11 +272,9 @@ async def process_group(self, group_id: str) -> None: if trace.id in samples_by_trace: self.episode_by_trace[trace.id] = episode self._drop_stale(samples_by_trace) - # A group's traces share one dispatch version, so the insertion sweep - # voids all or none of them. A fully-voided group shipped nothing — - # advance the zero-output tally instead of resetting it, or a stalled - # trainer plus a tight bound could void groups forever without ever - # surfacing the warning. + # A fully-voided group shipped nothing — advance the zero-output tally + # instead of resetting it, or a stalled trainer plus a tight bound + # could void groups forever without ever surfacing the warning. if not any(trace_id in self.pending_batch for trace_id in samples_by_trace): self._record_zero_output(group, [], n_owed) return diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 38fddcf7d5b..f461499f3bf 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -40,8 +40,8 @@ class Progress: @dataclass class GroupCancellation: - """Terminal marker for a dropped group: one message covering every episode - the group still owed the sink (in-flight and never-dispatched), so + """Cancelled episodes of one group (in flight and never dispatched): ``count`` + stands in for the episodes the group will never deliver, so count-to-``group_size`` finalization still fires. ``reason`` distinguishes pipeline decisions (staleness, overload cut, superseded eval) from episode errors.""" diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index e625eec3ce9..e60c0b91b59 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -39,23 +39,23 @@ def eval_work(episode: vf.Episode[Any, Any, Any]) -> vf.EvalWorkInfo: return run.work -def min_fresh_version(step: int, max_off_policy_steps: int) -> int: - """Oldest dispatch version whose episodes may still train in batch - ``step`` — anything older would ship past ``max_off_policy_steps``.""" - return (step - 1) - max_off_policy_steps +def rollout_age(start_version: int, step: int) -> int: + """Rollout staleness: the policy version batch ``step`` trains on, + v{step-1}, minus the version the rollout was dispatched at. Queue time + counts.""" + return (step - 1) - start_version def episode_staleness(episode: vf.Episode[Any, Any, Any], training_step: int) -> tuple[int, int, int]: - """``(total, in_flight, in_queue)`` staleness of one train episode when - consumed by batch ``training_step``: the version the batch trains on - (v{step-1}) minus the version that generated the episode. ``in_flight`` - is the span's share (weight updates during generation); ``in_queue`` is - time spent buffered between completion and ship. Frozen-sourced episodes - (no policy span) are never stale.""" + """``(total, in_flight, in_queue)`` split of ``rollout_age`` for one train + episode consumed by batch ``training_step``. ``in_flight`` is the span's + share (weight updates during generation); ``in_queue`` is time spent + buffered between completion and ship. Frozen-sourced episodes (no policy + span) are never stale.""" policy = train_work(episode).policy if policy is None: return 0, 0, 0 - total = max(0, (training_step - 1) - policy.start) + total = max(0, rollout_age(policy.start, training_step)) in_flight = min(total, max(0, policy.end - policy.start)) in_queue = total - in_flight return total, in_flight, in_queue