refactor RPC param

Signed-off-by: Superjomn <[email protected]>

Author: Superjomn

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"
 ]

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)

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.

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)

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()

tests/unittest/executor/test_rpc.py

@@ -4,7 +4,8 @@
 import pytest
 
 from tensorrt_llm.executor.rpc import (RPCCancelled, RPCClient, RPCError,
-                                       RPCServer, RPCStreamingError, RPCTimeout)
+                                       RPCParams, RPCServer, RPCStreamingError,
+                                       RPCTimeout)
 
 
 class RpcServerWrapper(RPCServer):
@@ -126,7 +127,7 @@ def get_task_submitted(self) -> bool:
         with RpcServerWrapper(App(),
                               addr="ipc:///tmp/rpc_test_no_wait") as server:
             with RPCClient("ipc:///tmp/rpc_test_no_wait") as client:
-                client.send_task(__rpc_need_response=False)
+                client.send_task(__rpc_params=RPCParams(need_response=False))
                 time.sleep(
                     0.1
                 )  # wait for some time to make sure the task is submitted
@@ -209,7 +210,8 @@ def task(self):
         with RPCClient(addr) as client:
             client.shutdown_server()
             pending_futures = [
-                client.task(__rpc_mode="future") for _ in range(10)
+                client.task(__rpc_params=RPCParams(mode="future"))
+                for _ in range(10)
             ]
 
             for future in pending_futures:
@@ -238,7 +240,7 @@ def slow_method(self):
             with RPCClient("ipc:///tmp/rpc_test_timeout",
                            timeout=0.5) as client:
                 with pytest.raises(RPCError) as exc_info:
-                    client.slow_method(__rpc_timeout=0.5)
+                    client.slow_method(__rpc_params=RPCParams(timeout=0.5))
 
                     error = exc_info.value
                     # Should be either a timeout error or RPC error indicating timeout
@@ -306,14 +308,14 @@ def send_task(self) -> None:
         with RPCClient("ipc:///tmp/rpc_test_no_wait") as client:
             time_start = time.time()
             for i in range(100):
-                client.send_task(__rpc_need_response=False)
+                client.send_task(__rpc_params=RPCParams(need_response=False))
             time_end = time.time()
 
             no_wait_time = time_end - time_start
 
             time_start = time.time()
             for i in range(100):
-                client.send_task(__rpc_need_response=True)
+                client.send_task(__rpc_params=RPCParams(need_response=True))
             time_end = time.time()
             wait_time = time_end - time_start
 
@@ -340,7 +342,8 @@ def cal(self, n: int):
 
             time_start = time.time()
             for i in range(100):
-                ret = client.cal(i, __rpc_timeout=10)  # sync call
+                ret = client.cal(
+                    i, __rpc_params=RPCParams(timeout=10))  # sync call
                 assert ret == i * 2, f"{ret} != {i * 2}"
             time_end = time.time()
             print(
@@ -378,7 +381,7 @@ def teardown_method(self):
 
     def run_sync_timeout_test(self):
         with pytest.raises(RPCTimeout) as exc_info:
-            self.client.slow_operation(2.0, __rpc_timeout=0.1)
+            self.client.slow_operation(2.0, __rpc_params=RPCParams(timeout=0.1))
         assert "timed out" in str(
             exc_info.value), f"Timeout message not found: {exc_info.value}"
 
@@ -387,26 +390,25 @@ def run_async_timeout_test(self):
 
         async def async_timeout():
             with pytest.raises(RPCTimeout) as exc_info:
-                await self.client.call_async('slow_operation',
-                                             2.0,
-                                             __rpc_timeout=0.1)
+                await self.client.call_async(
+                    'slow_operation', 2.0, __rpc_params=RPCParams(timeout=0.1))
             assert "timed out" in str(
                 exc_info.value), f"Timeout message not found: {exc_info.value}"
 
         asyncio.run(async_timeout())
 
     def run_sync_success_test(self):
-        result = self.client.slow_operation(0.1, __rpc_timeout=10.0)
+        result = self.client.slow_operation(
+            0.1, __rpc_params=RPCParams(timeout=10.0))
         assert result == "completed"
         print(f"final result: {result}")
 
     def run_async_success_test(self):
         import asyncio
 
         async def async_success():
-            result = await self.client.call_async('slow_operation',
-                                                  0.1,
-                                                  __rpc_timeout=10.0)
+            result = await self.client.call_async(
+                'slow_operation', 0.1, __rpc_params=RPCParams(timeout=10.0))
             assert result == "completed"
             print(f"final result: {result}")
             return result
@@ -457,7 +459,8 @@ def foo(self, delay: int):
         time.sleep(0.1)
         with RPCClient("ipc:///tmp/rpc_test_shutdown") as client:
             # This task should be continued after server shutdown
-            res = client.foo(10, __rpc_timeout=12, __rpc_mode="future")
+            res = client.foo(10,
+                             __rpc_params=RPCParams(timeout=12, mode="future"))
 
             # The shutdown will block until all pending requests are finished
             server.shutdown()
@@ -589,7 +592,7 @@ async def test_streaming_timeout(self):
         # Set short timeout
         with pytest.raises(RPCTimeout):
             async for value in client.streaming_timeout.call_streaming(
-                    delay=2.0, __rpc_timeout=0.5):
+                    delay=2.0, __rpc_params=RPCParams(timeout=0.5)):
                 pass  # Should timeout before first yield
 
     @pytest.mark.asyncio

tests/unittest/executor/test_rpc_worker.py

@@ -7,7 +7,7 @@
 from test_worker_base import create_fake_executor_config
 
 from tensorrt_llm.executor.request import GenerationRequest
-from tensorrt_llm.executor.rpc import RPCClient
+from tensorrt_llm.executor.rpc import RPCClient, RPCParams
 from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy
 from tensorrt_llm.executor.rpc_worker import RpcWorker
 from tensorrt_llm.sampling_params import SamplingParams
@@ -43,14 +43,14 @@ def create_rpc_client(self, addr: str):
     def test_main_loop(self):
         pool, addr = self.create_tp1_worker_process()
         client = self.create_rpc_client(addr)
-        client.setup_engine(__rpc_timeout=120)
+        client.setup_engine(__rpc_params=RPCParams(timeout=120))
         time.sleep(1)
 
         def process_request():
             ret = client.submit(GenerationRequest(
                 prompt_token_ids=[3, 4, 5],
                 sampling_params=SamplingParams(max_tokens=10)),
-                                __rpc_need_response=False)
+                                __rpc_params=RPCParams(need_response=False))
             assert ret is None  # need_response = False
 
             print(f"submit result: {ret}")
@@ -69,7 +69,7 @@ def process_request_streaming():
                 prompt_token_ids=[3, 4, 5],
                 sampling_params=SamplingParams(max_tokens=10),
                 streaming=True),
-                                __rpc_need_response=False)
+                                __rpc_params=RPCParams(need_response=False))
             assert ret is None
             print("submit result: ", ret)
 
@@ -80,7 +80,8 @@ def process_request_streaming():
 
             while not results:
                 time.sleep(1)
-                results.extend(client.fetch_responses(__rpc_timeout=10))
+                results.extend(
+                    client.fetch_responses(__rpc_params=RPCParams(timeout=10)))
                 print(f"try fetch_responses result: {results}")
             print(f"fetch_responses result: {results}")
             assert results
@@ -90,7 +91,7 @@ def process_request_streaming():
         process_request_streaming()
 
         print("call shutdown")
-        client.shutdown(__rpc_timeout=10)
+        client.shutdown(__rpc_params=RPCParams(timeout=10))
         pool.shutdown()
         client.close()