Offload finalized training publication calculations from the event loop

src/prime_rl/orchestrator/orchestrator.py

@@ -64,6 +64,7 @@
     episode_staleness,
     eval_work,
     intercept_vf_logging,
+    release_and_trim,
     set_default_executor,
     trim_process_memory,
 )
@@ -155,6 +156,7 @@ def __init__(self, config: OrchestratorConfig) -> None:
         self.wait_for_policy_time = 0.0
         self.eval_triggered_steps: set[int] = set()
         self.component_tasks = []
+        self.release_future: asyncio.Future | None = None
 
         # Always assigned by ``setup()``; None-initialized so teardown can run
         # on a partially completed setup with plain attribute checks
@@ -443,7 +445,7 @@ async def start(self) -> None:
                 get_logger().success("Orchestrator finished")
             else:
                 get_logger().warning("Orchestrator cleanup complete (forced)")
-            trim_process_memory()
+            await asyncio.to_thread(trim_process_memory)
 
     async def wait_for_version(self, version: int, reason: str) -> None:
         """Bounded wait until the watcher has applied v{version}."""
@@ -548,7 +550,7 @@ async def handle_train_result(self, batch: TrainBatch | None) -> None:
             await self.finalize_train_batch(batch)
             if self.draining or self.stopped.is_set():
                 return
-            batch = self.train_sink.take_batch()
+            batch = await self.train_sink.take_batch()
         if self.draining or self.stopped.is_set():
             if batch is not None:
                 self.train_sink.rollout_window.metrics.count("rollout/queued/cancelled", float(batch.cohort.num_traces))
@@ -629,7 +631,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
         self.last_batch_at = now
 
         effective = batch.cohort.effective
-        n_trainable = sum(is_trainable(record.trace) for record in effective.records)
+        n_trainable = await asyncio.to_thread(lambda: sum(is_trainable(record.trace) for record in effective.records))
         if effective.num_traces and n_trainable / effective.num_traces <= 0.1:
             get_logger().warning(
                 f"Only {n_trainable}/{effective.num_traces} effective traces are trainable "
@@ -663,8 +665,11 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
         # The effective (clean, trained-on) subset is logged at ship time as annotation
         # records against each trace's arrival record - membership, advantages, the step
         # it shipped at - never a second episode copy.
-        await monitors.log(effective.vf_episodes, step, "train", "effective")
-        await monitors.log_annotations(stamp_batch(effective.vf_episodes, step))
+        effective_episodes = await asyncio.to_thread(lambda: effective.vf_episodes)
+        # Monitors may retain the list for asynchronous upload after publication.
+        await monitors.log(effective_episodes.copy(), step, "train", "effective")
+        annotations = await asyncio.to_thread(stamp_batch, effective_episodes, step)
+        await monitors.log_annotations(annotations)
 
         pack_start_time = time.perf_counter()
         micro_batch_grid = await asyncio.to_thread(self.packer.pack, batch.samples)
@@ -674,54 +679,58 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
         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)
-        trim_process_memory()
 
-        metrics = rollout_metrics | effective.metrics.to_wandb(prefix="train/agg", subset="effective")
-        for env_name, env_pool in effective.by_env().items():
-            metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset="effective")
-        num_input = sum(record.trace.num_input_tokens for record in effective.records)
-        num_output = sum(record.trace.num_output_tokens for record in effective.records)
-        metrics |= {
-            "progress/input_tokens": num_input,
-            "progress/output_tokens": num_output,
-            "time/step": step_time,
-            "time/pack": pack_time,
-            "time/save_ckpt": save_ckpt_time,
-            "time/wait_for_policy": self.wait_for_policy_time,
-            "step": step,
-        }
-        # Staleness of the shipped cohort, decomposed into its in-flight and
-        # in-queue shares. Queue drops are reported with rollout windows.
-        staleness = [episode_staleness(episode, step) for episode in effective]
-        if staleness:
-            totals, in_flight, in_queue = (list(values) for values in zip(*staleness))
+        def shipped_metrics() -> dict[str, float]:
+            metrics = effective.metrics.to_wandb(prefix="train/agg", subset="effective")
+            for env_name, env_pool in effective.by_env().items():
+                metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset="effective")
+            num_input = sum(record.trace.num_input_tokens for record in effective.records)
+            num_output = sum(record.trace.num_output_tokens for record in effective.records)
             metrics |= {
-                "off_policy/mean": sum(totals) / len(totals),
-                "off_policy/max": float(max(totals)),
-                "off_policy/in_flight/mean": sum(in_flight) / len(in_flight),
-                "off_policy/in_flight/max": float(max(in_flight)),
-                "off_policy/in_queue/mean": sum(in_queue) / len(in_queue),
-                "off_policy/in_queue/max": float(max(in_queue)),
+                "progress/input_tokens": num_input,
+                "progress/output_tokens": num_output,
+                "time/step": step_time,
+                "time/pack": pack_time,
+                "time/save_ckpt": save_ckpt_time,
+                "step": step,
             }
-        for env_name, env_pool in effective.by_env().items():
-            metrics[f"batch/{env_name}"] = env_pool.num_traces / effective.num_traces
-        metrics |= self.train_source.metrics()
+            # Staleness of the shipped cohort, decomposed into its in-flight and
+            # in-queue shares. Queue drops are reported with rollout windows.
+            staleness = [episode_staleness(episode, step) for episode in effective]
+            if staleness:
+                totals, in_flight, in_queue = (list(values) for values in zip(*staleness))
+                metrics |= {
+                    "off_policy/mean": sum(totals) / len(totals),
+                    "off_policy/max": float(max(totals)),
+                    "off_policy/in_flight/mean": sum(in_flight) / len(in_flight),
+                    "off_policy/in_flight/max": float(max(in_flight)),
+                    "off_policy/in_queue/mean": sum(in_queue) / len(in_queue),
+                    "off_policy/in_queue/max": float(max(in_queue)),
+                }
+            for env_name, env_pool in effective.by_env().items():
+                metrics[f"batch/{env_name}"] = env_pool.num_traces / effective.num_traces
+            return metrics
+
+        source_metrics = self.train_source.metrics()
+        metrics = await asyncio.to_thread(shipped_metrics)
+        wait_for_policy_time = self.wait_for_policy_time
+        self.wait_for_policy_time = 0.0
+        metrics = rollout_metrics | metrics | source_metrics
+        metrics["time/wait_for_policy"] = wait_for_policy_time
         await monitors.log(metrics, step=step)
 
-        active_step_time = max(step_time - self.wait_for_policy_time, 0.0)
-        if step_time > 0 and self.wait_for_policy_time >= active_step_time:
+        active_step_time = max(step_time - wait_for_policy_time, 0.0)
+        if step_time > 0 and wait_for_policy_time >= active_step_time:
             get_logger().warning(
-                f"Orchestrator waited {format_time(self.wait_for_policy_time)} for policy updates, at least as long "
+                f"Orchestrator waited {format_time(wait_for_policy_time)} for policy updates, at least as long "
                 f"as its {format_time(active_step_time)} active step time. Train-inference compute is imbalanced; "
                 "add more trainer nodes."
             )
 
-        self.wait_for_policy_time = 0.0
-
         if self.heart is not None:
             self.heart.beat()
 
-        self.log_train_batch(batch, step=step, step_time=step_time)
+        await asyncio.to_thread(self.log_train_batch, batch, step=step, step_time=step_time)
 
         if config.max_steps is not None and step >= config.max_steps:
             await self.wait_for_version(step, reason="before shutdown")
@@ -731,7 +740,11 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
         # versions it would need are never broadcast).
         if config.max_steps is not None and step >= config.max_steps:
             await self.start_draining("Shipped the final batch")
-        trim_process_memory()
+        # Shielded so cancelling the loop does not cancel the release, which teardown still awaits.
+        if self.release_future is not None:
+            await asyncio.shield(self.release_future)
+        payload = (batch.samples, batch.cohort.episodes, effective_episodes, *micro_batch_grid)
+        self.release_future = asyncio.get_running_loop().run_in_executor(None, release_and_trim, *payload)
 
     async def start_draining(self, reason: str) -> None:
         """Stop scheduling train work and let the pipeline empty; triggered
@@ -978,6 +991,8 @@ async def stop(self) -> None:
         async def teardown() -> None:
             get_logger().debug("Closing micro batch sender")
             self.sender.close()
+            if self.release_future is not None:
+                await self.release_future
             if self.dispatcher is not None:
                 get_logger().debug("Stopping dispatcher")
                 await self.dispatcher.stop()

src/prime_rl/orchestrator/train_sink.py

@@ -117,7 +117,7 @@ async def add(self, episode: vf.Episode) -> TrainBatch | None:
         if not self._group_complete(group_id, env_name):
             return None
         await self.process_group(group_id)
-        return self.take_batch()
+        return await self.take_batch()
 
     async def cancel(self, cancellation: GroupCancellation) -> TrainBatch | None:
         """Process a dropped group's terminal marker: its ``count`` completes
@@ -129,7 +129,7 @@ async def cancel(self, cancellation: GroupCancellation) -> TrainBatch | None:
         if not self._group_complete(cancellation.group_id, cancellation.env_name):
             return None
         await self.process_group(cancellation.group_id)
-        return self.take_batch()
+        return await self.take_batch()
 
     async def fail(self, failure: DispatchFailure) -> TrainBatch | None:
         """Count a request failure toward its group without presenting it as
@@ -140,19 +140,19 @@ async def fail(self, failure: DispatchFailure) -> TrainBatch | None:
         if not self._group_complete(failure.group_id, failure.env_name):
             return None
         await self.process_group(failure.group_id)
-        return self.take_batch()
+        return await self.take_batch()
 
     def _group_complete(self, group_id: str, env_name: str) -> bool:
         cancellation = self.pending_group_cancellations.get(group_id)
         cancelled = cancellation.count if cancellation is not None else 0
         failed = len(self.pending_group_failures[group_id])
         return len(self.pending_groups[group_id]) + failed + cancelled >= self.group_size_for(env_name)
 
-    def take_batch(self) -> TrainBatch | None:
+    async def take_batch(self) -> TrainBatch | None:
         """Sweep stale queued traces, then cut a batch if the survivors still
         meet the threshold."""
         self._drop_stale()
-        return self.process_batch() if len(self.pending_batch) >= self.batch_size else None
+        return await self.process_batch() if len(self.pending_batch) >= self.batch_size else None
 
     def _drop_stale(self, trace_ids: Iterable[str] | None = None) -> None:
         """Void queued traces past ``max_off_policy_steps``. The batch being
@@ -345,7 +345,7 @@ def _warn_zero_output(self) -> None:
             f"({windows} zero-output batch equivalents)"
         )
 
-    def process_batch(self) -> TrainBatch:
+    async def process_batch(self) -> TrainBatch:
         selected = list(self.pending_batch.items())[: self.batch_size]
 
         selected_by_trace = dict(selected)
@@ -354,11 +354,15 @@ def process_batch(self) -> TrainBatch:
             del self.pending_batch[trace_id]
 
         if not self.config.constant_trainer_batch_size:
-            selected_by_trace = {
-                trace_id: [sample for sample in samples if _prune_zero_advantages(sample)]
-                for trace_id, samples in selected
-            }
-            selected_by_trace = {trace_id: samples for trace_id, samples in selected_by_trace.items() if samples}
+
+            def prune_selected() -> dict[str, list[TrainingSample]]:
+                pruned = {
+                    trace_id: [sample for sample in samples if _prune_zero_advantages(sample)]
+                    for trace_id, samples in selected
+                }
+                return {trace_id: samples for trace_id, samples in pruned.items() if samples}
+
+            selected_by_trace = await asyncio.to_thread(prune_selected)
         samples = [sample for trace_samples in selected_by_trace.values() for sample in trace_samples]
 
         shipped_ids = set(selected_by_trace)

src/prime_rl/orchestrator/utils.py

@@ -1,6 +1,5 @@
 import asyncio
 import ctypes
-import gc
 import logging
 import math
 from concurrent.futures import ThreadPoolExecutor
@@ -85,9 +84,16 @@ def set_default_executor(max_workers: int = 64) -> None:
     asyncio.get_event_loop().set_default_executor(ThreadPoolExecutor(max_workers=max_workers))
 
 
+def release_and_trim(*containers: list) -> None:
+    """Release shipped lists incrementally, then trim process memory."""
+    for items in containers:
+        while items:
+            items.pop()
+    trim_process_memory()
+
+
 def trim_process_memory() -> None:
     """Return freed heap pages to the OS on glibc systems."""
-    gc.collect()
     try:
         ctypes.CDLL("libc.so.6").malloc_trim(0)
     except Exception as exc:

src/prime_rl/transports/batch/base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from io import BytesIO
 from pathlib import Path
 
 import msgspec
@@ -21,6 +22,24 @@ async def send(self, micro_batch_grid: list[list[MicroBatch]]) -> None:
         """Send grid of micro batches to the trainers."""
         pass
 
+    def encode(self, micro_batches: list[MicroBatch]) -> bytes:
+        """Byte-identical to ``self.encoder.encode(micro_batches)``, but encodes one
+        micro batch at a time so a worker thread releases the GIL between them."""
+        n = len(micro_batches)
+        if n <= 1:
+            return self.encoder.encode(micro_batches)
+        if n < 16:
+            header = bytes([0x90 | n])
+        elif n < 2**16:
+            header = b"\xdc" + n.to_bytes(2, "big")
+        else:
+            header = b"\xdd" + n.to_bytes(4, "big")
+        with BytesIO() as payload:
+            payload.write(header)
+            for micro_batch in micro_batches:
+                payload.write(self.encoder.encode(micro_batch))
+            return payload.getvalue()
+
     def close(self) -> None:
         """Clean up any resources. Override if needed."""
         pass

src/prime_rl/transports/batch/filesystem.py

@@ -28,7 +28,7 @@ async def send(self, micro_batch_grid: list[list[MicroBatch]]) -> None:
 
     def _encode_and_write(self, micro_batch_grid: list[list[MicroBatch]], step_path: Path) -> None:
         for data_rank in range(self.data_world_size):
-            buffer = self.encoder.encode(micro_batch_grid[data_rank])
+            buffer = self.encode(micro_batch_grid[data_rank])
             tmp_path = step_path / f"rank_{data_rank}.bin.tmp"
             with open(tmp_path, "wb") as f:
                 f.write(buffer)

src/prime_rl/transports/batch/zmq.py

@@ -1,3 +1,4 @@
+import asyncio
 from pathlib import Path
 
 import zmq
@@ -74,7 +75,7 @@ async def send(self, micro_batch_grid: list[list[MicroBatch]]) -> None:
 
         self.logger.debug(f"Sending micro batch grid for step {self._current_step}")
         for data_rank in range(self.data_world_size):
-            buffer = self.encoder.encode(micro_batch_grid[data_rank])
+            buffer = await asyncio.to_thread(self.encode, micro_batch_grid[data_rank])
             topic = self._topic_prefix + str(data_rank).encode("utf-8") + b"|"
             await self.socket.send_multipart([topic, buffer], copy=False)
         self._current_step += 1

tests/unit/orchestrator/test_train_sink.py

@@ -151,7 +151,7 @@ async def run():
         assert s.take_rollout_window(force=True).attempts == 1
         del old
         s.progress.step = 4
-        assert s.take_batch() is None
+        assert await s.take_batch() is None
         gc.collect()
         assert ref() is None
         window = s.take_rollout_window(force=True)
@@ -173,7 +173,7 @@ async def run():
         assert len(batch.cohort) == 1 and batch.cohort.episodes[0].id == ep_id
         assert sibling.id in s.pending_batch
         assert zero.nodes[2].token_ids == [3, 4]
-        batch = s.take_batch()
+        batch = await s.take_batch()
         assert len(batch.samples) == 2
         assert len(batch.cohort) == 1 and batch.cohort.num_traces == 2
         assert batch.cohort.episodes[0].id == ep_id

tests/unit/transports/test_batch_encode.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+
+import msgspec
+import pytest
+
+from prime_rl.transports.batch.base import BatchSender
+from prime_rl.transports.batch.types import MicroBatch, RoutedExperts
+
+
+@pytest.fixture
+def sender(tmp_path):
+    class Sender(BatchSender):
+        async def send(self, micro_batch_grid: list[list[MicroBatch]]) -> None:
+            raise NotImplementedError
+
+    return Sender(Path(tmp_path), data_world_size=1)
+
+
+@pytest.mark.parametrize("count", [0, 1, 2, 15, 16, 255, 65535, 65536])
+@pytest.mark.parametrize("with_routing", [False, True])
+def test_batch_encode_preserves_wire_bytes(sender, count, with_routing):
+    batch = MicroBatch(
+        input_ids=[1, 2],
+        loss_mask=[False, True],
+        advantages=[0.0, 1.0],
+        inference_logprobs=[-0.1, -0.2],
+        position_ids=[0, 1],
+        sequence_lengths=[2],
+        temperatures=[1.0, 1.0],
+        env_names=["travel"],
+        seq_lens=[2],
+        routed_experts=RoutedExperts(b"\x00\x01\x02\x03", [2, 1, 2], "uint8") if with_routing else None,
+    )
+    micro_batches = [batch] * count
+    encoded = sender.encode(micro_batches)
+
+    assert isinstance(encoded, bytes)
+    assert encoded == msgspec.msgpack.encode(micro_batches)
+    assert msgspec.msgpack.decode(encoded, type=list[MicroBatch]) == micro_batches