Publish complete weight revisions and drain accepted rollout work

engine/megatron_core_train_engine.py

@@ -104,6 +104,7 @@ def __init__(
         self.world_size = int(os.getenv("WORLD_SIZE", "1"))
         self.is_main_process = self.rank == 0
         self.current_version = 0
+        self.optimizer_revision = 0
 
         self.bridge = None
         self.provider = None
@@ -533,6 +534,7 @@ def grpo_update(
             update_successful, grad_norm, _ = self.optimizer.step()
             if not update_successful:
                 raise FloatingPointError("Megatron optimizer rejected the update")
+            self.optimizer_revision += 1
             merged = self._merge_stats(local_stats)
             if isinstance(grad_norm, torch.Tensor):
                 merged["grad_norm"] = float(grad_norm.detach().item())

infra/execution/async_runner.py

@@ -36,6 +36,8 @@ def __init__(
 
         self.exiting = threading.Event()
         self.paused = threading.Event()
+        self._draining = threading.Event()
+        self._admission_lock = threading.Lock()
 
 
         self.input_queue: queue.Queue = queue.Queue(maxsize=max_queue_size)
@@ -62,6 +64,7 @@ def initialize(self):
         """Initialize."""
         self.exiting.clear()
         self.paused.clear()
+        self._draining.clear()
         self._loop_ready = threading.Event()
 
         self.thread = threading.Thread(
@@ -183,11 +186,9 @@ async def _run_async_loop(
 
         try:
             while not exit_event.is_set():
-                # Pausing blocks new dispatcher submissions. Tasks already
-                # accepted by this runner must still be started and reaped so
-                # callers can drain the rollout pipeline before restarting
-                # serving processes for a weight update.
-                if not pause_event.is_set():
+                # Explicit drain starts accepted inputs. Ordinary pause keeps
+                # them available for cancel_queued().
+                if not pause_event.is_set() or self._draining.is_set():
                     self._drain_pending_inputs(running_tasks, input_queue)
 
 
@@ -316,33 +317,33 @@ def submit(
     ) -> int:
         """Submit."""
 
-        with self._thread_exception_lock:
-            if self._thread_exception is not None:
-                raise RuntimeError(f"AsyncTaskRunner is unusable: {self._thread_exception}")
-
-
-        with self._active_task_ids_lock:
-            if task_id in self._active_task_ids:
-                raise ValueError(f"duplicate task ID: {task_id}")
-            self._active_task_ids.add(task_id)
-
-
-        fn_kwargs = {k: v for k, v in kwargs.items() if k != 'task_id'}
-
-        task_input = {
-            "async_fn": async_fn,
-            "args": args,
-            "kwargs": fn_kwargs,
-            "task_id": task_id,
-            "create_time": time.monotonic_ns(),
-        }
+        with self._admission_lock:
+            if self.paused.is_set():
+                raise queue.Full("runner admission is paused")
+            with self._thread_exception_lock:
+                if self._thread_exception is not None:
+                    raise RuntimeError(f"AsyncTaskRunner is unusable: {self._thread_exception}")
 
-        try:
-            self.input_queue.put_nowait(task_input)
-        except queue.Full:
             with self._active_task_ids_lock:
-                self._active_task_ids.discard(task_id)
-            raise queue.Full("input queue is full")
+                if task_id in self._active_task_ids:
+                    raise ValueError(f"duplicate task ID: {task_id}")
+                self._active_task_ids.add(task_id)
+
+            fn_kwargs = {k: v for k, v in kwargs.items() if k != 'task_id'}
+            task_input = {
+                "async_fn": async_fn,
+                "args": args,
+                "kwargs": fn_kwargs,
+                "task_id": task_id,
+                "create_time": time.monotonic_ns(),
+            }
+
+            try:
+                self.input_queue.put_nowait(task_input)
+            except queue.Full:
+                with self._active_task_ids_lock:
+                    self._active_task_ids.discard(task_id)
+                raise queue.Full("input queue is full")
 
         self._signal_new_input()
         return task_id
@@ -387,11 +388,13 @@ def wait(
 
     def pause(self):
         """Pause."""
-        self.paused.set()
+        with self._admission_lock:
+            self.paused.set()
 
     def resume(self):
         """Resume."""
-        self.paused.clear()
+        with self._admission_lock:
+            self.paused.clear()
         self._signal_new_input()
 
     def cancel_queued(self, timeout: float = 2.0) -> list[int]:
@@ -426,22 +429,27 @@ def drain_input_queue() -> None:
     def wait_until_idle(self, timeout: float = 3600.0) -> None:
         """Wait until every task accepted before pause has completed."""
         deadline = time.monotonic() + timeout
-        while True:
-            with self._thread_exception_lock:
-                if self._thread_exception is not None:
-                    raise RuntimeError(
-                        f"AsyncTaskRunner error: {self._thread_exception}"
+        self._draining.set()
+        self._signal_new_input()
+        try:
+            while True:
+                with self._thread_exception_lock:
+                    if self._thread_exception is not None:
+                        raise RuntimeError(
+                            f"AsyncTaskRunner error: {self._thread_exception}"
+                        )
+                with self._active_task_ids_lock:
+                    active_count = len(self._active_task_ids)
+                if active_count == 0 and self.input_queue.empty():
+                    return
+                if time.monotonic() >= deadline:
+                    raise TimeoutError(
+                        "Timed out draining rollout tasks before weight reload: "
+                        f"active={active_count}, queued={self.input_queue.qsize()}"
                     )
-            with self._active_task_ids_lock:
-                active_count = len(self._active_task_ids)
-            if active_count == 0 and self.input_queue.empty():
-                return
-            if time.monotonic() >= deadline:
-                raise TimeoutError(
-                    "Timed out draining rollout tasks before weight reload: "
-                    f"active={active_count}, queued={self.input_queue.qsize()}"
-                )
-            time.sleep(0.05)
+                time.sleep(0.05)
+        finally:
+            self._draining.clear()
 
     def get_queue_sizes(self) -> tuple[int, int]:
         """Get queue sizes."""

infra/execution/batch_dispatcher.py

@@ -147,16 +147,17 @@ def _commit_loop(self) -> None:
 
 
                 with self._reconfigure_lock:
-                    # pause() can race the small window between dequeue and
-                    # submit. Recheck under the same lock used by
-                    # cancel_queued() so an unsent request never crosses the
-                    # weight-sync boundary.
-                    if self.runner.paused.is_set():
-                        self.staleness_manager.on_rollouts_cancelled(enqueued=1)
-                        with self._result_cv:
-                            self._active_task_ids.discard(task_input.task_id)
-                            self._result_cv.notify_all()
-                        continue
+                    # Keep the selected task owned by the pending queue until
+                    # commit, so cancel_queued() can still return its identity
+                    # if pause wins the race for this lock.
+                    with self._input_cv:
+                        if (
+                            not self._pending_inputs
+                            or self._pending_inputs[0] is not task_input
+                            or self.runner.paused.is_set()
+                        ):
+                            continue
+                        self._pending_inputs.popleft()
                     task_fn = self.task_factory(task_input)
 
                     try:
@@ -191,7 +192,7 @@ def _commit_loop(self) -> None:
                 break
 
     def _get_next_task_for_submission(self) -> TInput | None:
-        """Get next task for submission."""
+        """Wait for a candidate, retaining queue ownership until commit."""
         with self._input_cv:
             while not self._shutdown_event.is_set():
                 self._check_thread_exception()
@@ -202,7 +203,7 @@ def _get_next_task_for_submission(self) -> TInput | None:
                     and self.staleness_manager.get_capacity() > 0
                     and self._pending_inputs
                 ):
-                    return self._pending_inputs.popleft()
+                    return self._pending_inputs[0]
 
 
                 self._input_cv.wait(timeout=0.5)

infra/sync/collective_export.py

@@ -0,0 +1,98 @@
+"""Keep every training rank in the same weight-export iterator protocol."""
+from __future__ import annotations
+
+import torch.distributed as dist
+
+
+def _gather(value, group):
+    if not dist.is_available() or not dist.is_initialized():
+        return [value]
+    values = [None] * dist.get_world_size(group)
+    dist.all_gather_object(values, value, group=group)
+    return values
+
+
+class _ExportIterator:
+    """One control checkpoint per yielded tensor, before the next TP export.
+
+    An iterator exception aborts this publication on every rank. No rank may
+    advance to the next tensor after a peer failed. An exception *inside* a
+    tensor's distributed collective remains a process-group failure and is
+    subject to that group's timeout/abort; it cannot be repaired by draining.
+    """
+
+    def __init__(self, factory, group, prepare):
+        self.group, self.prepare, self.finished = group, prepare, False
+        self.iterator = None
+        error = ""
+        try:
+            self.iterator = iter(factory())
+        except BaseException as exc:
+            error = f"{type(exc).__name__}: {exc}"
+        errors = _gather(error, group)
+        if any(errors):
+            self.finished = True
+            raise RuntimeError(f"collective weight export initialization failed: {errors}")
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.finished:
+            raise StopIteration
+        item, status = None, {"done": False, "error": "", "name": None}
+        try:
+            item = next(self.iterator)
+            if self.prepare is not None:
+                item = self.prepare(item)
+            status["name"] = str(item[0])
+        except StopIteration:
+            status["done"] = True
+        except BaseException as exc:
+            status["error"] = f"{type(exc).__name__}: {exc}"
+        statuses = _gather(status, self.group)
+        errors = [state["error"] for state in statuses if state["error"]]
+        if errors:
+            self.finished = True
+            raise RuntimeError(f"collective weight export iterator aborted: {errors}")
+        if any(state["done"] for state in statuses):
+            self.finished = True
+            if not all(state["done"] for state in statuses):
+                raise RuntimeError("collective weight export ended at different tensor positions")
+            raise StopIteration
+        if len({state["name"] for state in statuses}) != 1:
+            self.finished = True
+            raise RuntimeError("collective weight export tensor identities differ across ranks")
+        return item
+
+
+def stream_export_with_consensus(factory, *, is_sender, send, group=None, prepare=None):
+    """Drain training collectives after a transport failure, then fail together.
+
+    ``send`` runs only on the sender and must consume the supplied iterator.
+    A transport error leaves that iterator alive; the sender drains the same
+    stream as its peers without sending any further weights. Initialization,
+    preparation, iteration and final transport status all have common control
+    positions, so an iterator error never silently becomes successful EOF.
+    """
+    weights = _ExportIterator(factory, group, prepare)
+    error = ""
+    try:
+        if is_sender:
+            send(weights)
+        else:
+            for _name, tensor in weights:
+                del tensor
+    except BaseException as exc:
+        error = f"{type(exc).__name__}: {exc}"
+    # Also handle a sender which returned without consuming the entire stream.
+    # Draining is safe after transport errors, and a failed iterator marks
+    # itself finished so it cannot enter another TP collective.
+    try:
+        for _name, tensor in weights:
+            del tensor
+    except BaseException as exc:
+        error += f"; export drain failed: {type(exc).__name__}: {exc}"
+    errors = _gather(error, group)
+    if any(errors):
+        raise RuntimeError(f"collective weight publication failed: {errors}")