fix(rollout): hold payload egress for a whole weight sync

cosmos_rl/rollout/worker/rollout_control.py

@@ -81,6 +81,7 @@
     AsyncR2RSyncMode,
     get_async_r2r_sync_mode,
     get_broadcast_all_params,
+    holds_payload_egress,
     ensure_wst,
     sync_buffer_to_live,
     process_wst_deferred_actions,
@@ -1212,6 +1213,7 @@ def lazy_initialize_rollout_engine(self, load_format):
                 ensure_wst(self)
 
     @RolloutWorkerBase.register_rollout_command_handler(PolicyToRolloutUnicastCommand)
+    @holds_payload_egress
     @torch.no_grad()
     def policy_to_rollout_unicast(self, command: PolicyToRolloutUnicastCommand):
         """Sync the weight from policy to rollout.
@@ -1415,6 +1417,7 @@ def flush_completions(pending_bytes, pending_completions):
     @RolloutWorkerBase.register_rollout_command_handler(
         RolloutToRolloutBroadcastCommand
     )
+    @holds_payload_egress
     def broadcast_to_all_rollout_replica(
         self, broadcast_command: RolloutToRolloutBroadcastCommand
     ) -> None:
@@ -1432,13 +1435,6 @@ def broadcast_to_all_rollout_replica(
         src_replica_name: str = broadcast_command.src_replica_name
         dst_replica_names: List[str] = broadcast_command.dst_replica_names
 
-        # Forward-compat: flush any pending async NCCL sends (e.g. from data
-        # packers) so they complete before weight sync reuses the communicator.
-        if hasattr(self, "data_packer") and hasattr(
-            self.data_packer, "flush_pending_sends"
-        ):
-            self.data_packer.flush_pending_sends()
-
         # lazy initialization of the rollout engine.
         if self.replica_name != src_replica_name:
             # for replicas that needs to be broadcasted, use dummy format.

cosmos_rl/rollout/worker/weight_sync.py

@@ -51,12 +51,15 @@
 
 from __future__ import annotations
 
+import functools
 import os
 import queue
 import threading
 import time
+from collections.abc import Iterator
+from contextlib import contextmanager
 from enum import Enum
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 import torch
 from torch.distributed.tensor import DTensor
@@ -125,6 +128,53 @@ def get_broadcast_all_params(worker) -> bool:
     return worker.config.rollout.broadcast_all_params
 
 
+# ---------------------------------------------------------------------------
+# Payload egress
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def payload_egress_held(worker) -> Iterator[None]:
+    """Keep a data packer's payload egress off this device's NCCL for a sync.
+
+    NCCL gives no guarantee for two communicators at once on one device, so a
+    packer that ships payloads over NCCL must have no send in flight while
+    weight sync uses the device.  ``flush_pending_sends`` drains what is
+    already in flight, but nothing stops the packer claiming the next payload
+    the moment it returns, and a sync spends most of its wall time between that
+    drain and its transfer, waiting on a barrier for its peers.  A packer that
+    implements ``hold_sends`` gets the guarantee for the whole sync instead:
+    egress resumes when it is over.
+
+    Both are optional, so a packer that ships nothing over this device's NCCL
+    needs neither.
+
+    Yields:
+        Control, with payload egress held for as long as the block runs.
+    """
+    packer = getattr(worker, "data_packer", None)
+    hold_sends = getattr(packer, "hold_sends", None)
+    if hold_sends is not None:
+        with hold_sends():
+            yield
+        return
+    flush_pending_sends = getattr(packer, "flush_pending_sends", None)
+    if flush_pending_sends is not None:
+        flush_pending_sends()
+    yield
+
+
+def holds_payload_egress(handler: Callable) -> Callable:
+    """Run a weight-sync command handler with payload egress held."""
+
+    @functools.wraps(handler)
+    def handler_with_egress_held(self, command: Any, *args, **kwargs):
+        with payload_egress_held(self):
+            return handler(self, command, *args, **kwargs)
+
+    return handler_with_egress_held
+
+
 # ---------------------------------------------------------------------------
 # Buffer model helpers
 # ---------------------------------------------------------------------------
@@ -558,10 +608,11 @@ def _run(self) -> None:
                 self._idle.set()
                 continue
             try:
-                if cmd_type == "p2r":
-                    self._execute_p2r(command)
-                elif cmd_type == "r2r":
-                    self._execute_r2r(command)
+                with payload_egress_held(self._worker):
+                    if cmd_type == "p2r":
+                        self._execute_p2r(command)
+                    elif cmd_type == "r2r":
+                        self._execute_r2r(command)
             except Exception:
                 self._task_failed = True
                 logger.exception(
@@ -599,16 +650,11 @@ def _execute_r2r(self, command) -> None:
         When commands are routed directly from the background command
         thread (bypassing the main-thread handler), the WST is
         responsible for bookkeeping that would normally be done in the
-        handler: ``flush_pending_sends``, ``set_weight_synced``.
+        handler: ``set_weight_synced``.  Payload egress is held for it by
+        the run loop, around this whole command.
         """
         worker = self._worker
 
-        # Flush any pending async NCCL sends before reusing the communicator.
-        if hasattr(worker, "data_packer") and hasattr(
-            worker.data_packer, "flush_pending_sends"
-        ):
-            worker.data_packer.flush_pending_sends()
-
         weight_step = command.weight_step
         # Use the controller's authoritative recipient set for this round as the
         # barrier participant count so it stays in lockstep as replicas finish.

tests/run_test.sh

@@ -151,7 +151,7 @@ run python tests/test_put_rollouts.py
 run python tests/test_trajectory_iteration.py
 run python tests/test_gym_example.py
 # Pytest-style CPU suites; install pytest in case the image lacks it.
-run /bin/bash -c "python -m pip install --quiet pytest && python -m pytest -q tests/test_weight_sync.py tests/test_checkpoint.py tests/test_ranked_rollout_end_and_wst_fence.py tests/test_terminal_checkpoint_trainer_hooks.py tests/test_terminal_drain_protocol.py tests/test_training_complete_checkpoint.py"
+run /bin/bash -c "python -m pip install --quiet pytest && python -m pytest -q tests/test_weight_sync.py tests/test_weight_sync_payload_egress.py tests/test_checkpoint.py tests/test_ranked_rollout_end_and_wst_fence.py tests/test_terminal_checkpoint_trainer_hooks.py tests/test_terminal_drain_protocol.py tests/test_training_complete_checkpoint.py"
 run python -m unittest -v tests.contracts.test_trainer_metrics_contract
 run python -m unittest -v tests.contracts.test_config_routing_contract
 run python -m unittest -v tests.contracts.test_model_registry_contract

tests/test_weight_sync_payload_egress.py

@@ -0,0 +1,174 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests for holding payload egress across a weight sync (CPU-only).
+
+The packers here record when they are held and released; the transport calls
+are stubbed, so the ordering of those records against the broadcast is what is
+being checked.
+"""
+
+import queue
+import threading
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from cosmos_rl.dispatcher.command import (
+    PolicyToRolloutUnicastCommand,
+    RolloutToRolloutBroadcastCommand,
+)
+from cosmos_rl.rollout.worker import weight_sync
+from cosmos_rl.rollout.worker.rollout_control import DisaggregatedRolloutControlWorker
+from cosmos_rl.rollout.worker.weight_sync import (
+    AsyncR2RSyncMode,
+    WeightSyncThread,
+    payload_egress_held,
+)
+
+
+class _HoldingPacker:
+    """A packer that can keep its egress off the device for a whole sync."""
+
+    def __init__(self, events):
+        self._events = events
+
+    @contextmanager
+    def hold_sends(self):
+        self._events.append("hold")
+        try:
+            yield
+        finally:
+            self._events.append("release")
+
+
+class _FlushingPacker:
+    """A packer from before ``hold_sends``, which can only drain on demand."""
+
+    def __init__(self, events):
+        self._events = events
+
+    def flush_pending_sends(self):
+        self._events.append("flush")
+
+
+def _make_worker(events, packer):
+    worker = object.__new__(DisaggregatedRolloutControlWorker)
+    worker.data_packer = packer
+    worker.replica_name = "rollout-0"
+    worker.rank_in_rollout_repicas = 0
+    worker.replica_name_to_rank = {"rollout-0": 0, "rollout-1": 1}
+    worker.global_commnicator_idex = 7
+    worker.inference_stream = None
+    worker.weight_mapper = None
+    worker.trainable_params = {"a.weight"}
+    worker.non_trainable_params_received = True
+    worker.current_weight_version = 3
+    worker.prepare_trainable_params = lambda: None
+    worker.rollout = SimpleNamespace(
+        model_param_map=lambda _mapper: {"a.weight": torch.zeros(2)}
+    )
+    worker.state = SimpleNamespace(
+        weight_synced=lambda: True, set_weight_synced=lambda: None
+    )
+    worker.config = SimpleNamespace(
+        validation=SimpleNamespace(enable=False, val_before_train=False, freq=1)
+    )
+    worker.lazy_initialize_rollout_engine = lambda _load_format: events.append(
+        "lazy init"
+    )
+    return worker
+
+
+def _broadcast(worker, events):
+    command = RolloutToRolloutBroadcastCommand(
+        src_replica_name="rollout-0",
+        dst_replica_names=["rollout-0", "rollout-1"],
+        weight_step=4,
+        total_steps=10,
+        trainable_only=True,
+    )
+    with (
+        patch(
+            "cosmos_rl.rollout.worker.rollout_control.nccl_broadcast",
+            lambda *_args: events.append("broadcast"),
+        ),
+        patch(
+            "cosmos_rl.rollout.worker.rollout_control.get_async_r2r_sync_mode",
+            lambda _worker: AsyncR2RSyncMode.DISABLED,
+        ),
+        patch(
+            "cosmos_rl.rollout.worker.rollout_control.get_broadcast_all_params",
+            lambda _worker: False,
+        ),
+    ):
+        worker.broadcast_to_all_rollout_replica(command)
+
+
+def test_r2r_holds_egress_for_the_whole_broadcast():
+    events = []
+    worker = _make_worker(events, _HoldingPacker(events))
+
+    _broadcast(worker, events)
+
+    assert events == ["hold", "broadcast", "release"]
+
+
+def test_r2r_still_drains_a_packer_that_cannot_hold():
+    events = []
+    worker = _make_worker(events, _FlushingPacker(events))
+
+    _broadcast(worker, events)
+
+    assert events == ["flush", "broadcast"]
+
+
+def test_p2r_holds_egress_too():
+    events = []
+    worker = _make_worker(events, _HoldingPacker(events))
+    # Addressed to a peer, so the handler initializes the engine and returns
+    # without a transfer of its own.
+    command = PolicyToRolloutUnicastCommand(
+        src_replica_name="policy-0",
+        dst_replica_name="rollout-1",
+        src_replica_size=1,
+        dst_replica_size=1,
+    )
+
+    worker.policy_to_rollout_unicast(command)
+
+    assert events == ["hold", "lazy init", "release"]
+
+
+def test_the_weight_sync_thread_holds_egress_around_a_command():
+    events = []
+    wst = object.__new__(WeightSyncThread)
+    wst._queue = queue.PriorityQueue()
+    wst._stop = threading.Event()
+    wst._idle = threading.Event()
+    wst._worker = SimpleNamespace(device="cpu", data_packer=_HoldingPacker(events))
+    wst._queue.put((0, 0, ("r2r", object())))
+
+    def execute_r2r(_command):
+        events.append("r2r")
+        wst._stop.set()
+
+    wst._execute_r2r = execute_r2r
+    with patch.object(weight_sync.torch.cuda, "set_device"):
+        wst._run()
+
+    assert events == ["hold", "r2r", "release"]
+
+
+def test_a_packer_that_ships_nothing_over_nccl_needs_neither_hook():
+    worker = SimpleNamespace(data_packer=SimpleNamespace())
+
+    with payload_egress_held(worker):
+        pass
+
+
+def test_a_worker_without_a_packer_is_left_alone():
+    with payload_egress_held(SimpleNamespace()):
+        pass