Class RemoteMpiCommSessionClient

tensorrt_llm/llmapi/mpi_session.py:299–403

RemoteMpiCommSessionClient is a variant of MpiCommSession that is used to connect to a remote MPI pool. Note: This class uses a global singleton pattern because ZeroMQ PAIR sockets only support one connection at a time. Multiple LLM instances will reuse the same client connection.

class RemoteMpiCommSessionClient(MpiSession):
    '''
    RemoteMpiCommSessionClient is a variant of MpiCommSession that is used to connect to a remote MPI pool.

    Note: This class uses a global singleton pattern because ZeroMQ PAIR sockets only support
    one connection at a time. Multiple LLM instances will reuse the same client connection.
    '''
    _global_instance = None
    _global_instance_lock = threading.Lock()

    def __new__(cls, addr: str, hmac_key: Optional[bytes] = None):
        # Implement singleton pattern to reuse the same client connection
        # for multiple LLM instances, since PAIR sockets only support one connection
        with cls._global_instance_lock:
            if cls._global_instance is None or cls._global_instance.addr != addr:
                logger_debug(
                    f"Creating new global RemoteMpiCommSessionClient for {addr}\n",
                    "yellow")
                instance = super().__new__(cls)
                cls._global_instance = instance
                instance._initialized = False
            else:
                logger_debug(
                    f"Reusing existing global RemoteMpiCommSessionClient for {addr}\n",
                    "yellow")
            return cls._global_instance

    def __init__(self, addr: str, hmac_key: Optional[bytes] = None):
        # Only initialize once
        if self._initialized:
            return

        # FIXME: this is a hack to avoid circular import, resolve later
        from tensorrt_llm.executor.ipc import ZeroMqQueue
        self.addr = addr
        logger_debug(f"RemoteMpiCommSessionClient connecting to {addr}\n",
                     "yellow")
        self.queue = ZeroMqQueue((addr, hmac_key),
                                 is_server=False,
                                 socket_type=zmq.PAIR,
                                 use_hmac_encryption=bool(hmac_key))
        self._is_shutdown = False
        self._initialized = True

    def submit(self,
               task: Callable[..., T],
               *args,
               sync: bool = False,
               **kwargs) -> list:
        ''' Submit a task to the remote MPI pool. '''
        if self._is_shutdown:
            logger_debug("RemoteMpiCommSessionClient is already shut down\n",
                         "yellow")
            return []
        logger_debug(
            f"RemoteMpiCommSessionClient [rank{global_mpi_rank()}] sending task {task} to {self.addr}\n",
            "yellow")
        self.queue.put(RemoteTask(task, args, kwargs, sync=sync))