refactor RPC param

Superjomn · Superjomn · commit 34aa759fc585 · 2025-09-03T09:28:42.000+08:00
Signed-off-by: Superjomn &lt;328693+Superjomn@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/executor/proxy_base.py b/tensorrt_llm/executor/proxy_base.py
@@ -0,0 +1,84 @@
+import time
+from typing import Optional, Union
+
+from tensorrt_llm.logger import logger
+
+from ..llmapi.utils import AsyncQueue, _SyncQueue
+from .executor import GenerationExecutor
+from .ipc import FusedIpcQueue
+from .result import IterationResult
+from .utils import IntraProcessQueue
+
+
+class ProxyBase(GenerationExecutor):
+
+    def __init__(self,
+                 num_postprocess_workers: int = 0,
+                 postprocess_tokenizer_dir: Optional[str] = None,
+                 is_llm_executor: Optional[bool] = None):
+        super().__init__(num_postprocess_workers, postprocess_tokenizer_dir,
+                         is_llm_executor)
+
+    def _maybe_initialize_iteration_results(self):
+        if self._is_llm_executor:
+            if self._iter_stats_result is None:
+                # singleton to store cpp runtime stats
+                self._iter_stats_result = IterationResult()
+            else:
+                # expect more engine stats whenever new prompts are submitted
+                self._iter_stats_result.mark_undone()
+
+            if self._iter_kv_events_result is None:
+                self._iter_kv_events_result = IterationResult()
+            else:
+                self._iter_kv_events_result.mark_undone()
+
+    def _iteration_result_task(self, queue: Union[FusedIpcQueue,
+                                                  IntraProcessQueue],
+                               result_singleton: IterationResult) -> bool:
+        # iteration result is not urgent, so we can sleep a bit
+        time.sleep(0.2)
+
+        try:
+            data = queue.get()
+        except:
+            logger.debug(
+                "proxy.py: Error in _iteration_result_task: queue.get()")
+            return False
+
+        if data is None:
+            logger.debug("proxy.py: _iteration_result_task: data is None")
+            return False  # shutdown the thread
+
+        data = data if isinstance(data, list) else [data]
+        queue = result_singleton.queue
+        async_queues = []
+
+        while queue.full():
+            queue.get()
+
+        try:
+            for d in data:
+                if d is None:
+                    logger.debug("proxy.py: _iteration_result_task: d is None")
+                    return False
+
+                if isinstance(queue, _SyncQueue):
+                    queue.put_nowait(d)
+                    async_queues.append(queue)
+                else:
+                    queue.put(d)
+
+            if async_queues:
+                _SyncQueue.notify_many(queue.loop, async_queues)
+
+        except AsyncQueue.EventLoopShutdownError:
+            # This happens in the last loop while the generate workflow is
+            # stopped, or when get_stats() or aget_stats() are not called by users
+            # and therefore event loop can already be closed.
+            logger.debug("proxy.py: EventLoopShutdownError")
+        except Exception as e:
+            logger.debug(f"proxy.py: Error in _iteration_result_task: {e}")
+            raise e
+
+        return True  # success
diff --git a/tensorrt_llm/executor/rpc/__init__.py b/tensorrt_llm/executor/rpc/__init__.py
@@ -1,9 +1,10 @@
 from .rpc_client import RPCClient
-from .rpc_common import (RPCCancelled, RPCError, RPCRequest, RPCResponse,
-                         RPCStreamingError, RPCTimeout)
+from .rpc_common import (RPCCancelled, RPCError, RPCParams, RPCRequest,
+                         RPCResponse, RPCStreamingError, RPCTimeout)
 from .rpc_server import RPCServer, Server
 
 __all__ = [
     "RPCClient", "RPCServer", "Server", "RPCError", "RPCTimeout",
-    "RPCCancelled", "RPCStreamingError", "RPCRequest", "RPCResponse"
+    "RPCCancelled", "RPCStreamingError", "RPCRequest", "RPCResponse",
+    "RPCParams"
 ]
diff --git a/tensorrt_llm/executor/rpc/rpc_client.py b/tensorrt_llm/executor/rpc/rpc_client.py
@@ -6,7 +6,7 @@
 
 from ...logger import logger
 from ..ipc import ZeroMqQueue
-from .rpc_common import (RPCCancelled, RPCRequest, RPCResponse,
+from .rpc_common import (RPCCancelled, RPCParams, RPCRequest, RPCResponse,
                          RPCStreamingError, RPCTimeout)
 
 
@@ -164,33 +164,32 @@ def _start_response_reader_lazily(self):
             # Store the concurrent.futures.Future
             self._reader_task = future
 
-    async def _call_async(self, __rpc_method_name, *args, **kwargs):
+    async def _call_async(self, method_name, *args, **kwargs):
         """Async version of RPC call.
         Args:
-            __rpc_method_name: Method name to call
+            method_name: Method name to call
             *args: Positional arguments
             **kwargs: Keyword arguments
-            __rpc_timeout: The timeout (seconds) for the RPC call.
-            __rpc_need_response: Whether the RPC call needs a response.
-                If set to False, the remote call will return immediately.
+            __rpc_params: RPCParams object containing RPC parameters.
 
         Returns:
             The result of the remote method call
         """
         logger.debug(
-            f"RPC client calling method: {__rpc_method_name} with args: {args} and kwargs: {kwargs}"
+            f"RPC client calling method: {method_name} with args: {args} and kwargs: {kwargs}"
         )
         if self._server_stopped:
             raise RPCCancelled("Server is shutting down, request cancelled")
 
         self._start_response_reader_lazily()
-        need_response = kwargs.pop("__rpc_need_response", True)
-        timeout = kwargs.pop("__rpc_timeout", self._timeout)
+        rpc_params = kwargs.pop("__rpc_params", RPCParams())
+        need_response = rpc_params.need_response
+        timeout = rpc_params.timeout if rpc_params.timeout is not None else self._timeout
 
         request_id = uuid.uuid4().hex
         logger.debug(f"RPC client sending request: {request_id}")
         request = RPCRequest(request_id,
-                             __rpc_method_name,
+                             method_name,
                              args,
                              kwargs,
                              need_response,
@@ -216,7 +215,7 @@ async def _call_async(self, __rpc_method_name, *args, **kwargs):
             raise
         except asyncio.TimeoutError:
             raise RPCTimeout(
-                f"Request '{__rpc_method_name}' timed out after {timeout}s")
+                f"Request '{method_name}' timed out after {timeout}s")
         except Exception as e:
             raise e
         finally:
@@ -241,11 +240,11 @@ def run_loop():
             import time
             time.sleep(0.1)
 
-    def _call_sync(self, __rpc_method_name, *args, **kwargs):
+    def _call_sync(self, method_name, *args, **kwargs):
         """Synchronous version of RPC call."""
         self._ensure_event_loop()
         future = asyncio.run_coroutine_threadsafe(
-            self._call_async(__rpc_method_name, *args, **kwargs), self._loop)
+            self._call_async(method_name, *args, **kwargs), self._loop)
         return future.result()
 
     def call_async(self, name: str, *args, **kwargs):
@@ -263,7 +262,9 @@ def call_async(self, name: str, *args, **kwargs):
         Example:
             result = await client.call_async('remote_method', arg1, arg2, key=value)
         """
-        return self._call_async(name, *args, **kwargs, __rpc_need_response=True)
+        if "__rpc_params" not in kwargs:
+            kwargs["__rpc_params"] = RPCParams(need_response=True)
+        return self._call_async(name, *args, **kwargs)
 
     def call_future(self, name: str, *args,
                     **kwargs) -> concurrent.futures.Future:
@@ -331,7 +332,8 @@ async def call_streaming(self, name: str, *args,
             raise RPCCancelled("Server is shutting down, request cancelled")
 
         self._start_response_reader_lazily()
-        timeout = kwargs.pop("__rpc_timeout", self._timeout)
+        rpc_params = kwargs.pop("__rpc_params", RPCParams())
+        timeout = rpc_params.timeout if rpc_params.timeout is not None else self._timeout
 
         request_id = uuid.uuid4().hex
         queue = asyncio.Queue()
@@ -379,7 +381,9 @@ async def call_streaming(self, name: str, *args,
     def get_server_attr(self, name: str):
         """ Get the attribute of the RPC server.
         This is mainly used for testing. """
-        return self._call_sync("__rpc_get_attr", name, __rpc_timeout=10)
+        return self._call_sync("__rpc_get_attr",
+                               name,
+                               __rpc_params=RPCParams(timeout=10))
 
     def __getattr__(self, name):
         """
@@ -395,7 +399,8 @@ def __init__(self, client, method_name):
 
             def __call__(self, *args, **kwargs):
                 """Default synchronous call"""
-                mode = kwargs.pop("__rpc_mode", "sync")
+                rpc_params = kwargs.get("__rpc_params", RPCParams())
+                mode = rpc_params.mode
                 if mode == "sync":
                     return self.client._call_sync(self.method_name, *args,
                                                   **kwargs)
diff --git a/tensorrt_llm/executor/rpc/rpc_common.py b/tensorrt_llm/executor/rpc/rpc_common.py
@@ -1,6 +1,19 @@
 from typing import Any, Literal, NamedTuple, Optional
 
 
+class RPCParams(NamedTuple):
+    """ Parameters for RPC calls. """
+
+    # seconds to wait for the response
+    timeout: Optional[float] = None
+
+    # whether the client needs the response, if False, it will return immediately
+    need_response: bool = True
+
+    # mode for RPC calls: "sync", "async", or "future"
+    mode: str = "sync"
+
+
 # --- Custom Exceptions ---
 class RPCError(Exception):
     """Custom exception for RPC-related errors raised on the client side.
diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py
@@ -14,6 +14,7 @@
 from .request import GenerationRequest
 from .result import GenerationResult
 from .rpc import RPCClient
+from .rpc.rpc_common import RPCParams
 from .rpc_worker import RpcWorker
 from .utils import (ErrorResponse, create_mpi_comm_session,
                     get_spawn_proxy_process_env, is_llm_response)
@@ -90,7 +91,8 @@ def main_loop_task(self):
         clock = 0
         while not self._shutdown_event.is_set():
             if clock % 1 == 0:
-                responses = self.fetch_responses_remote()
+                responses = self.fetch_responses_remote(
+                )  # RPC request => RPC server; result => RPC client
                 self.handle_responses(responses)
             if clock % 10 == 0:
                 stats = self.fetch_stats_remote()  # TODO
@@ -144,7 +146,8 @@ def submit(self, request: GenerationRequest) -> GenerationResult:
         logprob_params = self._get_logprob_params(request)
 
         # submit is a fire-and-forget operation, don't need to wait for response
-        self.rpc_client.submit(request, __rpc_need_response=False)
+        self.rpc_client.submit(request,
+                               __rpc_params=RPCParams(need_response=False))
 
         result = GenerationResult(
             request,
@@ -157,16 +160,19 @@ def submit(self, request: GenerationRequest) -> GenerationResult:
         return result
 
     def fetch_responses_remote(self):
-        return self.rpc_client.fetch_responses(__rpc_timeout=20)
+        return self.rpc_client.fetch_responses(__rpc_params=RPCParams(
+            timeout=20))
 
     def fetch_stats_remote(self):
         return self.rpc_client.fetch_stats()
 
     def setup_engine_remote(self):
-        return self.rpc_client.setup_engine(__rpc_timeout=60 * 20)  # 20 min
+        return self.rpc_client.setup_engine(
+            __rpc_params=RPCParams(timeout=60 * 20))  # 20 min
 
     def shutdown_remote(self):
-        self.rpc_client.shutdown(__rpc_timeout=60 * 20)  # 20 min
+        self.rpc_client.shutdown(__rpc_params=RPCParams(timeout=60 *
+                                                        20))  # 20 min
 
     def abort_request(self, request_id: int) -> None:
         return self.rpc_client.abort_request(request_id)
diff --git a/tensorrt_llm/executor/rpc_worker.py b/tensorrt_llm/executor/rpc_worker.py
@@ -58,6 +58,12 @@ def fetch_responses(self) -> list:
         qsize = self._response_queue.qsize()
         return [self._response_queue.get() for _ in range(qsize)]
 
+    # for streaming performance
+    async def fetch_responses_async(self) -> list:
+        while not self.shutdown_event.is_set():
+            responses = self.fetch_responses()  # will block
+            yield responses  # batching the responses to opt IPC performance
+
     def shutdown(self):
         logger.debug(f"RPC worker {mpi_rank()} is shutting down")
         self.shutdown_event.set()
diff --git a/tests/unittest/executor/test_rpc.py b/tests/unittest/executor/test_rpc.py
diff --git a/tests/unittest/executor/test_rpc_worker.py b/tests/unittest/executor/test_rpc_worker.py
diff --git a/tests/unittest/llmapi/test_llm_pytorch_rpc.py b/tests/unittest/llmapi/test_llm_pytorch_rpc.py
