[DEP-950][feat] add multi-rank sleep/wakeup support to MPI executor path

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -100,6 +100,13 @@ class PPCommTag(IntEnum):
     SAMPLE_STATE = 20003
 
 
+# Tags for the dedicated control communicator used by multi-rank sleep/wakeup.
+# Because _control_comm is a duplicated communicator isolated from all other
+# MPI traffic, small tag values are safe and do not conflict with PPCommTag.
+_CONTROL_ACTION_TAG = 0  # rank-0 -> non-rank-0: sleep/wakeup/shutdown msg
+_CONTROL_ACK_TAG = 1     # non-rank-0 -> rank-0: ACK after action completes
+
+
 @functools.cache
 def _load_iteration_indexes(env_var: str):
     spans = os.environ.get(env_var, None)
@@ -628,6 +635,16 @@ def on_detected():
         self.worker_started = False
         self.worker_lock = threading.Lock()
         self._broadcast_mpi_comm = None
+        # Secondary MPI communicator and listener thread for multi-rank
+        # sleep/wakeup control messages.  Both are None until start_worker()
+        # calls Dup() (a collective) on the main thread.
+        self._control_comm = None
+        self._control_listener_thread: Optional[threading.Thread] = None
+        # Serialises concurrent sleep/wakeup calls so that the control_action
+        # + _control_comm send/recv sequence is never interleaved.  Without
+        # this, two concurrent RPC calls could both pass control_request_barrier
+        # and race on the shared communicator, consuming the wrong ACKs.
+        self._control_action_lock = threading.Lock()
 
         self.kv_connector_manager = kv_connector_manager
 
@@ -813,6 +830,22 @@ def start_worker(self):
                         name="broadcast_sample_state_handler",
                     )
                     self.broadcast_sample_state_handler.start()
+                # Duplicate the communicator on the main thread for
+                # sleep/wakeup control messages.  MPI_Comm_dup is collective
+                # across all ranks and must run before any worker thread
+                # starts to avoid racing with the worker's MPI traffic.
+                if self.dist.world_size > 1:
+                    logger.info(
+                        "Create new MPI comm for sleep/wakeup control listener."
+                    )
+                    self._control_comm = mpi_comm().Dup()
+                    if self.dist.rank != 0:
+                        self._control_listener_thread = threading.Thread(
+                            target=self._control_listener_loop,
+                            daemon=True,
+                            name="control_listener_thread",
+                        )
+                        self._control_listener_thread.start()
                 self.worker_thread = threading.Thread(
                     target=self._event_loop_wrapper, daemon=True)
                 self.worker_thread.start()
@@ -914,6 +947,19 @@ def shutdown(self):
         if self.dist.pp_size > 1:
             self.executed_batch_queue.put(None)
             self.broadcast_sample_state_handler.join()
+        # Signal non-rank-0 control-listener threads to exit.  This runs after
+        # the worker thread has joined, which guarantees that the non-rank-0
+        # executor loops have already processed the shutdown broadcast and are
+        # no longer driving NCCL, so the control-comm send cannot deadlock.
+        if (self.dist.rank == 0
+                and self._control_comm is not None
+                and self.dist.world_size > 1):
+            for dest in range(1, self.dist.world_size):
+                self._control_comm.send(
+                    {"action": "shutdown", "tags": []},
+                    dest=dest,
+                    tag=_CONTROL_ACTION_TAG,
+                )
         self.worker_started = False
         # Release CUDA graphs before resource managers free their GPU memory.
         # Resource managers (e.g. SuffixAutomatonManager) allocate GPU workspace
@@ -2036,6 +2082,78 @@ def _broadcast_sample_state_loop(self):
             # executor), and MPI reclaims the communicators on MPI_Finalize.
             set_thread_local_mpi_comm(None)
 
+    def _control_listener_loop(self) -> None:
+        """Listener loop for sleep/wakeup control messages on non-rank-0 workers.
+
+        Runs in a dedicated daemon thread on every non-rank-0 rank.  Blocks on
+        ``_control_comm.recv`` waiting for a control message from rank-0, then
+        executes the requested VMM operation (``release_with_tag`` or
+        ``materialize_with_tag``) and sends an ACK back.  A ``"shutdown"``
+        message breaks the loop cleanly.
+
+        Safety notes:
+        - ``_control_comm`` is a dedicated duplicated MPI communicator, isolated
+          from all regular executor MPI traffic.
+        - ``torch.cuda.synchronize()`` is called before each VMM operation so
+          that in-flight CUDA kernels from the event-loop thread complete before
+          memory mappings are changed.  The event-loop thread is effectively
+          idle at this point (it is blocked in a collective waiting for rank-0,
+          whose event loop is paused via ``control_action()``).
+        - ``gc.collect()`` and ``torch.cuda.empty_cache()`` are safe to call
+          from a non-main thread.
+        """
+        import gc
+
+        from tensorrt_llm._torch.virtual_memory import (
+            materialize_with_tag, release_with_tag)
+        from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType
+
+        torch.cuda.set_device(self.device_id)
+        CUASSERT(cudart.cudaSetDevice(self.device_id))
+        set_thread_local_mpi_comm(self._control_comm)
+        try:
+            while True:
+                msg = self._control_comm.recv(
+                    source=0, tag=_CONTROL_ACTION_TAG)
+                action: str = msg["action"]
+                if action == "shutdown":
+                    break
+                tag_strings: list = msg["tags"]
+                tags = [ExecutorMemoryType(t) for t in tag_strings]
+                error_msg: Optional[str] = None
+                try:
+                    torch.cuda.synchronize()
+                    if action == "sleep":
+                        release_with_tag(*tags)
+                        torch.cuda.synchronize()
+                        gc.collect()
+                        torch.cuda.empty_cache()
+                    elif action == "wakeup":
+                        materialize_with_tag(*tags)
+                        torch.cuda.synchronize()
+                    else:
+                        error_msg = f"unknown action '{action}'"
+                        logger.warning(
+                            f"Control listener: {error_msg}, ignoring.")
+                except Exception as exc:
+                    error_msg = (
+                        f"rank {self.dist.rank} '{action}' failed: "
+                        f"{exc}\n{traceback.format_exc()}")
+                    logger.error(
+                        f"Control listener: error executing '{action}':",
+                        exc_info=True)
+                finally:
+                    # Always ACK so rank-0 does not deadlock; carry error
+                    # details so rank-0 can raise after all ranks respond.
+                    self._control_comm.send(
+                        {"status": "ok" if error_msg is None else "error",
+                         "error": error_msg},
+                        dest=0,
+                        tag=_CONTROL_ACK_TAG,
+                    )
+        finally:
+            set_thread_local_mpi_comm(None)
+
     def _ring_broadcast_sample_state(
         self,
         executed_batch: Optional[BatchStatePP],

tensorrt_llm/executor/base_worker.py

@@ -18,6 +18,7 @@
 import gc
 import json
 import os
+import traceback
 import weakref
 from pathlib import Path
 from queue import Queue
@@ -701,7 +702,6 @@ def _check_sleep_wakeup_preconditions(self, method: str) -> None:
         Raises:
             ValueError: If the backend is not ``"pytorch"`` or
                 ``sleep_config`` is not set.
-            NotImplementedError: If ``parallel_config.world_size > 1``.
         """
         # _autodeploy is intentionally excluded: its allocations are not tagged
         # under sleep_config VMM scopes, so release_with_tag would silently
@@ -715,20 +715,112 @@ def _check_sleep_wakeup_preconditions(self, method: str) -> None:
             raise ValueError(
                 "Sleep feature is not enabled, please set sleep_config in "
                 "the LLM arguments.")
-        # Non-rank-0 processes block on their local control_action_done
-        # threading.Event with no Python caller to release it — deadlock.
-        if self.llm_args.parallel_config.world_size > 1:
-            raise NotImplementedError(
-                f"{method}() requires parallel_config.world_size == 1; "
-                "use the Ray executor for multi-rank deployments.")
+
+    def _multi_rank_sleep_wakeup(
+        self,
+        action: str,
+        tags: List[ExecutorMemoryType],
+    ) -> None:
+        """Coordinate a sleep or wakeup operation across all MPI ranks.
+
+        Called on rank-0 only for ``world_size > 1`` deployments.
+
+        Sequence:
+        1. Enter ``control_action()`` to drain in-flight requests and pause
+           rank-0's event loop.  Non-rank-0 event loops become idle (starved
+           of NCCL collectives from rank-0) once the current iteration drains.
+        2. Broadcast the control message to every non-rank-0 rank via the
+           dedicated ``_control_comm`` communicator.
+        3. Execute the VMM operation (``release_with_tag`` or
+           ``materialize_with_tag``) locally on rank-0.
+        4. Collect ACKs from all non-rank-0 ranks to confirm they have
+           finished their local VMM operations.
+        5. Exit ``control_action()``, resuming rank-0's event loop.
+
+        Args:
+            action: ``"sleep"`` or ``"wakeup"``.
+            tags: Parsed :class:`~tensorrt_llm.llmapi.llm_args.ExecutorMemoryType`
+                values; forwarded verbatim to each rank's VMM call.
+        """
+        from tensorrt_llm._torch.pyexecutor.py_executor import (
+            _CONTROL_ACK_TAG, _CONTROL_ACTION_TAG)
+        from tensorrt_llm._torch.virtual_memory import (materialize_with_tag,
+                                                         release_with_tag)
+
+        assert self.rank == 0, (
+            "_multi_rank_sleep_wakeup must only be called on rank 0")
+
+        control_comm = self.engine._control_comm
+        assert control_comm is not None, (
+            "_control_comm not initialised; was start_worker() called?")
+
+        world_size = self.llm_args.parallel_config.world_size
+        tag_strings = [t.value for t in tags]
+        msg = {"action": action, "tags": tag_strings}
+
+        # Serialise concurrent control actions.  control_action() uses an
+        # Event-based barrier, not a mutex, so two concurrent callers can both
+        # pass the barrier and then interleave sends/recvs on _control_comm,
+        # consuming the wrong ACKs or resuming the event loop prematurely.
+        # _control_action_lock turns the whole sequence into a critical section.
+        with self.engine._control_action_lock, self.engine.control_action():
+            # Broadcast control message to non-rank-0 listeners.
+            for dest in range(1, world_size):
+                control_comm.send(msg, dest=dest, tag=_CONTROL_ACTION_TAG)
+
+            # Execute locally on rank-0; capture error so ACKs are drained
+            # before raising.
+            local_error = None
+            try:
+                torch.cuda.synchronize()
+                if action == "sleep":
+                    release_with_tag(*tags)
+                    torch.cuda.synchronize()
+                    gc.collect()
+                    torch.cuda.empty_cache()
+                else:
+                    materialize_with_tag(*tags)
+                    torch.cuda.synchronize()
+            except Exception as exc:
+                local_error = (
+                    f"rank 0 '{action}' failed: {exc}\n"
+                    f"{traceback.format_exc()}"
+                )
+                logger.error(
+                    f"_multi_rank_sleep_wakeup: rank-0 local {action} failed:",
+                    exc_info=True,
+                )
+
+            # Always drain all ACKs to keep the communicator clean even if
+            # rank-0's local op failed.
+            errors = []
+            if local_error:
+                errors.append(local_error)
+            for src in range(1, world_size):
+                ack = control_comm.recv(source=src, tag=_CONTROL_ACK_TAG)
+                if ack.get("status") != "ok":
+                    errors.append(
+                        ack.get("error")
+                        or f"rank {src} returned unknown ACK"
+                    )
+            if errors:
+                raise RuntimeError(
+                    f"{action}() failed on {len(errors)} rank(s):\n"
+                    + "\n".join(errors)
+                )
 
     def sleep(self, sleep_tags: List[str]) -> None:
         """Release GPU virtual memory for the specified memory type tags.
 
-        Single-rank (``world_size == 1``) only.  Uses
-        ``PyExecutor.control_action()`` to drain in-flight requests and pause
-        the event loop before calling ``release_with_tag()``, matching the
-        ``@control_action_decorator`` behaviour used in Ray.
+        Supports both single-rank (``world_size == 1``) and multi-rank
+        (TP/PP > 1) MPI executor deployments.  For multi-rank, a lightweight
+        control-listener thread on each non-rank-0 worker handles the local
+        ``release_with_tag()`` call while rank-0 coordinates via a dedicated
+        secondary MPI communicator.
+
+        Uses ``PyExecutor.control_action()`` to drain in-flight requests and
+        pause rank-0's event loop, matching the ``@control_action_decorator``
+        behaviour used in the Ray executor.
 
         Only allocations backed by virtual memory (VMM) and registered under
         the active :class:`~tensorrt_llm.llmapi.llm_args.SleepConfig` are
@@ -743,32 +835,36 @@ def sleep(self, sleep_tags: List[str]) -> None:
 
         Returns:
             None.  The call is synchronous; when it returns all requested
-            VMM-tagged allocations have been released and the event loop
-            has been resumed.
+            VMM-tagged allocations have been released on every rank and the
+            event loop has been resumed.
 
         Raises:
             ValueError: If the backend is not ``"pytorch"`` or
                 ``sleep_config`` is not set.
-            NotImplementedError: If ``parallel_config.world_size > 1``.
         """
         self._check_sleep_wakeup_preconditions("sleep")
 
         from tensorrt_llm._torch.virtual_memory import release_with_tag
 
         tags = [ExecutorMemoryType(tag) for tag in sleep_tags]
         logger.info(f"Sleep: {tags}")
-        with self.engine.control_action():
-            torch.cuda.synchronize()
-            release_with_tag(*tags)
-            torch.cuda.synchronize()
-            gc.collect()
-            torch.cuda.empty_cache()
+        if self.llm_args.parallel_config.world_size > 1:
+            self._multi_rank_sleep_wakeup("sleep", tags)
+        else:
+            with self.engine._control_action_lock, self.engine.control_action():
+                torch.cuda.synchronize()
+                release_with_tag(*tags)
+                torch.cuda.synchronize()
+                gc.collect()
+                torch.cuda.empty_cache()
 
     def wakeup(self, wakeup_tags: List[str]) -> None:
         """Materialize GPU virtual memory for the specified memory type tags.
 
-        Single-rank (``world_size == 1``) only.  See :meth:`sleep` for
-        details on VMM scope restrictions and backend prerequisites.
+        Supports both single-rank (``world_size == 1``) and multi-rank
+        (TP/PP > 1) MPI executor deployments.  See :meth:`sleep` for details
+        on VMM scope restrictions, backend prerequisites, and multi-rank
+        coordination.
 
         Args:
             wakeup_tags: List of
@@ -777,24 +873,26 @@ def wakeup(self, wakeup_tags: List[str]) -> None:
 
         Returns:
             None.  The call is synchronous; when it returns all requested
-            VMM-tagged allocations have been materialized and the event loop
-            has been resumed.
+            VMM-tagged allocations have been materialized on every rank and the
+            event loop has been resumed.
 
         Raises:
             ValueError: If the backend is not ``"pytorch"`` or
                 ``sleep_config`` is not set.
-            NotImplementedError: If ``parallel_config.world_size > 1``.
         """
         self._check_sleep_wakeup_preconditions("wakeup")
 
         from tensorrt_llm._torch.virtual_memory import materialize_with_tag
 
         tags = [ExecutorMemoryType(tag) for tag in wakeup_tags]
         logger.info(f"Wakeup: {tags}")
-        with self.engine.control_action():
-            torch.cuda.synchronize()
-            materialize_with_tag(*tags)
-            torch.cuda.synchronize()
+        if self.llm_args.parallel_config.world_size > 1:
+            self._multi_rank_sleep_wakeup("wakeup", tags)
+        else:
+            with self.engine._control_action_lock, self.engine.control_action():
+                torch.cuda.synchronize()
+                materialize_with_tag(*tags)
+                torch.cuda.synchronize()
 
     def shutdown(self):
         if self.doing_shutdown: