better reuse code between remote methods

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

Author: Superjomn

tensorrt_llm/executor/rpc/rpc_client.py

@@ -20,48 +20,52 @@ def __init__(self, client: 'RPCClient', method_name: str, *args, **kwargs):
         self.args = args
         self.kwargs = kwargs
 
+    def _prepare_and_call(self, timeout: Optional[float], need_response: bool,
+                          mode: str, call_method: str) -> Any:
+        """Common method to prepare RPC params and make the call.
+
+        Args:
+            timeout: Timeout for the RPC call
+            need_response: Whether a response is expected
+            mode: The RPC mode ("sync", "async", "future")
+            call_method: The method name to call on the client
+
+        Returns:
+            The result of the client method call
+        """
+        rpc_params = RPCParams(timeout=timeout,
+                               need_response=need_response,
+                               mode=mode)
+        self.kwargs["__rpc_params"] = rpc_params
+        client_method = getattr(self.client, call_method)
+        return client_method(self.method_name, *self.args, **self.kwargs)
+
     def remote(self,
                timeout: Optional[float] = None,
                need_response: bool = True) -> Any:
         """Synchronous remote call with optional RPC parameters."""
-        rpc_params = RPCParams(timeout=timeout,
-                               need_response=need_response,
-                               mode="sync")
-        self.kwargs["__rpc_params"] = rpc_params
-        return self.client._call_sync(self.method_name, *self.args,
-                                      **self.kwargs)
+        return self._prepare_and_call(timeout, need_response, "sync",
+                                      "_call_sync")
 
     def remote_async(self,
                      timeout: Optional[float] = None,
                      need_response: bool = True):
         """Asynchronous remote call that returns a coroutine."""
-        rpc_params = RPCParams(timeout=timeout,
-                               need_response=need_response,
-                               mode="async")
-        self.kwargs["__rpc_params"] = rpc_params
-        return self.client._call_async(self.method_name, *self.args,
-                                       **self.kwargs)
+        return self._prepare_and_call(timeout, need_response, "async",
+                                      "_call_async")
 
     def remote_future(self,
                       timeout: Optional[float] = None,
                       need_response: bool = True) -> concurrent.futures.Future:
         """Remote call that returns a Future object."""
-        rpc_params = RPCParams(timeout=timeout,
-                               need_response=need_response,
-                               mode="future")
-        self.kwargs["__rpc_params"] = rpc_params
-        return self.client.call_future(self.method_name, *self.args,
-                                       **self.kwargs)
+        return self._prepare_and_call(timeout, need_response, "future",
+                                      "call_future")
 
     def remote_streaming(self,
                          timeout: Optional[float] = None) -> AsyncIterator[Any]:
         """Remote call for streaming results."""
-        rpc_params = RPCParams(timeout=timeout,
-                               need_response=True,
-                               mode="async")
-        self.kwargs["__rpc_params"] = rpc_params
-        return self.client.call_streaming(self.method_name, *self.args,
-                                          **self.kwargs)
+        # Streaming always needs a response
+        return self._prepare_and_call(timeout, True, "async", "call_streaming")
 
 
 class RPCClient:
@@ -309,7 +313,7 @@ async def _call_async(self, method_name, *args, **kwargs):
                 res = await future
             else:
                 # Add 1 second to the timeout to ensure the client can get
-                res = await asyncio.wait_for(future, timeout + 1)
+                res = await asyncio.wait_for(future, timeout)
             logger_debug(
                 f"RPC Client _call_async: Got result for request_id: {request_id}: {res}"
             )
@@ -361,7 +365,7 @@ def _call_sync(self, method_name, *args, **kwargs):
             f"RPC Client _call_sync: Got result for {method_name}: {result}")
         return result
 
-    def call_async(self, name: str, *args, **kwargs):
+    def call_async(self, name: str, *args, **kwargs) -> Any:
         """
         Call a remote method asynchronously.
 
@@ -408,7 +412,7 @@ def _async_to_sync():
 
         return self._executor.submit(_async_to_sync)
 
-    def call_sync(self, name: str, *args, **kwargs):
+    def call_sync(self, name: str, *args, **kwargs) -> Any:
         """
         Call a remote method synchronously (blocking).
 
@@ -476,7 +480,7 @@ async def call_streaming(self, name: str, *args,
                     response = await queue.get()
                 else:
                     response = await asyncio.wait_for(queue.get(),
-                                                      timeout=timeout + 1)
+                                                      timeout=timeout)
 
                 logger_debug(
                     f"RPC Client call_streaming received [{response.stream_status}] response: {response}",

tensorrt_llm/executor/rpc/rpc_common.py

@@ -1,3 +1,5 @@
+import time
+from dataclasses import dataclass
 from typing import Any, Literal, NamedTuple, Optional
 
 
@@ -48,14 +50,22 @@ class RPCStreamingError(RPCError):
     """Exception for streaming-related errors."""
 
 
-class RPCRequest(NamedTuple):
+@dataclass
+class RPCRequest:
     request_id: str
     method_name: str
     args: tuple
     kwargs: dict
     need_response: bool = True
     timeout: float = 0.5
     is_streaming: bool = False
+    creation_timestamp: Optional[
+        float] = None  # Unix timestamp when request was created
+
+    def __post_init__(self):
+        """Initialize creation_timestamp if not provided."""
+        if self.creation_timestamp is None:
+            self.creation_timestamp = time.time()
 
 
 class RPCResponse(NamedTuple):

tensorrt_llm/executor/rpc/rpc_server.py

@@ -292,18 +292,46 @@ async def _worker_routine(self, stop_event: threading.Event):
             if req.method_name not in ["_rpc_shutdown", "shutdown"]:
                 self._num_pending_requests -= 1
 
+    def _calculate_adjusted_timeout(self,
+                                    req: RPCRequest,
+                                    is_streaming: bool = False) -> float:
+        """Calculate adjusted timeout based on pending overhead.
+
+        Args:
+            req: The RPC request
+            is_streaming: Whether this is for a streaming request
+
+        Returns:
+            The adjusted timeout value
+        """
+        adjusted_timeout = req.timeout
+        if req.creation_timestamp is not None and req.timeout is not None and req.timeout > 0:
+            pending_time = time.time() - req.creation_timestamp
+            adjusted_timeout = max(0.1, req.timeout -
+                                   pending_time)  # Keep at least 0.1s timeout
+            if pending_time > 0.1:  # Only log if significant pending time
+                method_type = "streaming " if is_streaming else ""
+                logger_debug(
+                    f"RPC Server adjusted timeout for {method_type}{req.method_name}: "
+                    f"original={req.timeout}s, pending={pending_time:.3f}s, adjusted={adjusted_timeout:.3f}s"
+                )
+        return adjusted_timeout
+
     async def _process_request(self, req: RPCRequest) -> Optional[RPCResponse]:
         """Process a request. Returns None for streaming requests (handled separately)."""
         func = self._functions[req.method_name]
 
+        # Calculate adjusted timeout based on pending overhead
+        adjusted_timeout = self._calculate_adjusted_timeout(req)
+
         try:
             if inspect.iscoroutinefunction(func):
                 # Execute async function directly in event loop, no need to run in executor due to the GIL
                 logger_debug(
                     f"RPC Server running async task {req.method_name} in dispatcher"
                 )
                 result = await asyncio.wait_for(func(*req.args, **req.kwargs),
-                                                timeout=req.timeout)
+                                                timeout=adjusted_timeout)
             else:
                 # Execute sync function in thread executor
                 loop = asyncio.get_running_loop()
@@ -317,7 +345,7 @@ def call_with_kwargs():
                 # TODO: let num worker control the pool size
                 result = await asyncio.wait_for(loop.run_in_executor(
                     self._executor, call_with_kwargs),
-                                                timeout=req.timeout)
+                                                timeout=adjusted_timeout)
 
             logger_debug(f"RPC Server returned result for request {req}")
             response = RPCResponse(req.request_id, result)
@@ -354,6 +382,10 @@ async def _process_streaming_request(self, req: RPCRequest):
 
         sequence_number = 0
 
+        # Calculate adjusted timeout based on pending overhead
+        adjusted_timeout = self._calculate_adjusted_timeout(req,
+                                                            is_streaming=True)
+
         try:
             logger_debug(f"RPC Server running streaming task {req.method_name}")
             # Send start signal
@@ -362,16 +394,37 @@ async def _process_streaming_request(self, req: RPCRequest):
                             'start'))
             sequence_number += 1
 
-            # Stream the results
-            async for result in func(*req.args, **req.kwargs):
-                logger_debug(
-                    f"RPC Server got data and ready to send result {result}")
-                response = RPCResponse(req.request_id, result, None, True,
-                                       sequence_number, 'data')
-                if not await self._send_response(req, response):
-                    # Stop streaming after a pickle error
-                    return
-                sequence_number += 1
+            # Apply timeout to the entire streaming operation if specified
+            if adjusted_timeout is not None and adjusted_timeout > 0:
+                # Create a task for the async generator with timeout
+                async def stream_with_timeout():
+                    nonlocal sequence_number
+                    async for result in func(*req.args, **req.kwargs):
+                        logger_debug(
+                            f"RPC Server got data and ready to send result {result}"
+                        )
+                        response = RPCResponse(req.request_id, result, None,
+                                               True, sequence_number, 'data')
+                        if not await self._send_response(req, response):
+                            # Stop streaming after a pickle error
+                            return
+                        sequence_number += 1
+
+                # Use wait_for for timeout handling
+                await asyncio.wait_for(stream_with_timeout(),
+                                       timeout=adjusted_timeout)
+            else:
+                # No timeout specified, stream normally
+                async for result in func(*req.args, **req.kwargs):
+                    logger_debug(
+                        f"RPC Server got data and ready to send result {result}"
+                    )
+                    response = RPCResponse(req.request_id, result, None, True,
+                                           sequence_number, 'data')
+                    if not await self._send_response(req, response):
+                        # Stop streaming after a pickle error
+                        return
+                    sequence_number += 1
 
             # Send end signal
             await self._client_socket.put_async(