refine rpc timeout

TODO: The timeout error may need a specific Error type

Author: Superjomn

tensorrt_llm/executor/rpc.py

@@ -27,6 +27,7 @@ class RPCRequest(NamedTuple):
     args: tuple
     kwargs: dict
     need_response: bool = True
+    timeout: float = 0.5
 
 
 class RPCResponse(NamedTuple):
@@ -176,15 +177,35 @@ async def _worker_routine(self, stop_event: threading.Event):
             if req.method_name in self._functions:
                 try:
                     if self._executor is not None:
-                        # Dispatch to worker thread and await result
+                        # Dispatch to worker thread and await result with timeout
                         loop = asyncio.get_running_loop()
-                        result = await loop.run_in_executor(
-                            self._executor, self._functions[req.method_name],
-                            *req.args, **req.kwargs)
+
+                        # Create a wrapper function to handle keyword arguments
+                        def call_with_kwargs():
+                            return self._functions[req.method_name](
+                                *req.args, **req.kwargs)
+
+                        result = await asyncio.wait_for(loop.run_in_executor(
+                            self._executor, call_with_kwargs),
+                                                        timeout=req.timeout)
                     else:
-                        result = self._functions[req.method_name](*req.args,
-                                                                  **req.kwargs)
+                        # For synchronous execution, we need to run in executor to support timeout
+                        loop = asyncio.get_running_loop()
+
+                        # Create a wrapper function to handle keyword arguments
+                        def call_with_kwargs():
+                            return self._functions[req.method_name](
+                                *req.args, **req.kwargs)
+
+                        result = await asyncio.wait_for(loop.run_in_executor(
+                            None, call_with_kwargs),
+                                                        timeout=req.timeout)
                     response = RPCResponse(req.request_id, 'OK', result)
+                except asyncio.TimeoutError:
+                    response = RPCResponse(
+                        req.request_id, 'ERROR',
+                        f"Method '{req.method_name}' timed out after {req.timeout} seconds"
+                    )
                 except Exception:
                     tb = traceback.format_exc()
                     response = RPCResponse(req.request_id, 'ERROR', tb)
@@ -313,13 +334,28 @@ async def _start_reader_if_needed(self):
             self._reader_task = loop.create_task(self._response_reader())
 
     async def _call_async(self, name, *args, **kwargs):
-        """Async version of RPC call."""
+        """Async version of RPC call.
+        Args:
+            name: Method name to call
+            *args: Positional arguments
+            **kwargs: Keyword arguments
+            __rpc_timeout: The timeout (seconds) for the RPC call.
+
+        Returns:
+            The result of the remote method call
+        """
         await self._start_reader_if_needed()
         need_response = kwargs.pop("need_response", True)
 
         request_id = uuid.uuid4().hex
         logger.debug(f"RPC client sending request: {request_id}")
-        request = RPCRequest(request_id, name, args, kwargs, need_response)
+        timeout = kwargs.pop("__rpc_timeout", self._timeout)
+        request = RPCRequest(request_id,
+                             name,
+                             args,
+                             kwargs,
+                             need_response,
+                             timeout=timeout)
         logger.debug(f"RPC client sending request: {request}")
         await self._client_socket.put_async(request)
 
@@ -331,9 +367,12 @@ async def _call_async(self, name, *args, **kwargs):
         self._pending_futures[request_id] = future
 
         try:
-            return await asyncio.wait_for(future, self._timeout)
+            # If timeout, the remote call should return a timeout error timely,
+            # so we add 1 second to the timeout to ensure the client can get
+            # that result.
+            return await asyncio.wait_for(future, timeout + 1)
         except asyncio.TimeoutError:
-            raise RPCError(f"Request '{name}' timed out after {self._timeout}s")
+            raise RPCTimeout(f"Request '{name}' timed out after {timeout}s")
         finally:
             self._pending_futures.pop(request_id, None)
 

tests/unittest/executor/test_rpc.py

@@ -183,9 +183,66 @@ def cal(self, n: int):
 
         time_start = time.time()
         for i in range(10000):
-            ret = client.cal(i)  # sync call
+            ret = client.cal(i, __rpc_timeout=10)  # sync call
             assert ret == i * 2, f"{ret} != {i * 2}"
         time_end = time.time()
         print(
             f"Time taken: {time_end - time_start} seconds, {10000 / (time_end - time_start)} calls/second"
         )
+
+
+@pytest.mark.parametrize("use_async", [True, False])
+def test_rpc_timeout(use_async: bool):
+    """Test RPC timeout functionality.
+
+    Args:
+        use_async: Whether to test async RPC calls or sync RPC calls
+    """
+
+    class App:
+
+        def slow_operation(self, delay: float):
+            """A method that takes a long time to complete."""
+            time.sleep(delay)
+            return "completed"
+
+    with RPCServer(App()) as server:
+        server.bind("ipc:///tmp/rpc_test_timeout")
+        server.start()
+        time.sleep(0.1)
+        client = RPCClient("ipc:///tmp/rpc_test_timeout")
+
+        # Test that a short timeout causes RPCTimeout exception
+        with pytest.raises(RPCError) as exc_info:
+            if use_async:
+                # Test async call with timeout
+                import asyncio
+
+                async def test_async_timeout():
+                    return await client.call_async('slow_operation',
+                                                   2.0,
+                                                   __rpc_timeout=0.1)
+
+                asyncio.run(test_async_timeout())
+            else:
+                # Test sync call with timeout
+                client.slow_operation(2.0, __rpc_timeout=0.1)
+
+            assert "timed out" in str(
+                exc_info.value), f"Timeout message not found: {exc_info.value}"
+
+        # Test that a long timeout allows the operation to complete
+        if use_async:
+            # Test async call with sufficient timeout
+            import asyncio
+
+            async def test_async_success():
+                return await client.call_async('slow_operation',
+                                               0.1,
+                                               __rpc_timeout=1.0)
+
+            result = asyncio.run(test_async_success())
+        else:
+            result = client.slow_operation(0.1, __rpc_timeout=1.0)
+
+        assert result == "completed"