[None][feat] Support externally provided MPI sessions with explicit ownership (NVIDIA#16053)

Signed-off-by: qgai <qgai@nvidia.com>

Author: sunnyqgg

tensorrt_llm/_torch/distributed/communicator.py

@@ -1240,7 +1240,28 @@ def init_pp_comm(mapping):
     global _pp_comm
     if mpi_disabled():
         _pp_comm = PPCommTorch(mapping)
+    elif isinstance(_pp_comm, PPCommNCCL) and \
+            _pp_comm.mapping.world_size == mapping.world_size:
+        # Reuse the existing world NCCL communicator across LLM instances that
+        # share the same worker processes (e.g. a reused MpiPoolSession). The
+        # underlying comm depends only on (world_size, rank) -- it is a world
+        # communicator, independent of the pp/tp/ep layout -- so only the
+        # routing mapping needs refreshing. Recreating it would drop the old
+        # comm and trigger a collective ncclCommDestroy at an unsynchronized
+        # point during the next model build, which can deadlock on reused
+        # workers. Single-LLM (production) runs are unaffected: _pp_comm starts
+        # as None, so the first call still constructs a fresh PPCommNCCL.
+        _pp_comm.mapping = mapping
     else:
+        if _pp_comm is not None:
+            # Rebinding drops the old comm; its ncclCommDestroy runs at an
+            # unsynchronized point and can deadlock on reused worker processes
+            # (see the reuse branch above). Surface it instead of hanging
+            # silently -- pools sharing workers must keep one world_size.
+            logger.warning(
+                "init_pp_comm: replacing existing PP comm (world_size "
+                f"{_pp_comm.mapping.world_size} -> {mapping.world_size}) on a "
+                "live process; this can deadlock on reused MPI workers.")
         _pp_comm = PPCommNCCL(mapping)
     init_helix_cp_comm(mapping)
 

tensorrt_llm/executor/executor.py

@@ -652,7 +652,7 @@ def create(
             return GenerationExecutor._create_ipc_executor(
                 worker_kwargs,
                 model_world_size=model_world_size,
-                mpi_session=None,  # use mpi4py
+                mpi_session=mpi_session,
                 postproc_worker_config=postproc_worker_config,
                 is_llm_executor=is_llm_executor,
                 use_worker=False)
@@ -662,13 +662,17 @@ def create(
             mpi_session = ProcessPoolExecutorSession(n_workers=1,
                                                      mp_context=ctx)
             # TODO: add rpc worker here
-            return GenerationExecutor._create_ipc_executor(
+            executor = GenerationExecutor._create_ipc_executor(
                 worker_kwargs,
                 model_world_size=model_world_size,
                 mpi_session=mpi_session,
                 postproc_worker_config=postproc_worker_config,
                 is_llm_executor=is_llm_executor,
                 use_worker=False)
+            # The session was created right here with no outer owner, so the
+            # proxy must shut it down despite it arriving as "external".
+            executor._owns_mpi_session = True
+            return executor
 
     def wait_first_completed(
         self, futures: List[GenerationResult]

tensorrt_llm/executor/proxy.py

@@ -29,7 +29,8 @@
 
 from .._utils import customized_gc_thresholds, mpi_rank, nvtx_range_debug
 from ..llmapi.mpi_session import (MpiCommSession, MpiPoolSession, MpiSession,
-                                  RemoteMpiCommSessionClient)
+                                  RemoteMpiCommSessionClient,
+                                  validate_session_world_size)
 from ..llmapi.tracer import enable_llm_tracer, get_tracer, global_tracer
 from ..llmapi.utils import (AsyncQueue, ManagedThread, _SyncQueue,
                             enable_llm_debug, logger_debug, print_colored)
@@ -117,6 +118,7 @@ def __init__(
         self.worker_cls = worker_cls
 
         mpi_process_pre_spawned: bool = get_spawn_proxy_process_env()
+        self._owns_mpi_session = mpi_session is None
 
         if mpi_session is None:
             if mpi_process_pre_spawned:
@@ -126,6 +128,10 @@ def __init__(
                 logger_debug('create pool session ...\n', "yellow")
                 self.mpi_session = MpiPoolSession(n_workers=model_world_size)
         else:
+            # submit() launches one worker task per pool worker, so an
+            # external session must match the model's world size exactly;
+            # fail loudly instead of starting the wrong number of executors.
+            validate_session_world_size(mpi_session, model_world_size)
             logger_debug('using external mpi session ...\n', "yellow")
             self.mpi_session = mpi_session
 
@@ -406,11 +412,11 @@ def resource_governor_queue(self):
         return self._resource_governor_queue
 
     def abort_request(self, request_id: int) -> None:
-        ''' Abort a request by sending a cancelling request to the request queue.
+        """Abort a request by sending a cancelling request to the request queue.
 
         Args:
             request_id (int): The id of the request to abort.
-        '''
+        """
         # NOTE, it just sends a cancelling request to the request queue, but it
         # may take a while for the request to be cancelled in the worker and
         # send back a finished result.
@@ -543,7 +549,10 @@ def mpi_done_callback(future: concurrent.futures.Future):
 
         if ready_signal != GenerationExecutorProxy.READY_SIGNAL:
             logger.error(f"Executor worker initialization error: {error_trace}")
-            self.mpi_session.shutdown_abort(reason=ready_signal)
+            # Only abort a session this proxy created; an externally owned
+            # (shared) session must stay alive for its owner to tear down.
+            if self._owns_mpi_session:
+                self.mpi_session.shutdown_abort(reason=ready_signal)
             raise RuntimeError(
                 "Executor worker returned error") from ready_signal
 
@@ -630,7 +639,8 @@ def shutdown(self):
             self._resource_governor_queue.close()
 
         self.workers_started = False
-        self.mpi_session.shutdown()
+        if self._owns_mpi_session:
+            self.mpi_session.shutdown()
 
         # Process the errors in-case error during shutting down the threads
         self._handle_background_error()

tensorrt_llm/executor/rpc_proxy.py

@@ -16,7 +16,8 @@
 import threading
 from typing import List, Optional, Union
 
-from ..llmapi.mpi_session import MpiPoolSession, MpiSession
+from ..llmapi.mpi_session import (MpiPoolSession, MpiSession,
+                                  validate_session_world_size)
 from ..llmapi.utils import logger_debug, print_colored
 from ..logger import logger
 from .executor import GenerationExecutor
@@ -42,13 +43,12 @@ def __init__(
         postproc_worker_config: Optional[PostprocWorkerConfig] = None,
         is_llm_executor: Optional[bool] = None,
     ):
-        """
-        Args:
-            worker_kwargs: kwargs for the rpc worker
-            model_world_size: the world size of the model
-            mpi_session: the mpi session to use
-            postproc_worker_config: the postproc worker config
-            is_llm_executor: whether this is an llm executor
+        """Args:
+        worker_kwargs: kwargs for the rpc worker
+        model_world_size: the world size of the model
+        mpi_session: the mpi session to use
+        postproc_worker_config: the postproc worker config
+        is_llm_executor: whether this is an llm executor
         """
         GenerationExecutorRpcProxy.INSTANCE_COUNTER += 1
         self.init_rpc_executor()
@@ -81,11 +81,13 @@ def __init__(
         self._setup_mainloop_with_tasks()
 
     def launch_workers(self):
-        logger.debug(f"Launching workers")
+        logger.debug("Launching workers")
         assert self.mpi_session is not None
-        self.mpi_session.submit(RpcWorker.main_task,
-                                rpc_addr=self.rpc_addr,
-                                **self.worker_kwargs)
+        # Keep the futures: on an externally owned (shared) session, shutdown
+        # waits on them so worker teardown finishes before the pool is reused.
+        self.worker_futures = self.mpi_session.submit(RpcWorker.main_task,
+                                                      rpc_addr=self.rpc_addr,
+                                                      **self.worker_kwargs)
 
     def _setup_mainloop_with_tasks(self):
         """Setup mainloop with tasks needed for RpcProxy.
@@ -256,7 +258,7 @@ def setup_engine_remote(self):
         return self.rpc_client.setup_engine().remote(need_response=True)
 
     def shutdown_remote(self):
-        logger_debug(f"Shutting down rpc remote", color="yellow")
+        logger_debug("Shutting down rpc remote", color="yellow")
         self.rpc_client.shutdown().remote(need_response=False)
 
     def abort_request(self, request_id: int) -> None:
@@ -266,8 +268,7 @@ def shutdown(self):
         if self._shutdown_event.is_set():
             return
         self._shutdown_event.set()
-        logger_debug(f"Shutting down GenerationExecutorRpcProxy",
-                     color="yellow")
+        logger_debug("Shutting down GenerationExecutorRpcProxy", color="yellow")
 
         # 1. shutdown the rpc server (PyExecutor Rank 0 + RPC server)
         self.shutdown_remote()
@@ -294,9 +295,24 @@ def shutdown(self):
         # 3. shutdown the mpi session, this should wait until all the PyExecutor
         # processes are shutdown
         if self.mpi_session is not None:
-            logger_debug(f"Shutting down mpi session", color="yellow")
-            self.mpi_session.shutdown()
-            logger_debug(f"Mpi session shutdown", color="yellow")
+            if self._owns_mpi_session:
+                logger_debug("Shutting down mpi session", color="yellow")
+                self.mpi_session.shutdown()
+            else:
+                # Externally owned (shared) session: leave the pool alive, but
+                # wait for this executor's worker tasks to finish so the next
+                # LLM on the pool doesn't race with PyExecutor teardown. Block
+                # without a timeout, mirroring the owned path above
+                # (mpi_session.shutdown() also waits indefinitely); a timeout
+                # that expires would just reintroduce the race it prevents.
+                for future in getattr(self, "worker_futures", []):
+                    try:
+                        future.result()
+                    except Exception as e:
+                        logger.warning(
+                            f"RPC worker task raised during shutdown on "
+                            f"shared MPI session: {e}")
+            logger_debug("Mpi session shutdown", color="yellow")
             self.mpi_session = None
 
         self.rpc_client.close()
@@ -309,14 +325,20 @@ def __exit__(self, exc_type, exc_value, traceback):
 
     def _create_mpi_session(self, model_world_size: int,
                             mpi_session: Optional[MpiSession]):
+        # Ownership is decided here, next to the create-vs-adopt branch: a
+        # session this proxy created is shut down by it; an external one is
+        # owned (and shut down) by the caller.
         mpi_process_pre_spawned: bool = get_spawn_proxy_process_env()
         if mpi_session is None:
+            self._owns_mpi_session = True
             if mpi_process_pre_spawned:
                 logger_debug('[proxy] create comm session ...\n', "yellow")
                 self.mpi_session = create_mpi_comm_session(model_world_size)
             else:
                 logger_debug('[proxy] create pool session ...\n', "yellow")
                 self.mpi_session = MpiPoolSession(n_workers=model_world_size)
         else:
+            validate_session_world_size(mpi_session, model_world_size)
+            self._owns_mpi_session = False
             logger_debug('[proxy] using external mpi session ...\n', "yellow")
             self.mpi_session = mpi_session

tensorrt_llm/llmapi/llm.py

@@ -314,6 +314,10 @@ def __init__(self,
         logger_debug(f"LLM.args.mpi_session: {self.args.mpi_session}\n",
                      "yellow")
         self.mpi_session = self.args.mpi_session
+        # Keep the live session on LLM only. LLM args are passed to model-build
+        # tasks and executor workers, and MpiSession objects are not pickleable.
+        self.args.mpi_session = None
+        self._owns_mpi_session = self.mpi_session is None
 
         # Build this LLM's post-processing hook for the in-proxy detok path (each
         # postproc worker builds its own). Resolving here fails fast on a bad
@@ -333,6 +337,8 @@ def __init__(self,
             logger.info(
                 f'start MpiSession with {self.args.parallel_config.world_size} workers'
             )
+            # _owns_mpi_session is already True here: this branch only runs
+            # when no external session was supplied.
             if not self.mpi_session:
                 mpi_process_pre_spawned: bool = get_spawn_proxy_process_env()
                 if not mpi_process_pre_spawned:
@@ -365,7 +371,9 @@ def __init__(self,
             self._build_model()
 
         except Exception:
-            if self.mpi_session is not None:
+            # _owns_mpi_session is assigned before this try block, so it is
+            # always present here.
+            if self.mpi_session is not None and self._owns_mpi_session:
                 self.mpi_session.shutdown()
             raise
 
@@ -1500,7 +1508,8 @@ def shutdown(self) -> None:
             self._encoder_executor.shutdown()
             self._encoder_executor = None
 
-        if hasattr(self, 'mpi_session') and self.mpi_session is not None:
+        if (hasattr(self, 'mpi_session') and self.mpi_session is not None
+                and getattr(self, "_owns_mpi_session", True)):
             self.mpi_session.shutdown()
             self.mpi_session = None