diff --git a/megatron/rl/agent/api.py b/megatron/rl/agent/api.py index ffa38be06bc..b4ffb35048f 100644 --- a/megatron/rl/agent/api.py +++ b/megatron/rl/agent/api.py @@ -1,23 +1,16 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -import asyncio -import time from abc import ABC, abstractmethod -from typing import AsyncIterator, Awaitable, Callable, Generic, NamedTuple, TypeVar +from typing import Awaitable, Callable, Generic, NamedTuple, TypeVar -import numpy as np from pydantic import BaseModel -from megatron.core.inference.utils import asyncio_Queue, asyncio_QueueShutDown -from megatron.core.utils import trace_async_exceptions - from ..__init__ import Request, TypeLookupable from ..inference import ( InferenceInterface, InferenceRequest, InferenceResponse, LLMChatMessage, - ReturnsRaw, ) from ..rollout_granularity import ConsumptionGranularity, SubmissionGranularity @@ -200,359 +193,31 @@ class TokenizedRolloutGenerator(Agent, ABC): async def get_reward_rollouts(self, request: RolloutRequest) -> list[TokenRollout]: ... -class _GranularityConfig(NamedTuple): - submission: SubmissionGranularity - consumption: ConsumptionGranularity - num_groups_per_batch: int - - @classmethod - def from_request(cls, request: GroupedRolloutRequest) -> "_GranularityConfig": - cls._validate(request) - return cls( - submission=request.submission_granularity, - consumption=request.consumption_granularity, - num_groups_per_batch=request.num_groups, - ) - - @property - def prevent_dataset_reorder(self) -> bool: - return self.consumption == "B" - - @staticmethod - def _validate(request: GroupedRolloutRequest) -> None: - assert not ( - request.submission_granularity == "B" and request.consumption_granularity == "G" - ), "Batch submission with group consumption is not supported." - assert not request.filter_groups_with_same_reward, ( - "filter_groups_with_same_reward is not currently supported: dropped groups " - "are not regenerated, so non-streaming callers receive fewer groups than " - "requested and batch-order consumers stall on incomplete batches." - ) - - -class _SubmissionGate: - """Gate capacity is measured in units of the configured submission granularity. - - Each granularity has a single release point: R slots free when inference - completes, so the gate bounds engine concurrency in rollouts. G and B - slots free when the trainer consumes the group/batch, so the gate - enforces the --rl-generation-lag run-ahead cap in groups/batches - respectively. - """ - - def __init__( - self, - *, - capacity: int, - submission: SubmissionGranularity, - ) -> None: - self._sem = asyncio.Semaphore(capacity) - self._submission = submission - self.capacity = capacity - # Observability counters, updated only on the configured submission - # granularity (the only path that touches the semaphore). `held` - # counts slots currently held; `prepare_blocked_seconds` accumulates - # time stage_prepare spent waiting on the semaphore. - self.held = 0 - self.prepare_blocked_seconds = 0.0 - self.acquire_calls = 0 - self.release_calls = 0 - - async def acquire_for(self, granularity: SubmissionGranularity) -> None: - if self._submission == granularity: - start = time.monotonic() - await self._sem.acquire() - self.prepare_blocked_seconds += time.monotonic() - start - self.held += 1 - self.acquire_calls += 1 - - def release_for(self, granularity: SubmissionGranularity) -> None: - if self._submission == granularity: - self._sem.release() - self.held -= 1 - self.release_calls += 1 - - -class _InferWorkItem(NamedTuple): - """One rollout's worth of work flowing from prepare to infer. - - Timestamps are wall-clock monotonic seconds: `prepared_at` is stamped at - construction and `infer_dequeued_at` is filled in via `_replace` when an - infer worker dequeues the item. Zero means "not yet reached". - """ - - group_id: int - rollout_idx: int - batch_id: int - index_in_batch: int - params: GroupRolloutParams - prepared_at: float = 0.0 - infer_dequeued_at: float = 0.0 - - -class _InferredItem(NamedTuple): - """One rollout post-inference, flowing from infer to assemble.""" - - item: _InferWorkItem - episode: EpisodeResult - inferred_at: float = 0.0 - - -class _RolloutPipeline: - """Per-call orchestrator for grouped rollout generation.""" +class EnvAllocation(NamedTuple): + """One env's constant share of every trainer batch.""" - def __init__( - self, - agent: "GroupedRolloutGenerator", - request: GroupedRolloutRequest, - parallel_generation_tasks: int, - ) -> None: - self.agent = agent - self.request = request - self.gran_policy = _GranularityConfig.from_request(request) - self.gate = _SubmissionGate( - capacity=parallel_generation_tasks, - submission=self.gran_policy.submission, - ) - rollouts_per_submission_unit = { - "R": 1, - "G": request.rollouts_per_group, - "B": self.gran_policy.num_groups_per_batch * request.rollouts_per_group, - }[self.gran_policy.submission] - self.num_infer_workers = parallel_generation_tasks * rollouts_per_submission_unit - if not request.streaming: - self.num_infer_workers = min( - self.num_infer_workers, request.num_groups * request.rollouts_per_group - ) - self.infer_queue = asyncio_Queue() - self.assemble_queue = asyncio_Queue() - # Unbounded: flow control is owned entirely by the submission gate. - # Bounding this queue would add a second backpressure that silently - # clamps the run-ahead configured via --rl-generation-lag. - self.output_queue = asyncio_Queue() - # Buffer of pending groups (incomplete groups being filled by - # stage_assemble). Held here so metric collection can report its size. - self._assemble_pending: dict[int, list[_InferredItem]] = {} - # Pending groups waiting for their batch to fill in stage_consume - # (only populated when prevent_dataset_reorder is True). - self._consume_pending: dict[int, list[RolloutGroup]] = {} - # Per-group "output entry" times, keyed by (batch_id, index_in_batch), - # so stage_consume can compute output_queue_dwell when yielding. - self._output_enqueued_at: dict[tuple[int, int], float] = {} - # Observability accumulators. Measured here; snapshot/reset and - # wandb formatting happen in rl_utils during metric logging. - self.infer_queue_dwell: list[float] = [] - self.engine_dwell: list[float] = [] - self.assemble_queue_dwell: list[float] = [] - self.output_queue_dwell: list[float] = [] - self.prepared_count = 0 - self.inferred_count = 0 - self.assembled_count = 0 - self.yielded_count = 0 - - async def stage_prepare(self) -> None: - """Generate gated inference work items.""" - assert ( - self.request.streaming - or self.request.num_groups % self.gran_policy.num_groups_per_batch == 0 - ), "non-streaming requires num_groups to be a multiple of num_groups_per_batch" - group_id = 0 - try: - while self.request.streaming or group_id < self.request.num_groups: - await self.gate.acquire_for("B") - batch_id = group_id // self.gran_policy.num_groups_per_batch - - for index_in_batch in range(self.gran_policy.num_groups_per_batch): - await self.gate.acquire_for("G") - params: GroupRolloutParams = await self.agent.prepare_group_rollout(self.request) - - for rollout_idx in range(self.request.rollouts_per_group): - await self.gate.acquire_for("R") - item = _InferWorkItem( - group_id=group_id, - rollout_idx=rollout_idx, - batch_id=batch_id, - index_in_batch=index_in_batch, - params=params, - prepared_at=time.monotonic(), - ) - await self.infer_queue.put(item) - self.prepared_count += 1 - group_id += 1 - finally: - self.infer_queue.shutdown() - - async def stage_infer(self) -> None: - """Run a persistent pool of inference workers, spawned once per pipeline.""" - workers = [ - asyncio.create_task(self._infer_worker()) for _ in range(self.num_infer_workers) - ] - try: - await asyncio.gather(*workers, return_exceptions=True) - finally: - for worker in workers: - worker.cancel() - self.assemble_queue.shutdown() - - async def _infer_worker(self) -> None: - while True: - try: - item = await self.infer_queue.get() - except asyncio_QueueShutDown: - return - item = item._replace(infer_dequeued_at=time.monotonic()) - if item.prepared_at: - self.infer_queue_dwell.append(item.infer_dequeued_at - item.prepared_at) - await self._infer_one(item) - - @trace_async_exceptions(verbose=True) - async def _infer_one(self, item: _InferWorkItem) -> None: - episode = await item.params.run_episode() - inferred_at = time.monotonic() - self.gate.release_for("R") - if item.infer_dequeued_at: - self.engine_dwell.append(inferred_at - item.infer_dequeued_at) - self.inferred_count += 1 - await self.assemble_queue.put( - _InferredItem(item=item, episode=episode, inferred_at=inferred_at) - ) - - async def stage_assemble(self) -> None: - """Build complete rollout groups from inferred items.""" - pending = self._assemble_pending - try: - while True: - try: - inferred = await self.assemble_queue.get() - except asyncio_QueueShutDown: - break - dequeued_at = time.monotonic() - if inferred.inferred_at: - self.assemble_queue_dwell.append(dequeued_at - inferred.inferred_at) - bucket = pending.setdefault(inferred.item.group_id, []) - bucket.append(inferred) - if len(bucket) < self.request.rollouts_per_group: - continue - completed = pending.pop(inferred.item.group_id) - completed.sort(key=lambda item: item.item.rollout_idx) - rollouts = await asyncio.gather( - *[item.item.params.build_rollout(item.episode) for item in completed] - ) - self.assembled_count += 1 - # NOTE: this filter is currently non-functional dead code: - # _GranularityConfig._validate rejects filter_groups_with_same_reward - # at pipeline construction, so `keep` is always True. Kept for a - # future PR that regenerates dropped groups instead of - # under-delivering to the caller. That PR must also release the - # gate slot on the drop path: G/B slots free on consumption, and - # a dropped group never reaches stage_consume, so its slot (and - # eventually its batch's) would leak permanently. - keep = ( - not self.request.filter_groups_with_same_reward - or np.std([rollout.reward for rollout in rollouts]) > 1e-6 - ) - if keep: - first = completed[0] - output_enqueued_at = time.monotonic() - self._output_enqueued_at[ - (first.item.batch_id, first.item.index_in_batch) - ] = output_enqueued_at - await self.output_queue.put( - RolloutGroup( - rollouts=rollouts, - batch_id=first.item.batch_id, - index_in_batch=first.item.index_in_batch, - ) - ) - finally: - self.output_queue.shutdown() - - def _record_output_dwell(self, group: RolloutGroup) -> None: - """Record how long a group sat in output_queue before being yielded.""" - key = (group.batch_id, group.index_in_batch) - enqueued_at = self._output_enqueued_at.pop(key, 0.0) - if enqueued_at: - self.output_queue_dwell.append(time.monotonic() - enqueued_at) - self.yielded_count += 1 - - async def stage_consume(self) -> AsyncIterator[RolloutGroup]: - if not self.gran_policy.prevent_dataset_reorder: - while True: - try: - group = await self.output_queue.get() - except asyncio_QueueShutDown: - return - self._record_output_dwell(group) - yield group - self.gate.release_for("G") - - next_batch_id = 0 - pending = self._consume_pending - while True: - try: - group = await self.output_queue.get() - except asyncio_QueueShutDown: - return - self._record_output_dwell(group) - pending.setdefault(group.batch_id, []).append(group) - while ( - len(pending.get(next_batch_id, [])) - >= self.gran_policy.num_groups_per_batch - ): - batch = pending.pop(next_batch_id) - batch.sort(key=lambda group: group.index_in_batch) - next_batch_id += 1 - for group in batch: - yield group - self.gate.release_for("G") - self.gate.release_for("B") + agent: "GroupedRolloutGenerator" + env_id: str + num_groups: int class GroupedRolloutGenerator(Agent, ABC): - """An interface to return grouped Rollout objects to support algorithms like GRPO.""" - - parallel_generation_tasks: int = 512 - - def __init__(self, *, parallel_generation_tasks: int | None = None, **kwargs): - super().__init__(**kwargs) - if parallel_generation_tasks is not None: - self.parallel_generation_tasks = parallel_generation_tasks + """Agent contract consumed by RolloutPipeline to generate grouped rollouts (e.g. GRPO).""" @abstractmethod - async def prepare_group_rollout( - self, - request: GroupedRolloutRequest, - ) -> GroupRolloutParams: + async def prepare_group_rollout(self, request: GroupedRolloutRequest) -> GroupRolloutParams: """Return the params for one group's rollouts.""" ... - async def get_grouped_rollouts( - self, request: GroupedRolloutRequest - ) -> AsyncIterator[RolloutGroup]: - assert isinstance( - request.inference_interface, ReturnsRaw - ), "InferenceInterface must support raw_text return to provide rollouts." - pipeline = _RolloutPipeline( - agent=self, - request=request, - parallel_generation_tasks=self.parallel_generation_tasks, - ) - # Expose the live pipeline for observability; rl_utils reads its - # queue sizes, gate state, and timing accumulators during logging. - self._active_pipeline = pipeline - stage_prepare_task = asyncio.create_task(pipeline.stage_prepare()) - infer_task = asyncio.create_task(pipeline.stage_infer()) - assemble_task = asyncio.create_task(pipeline.stage_assemble()) - tasks = (stage_prepare_task, infer_task, assemble_task) - - try: - async for group in pipeline.stage_consume(): - yield group - finally: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - self._active_pipeline = None + def rollout_allocations(self, num_groups: int) -> list[EnvAllocation]: + """Returns each env's per-trainer-batch allocation, in env order.""" + return [ + EnvAllocation( + agent=self, + env_id=getattr(self, "env_id", None) or "rollout", + num_groups=num_groups, + ) + ] class EvaluationAgent(Agent, ABC): diff --git a/megatron/rl/agent/rollout_pipeline.py b/megatron/rl/agent/rollout_pipeline.py new file mode 100644 index 00000000000..35712ce6992 --- /dev/null +++ b/megatron/rl/agent/rollout_pipeline.py @@ -0,0 +1,506 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Caller-owned orchestration of grouped rollout generation over an agent.""" + +import asyncio +import time +from collections import deque +from typing import TYPE_CHECKING, AsyncIterator, NamedTuple + +import numpy as np + +from megatron.core.inference.utils import asyncio_Queue, asyncio_QueueShutDown +from megatron.core.utils import trace_async_exceptions + +from ..inference import ReturnsRaw +from ..rollout_granularity import ( + GRANULARITY_RANK, + ConsumptionGranularity, + SubmissionGranularity, +) +from .api import EpisodeResult, GroupedRolloutRequest, GroupRolloutParams, RolloutGroup + +if TYPE_CHECKING: + from .api import GroupedRolloutGenerator + + +class _GranularityConfig(NamedTuple): + submission: SubmissionGranularity + consumption: ConsumptionGranularity + num_groups_per_batch: int + rollouts_per_group: int + num_groups_per_env: tuple[int, ...] + + @classmethod + def from_request( + cls, request: GroupedRolloutRequest, num_groups_per_env: list[int] + ) -> "_GranularityConfig": + """Build the per-request granularity policy. + + Args: + request: Grouped rollout request carrying the granularity choices. + num_groups_per_env: Groups each env contributes to one batch, in env order. + + Returns: + A validated _GranularityConfig. + """ + cls._validate(request, num_groups_per_env) + return cls( + submission=request.submission_granularity, + consumption=request.consumption_granularity, + num_groups_per_batch=request.num_groups, + rollouts_per_group=request.rollouts_per_group, + num_groups_per_env=tuple(num_groups_per_env), + ) + + def env_of_index(self, index_in_batch: int) -> int: + """Map a batch slot to the env owning it (slots are env-blocked, in env order). + + Args: + index_in_batch: Slot index within one trainer batch. + + Returns: + The env_index owning the slot. + """ + boundary = 0 + for env_index, groups in enumerate(self.num_groups_per_env): + boundary += groups + if index_in_batch < boundary: + return env_index + raise IndexError( + f"index_in_batch {index_in_batch} outside batch of {self.num_groups_per_batch}" + ) + + def units_per_batch(self) -> int: + """Submission units in one batch; gate capacity = depth-in-batches x this. + + Returns: + The number of submission units one trainer batch contains. + """ + return { + "R": self.num_groups_per_batch * self.rollouts_per_group, + "G": self.num_groups_per_batch, + "B": 1, + }[self.submission] + + @staticmethod + def _validate(request: GroupedRolloutRequest, num_groups_per_env: list[int]) -> None: + """Reject invalid granularity, layout, and filter combinations. + + Args: + request: Grouped rollout request to check. + num_groups_per_env: Proposed per-env group layout. + """ + assert ( + GRANULARITY_RANK[request.consumption_granularity] + >= GRANULARITY_RANK[request.submission_granularity] + ), ( + f"Consumption granularity ({request.consumption_granularity}) must be no finer " + f"than submission granularity ({request.submission_granularity})." + ) + assert all( + groups > 0 for groups in num_groups_per_env + ), "Each environment must request at least one group per batch." + assert ( + sum(num_groups_per_env) == request.num_groups + ), "The sum of groups per environment must equal the total number of groups requested." + assert not request.filter_groups_with_same_reward, ( + "filter_groups_with_same_reward is not currently supported: dropped groups " + "are not regenerated, so non-streaming callers receive fewer groups than " + "requested and batch-order consumers stall on incomplete batches." + ) + + +class _SubmissionGate: + """Gate capacity is measured in units of the configured submission granularity. + + Each granularity has a single release point: R slots free when inference + completes, so the gate bounds engine concurrency in rollouts. G and B + slots free when the trainer consumes the group/batch, so the gate + enforces the --rl-generation-lag run-ahead cap in groups/batches + respectively. + """ + + def __init__( + self, + *, + capacity: int, + submission: SubmissionGranularity, + ) -> None: + """Create a gate with `capacity` slots counted at `submission` granularity. + + Args: + capacity: Maximum submission units in flight. + submission: Configured submission granularity. + """ + self._sem = asyncio.Semaphore(capacity) + self._submission = submission + self.capacity = capacity + # Observability counters, updated only on the configured submission + # granularity (the only path that touches the semaphore). `held` + # counts slots currently held; `prepare_blocked_seconds` accumulates + # time stage_prepare spent waiting on the semaphore. + self.held = 0 + self.prepare_blocked_seconds = 0.0 + self.acquire_calls = 0 + self.release_calls = 0 + + async def acquire_for(self, granularity: SubmissionGranularity) -> None: + """Take one slot when crossing a boundary of the configured granularity. + + Args: + granularity: The dispatch boundary being crossed. + """ + if self._submission == granularity: + start = time.monotonic() + await self._sem.acquire() + self.prepare_blocked_seconds += time.monotonic() - start + self.held += 1 + self.acquire_calls += 1 + + def release_for(self, granularity: SubmissionGranularity) -> None: + """Release one slot when work at the given granularity reaches its release point. + + Args: + granularity: The granularity whose release point was just reached. + """ + if self._submission == granularity: + self._sem.release() + self.held -= 1 + self.release_calls += 1 + + +class _InferWorkItem(NamedTuple): + """One rollout's worth of work flowing from prepare to infer. + + Timestamps are wall-clock monotonic seconds: `prepared_at` is stamped at + construction and `infer_dequeued_at` is filled in via `_replace` when an + infer worker dequeues the item. Zero means "not yet reached". + """ + + group_id: int + rollout_idx: int + batch_id: int + index_in_batch: int + params: GroupRolloutParams + env_index: int = 0 + prepared_at: float = 0.0 + infer_dequeued_at: float = 0.0 + + +class _InferredItem(NamedTuple): + """One rollout post-inference, flowing from infer to assemble.""" + + item: _InferWorkItem + episode: EpisodeResult + inferred_at: float = 0.0 + + +class RolloutPipeline: + """Orchestrates grouped rollout generation over an agent, one instance per request. + + Constructed and driven by the caller (e.g. the trainer via run()); + the agent only supplies the env allocations, per-group preparation, and inference calls. + """ + + def __init__( + self, + agent: "GroupedRolloutGenerator", + request: GroupedRolloutRequest, + parallel_generation_tasks: int, + ) -> None: + """Validate the request and size the gate, queues, and worker pool. + + Args: + agent: Agent supplying the env layout, preparation, and inference. + request: Grouped rollout request to serve; one pipeline per request. + parallel_generation_tasks: Submission gate depth in trainer batches. + """ + assert isinstance( + request.inference_interface, ReturnsRaw + ), "InferenceInterface must support raw_text return to provide rollouts." + self.agent = agent + self.request = request + self.allocations = agent.rollout_allocations(request.num_groups) + self.gran_policy = _GranularityConfig.from_request( + request, [allocation.num_groups for allocation in self.allocations] + ) + self.gate = _SubmissionGate( + capacity=parallel_generation_tasks + * self.gran_policy.units_per_batch(), + submission=self.gran_policy.submission, + ) + self.num_infer_workers = ( + parallel_generation_tasks + * self.gran_policy.num_groups_per_batch + * request.rollouts_per_group + ) + if not request.streaming: + self.num_infer_workers = min( + self.num_infer_workers, request.num_groups * request.rollouts_per_group + ) + + # Core queues. + self.infer_queue = asyncio_Queue() + self.assemble_queue = asyncio_Queue() + self.output_queue = asyncio_Queue() + + # Buffers of partial results. + self._assemble_pending: dict[int, list[_InferredItem]] = {} + self._consume_pending: dict[int, list[RolloutGroup]] = {} + self._output_enqueued_at: dict[tuple[int, int], float] = {} + + # Observability accumulators. + self.infer_queue_dwell: list[float] = [] + self.engine_dwell: list[float] = [] + self.assemble_queue_dwell: list[float] = [] + self.output_queue_dwell: list[float] = [] + self.prepared_count = 0 + self.inferred_count = 0 + self.assembled_count = 0 + self.yielded_count = 0 + self.prepared_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) + self.assembled_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) + self.yielded_groups_per_env = [0] * len(self.gran_policy.num_groups_per_env) + + async def run(self) -> AsyncIterator[RolloutGroup]: + """Run the pipeline stages; cancels them when the iterator is closed. + + Yields: + RolloutGroup: Groups in consumption-granularity order. + """ + tasks = ( + asyncio.create_task(self.stage_prepare()), + asyncio.create_task(self.stage_infer()), + asyncio.create_task(self.stage_assemble()), + ) + try: + async for group in self.stage_consume(): + yield group + for task in tasks: + task.cancel() + results = await asyncio.gather(*tasks, return_exceptions=True) + failure = next( + ( + result + for result in results + if isinstance(result, BaseException) + and not isinstance(result, asyncio.CancelledError) + ), + None, + ) + expected_end = ( + not self.request.streaming + and self.yielded_count == self.request.num_groups + ) + if failure is not None or not expected_end: + raise RuntimeError( + "RolloutPipeline output stream ended: a pipeline stage died" + + ("" if failure is not None else " (no stage exception was recovered)") + ) from failure + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + async def stage_prepare(self) -> None: + """Generate gated inference work items.""" + assert ( + self.request.streaming + or self.request.num_groups % self.gran_policy.num_groups_per_batch == 0 + ), "non-streaming requires num_groups to be a multiple of num_groups_per_batch" + group_id = 0 + try: + while self.request.streaming or group_id < self.request.num_groups: + await self.gate.acquire_for("B") + batch_id = group_id // self.gran_policy.num_groups_per_batch + + for index_in_batch in range(self.gran_policy.num_groups_per_batch): + env_index = self.gran_policy.env_of_index(index_in_batch) + await self.gate.acquire_for("G") + agent = self.allocations[env_index].agent + params: GroupRolloutParams = await agent.prepare_group_rollout(self.request) + self.prepared_groups_per_env[env_index] += 1 + + for rollout_idx in range(self.request.rollouts_per_group): + await self.gate.acquire_for("R") + item = _InferWorkItem( + group_id=group_id, + rollout_idx=rollout_idx, + batch_id=batch_id, + index_in_batch=index_in_batch, + params=params, + env_index=env_index, + prepared_at=time.monotonic(), + ) + await self.infer_queue.put(item) + self.prepared_count += 1 + group_id += 1 + finally: + self.infer_queue.shutdown() + + async def stage_infer(self) -> None: + """Run a persistent pool of inference workers, spawned once per pipeline.""" + workers = [ + asyncio.create_task(self._infer_worker()) for _ in range(self.num_infer_workers) + ] + try: + await asyncio.gather(*workers, return_exceptions=True) + finally: + for worker in workers: + worker.cancel() + self.assemble_queue.shutdown() + + async def _infer_worker(self) -> None: + while True: + try: + item = await self.infer_queue.get() + except asyncio_QueueShutDown: + return + item = item._replace(infer_dequeued_at=time.monotonic()) + if item.prepared_at: + self.infer_queue_dwell.append(item.infer_dequeued_at - item.prepared_at) + await self._infer_one(item) + + @trace_async_exceptions(verbose=True) + async def _infer_one(self, item: _InferWorkItem) -> None: + """Run one episode for one work item and hand the result to assemble. + + Args: + item: The dequeued work item; its params carry the episode closure. + """ + episode = await item.params.run_episode() + inferred_at = time.monotonic() + self.gate.release_for("R") + if item.infer_dequeued_at: + self.engine_dwell.append(inferred_at - item.infer_dequeued_at) + self.inferred_count += 1 + await self.assemble_queue.put( + _InferredItem(item=item, episode=episode, inferred_at=inferred_at) + ) + + async def stage_assemble(self) -> None: + """Build complete rollout groups from inferred items.""" + pending = self._assemble_pending + try: + while True: + try: + inferred = await self.assemble_queue.get() + except asyncio_QueueShutDown: + break + dequeued_at = time.monotonic() + if inferred.inferred_at: + self.assemble_queue_dwell.append(dequeued_at - inferred.inferred_at) + bucket = pending.setdefault(inferred.item.group_id, []) + bucket.append(inferred) + if len(bucket) < self.request.rollouts_per_group: + continue + completed = pending.pop(inferred.item.group_id) + completed.sort(key=lambda item: item.item.rollout_idx) + rollouts = await asyncio.gather( + *[item.item.params.build_rollout(item.episode) for item in completed] + ) + first = completed[0] + self.assembled_count += 1 + self.assembled_groups_per_env[first.item.env_index] += 1 + # NOTE: this filter is currently non-functional dead code: + # _GranularityConfig._validate rejects filter_groups_with_same_reward + # at pipeline construction, so `keep` is always True. Kept for a + # future PR that regenerates dropped groups instead of + # under-delivering to the caller. That PR must also release the + # gate slot on the drop path: G/B slots free on consumption, + # and a dropped group never reaches stage_consume, so its slot + # (and eventually its batch's) would leak permanently. + keep = ( + not self.request.filter_groups_with_same_reward + or np.std([rollout.reward for rollout in rollouts]) > 1e-6 + ) + if keep: + output_enqueued_at = time.monotonic() + self._output_enqueued_at[ + (first.item.batch_id, first.item.index_in_batch) + ] = output_enqueued_at + await self.output_queue.put( + RolloutGroup( + rollouts=rollouts, + batch_id=first.item.batch_id, + index_in_batch=first.item.index_in_batch, + ) + ) + finally: + self.output_queue.shutdown() + + def _record_output_dwell(self, group: RolloutGroup) -> None: + """Record how long a group sat in output_queue before being yielded.""" + key = (group.batch_id, group.index_in_batch) + enqueued_at = self._output_enqueued_at.pop(key, 0.0) + if enqueued_at: + self.output_queue_dwell.append(time.monotonic() - enqueued_at) + self.yielded_count += 1 + self.yielded_groups_per_env[self.gran_policy.env_of_index(group.index_in_batch)] += 1 + + async def _next_complete_group(self) -> RolloutGroup | None: + """Pop the next group off output_queue and record its dwell.""" + try: + group = await self.output_queue.get() + except asyncio_QueueShutDown: + return None + self._record_output_dwell(group) + return group + + async def stage_consume(self) -> AsyncIterator[RolloutGroup]: + """Deliver groups in the order defined by the consumption granularity.""" + consume = { + "G": self._consume_completion_order, + "B": self._consume_batch_order, + }[self.gran_policy.consumption] + async for group in consume(): + yield group + + async def _consume_completion_order(self) -> AsyncIterator[RolloutGroup]: + """G consumption: deliver groups in completion order, balanced across envs.""" + groups_per_env_per_batch = self.gran_policy.num_groups_per_env + pending_groups_by_env: list[deque[RolloutGroup]] = [ + deque() for _ in groups_per_env_per_batch + ] + delivered_groups_by_env = [0] * len(groups_per_env_per_batch) + while (group := await self._next_complete_group()) is not None: + env_index = self.gran_policy.env_of_index(group.index_in_batch) + pending_groups_by_env[env_index].append(group) + yielded_any = True + while yielded_any: + yielded_any = False + for env, queue in enumerate(pending_groups_by_env): + if queue and delivered_groups_by_env[env] < groups_per_env_per_batch[env]: + yield queue.popleft() + self.gate.release_for("G") + delivered_groups_by_env[env] += 1 + yielded_any = True + if all( + count == quota + for count, quota in zip(delivered_groups_by_env, groups_per_env_per_batch) + ): + delivered_groups_by_env = [0] * len(groups_per_env_per_batch) + # The stream is over; nothing is left to balance against, so drain any pending groups. + for queue in pending_groups_by_env: + while queue: + yield queue.popleft() + self.gate.release_for("G") + + async def _consume_batch_order(self) -> AsyncIterator[RolloutGroup]: + """B consumption: deliver whole batches in dataset order.""" + next_batch_id = 0 + pending = self._consume_pending + while (group := await self._next_complete_group()) is not None: + pending.setdefault(group.batch_id, []).append(group) + while ( + len(pending.get(next_batch_id, [])) + >= self.gran_policy.num_groups_per_batch + ): + batch = pending.pop(next_batch_id) + batch.sort(key=lambda group: group.index_in_batch) + next_batch_id += 1 + for group in batch: + yield group + self.gate.release_for("G") + self.gate.release_for("B") diff --git a/megatron/rl/agent/weighted_multi_task.py b/megatron/rl/agent/weighted_multi_task.py index 2c52784be1c..a11e93fd7a7 100644 --- a/megatron/rl/agent/weighted_multi_task.py +++ b/megatron/rl/agent/weighted_multi_task.py @@ -4,13 +4,12 @@ import logging from typing import Any, Optional, Type -import numpy as np - from .registry import get_agent_class from .api import ( AgentBaseModel, ContrastiveRollout, ContrastiveRolloutGenerator, + EnvAllocation, EvaluationAgent, EvaluationRequest, EvaluationResponse, @@ -70,9 +69,7 @@ def __init__(self, agent_configs: list[AgentConfig]): self.weights.append(config.weight / total_weight) @classmethod - def from_config( - cls, config: list[dict[str, Any]], *, parallel_generation_tasks: int | None = None - ) -> 'WeightedMultiTask': + def from_config(cls, config: list[dict[str, Any]]) -> 'WeightedMultiTask': """Create a WeightedMultiTask from a config list. Args: @@ -89,8 +86,6 @@ def from_config( if not all(k in entry for k in ['agent_type', 'agent_args', 'weight']): raise ValueError(f"Missing required keys in config entry: {entry}") agent_args = entry.get('agent_args', {}) - agent_args['parallel_generation_tasks'] = parallel_generation_tasks - agent_type = get_agent_class(entry['agent_type']) agent_configs.append( AgentConfig( @@ -101,10 +96,7 @@ def from_config( ) ) - instance = cls(agent_configs) - if parallel_generation_tasks is not None: - instance.parallel_generation_tasks = parallel_generation_tasks - return instance + return cls(agent_configs) def _distribute_counts(self, total_count: int, distribute_remainder: bool = True) -> list[int]: """Helper method to distribute counts according to weights. @@ -156,12 +148,52 @@ def _distribute_counts(self, total_count: int, distribute_remainder: bool = True return final_counts + def rollout_allocations(self, num_groups: int) -> list[EnvAllocation]: + """Constant per-batch allocation for each weighted env, in env order.""" + counts = self._distribute_counts(num_groups) + env_ids = [ + getattr(agent, "env_id", None) or f"agent_{idx}" + for idx, agent in enumerate(self.agents) + ] + starved = [ + env_ids[idx] + for idx, count in enumerate(counts) + if count == 0 and self.weights[idx] > 0 + ] + if starved: + raise ValueError( + f"num_groups={num_groups} is too small to give every weighted env a group " + f"per batch (starved envs: {starved}); increase the trainer batch size." + ) + for agent, count in zip(self.agents, counts): + if count > 0 and not isinstance(agent, GroupedRolloutGenerator): + raise TypeError( + f"Agent of type {type(agent)} does not support grouped rollouts" + ) + # Snapshot for metric logging; read back by rl_utils. + self.latest_distribution = { + "env_ids": env_ids, + "agent_groups": list(counts), + "num_groups": num_groups, + } + logger.info( + "WeightedMultiTask layout: num_groups=%d per_agent=%s", + num_groups, + ", ".join(f"{eid}(groups={c})" for eid, c in zip(env_ids, counts)), + ) + return [ + EnvAllocation(agent=agent, env_id=env_id, num_groups=count) + for agent, env_id, count in zip(self.agents, env_ids, counts) + if count > 0 + ] + async def prepare_group_rollout( self, request: GroupedRolloutRequest, ) -> GroupRolloutParams: raise NotImplementedError( - "WeightedMultiTask is a collection of tasks and therefore doesn't implement this method directly. Use get_grouped_rollouts instead to generate grouped rollouts." + "WeightedMultiTask only routes; the pipeline prepares each group via the " + "agent in the matching rollout_allocations entry." ) async def get_rollout_response(self, request, inference_request): @@ -189,104 +221,6 @@ async def get_reward_rollouts(self, request: RolloutRequest) -> list[Rollout]: all_rollouts_lists = await asyncio.gather(*tasks) return [rollout for rollouts in all_rollouts_lists for rollout in rollouts] - async def get_grouped_rollouts(self, request: GroupedRolloutRequest): - """Distribute grouped rollouts across sub-agents according to weights.""" - agent_groups = self._distribute_counts(request.num_groups) - if request.submission_granularity == "B": - # In BATCH mode, pgt counts local batches in flight. agent_groups already - # splits each batch by weight, so copy pgt to every active agent. - agent_pgts = [ - self.parallel_generation_tasks if num_groups > 0 else 0 - for num_groups in agent_groups - ] - else: - # In GROUP/ROLLOUT mode, pgt counts fine-grained work units, so split it by weight. - agent_pgts = self._distribute_counts(self.parallel_generation_tasks) - agent_slots = self._distribute_counts(request.num_groups, distribute_remainder=False) - agent_slots = np.array(agent_slots) / np.gcd.reduce(agent_slots) - - # Snapshot the distribution for observability. Read back by rl_utils - # during per-iteration metric logging. - env_ids = [getattr(a, "env_id", f"agent_{i}") or f"agent_{i}" - for i, a in enumerate(self.agents)] - self.latest_distribution = { - "env_ids": env_ids, - "agent_groups": list(agent_groups), - "agent_pgts": list(agent_pgts), - "agent_slots": agent_slots.tolist(), - "total_pgt": int(sum(agent_pgts)), - "num_groups": request.num_groups, - } - logger.info( - "WeightedMultiTask distribution: sub=%s cons=%s num_groups=%d " - "rollouts_per_group=%d total_pgt=%d per_agent=" - + ", ".join( - f"{eid}(groups={g}, pgt={p}, slots={s:g})" - for eid, g, p, s in zip(env_ids, agent_groups, agent_pgts, agent_slots) - ), - request.submission_granularity, - request.consumption_granularity, - request.num_groups, - request.rollouts_per_group, - int(sum(agent_pgts)), - ) - - # Create tasks for each agent with non-zero groups - generators = [] - for agent, num_groups, pgt in zip( - self.agents, agent_groups, agent_pgts, strict=True - ): - if num_groups > 0: - if not isinstance(agent, GroupedRolloutGenerator): - raise TypeError( - f"Agent of type {type(agent)} does not support grouped rollouts" - ) - agent.parallel_generation_tasks = pgt - agent_request = GroupedRolloutRequest( - num_groups=num_groups, - streaming=request.streaming, - rollouts_per_group=request.rollouts_per_group, - inference_interface=request.inference_interface, - validation=request.validation, - generation_args=request.generation_args, - filter_groups_with_same_reward=request.filter_groups_with_same_reward, - submission_granularity=request.submission_granularity, - consumption_granularity=request.consumption_granularity, - ) - generators.append(agent.get_grouped_rollouts(agent_request)) - else: - generators.append(None) - - while any(generators): - balanced_rollouts = asyncio.Queue() - - async def get_balanced_rollouts_if_remaining(agent_id): - generated_rollouts = 0 - while generated_rollouts < agent_slots[agent_id]: - if generators[agent_id] is None: - return - try: - await balanced_rollouts.put(await anext(generators[agent_id])) - generated_rollouts += 1 - except StopAsyncIteration: - await balanced_rollouts.put(None) - generators[agent_id] = None - return - - tasks = [ - asyncio.create_task(get_balanced_rollouts_if_remaining(agent_id)) - for agent_id in range(len(generators)) - ] - - try: - while balanced_rollouts.qsize() > 0 or not all(task.done() for task in tasks): - rollout = await balanced_rollouts.get() - if rollout is not None: - yield rollout - finally: - for task in tasks: - task.cancel() - async def get_contrastive_rollouts(self, request: RolloutRequest) -> list[ContrastiveRollout]: """Distribute contrastive rollouts across sub-agents according to weights.""" agent_rollouts = self._distribute_counts(request.num_rollouts) diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index b2a6b320a97..1bbb2fb3f25 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -28,7 +28,6 @@ ReturnsRaw, ReturnsTokens, ) -from ..rollout_granularity import get_rl_parallel_generation_tasks from ..server.api import InferenceServer logger = logging.getLogger(__name__) @@ -140,7 +139,7 @@ async def launch(cls, model: GPTModel, **kwargs): concurrency_limit = ( args.grpo_prompts_per_step * args.grpo_group_size - * get_rl_parallel_generation_tasks(args) + * (args.rl_generation_lag + 1) ) custom_limits = httpx.Limits( max_connections=concurrency_limit, diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index e43e2f72ced..14551b685e9 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from functools import partial from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional import numpy as np import torch @@ -71,11 +71,13 @@ RolloutGroup, TokenRollout, ) +from megatron.rl.agent.rollout_pipeline import RolloutPipeline from megatron.rl.agent.weighted_multi_task import WeightedMultiTask +from megatron.rl.inference import ReturnsRaw from megatron.rl.inference.megatron import MegatronLocal from megatron.rl.logging import LOG_DIR as lang_rl_log_dir from megatron.rl.logging import log as lang_rl_log -from megatron.rl.rollout_granularity import get_rl_parallel_generation_tasks +from megatron.rl.rollout_granularity import ConsumptionGranularity, SubmissionGranularity from megatron.rl.sequence_packing_utils import ( compute_packed_inference_logprobs_stats, get_default_packed_seq_params, @@ -537,19 +539,12 @@ def align_unpacked_inference_logprobs( return padded_inference_logprobs -def get_agent(args, parallel_generation_tasks: int | None = None): - """Get an agent based on environment configuration. - - If args.langrl_env_config is provided, uses weighted environment selection. - Otherwise falls back to legacy single environment selection. - """ - with open(args.langrl_env_config, 'r') as f: +def get_agent(env_config_path): + """Build the rollout agent tree from the environment configuration.""" + with open(env_config_path, 'r') as f: config = yaml.safe_load(f) - return WeightedMultiTask.from_config( - config, - parallel_generation_tasks=parallel_generation_tasks, - ) + return WeightedMultiTask.from_config(config) _INFERENCE_INTERFACE = None @@ -569,33 +564,46 @@ def get_inference_interface(args, loop, model): _ROLLOUT_GENERATOR = None -_ROLLOUT_AGENT = None +_ROLLOUT_PIPELINE = None + + +def get_rollout_generator( + inference_interface: ReturnsRaw, + n_prompts: int, + samples_per_group: int, + *, + streaming: bool, + generation_args: dict[str, Any], + filter_groups_with_same_reward: bool, + submission_granularity: SubmissionGranularity, + consumption_granularity: ConsumptionGranularity, + generation_lag: int, + env_config_path: str, +) -> AsyncIterator[RolloutGroup]: + """Return the rollout group iterator for this step. - -def get_rollout_generator(args, inference_interface, n_prompts, samples_per_group): - global _ROLLOUT_GENERATOR, _ROLLOUT_AGENT - if not (streaming := args.rl_partial_rollouts) or _ROLLOUT_GENERATOR is None: - parallel_generation_tasks = get_rl_parallel_generation_tasks(args) - agent = get_agent(args, parallel_generation_tasks=parallel_generation_tasks) + Returns: + The async iterator produced by RolloutPipeline.run(). + """ + global _ROLLOUT_GENERATOR, _ROLLOUT_PIPELINE + if not streaming or _ROLLOUT_GENERATOR is None: request = GroupedRolloutRequest( num_groups=n_prompts, streaming=streaming, rollouts_per_group=samples_per_group, inference_interface=inference_interface, - generation_args={ - 'temperature': args.rl_default_temperature, - 'max_tokens': args.inference_max_seq_length, - 'top_p': args.rl_default_top_p, - 'top_k': args.rl_default_top_k, - }, - filter_groups_with_same_reward=args.grpo_filter_groups_with_same_reward, - submission_granularity=args.rl_submission_granularity, - consumption_granularity=args.rl_consumption_granularity, + generation_args=generation_args, + filter_groups_with_same_reward=filter_groups_with_same_reward, + submission_granularity=submission_granularity, + consumption_granularity=consumption_granularity, ) - # Keep the agent handle so metric logging can read the live rollout - # pipelines (see _collect_rollout_pipeline_metrics). - _ROLLOUT_AGENT = agent - _ROLLOUT_GENERATOR = agent.get_grouped_rollouts(request) + # Keep the pipeline handle so logging can read its queues, gate state, and per-env counters. + _ROLLOUT_PIPELINE = RolloutPipeline( + agent=get_agent(env_config_path), + request=request, + parallel_generation_tasks=generation_lag + 1, + ) + _ROLLOUT_GENERATOR = _ROLLOUT_PIPELINE.run() return _ROLLOUT_GENERATOR @@ -663,7 +671,21 @@ def get_environment_rollouts( with nvtx_range("rl/inference-setup", time=True): # Asyncronously run inference and rollout collection rollout_generator = get_rollout_generator( - args, inference_interface, n_prompts, samples_per_group + inference_interface, + n_prompts, + samples_per_group, + streaming=args.rl_partial_rollouts, + generation_args={ + 'temperature': args.rl_default_temperature, + 'max_tokens': args.inference_max_seq_length, + 'top_p': args.rl_default_top_p, + 'top_k': args.rl_default_top_k, + }, + filter_groups_with_same_reward=args.grpo_filter_groups_with_same_reward, + submission_granularity=args.rl_submission_granularity, + consumption_granularity=args.rl_consumption_granularity, + generation_lag=args.rl_generation_lag, + env_config_path=args.langrl_env_config, ) # NOTE(jbarker): we need to double check this when using PP>1 @@ -1174,91 +1196,96 @@ def _real(grouped): def _collect_rollout_pipeline_metrics() -> dict: - """Snapshot per-pipeline instrumentation into wandb-loggable scalars. + """Snapshot pipeline instrumentation into wandb-loggable scalars. - Walks the live rollout agent (set by get_rollout_generator) and, for each - sub-agent with an active _RolloutPipeline, reads queue sizes, gate state, - per-stage dwell times, and rate counters. Accumulators are reset after - reading; point-in-time values (queue sizes, gate held) are re-read next - call. Keys follow the existing f"{env_id}_{metric}" convention. + Reads the RolloutPipeline held by get_rollout_generator: queue sizes, gate state, + per-stage dwell times, and rate counters, plus per-env group counters for multi-env agents. + + Returns: + Metric name -> value dict; empty when no pipeline exists yet. """ - if _ROLLOUT_AGENT is None: + if _ROLLOUT_PIPELINE is None: return {} - sub_agents = ( - _ROLLOUT_AGENT.agents - if isinstance(_ROLLOUT_AGENT, WeightedMultiTask) - else [_ROLLOUT_AGENT] - ) + pipeline = _ROLLOUT_PIPELINE + dist = getattr(pipeline.agent, "latest_distribution", None) metrics: dict = {} - for sub_agent in sub_agents: - pipeline = getattr(sub_agent, "_active_pipeline", None) - if pipeline is None: - continue - env_id = getattr(sub_agent, "env_id", "") or "rollout" - gate = pipeline.gate - metrics.update({ - # Queue sizes and gate held are point-in-time reads. - f"{env_id}_pipeline_infer_queue_size": pipeline.infer_queue.qsize(), - f"{env_id}_pipeline_assemble_queue_size": pipeline.assemble_queue.qsize(), - f"{env_id}_pipeline_output_queue_size": pipeline.output_queue.qsize(), - f"{env_id}_pipeline_assemble_pending_groups": len(pipeline._assemble_pending), - f"{env_id}_pipeline_consume_pending_groups": len(pipeline._consume_pending), - f"{env_id}_pipeline_gate_capacity": gate.capacity, - f"{env_id}_pipeline_gate_held": gate.held, - f"{env_id}_pipeline_gate_utilization": ( - gate.held / gate.capacity if gate.capacity else 0.0 - ), - # Counters below accumulate since the previous collection. - f"{env_id}_pipeline_gate_prepare_blocked_seconds": gate.prepare_blocked_seconds, - f"{env_id}_pipeline_gate_acquire_calls": gate.acquire_calls, - f"{env_id}_pipeline_gate_release_calls": gate.release_calls, - f"{env_id}_pipeline_prepared_count": pipeline.prepared_count, - f"{env_id}_pipeline_inferred_count": pipeline.inferred_count, - f"{env_id}_pipeline_assembled_count": pipeline.assembled_count, - f"{env_id}_pipeline_yielded_count": pipeline.yielded_count, - }) - for name, samples in ( - ("infer_queue_dwell", pipeline.infer_queue_dwell), - ("engine_dwell", pipeline.engine_dwell), - ("assemble_queue_dwell", pipeline.assemble_queue_dwell), - ("output_queue_dwell", pipeline.output_queue_dwell), - ): - if samples: - arr = np.asarray(samples, dtype=np.float64) - metrics[f"{env_id}_pipeline_mean_{name}_s"] = float(arr.mean()) - metrics[f"{env_id}_pipeline_max_{name}_s"] = float(arr.max()) - metrics[f"{env_id}_pipeline_p50_{name}_s"] = float(np.percentile(arr, 50)) - metrics[f"{env_id}_pipeline_p99_{name}_s"] = float(np.percentile(arr, 99)) - # Reset accumulators; queue sizes and gate held are point-in-time. - pipeline.infer_queue_dwell = [] - pipeline.engine_dwell = [] - pipeline.assemble_queue_dwell = [] - pipeline.output_queue_dwell = [] - pipeline.prepared_count = 0 - pipeline.inferred_count = 0 - pipeline.assembled_count = 0 - pipeline.yielded_count = 0 - gate.prepare_blocked_seconds = 0.0 - gate.acquire_calls = 0 - gate.release_calls = 0 - - # WeightedMultiTask work distribution (agent_slots / agent_pgts). - dist = getattr(_ROLLOUT_AGENT, "latest_distribution", None) + gate = pipeline.gate + metrics.update({ + # Queue sizes and gate held are point-in-time reads. + "rollout_pipeline_infer_queue_size": pipeline.infer_queue.qsize(), + "rollout_pipeline_assemble_queue_size": pipeline.assemble_queue.qsize(), + "rollout_pipeline_output_queue_size": pipeline.output_queue.qsize(), + "rollout_pipeline_assemble_pending_groups": len(pipeline._assemble_pending), + "rollout_pipeline_consume_pending_groups": len(pipeline._consume_pending), + "rollout_pipeline_gate_capacity": gate.capacity, + "rollout_pipeline_gate_held": gate.held, + "rollout_pipeline_gate_utilization": ( + gate.held / gate.capacity if gate.capacity else 0.0 + ), + # Counters below accumulate since the previous collection. + "rollout_pipeline_gate_prepare_blocked_seconds": gate.prepare_blocked_seconds, + "rollout_pipeline_gate_acquire_calls": gate.acquire_calls, + "rollout_pipeline_gate_release_calls": gate.release_calls, + "rollout_pipeline_prepared_count": pipeline.prepared_count, + "rollout_pipeline_inferred_count": pipeline.inferred_count, + "rollout_pipeline_assembled_count": pipeline.assembled_count, + "rollout_pipeline_yielded_count": pipeline.yielded_count, + }) + for name, samples in ( + ("infer_queue_dwell", pipeline.infer_queue_dwell), + ("engine_dwell", pipeline.engine_dwell), + ("assemble_queue_dwell", pipeline.assemble_queue_dwell), + ("output_queue_dwell", pipeline.output_queue_dwell), + ): + if samples: + arr = np.asarray(samples, dtype=np.float64) + metrics[f"rollout_pipeline_mean_{name}_s"] = float(arr.mean()) + metrics[f"rollout_pipeline_max_{name}_s"] = float(arr.max()) + metrics[f"rollout_pipeline_p50_{name}_s"] = float(np.percentile(arr, 50)) + metrics[f"rollout_pipeline_p99_{name}_s"] = float(np.percentile(arr, 99)) + # Per-env group counters, mapped from env_index to env_id via the multi-env layout. + if dist: + active_env_ids = [ + env_id + for env_id, groups in zip(dist["env_ids"], dist["agent_groups"]) + if groups > 0 + ] + for env_index, env_id in enumerate(active_env_ids): + metrics[f"{env_id}_prepared_groups"] = ( + pipeline.prepared_groups_per_env[env_index] + ) + metrics[f"{env_id}_assembled_groups"] = ( + pipeline.assembled_groups_per_env[env_index] + ) + metrics[f"{env_id}_yielded_groups"] = ( + pipeline.yielded_groups_per_env[env_index] + ) + # Reset accumulators; queue sizes and gate held are point-in-time. + pipeline.infer_queue_dwell = [] + pipeline.engine_dwell = [] + pipeline.assemble_queue_dwell = [] + pipeline.output_queue_dwell = [] + pipeline.prepared_count = 0 + pipeline.inferred_count = 0 + pipeline.assembled_count = 0 + pipeline.yielded_count = 0 + pipeline.prepared_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) + pipeline.assembled_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) + pipeline.yielded_groups_per_env = [0] * len(pipeline.gran_policy.num_groups_per_env) + gate.prepare_blocked_seconds = 0.0 + gate.acquire_calls = 0 + gate.release_calls = 0 + + # WeightedMultiTask per-batch group distribution. if dist: # An env_id can appear more than once in the config (e.g. an active # entry plus an evaluation-only twin with zero weight). Sum per # env_id so the zero twin does not overwrite the active entry. per_env: dict = {} - for env_id, groups, pgt, slots in zip( - dist["env_ids"], dist["agent_groups"], dist["agent_pgts"], dist["agent_slots"] - ): - g, p, s = per_env.get(env_id, (0, 0, 0.0)) - per_env[env_id] = (g + groups, p + pgt, s + slots) - for env_id, (groups, pgt, slots) in per_env.items(): + for env_id, groups in zip(dist["env_ids"], dist["agent_groups"]): + per_env[env_id] = per_env.get(env_id, 0) + groups + for env_id, groups in per_env.items(): metrics[f"{env_id}_agent_groups"] = groups - metrics[f"{env_id}_agent_pgts"] = pgt - metrics[f"{env_id}_agent_slots"] = slots - metrics["multitask_total_pgt"] = dist["total_pgt"] return metrics @@ -2008,7 +2035,7 @@ def evaluate_and_print_results_rl( rank = torch.distributed.get_rank() if rank == 0: logger.info("Collecting evaluation results...") - agent = get_agent(args) + agent = get_agent(args.langrl_env_config) request = EvaluationRequest( inference_interface=inference_interface, num_prompts=args.rl_prompts_per_eval, @@ -2319,13 +2346,13 @@ def megatron_rl_inference_mode( def rl_inference_interface_shutdown(): global _INFERENCE_INTERFACE global _ROLLOUT_GENERATOR - global _ROLLOUT_AGENT + global _ROLLOUT_PIPELINE if _ROLLOUT_GENERATOR is not None: loop = get_asyncio_loop() loop.run_until_complete(_ROLLOUT_GENERATOR.aclose()) _ROLLOUT_GENERATOR = None - _ROLLOUT_AGENT = None + _ROLLOUT_PIPELINE = None if _INFERENCE_INTERFACE is not None: loop = get_asyncio_loop() diff --git a/megatron/rl/rollout_granularity.py b/megatron/rl/rollout_granularity.py index b9432f0bd4d..aa7ea545c54 100644 --- a/megatron/rl/rollout_granularity.py +++ b/megatron/rl/rollout_granularity.py @@ -7,12 +7,5 @@ SubmissionGranularity = Literal["R", "G", "B"] ConsumptionGranularity = Literal["G", "B"] - -def get_rl_parallel_generation_tasks(args) -> int: - """Return the number of generation slots implied by RL lag and submission granularity.""" - parallel_generation_tasks = args.rl_generation_lag + 1 - if args.rl_submission_granularity != "B": - parallel_generation_tasks *= args.grpo_prompts_per_step - if args.rl_submission_granularity == "R": - parallel_generation_tasks *= args.grpo_group_size - return parallel_generation_tasks +# Coarseness order of the granularity ladder (rollout < group < batch). +GRANULARITY_RANK: dict[str, int] = {"R": 0, "G": 1, "B": 2} diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index 72b65689ff5..781d38c378f 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -38,7 +38,6 @@ from megatron.rl import rl_utils from megatron.rl.agent.api import TokenRollout from megatron.rl.inference import ReturnsRaw -from megatron.rl.rollout_granularity import get_rl_parallel_generation_tasks from megatron.rl.sequence_packing_utils import get_default_packed_seq_params from megatron.training.arguments import parse_args, validate_args from megatron.training.global_vars import destroy_global_vars, set_global_variables @@ -214,30 +213,6 @@ def test_rl_granularity_defaults(self): assert args.rl_consumption_granularity == "B" assert args.rl_generation_lag == 0 assert not hasattr(args, "rl_parallel_generation_tasks") - assert get_rl_parallel_generation_tasks(args) == 1 - - @pytest.mark.parametrize( - "submission_granularity, generation_lag, expected_parallel_generation_tasks", - [ - pytest.param("B", 0, 1, id="batch"), - pytest.param("B", 2, 3, id="batch_with_lag"), - pytest.param("G", 0, 8, id="group"), - pytest.param("G", 2, 24, id="group_with_lag"), - pytest.param("R", 0, 32, id="rollout"), - pytest.param("R", 2, 96, id="rollout_with_lag"), - ], - ) - def test_get_rl_parallel_generation_tasks( - self, submission_granularity, generation_lag, expected_parallel_generation_tasks - ): - args = SimpleNamespace( - rl_submission_granularity=submission_granularity, - rl_generation_lag=generation_lag, - grpo_prompts_per_step=8, - grpo_group_size=4, - ) - - assert get_rl_parallel_generation_tasks(args) == expected_parallel_generation_tasks @pytest.mark.parametrize( "rl_partial_rollouts, submission_granularity", @@ -254,49 +229,52 @@ def test_get_rollout_generator_keeps_num_groups_at_trainer_batch_size( """Regression for the removed ``num_groups=1`` streaming override. Previously ``get_rollout_generator`` forced ``num_groups`` to 1 whenever it - streamed with a non-batch submission granularity. For a multi-environment - agent that collapses the per-env group distribution so some environments - receive zero groups (and a degenerate all-zero ``agent_slots``), stalling - ``get_grouped_rollouts``. ``num_groups`` must stay at the trainer batch size - (``n_prompts``) regardless of streaming or submission granularity. + streamed with a non-batch submission granularity, collapsing the per-env + group layout so some environments received zero groups (now a loud + ``ValueError`` from ``rollout_allocations``). ``num_groups`` must stay at + the trainer batch size (``n_prompts``) regardless of streaming or + submission granularity. """ n_prompts = 8 captured = {} rollout_generator = object() + agent = object() - class Agent: - def get_grouped_rollouts(self, request): + class FakePipeline: + def __init__(self, agent, request, parallel_generation_tasks): + captured["agent"] = agent captured["request"] = request - return rollout_generator + captured["parallel_generation_tasks"] = parallel_generation_tasks - def get_agent(_args, parallel_generation_tasks=None): - captured["parallel_generation_tasks"] = parallel_generation_tasks - return Agent() + def run(self): + return rollout_generator monkeypatch.setattr(rl_utils, "_ROLLOUT_GENERATOR", None) - monkeypatch.setattr(rl_utils, "get_agent", get_agent) - - args = SimpleNamespace( - rl_partial_rollouts=rl_partial_rollouts, - rl_submission_granularity=submission_granularity, - rl_consumption_granularity="B", - rl_generation_lag=0, - grpo_prompts_per_step=n_prompts, - grpo_group_size=4, - rl_default_temperature=1.0, - inference_max_seq_length=128, - rl_default_top_p=1.0, - rl_default_top_k=0, - grpo_filter_groups_with_same_reward=False, - ) + monkeypatch.setattr(rl_utils, "_ROLLOUT_PIPELINE", None) + monkeypatch.setattr(rl_utils, "get_agent", lambda _env_config_path: agent) + monkeypatch.setattr(rl_utils, "RolloutPipeline", FakePipeline) + generation_lag = 0 result = rl_utils.get_rollout_generator( - args, inference_interface=ReturnsRaw(), n_prompts=n_prompts, samples_per_group=4 + ReturnsRaw(), + n_prompts, + 4, + streaming=rl_partial_rollouts, + generation_args={'temperature': 1.0, 'max_tokens': 128, 'top_p': 1.0, 'top_k': 0}, + filter_groups_with_same_reward=False, + submission_granularity=submission_granularity, + consumption_granularity="B", + generation_lag=generation_lag, + env_config_path="unused.yaml", ) assert result is rollout_generator + assert captured["agent"] is agent assert captured["request"].num_groups == n_prompts assert captured["request"].streaming == rl_partial_rollouts + # The gate depth handed to the pipeline is lag + 1 trainer batches, + # independent of submission granularity. + assert captured["parallel_generation_tasks"] == generation_lag + 1 assert captured["request"].submission_granularity == submission_granularity @pytest.mark.parametrize( @@ -317,11 +295,6 @@ def get_agent(_args, parallel_generation_tasks=None): "--rl-consumption-granularity R is not currently supported", id="rollout_consumption_unsupported", ), - pytest.param( - {"rl_submission_granularity": "B", "rl_consumption_granularity": "G"}, - "--rl-submission-granularity B with --rl-consumption-granularity G", - id="batch_submit_group_consume_unsupported", - ), ], ) def test_rl_granularity_validation_rejects_unsupported_modes(self, overrides, match): diff --git a/tests/unit_tests/rl/test_rollout_generation.py b/tests/unit_tests/rl/test_rollout_generation.py index b9c26c7600e..ab1d960836c 100644 --- a/tests/unit_tests/rl/test_rollout_generation.py +++ b/tests/unit_tests/rl/test_rollout_generation.py @@ -1,9 +1,9 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import asyncio +from contextlib import aclosing from unittest.mock import MagicMock -import numpy as np import pytest from pydantic import Field, ValidationError @@ -16,9 +16,9 @@ RolloutGenerator, RolloutRequest, TokenRollout, - _SubmissionGate, ) from megatron.rl.agent.reward_only_agent import RewardOnlyAgent +from megatron.rl.agent.rollout_pipeline import RolloutPipeline, _SubmissionGate from megatron.rl.agent.weighted_multi_task import AgentConfig, WeightedMultiTask from megatron.rl.inference import InferenceResponse, LLMChatMessage, ReturnsRaw, ReturnsTokens @@ -60,11 +60,13 @@ def __init__(self, env_id="test", **kwargs): self.env_id = env_id self._call_count = 0 self.prepare_group_rollout_calls = 0 + self.get_rollout_response_calls = 0 async def get_reward_rollouts(self, request): raise NotImplementedError async def get_rollout_response(self, request, inference_request): + self.get_rollout_response_calls += 1 return await request.inference_interface.agenerate(inference_request) async def prepare_group_rollout(self, request): @@ -152,8 +154,10 @@ class TestConsumptionRelease: async def test_group_submission_stalls_until_consumption( self, consumption_granularity, num_groups ): + # Gate capacity in G-submission slots is parallel_generation_tasks + # (a depth in batches) x num_groups (groups per batch). capacity = 4 - gen = MockGenerator(parallel_generation_tasks=capacity) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=num_groups, rollouts_per_group=1, @@ -162,7 +166,8 @@ async def test_group_submission_stalls_until_consumption( submission_granularity="G", consumption_granularity=consumption_granularity, ) - it = gen.get_grouped_rollouts(request) + pipeline = RolloutPipeline(gen, request, parallel_generation_tasks=capacity // num_groups) + it = pipeline.run() try: for pulled in range(1, capacity + 3): # wait_for turns the deadlock failure mode (a slot never freed) @@ -181,7 +186,7 @@ async def test_group_submission_stalls_until_consumption( @pytest.mark.asyncio async def test_batch_submission_releases_once_per_batch(self): - gen = MockGenerator(parallel_generation_tasks=1) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=2, rollouts_per_group=1, @@ -190,12 +195,13 @@ async def test_batch_submission_releases_once_per_batch(self): submission_granularity="B", consumption_granularity="B", ) - it = gen.get_grouped_rollouts(request) + pipeline = RolloutPipeline(gen, request, parallel_generation_tasks=1) + it = pipeline.run() try: await asyncio.wait_for(anext(it), timeout=10) await asyncio.wait_for(anext(it), timeout=10) await _flush() - gate = gen._active_pipeline.gate + gate = pipeline.gate # Batch 0 fully yielded but the consumer hasn't come back yet: its # single batch slot is still held (a per-group release here would # show release_calls == 2 and prepared == 4). @@ -209,6 +215,74 @@ async def test_batch_submission_releases_once_per_batch(self): await it.aclose() +class TestStageFailurePropagation: + """A dead stage must fail run() loudly, never read as a clean end-of-stream. + + A stage that dies runs the queue-shutdown cascade, which reaches + stage_consume exactly like a clean end-of-stream; before run() reaped the + stage tasks, the caller saw StopAsyncIteration and waited forever for + rollouts nobody would ever generate (observed live 2026-07-30: a TypeError + in the first prepare_group_rollout idled a training job to its time limit). + """ + + @pytest.mark.asyncio + async def test_prepare_failure_raises_out_of_run(self): + class BrokenPrepareGenerator(MockGenerator): + async def prepare_group_rollout(self, request): + raise TypeError("agent/pipeline interface mismatch") + + request = GroupedRolloutRequest( + num_groups=2, + rollouts_per_group=2, + inference_interface=MockInferenceInterface(), + submission_granularity="R", + consumption_granularity="G", + ) + pipeline = RolloutPipeline(BrokenPrepareGenerator(), request, parallel_generation_tasks=1) + async with aclosing(pipeline.run()) as it: + with pytest.raises(RuntimeError, match="stage died") as excinfo: + # wait_for turns the pre-fix failure mode (an eternal hang once + # the cascade is mistaken for end-of-stream) into a test failure. + await asyncio.wait_for(anext(it), timeout=10) + assert isinstance(excinfo.value.__cause__, TypeError) + + @pytest.mark.asyncio + async def test_midstream_stage_failure_raises_after_delivered_groups(self): + class BrokenBuildGenerator(MockGenerator): + """First group builds normally; every later group's build_rollout raises.""" + + async def prepare_group_rollout(self, request): + idx = self._call_count + params = await super().prepare_group_rollout(request) + if idx < 1: + return params + + async def broken_build(episode): + raise ValueError("reward model exploded") + + return GroupRolloutParams( + run_episode=params.run_episode, build_rollout=broken_build + ) + + request = GroupedRolloutRequest( + num_groups=2, + rollouts_per_group=2, + inference_interface=MockInferenceInterface(), + submission_granularity="G", + consumption_granularity="G", + ) + pipeline = RolloutPipeline(BrokenBuildGenerator(), request, parallel_generation_tasks=1) + async with aclosing(pipeline.run()) as it: + # Group 0 is healthy and must still be delivered. + group = await asyncio.wait_for(anext(it), timeout=10) + assert len(group.rollouts) == 2 + # Group 1's build_rollout kills stage_assemble; the next pull must + # surface that failure instead of hanging on the drained stream. + with pytest.raises(RuntimeError, match="stage died") as excinfo: + await asyncio.wait_for(anext(it), timeout=10) + assert isinstance(excinfo.value.__cause__, ValueError) + + class TestRewardRollouts: @pytest.mark.asyncio async def test_get_reward_rollouts_matches_per_rollout_composition(self): @@ -246,7 +320,7 @@ def test_grouped_rollout_request_rejects_unknown_granularity(self, field): async def test_filter_groups_with_same_reward_rejected( self, num_groups, submission_granularity, consumption_granularity ): - gen = MockGenerator(parallel_generation_tasks=8) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=num_groups, rollouts_per_group=2, @@ -256,8 +330,7 @@ async def test_filter_groups_with_same_reward_rejected( consumption_granularity=consumption_granularity, ) with pytest.raises(AssertionError, match="filter_groups_with_same_reward"): - async for _ in gen.get_grouped_rollouts(request): - pass + RolloutPipeline(gen, request, parallel_generation_tasks=8) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -307,7 +380,7 @@ async def test_filter_groups_with_same_reward_rejected( ), ], ) - async def test_get_grouped_rollouts( + async def test_grouped_rollout_generation( self, num_slow_calls, streaming, @@ -318,7 +391,7 @@ async def test_get_grouped_rollouts( expected_batch_ids, expected_trajectories, ): - gen = MockGenerator(parallel_generation_tasks=8) + gen = MockGenerator() request = GroupedRolloutRequest( num_groups=num_groups, rollouts_per_group=1, @@ -329,7 +402,7 @@ async def test_get_grouped_rollouts( ) groups = [] - async for group in gen.get_grouped_rollouts(request): + async for group in RolloutPipeline(gen, request, parallel_generation_tasks=8).run(): groups.append(group) if request.streaming and len(groups) >= expected_count: break @@ -343,11 +416,14 @@ async def test_get_grouped_rollouts( @pytest.mark.asyncio async def test_rollout_submission_granularity_limits_inference_concurrency(self): - gen = MockGenerator(parallel_generation_tasks=2) + # parallel_generation_tasks is a depth in batches; the R gate admits at + # most depth x (num_groups x rollouts_per_group) rollouts at once. + parallel_generation_tasks = 1 + gen = MockGenerator() inference_interface = MockInferenceInterface(num_slow_calls=100) request = GroupedRolloutRequest( - num_groups=1, - rollouts_per_group=4, + num_groups=2, + rollouts_per_group=2, inference_interface=inference_interface, streaming=True, submission_granularity="R", @@ -355,42 +431,30 @@ async def test_rollout_submission_granularity_limits_inference_concurrency(self) ) groups = [] - async for group in gen.get_grouped_rollouts(request): + pipeline = RolloutPipeline( + gen, request, parallel_generation_tasks=parallel_generation_tasks + ) + async for group in pipeline.run(): groups.append(group) - break + if len(groups) >= 4: + break - assert len(groups) == 1 - assert len(groups[0]) == 4 - assert inference_interface.max_active_requests <= gen.parallel_generation_tasks + assert all(len(group) == 2 for group in groups) + assert inference_interface.max_active_requests <= ( + parallel_generation_tasks * request.num_groups * request.rollouts_per_group + ) @pytest.mark.asyncio @pytest.mark.parametrize( - "submission_granularity, consumption_granularity, expected_parallel_generation_tasks", - [ - pytest.param("B", "B", [4, 4], id="batch_submission"), - pytest.param("G", "G", [3, 1], id="group_submission"), - ], + "submission_granularity, consumption_granularity", + [pytest.param("B", "B", id="batch_batch"), pytest.param("G", "G", id="group_group")], ) - async def test_weighted_multi_task( - self, submission_granularity, consumption_granularity, expected_parallel_generation_tasks - ): + async def test_weighted_multi_task(self, submission_granularity, consumption_granularity): configs = [ AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), ] mt = WeightedMultiTask(configs) - mt.parallel_generation_tasks = 4 - - captured = [] - for agent in mt.agents: - original = agent.get_grouped_rollouts - - async def spy(req, orig=original): - captured.append(req) - async for group in orig(req): - yield group - - agent.get_grouped_rollouts = spy request = GroupedRolloutRequest( num_groups=4, @@ -401,59 +465,91 @@ async def spy(req, orig=original): consumption_granularity=consumption_granularity, ) groups = [] - async for group in mt.get_grouped_rollouts(request): + pipeline = RolloutPipeline(mt, request, parallel_generation_tasks=1) + async for group in pipeline.run(): groups.append(group) assert len(groups) == 4 - # Weights 3:1 → agent "a" produces 3 groups, agent "b" produces 1. + # Weights 3:1 → env "a" owns 3 batch slots, env "b" owns 1; the single + # pipeline routes preparation and generation to the owning sub-agent. env_ids = [g[0].env_id for g in groups] assert sorted(env_ids) == ["a", "a", "a", "b"] - for sub_req in captured: - assert sub_req.num_groups in (1, 3) # distributed proportionally by weight - assert sub_req.streaming == request.streaming - assert sub_req.submission_granularity == request.submission_granularity - assert sub_req.consumption_granularity == request.consumption_granularity - assert [agent.parallel_generation_tasks for agent in mt.agents] == ( - expected_parallel_generation_tasks + assert [agent.prepare_group_rollout_calls for agent in mt.agents] == [3, 1] + assert [agent.get_rollout_response_calls for agent in mt.agents] == [3, 1] + assert mt.latest_distribution["agent_groups"] == [3, 1] + # The pipeline drains fully: every gate slot is released at exhaustion. + assert pipeline.gate.held == 0 + + @pytest.mark.asyncio + async def test_group_consumption_balances_each_batch(self): + """Balanced G: every trainer-batch window holds each env's exact share.""" + configs = [ + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), + AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), + ] + mt = WeightedMultiTask(configs) + + request = GroupedRolloutRequest( + num_groups=4, + rollouts_per_group=1, + inference_interface=MockInferenceInterface(num_slow_calls=2), + streaming=True, + submission_granularity="G", + consumption_granularity="G", ) + groups = [] + async for group in RolloutPipeline(mt, request, parallel_generation_tasks=2).run(): + groups.append(group) + if len(groups) >= 12: + break + + for start in range(0, 12, 4): + env_ids = [g[0].env_id for g in groups[start : start + 4]] + assert sorted(env_ids) == ["a", "a", "a", "b"] + @pytest.mark.asyncio @pytest.mark.parametrize( - "num_groups, all_envs_active", - [ - pytest.param(1, False, id="num_groups_1_starves_an_env"), - pytest.param(8, True, id="trainer_batch_size_keeps_all_envs_active"), - ], + "submission_granularity, consumption_granularity", + [pytest.param("B", "G", id="batch_group")], ) - def test_multi_env_distribution_requires_num_groups_above_one( - self, num_groups, all_envs_active + async def test_consumption_finer_than_submission_rejected( + self, submission_granularity, consumption_granularity ): - """Regression for the removed ``num_groups=1`` streaming override. - - With multiple weighted environments, ``num_groups=1`` hands the single - group to one environment and leaves the other with zero groups. It also - collapses ``agent_slots`` (computed without remainder distribution) to all - zeros, so ``np.gcd.reduce`` is 0 and the per-agent slot counts become - ``nan`` -- which stalls ``get_grouped_rollouts``. Keeping ``num_groups`` at - the trainer batch size (> 1) keeps every environment active with a valid, - non-zero slot count. - """ + gen = MockGenerator() + request = GroupedRolloutRequest( + num_groups=2, + rollouts_per_group=1, + inference_interface=MockInferenceInterface(), + submission_granularity=submission_granularity, + consumption_granularity=consumption_granularity, + ) + with pytest.raises(AssertionError, match="no finer"): + RolloutPipeline(gen, request, parallel_generation_tasks=1) + + def test_multi_env_layout_rejects_starving_batch_size(self): + """The layout raises rather than silently starving a weighted env.""" configs = [ AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "a"}, weight=3.0), AgentConfig(agent_type=MockGenerator, agent_args={"env_id": "b"}, weight=1.0), ] mt = WeightedMultiTask(configs) - - agent_groups = mt._distribute_counts(num_groups) - agent_slots = mt._distribute_counts(num_groups, distribute_remainder=False) - - assert all(groups > 0 for groups in agent_groups) is all_envs_active - if all_envs_active: - assert min(agent_slots) > 0 - assert np.gcd.reduce(agent_slots) > 0 - else: - assert min(agent_groups) == 0 - assert all(slots == 0 for slots in agent_slots) - assert np.gcd.reduce(agent_slots) == 0 + with pytest.raises(ValueError, match="starved"): + mt.rollout_allocations(1) + assert [a.num_groups for a in mt.rollout_allocations(8)] == [6, 2] + + # Evaluation-only envs take no groups and never count as starved. + mt = WeightedMultiTask( + configs + + [ + AgentConfig( + agent_type=MockGenerator, + agent_args={"env_id": "c"}, + weight=1.0, + evaluation_only=True, + ) + ] + ) + assert [a.num_groups for a in mt.rollout_allocations(8)] == [6, 2] def make_response(epochs, prompt_length, total_len, content="resp", finish_reason="stop"): @@ -589,12 +685,15 @@ async def test_run_episode(self, driver, max_turns, done_at_turn, scripted, expe groups = [] async def _drain(): - async for group in agent.get_grouped_rollouts( - GroupedRolloutRequest( - num_groups=1, rollouts_per_group=1, inference_interface=iface - ) - ): - groups.append(group) + request = GroupedRolloutRequest( + num_groups=1, rollouts_per_group=1, inference_interface=iface + ) + async with aclosing( + RolloutPipeline(agent, request, parallel_generation_tasks=1).run() + ) as iterator: + async for group in iterator: + groups.append(group) + break # Bounded so a wedged pipeline fails fast instead of hanging. await asyncio.wait_for(_drain(), timeout=5.0)