feat(orchestrator): opt-in background group finalization

docs/algorithms.md

@@ -433,6 +433,8 @@ Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated
 
 Curriculum admission and metrics can inspect the resulting graph-native streams after group scoring.
 
+`score_group` normally runs inline in the orchestrator loop. An algorithm whose group scoring awaits slow work sets `finalize_in_background = True`: each complete group's `finalize_group` then runs as a background task, and the group is admitted and batched once it lands. These calls run concurrently and complete in any order, so a stateful algorithm must tolerate interleaved groups. Train dispatch pauses while an env has `max_finalizing_groups` (default 64) groups finalizing, or one batch of its groups (`batch_size // group_size`) if that is more. Groups that go stale while finalizing are cancelled when the step advances, and draining cancels the rest.
+
 ### Reference Scoring
 
 `OPDAlgorithm` / `OPSDAlgorithm` do their model I/O in `score_episode`: as each episode arrives they query a reference and attach sampled-token reference logprobs to its graph nodes:

src/prime_rl/orchestrator/algo/base.py

@@ -51,6 +51,12 @@ class Algorithm:
     """
 
     action_loss_type: ClassVar[ActionLossType] = "rl"
+    # Run ``finalize_group`` as a background task instead of inline in the
+    # orchestrator loop; for algorithms whose group scoring awaits slow work.
+    finalize_in_background: ClassVar[bool] = False
+    # Background finalizations per env before train dispatch pauses; the train
+    # sink raises it to one batch of the env's groups.
+    max_finalizing_groups: int = 64
 
     def __init__(self, config: AlgoConfig, clients: InferenceClient):
         self.clients = clients

src/prime_rl/orchestrator/orchestrator.py

@@ -501,11 +501,21 @@ async def main_loop(self) -> None:
         the result."""
         while not self.stopped.is_set():
             self._raise_if_component_stopped()
+            # Apply the background-finalization cap to the dispatch gate.
+            if self.train_sink.finalizing or not self.dispatcher.dispatch_allowed.is_set():
+                self.update_dispatch_gate()
             if self.draining and self.dispatcher.is_idle:
                 get_logger().info("Pipeline drained, exiting main loop")
                 self.stopped.set()
                 break
 
+            # Groups finalized in the background settle before new results.
+            if not self.train_sink.finalized_q.empty():
+                train_batch = await self.train_sink.finalized(self.train_sink.finalized_q.get_nowait())
+                if train_batch is not None and not self.draining and not self.stopped.is_set():
+                    await self.finalize_train_batch(train_batch)
+                continue
+
             try:
                 item = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5)
             except asyncio.TimeoutError:
@@ -558,7 +568,9 @@ async def main_loop(self) -> None:
                 await self.finalize_train_batch(train_batch)
 
     def _raise_if_component_stopped(self) -> None:
-        """Propagate unexpected background-component termination to the run."""
+        """Propagate unexpected background-component termination (or a failed
+        background group finalization) to the run."""
+        self.train_sink.raise_if_failed()
         for task in self.component_tasks:
             if not task.done():
                 continue
@@ -634,6 +646,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
         pack_time = time.perf_counter() - pack_start_time
         await self.sender.send(micro_batch_grid)
         self.progress.step += 1
+        await self.train_sink.cancel_stale_finalizing()
         self.update_dispatch_gate()
         # Checkpoint the step we just shipped (resume point: continue at step + 1).
         save_ckpt_time = await self.maybe_save_ckpt(step)
@@ -765,6 +778,7 @@ async def start_draining(self, reason: str) -> None:
         eval epochs still run to completion."""
         self.draining = True
         self.dispatcher.disable_train_scheduling()
+        await self.train_sink.stop()
         n_cancelled = await self.dispatcher.cancel_inflight_train_episodes()
         get_logger().info(
             f"{reason} — draining pipeline (cancelled {n_cancelled} in-flight "
@@ -821,7 +835,8 @@ def collect_pipeline_view(self) -> tuple[str, dict[str, float]]:
         multi_eval = self.eval_envs is not None and len(self.eval_envs) > 1
 
         # Train batch: finalized-group survivors only (0→target). Partial-group
-        # arrivals are surfaced as a separate ``(+N buffered)`` addendum
+        # arrivals and groups finalizing in the background are surfaced as a
+        # separate ``(+N buffered)`` addendum
         train_pct = train_batch / train_target if train_target else 0.0
         train_batch_part = f"Train batch {train_batch}/{train_target} ({train_pct:.1%})"
         if multi_train:
@@ -979,7 +994,8 @@ def update_dispatch_gate(self) -> None:
         collected — advanced right after shipping — so both call sites (ship time
         here, policy update in ``on_new_version``) share one lead formula. Steps
         are 1-indexed while policy versions stay 0-indexed, so the shipped-batch
-        count is ``progress.step - 1``."""
+        count is ``progress.step - 1``. Train dispatch also pauses while an env
+        has ``max_finalizing_groups`` groups finalizing in the background."""
         lead = (self.progress.step - 1) - self.policy.version
         gate = self.dispatcher.dispatch_allowed
         was_set = gate.is_set()
@@ -991,6 +1007,12 @@ def update_dispatch_gate(self) -> None:
                 )
                 self.gate_closed_at = time.perf_counter()
             gate.clear()
+        elif self.train_sink.finalizing_full:
+            if was_set:
+                get_logger().info(
+                    f"Pausing dispatcher while {len(self.train_sink.finalizing)} groups finalize in the background"
+                )
+            gate.clear()
         else:
             if not was_set:
                 get_logger().info(f"Resuming dispatcher (policy v{self.policy.version})")
@@ -1015,6 +1037,7 @@ async def teardown() -> None:
             if self.dispatcher is not None:
                 get_logger().debug("Stopping dispatcher")
                 await self.dispatcher.stop()
+            await self.train_sink.stop()
             if self.watcher is not None:
                 get_logger().debug("Stopping weight watcher")
                 await self.watcher.stop()

src/prime_rl/orchestrator/train_sink.py

@@ -1,31 +1,41 @@
 """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
+episode, and ``cancel()`` a dropped group's ``GroupCancellation``. Algorithms
+with ``finalize_in_background`` run ``finalize_group`` as a background task;
+the orchestrator settles the group via ``finalized()`` once it lands. 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."""
 
 from __future__ import annotations
 
 import asyncio
-from collections import defaultdict
+from collections import Counter, defaultdict
 from collections.abc import Callable, Iterable
+from dataclasses import dataclass
 
 import verifiers.v1 as vf
 
 from prime_rl.configs.orchestrator import OrchestratorConfig
-from prime_rl.orchestrator.algo.base import iter_trainable_traces
+from prime_rl.orchestrator.algo.base import Algorithm, iter_trainable_traces
 from prime_rl.orchestrator.algo.routing import stamp_loss_routing
 from prime_rl.orchestrator.envs import TrainEnvs
 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.transports.batch import TrainingSample
+from prime_rl.utils.async_utils import safe_cancel_all
 from prime_rl.utils.logger import get_logger
 
 
+def _is_stale(episode: vf.Episode, min_version: int) -> bool:
+    """Frozen-sourced episodes (no policy span) never go stale."""
+    policy = train_work(episode).policy
+    return policy is not None and policy.start < min_version
+
+
 def _prune_zero_advantages(sample: TrainingSample) -> bool:
     """Remove zero-advantage tokens from the RL component."""
     if sample.advantages is None:
@@ -54,6 +64,17 @@ def _prune_zero_advantages(sample: TrainingSample) -> bool:
     return has_rl or has_ce or has_ref_kl
 
 
+@dataclass
+class FinalizingGroup:
+    """A complete group whose background ``finalize_group`` is running."""
+
+    env_name: str
+    task: asyncio.Task
+    group: list[vf.Episode]
+    failures: list[DispatchFailure]
+    cancellation: GroupCancellation | None
+
+
 class TrainSink:
     """Score native episodes, admit groups, then compile trainer payloads."""
 
@@ -83,6 +104,11 @@ def __init__(
         # 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] = {}
+        self.finalizing: dict[str, FinalizingGroup] = {}
+        # Group ids whose background finalization completed, for the
+        # orchestrator loop to settle via ``finalized()``.
+        self.finalized_q: asyncio.Queue[str] = asyncio.Queue()
+        self.finalize_error: BaseException | None = None
         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;
@@ -103,7 +129,21 @@ def batch_progress(self) -> tuple[int, int]:
     def buffered_count(self) -> int:
         episodes = sum(len(group) for group in self.pending_groups.values())
         failures = sum(len(group) for group in self.pending_group_failures.values())
-        return episodes + failures
+        finalizing = sum(len(entry.group) + len(entry.failures) for entry in self.finalizing.values())
+        return episodes + failures + finalizing
+
+    @property
+    def finalizing_full(self) -> bool:
+        """Whether an env has as many groups finalizing as its algorithm's
+        ``max_finalizing_groups`` or, if more, as one batch of its groups: a
+        bigger batch needs more groups graded per step."""
+        counts = Counter(entry.env_name for entry in self.finalizing.values())
+        for name, n in counts.items():
+            env = self.train_envs.get(name)
+            batch_groups = self.batch_size // env.config.group_size
+            if n >= max(env.algorithm.max_finalizing_groups, batch_groups):
+                return True
+        return False
 
     def pending_batch_by_env(self) -> dict[str, int]:
         counts: dict[str, int] = defaultdict(int)
@@ -146,6 +186,48 @@ async def fail(self, failure: DispatchFailure) -> TrainBatch | None:
         await self.process_group(failure.group_id)
         return self._maybe_batch()
 
+    async def finalized(self, group_id: str) -> TrainBatch | None:
+        """Settle a group whose background ``finalize_group`` completed. Ids
+        already dropped (stale cancel, stop) are ignored."""
+        entry = self.finalizing.pop(group_id, None)
+        if entry is None:
+            return None
+        await self._settle_group(entry.group, entry.failures, entry.cancellation, finalized=True)
+        return self._maybe_batch()
+
+    async def cancel_stale_finalizing(self) -> None:
+        """Called when ``progress.step`` advances: cancel finalizing groups the
+        insertion sweep would void on arrival, accounted like a ``stale`` group
+        cancellation."""
+        min_version = min_fresh_version(self.progress.step, self.config.max_off_policy_steps)
+        if min_version <= 0:
+            return
+        stale = [
+            group_id
+            for group_id, entry in self.finalizing.items()
+            if not entry.task.done() and all(_is_stale(episode, min_version) for episode in entry.group)
+        ]
+        for group_id in stale:
+            entry = self.finalizing.pop(group_id)
+            entry.task.cancel()
+            await self._settle_group(entry.group, entry.failures, entry.cancellation, stale=True)
+        if stale:
+            get_logger().warning(
+                f"Cancelled finalization of {len(stale)} groups past "
+                f"max_off_policy_steps={self.config.max_off_policy_steps}. Consider increasing it to avoid this."
+            )
+
+    def raise_if_failed(self) -> None:
+        """Re-raise the first exception from a background ``finalize_group``."""
+        if self.finalize_error is not None:
+            raise self.finalize_error
+
+    async def stop(self) -> None:
+        """Cancel all background finalizations; their groups are dropped."""
+        tasks = [entry.task for entry in self.finalizing.values()]
+        self.finalizing.clear()
+        await safe_cancel_all(tasks)
+
     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
@@ -183,8 +265,7 @@ def _drop_stale(self, trace_ids: Iterable[str] | None = None) -> None:
         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 _is_stale(episode, min_version):
                 continue
             del self.pending_batch[trace_id]
             del self.episode_by_trace[trace_id]
@@ -208,7 +289,41 @@ async def process_group(self, group_id: str) -> None:
         cancellation = self.pending_group_cancellations.pop(group_id, None)
         if not group and not failures and cancellation is None:
             return
+        env_name = episode_env_name(group[0]) if group else None
+        algorithm = self.train_envs.get(env_name).algorithm if env_name is not None else None
+        stale = cancellation is not None and cancellation.reason == "stale"
+        if (
+            algorithm is not None
+            and algorithm.finalize_in_background
+            and not stale
+            and any(True for _ in iter_trainable_traces(group))
+        ):
+            task = asyncio.create_task(self._finalize_group(group_id, algorithm, group))
+            task.add_done_callback(self._on_finalize_done)
+            self.finalizing[group_id] = FinalizingGroup(env_name, task, group, failures, cancellation)
+            return
+        await self._settle_group(group, failures, cancellation)
 
+    async def _finalize_group(self, group_id: str, algorithm: Algorithm, group: list[vf.Episode]) -> None:
+        await algorithm.finalize_group(group)
+        self.finalized_q.put_nowait(group_id)
+
+    def _on_finalize_done(self, task: asyncio.Task) -> None:
+        if not task.cancelled() and task.exception() is not None and self.finalize_error is None:
+            self.finalize_error = task.exception()
+
+    async def _settle_group(
+        self,
+        group: list[vf.Episode],
+        failures: list[DispatchFailure],
+        cancellation: GroupCancellation | None,
+        *,
+        finalized: bool = False,
+        stale: bool = False,
+    ) -> None:
+        """Account, admit, and compile one complete group. ``finalized`` skips
+        ``finalize_group`` (it already ran in the background); ``stale`` voids
+        the group like a ``stale`` cancellation."""
         env_name = (
             episode_env_name(group[0]) if group else (failures[0].env_name if failures else cancellation.env_name)
         )
@@ -231,7 +346,7 @@ async def process_group(self, group_id: str) -> None:
         # 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":
+        if stale or (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(
@@ -241,7 +356,7 @@ async def process_group(self, group_id: str) -> None:
             return
 
         survivors = [trace for _, trace in iter_trainable_traces(group)]
-        if survivors:
+        if survivors and not finalized:
             await env.algorithm.finalize_group(group)
         admitted = self._admit(group) if group else False
         if not survivors or not admitted: