rl: add durable rollout bank for completed groups

megatron/rl/agent/api.py

@@ -6,7 +6,6 @@
 from typing import AsyncIterator, 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
@@ -19,11 +18,19 @@
     LLMChatMessage,
     ReturnsRaw,
 )
+from ..rollout_bank import RolloutBank
 from ..rollout_granularity import ConsumptionGranularity, SubmissionGranularity
-
-
-class AgentBaseModel(BaseModel, extra='allow'):
-    pass
+from ..types import (
+    AgentBaseModel,
+    EnvId,
+    GroupedRollouts,
+    GroupQueuesPerEnv,
+    GroupsPerEnv,
+    Rollout,
+    RolloutGroup,
+    Rollouts,
+    TokenRollout,
+)
 
 
 class RolloutRequest(Request):
@@ -45,56 +52,7 @@ class GroupedRolloutRequest(Request):
     streaming: bool = False
     submission_granularity: SubmissionGranularity = "B"
     consumption_granularity: ConsumptionGranularity = "B"
-
-
-class Rollout(AgentBaseModel):
-    """Data for language-based Rollout."""
-
-    trajectory: list[str]
-    prompt_length: list[int] | None = None
-    reward: float = None
-    env_id: str = ''
-    problem_id: str | None = None
-    policy_epoch: list[list[tuple[int, int]]]
-    kv_cache_epoch: list[list[tuple[int, int]]]
-    num_evictions: list[int]
-
-
-class TokenRollout(AgentBaseModel):
-    """Tokenized representation of a language-based Rollout."""
-
-    trajectory: list[list[int]]
-    reward: list[float] | float
-    generation_mask: list[list[bool]] | None = None
-    logprobs: list[list[float]] | None = None
-    env_id: str = ''
-    problem_id: str | None = None
-    policy_epoch: list[list[tuple[int, int]]]
-    kv_cache_epoch: list[list[tuple[int, int]]]
-    num_evictions: list[int]
-
-
-Rollouts = list[TokenRollout | Rollout]
-
-
-class RolloutGroup(AgentBaseModel):
-    """A group of rollouts (e.g. multiple completions for one prompt) with batch metadata."""
-
-    rollouts: Rollouts
-    batch_id: int = 0
-    index_in_batch: int = 0
-
-    def __iter__(self):
-        return iter(self.rollouts)
-
-    def __len__(self):
-        return len(self.rollouts)
-
-    def __getitem__(self, idx):
-        return self.rollouts[idx]
-
-
-GroupedRollouts = list[RolloutGroup]
+    num_groups_per_env: GroupsPerEnv | None = None
 
 
 class EpisodeResult(NamedTuple):
@@ -306,9 +264,13 @@ def __init__(
         agent: "GroupedRolloutGenerator",
         request: GroupedRolloutRequest,
         parallel_generation_tasks: int,
+        bank: "RolloutBank | None" = None,
     ) -> None:
         self.agent = agent
         self.request = request
+        # Optional durable rollout bank. Injected (never read from a global) so the
+        # pipeline stays testable; when None, all bank calls are skipped.
+        self.bank = bank
         self.gran_policy = _GranularityConfig.from_request(request)
         self.gate = _SubmissionGate(
             capacity=parallel_generation_tasks,
@@ -457,13 +419,18 @@ async def stage_assemble(self) -> None:
                     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,
-                        )
+                    group = RolloutGroup(
+                        rollouts=rollouts,
+                        batch_id=first.item.batch_id,
+                        index_in_batch=first.item.index_in_batch,
                     )
+                    # Write-through: the completed group hits durable storage the
+                    # instant it exists, before it is queued for the trainer. The
+                    # returned uid rides on the group so the consume side can mark
+                    # it consumed. Rank-0 single writer (only rank 0 runs the pipeline).
+                    if self.bank is not None:
+                        group.uid = self.bank.append(group)
+                    await self.output_queue.put(group)
         finally:
             self.output_queue.shutdown()
 
@@ -515,6 +482,10 @@ class GroupedRolloutGenerator(Agent, ABC):
 
     def __init__(self, *, parallel_generation_tasks: int | None = None, **kwargs):
         super().__init__(**kwargs)
+        # Durable rollout bank, wired in by rl_utils before generation. None when
+        # the feature is disabled (or on non-rank-0). Propagated to the pipeline
+        # (and, for WeightedMultiTask, to each sub-agent) on build.
+        self._rollout_bank = None
         if parallel_generation_tasks is not None:
             self.parallel_generation_tasks = parallel_generation_tasks
 
@@ -536,6 +507,7 @@ async def get_grouped_rollouts(
             agent=self,
             request=request,
             parallel_generation_tasks=self.parallel_generation_tasks,
+            bank=self._rollout_bank,
         )
         # Expose the live pipeline for observability; rl_utils reads its
         # queue sizes, gate state, and timing accumulators during logging.

megatron/rl/agent/weighted_multi_task.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from .registry import get_agent_class
+from ..types import GroupsPerEnv
 from .api import (
     AgentBaseModel,
     ContrastiveRollout,
@@ -21,6 +21,7 @@
     RolloutGenerator,
     RolloutRequest,
 )
+from .registry import get_agent_class
 
 logger = logging.getLogger(__name__)
 
@@ -189,9 +190,53 @@ 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]
 
+    def _env_ids(self) -> list[str]:
+        """Per-agent env_ids (aligned with ``self.agents``)."""
+        env_ids = []
+        for i, (a, config) in enumerate(zip(self.agents, self.agent_configs)):
+            env_id = getattr(a, "env_id", None)
+            if not env_id and not config.evaluation_only:
+                raise ValueError(
+                    f"Active agent {i} ({type(a).__name__}) has no env_id; it is "
+                    f"required to weight-balance restored rollout-bank groups by env. "
+                    f"Set env_id for every non-evaluation agent in the environment config."
+                )
+            env_ids.append(env_id or f"agent_{i}")
+        return env_ids
+
+    def env_group_targets(self, total_count: int) -> GroupsPerEnv:
+        """Per-env group counts for a full batch of ``total_count`` groups.
+
+        The weight split (``_distribute_counts``) keyed by env_id and summed over
+        any agents that share an env_id — the whole-batch target used to
+        weight-balance injected restored rollout-bank groups. Living here, beside
+        ``_distribute_counts``/``_env_ids``, keeps this target and the generator's
+        actual split identical. Counts always sum to ``total_count``
+        (weight-proportional, remainder to the largest fractional parts), and
+        agents that share an env_id are merged into one entry.
+        """
+        target: GroupsPerEnv = {}
+        for eid, c in zip(self._env_ids(), self._distribute_counts(total_count)):
+            target[eid] = target.get(eid, 0) + c
+        return target
+
     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)
+        override: GroupsPerEnv | None = request.num_groups_per_env
+        if override is not None:
+            # Explicit per-env counts (e.g. the residual after injecting restored
+            # rollout-bank groups). Route each env's count to its agent(s).
+            env_ids = self._env_ids()
+            unknown = set[Any](override) - set(env_ids)
+            if unknown:
+                raise ValueError(
+                    f"num_groups_per_env references unknown env_id(s) {sorted(unknown)}; "
+                    f"known env_ids: {sorted(set(env_ids))}. Check that the envs defined "
+                    f"in this run are consistent with the restored rollouts."
+                )
+            agent_groups = [override.get(eid, 0) for eid in env_ids]
+        else:
+            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.
@@ -202,8 +247,13 @@ async def get_grouped_rollouts(self, request: GroupedRolloutRequest):
         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)
+        if override is not None:
+            # With an explicit per-env override, slots must follow the residual
+            # shape, not a weight split of num_groups.
+            raw_slots = list(agent_groups)
+        else:
+            raw_slots = self._distribute_counts(request.num_groups, distribute_remainder=False)
+        agent_slots = np.array(raw_slots) / np.gcd.reduce(raw_slots)
 
         # Snapshot the distribution for observability. Read back by rl_utils
         # during per-iteration metric logging.
@@ -242,6 +292,7 @@ async def get_grouped_rollouts(self, request: GroupedRolloutRequest):
                         f"Agent of type {type(agent)} does not support grouped rollouts"
                     )
                 agent.parallel_generation_tasks = pgt
+                agent._rollout_bank = self._rollout_bank
                 agent_request = GroupedRolloutRequest(
                     num_groups=num_groups,
                     streaming=request.streaming,