add correctness tests

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

Author: Superjomn

tensorrt_llm/executor/rpc/rpc_client.py

@@ -2,7 +2,7 @@
 import concurrent.futures
 import threading
 import uuid
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncIterator, Callable, Optional
 
 from ...llmapi.utils import logger_debug
 from ...logger import logger
@@ -14,7 +14,8 @@
 class RemoteCall:
     """Helper class to enable chained remote call syntax like client.method().remote()"""
 
-    def __init__(self, client: 'RPCClient', method_name: str, *args, **kwargs):
+    def __init__(self, client: 'RPCClient', method_name: str, *args,
+                 **kwargs) -> None:
         self.client = client
         self.method_name = method_name
         self.args = args
@@ -57,7 +58,7 @@ def remote(self,
 
     def remote_async(self,
                      timeout: Optional[float] = None,
-                     need_response: bool = True):
+                     need_response: bool = True) -> Any:
         """Asynchronous remote call that returns a coroutine.
 
         Args:
@@ -106,9 +107,9 @@ class RPCClient:
 
     def __init__(self,
                  address: str,
-                 hmac_key=None,
+                 hmac_key: Optional[bytes] = None,
                  timeout: Optional[float] = None,
-                 num_workers: int = 4):
+                 num_workers: int = 4) -> None:
         '''
         Args:
             address: The ZMQ address to connect to.
@@ -137,7 +138,7 @@ def __init__(self,
 
         logger_debug(f"RPC Client initialized. Connected to {self._address}")
 
-    def shutdown_server(self):
+    def shutdown_server(self) -> None:
         """Shutdown the server."""
         if self._server_stopped:
             return
@@ -146,7 +147,7 @@ def shutdown_server(self):
 
         self._server_stopped = True
 
-    def close(self):
+    def close(self) -> None:
         """Gracefully close the client, cleaning up background tasks."""
         if self._closed:
             return
@@ -172,7 +173,7 @@ def close(self):
 
         logger_debug("RPC Client closed")
 
-    async def _response_reader(self):
+    async def _response_reader(self) -> None:
         """Task to read responses from the socket and set results on futures."""
         logger_debug("Response reader started")
 
@@ -242,7 +243,7 @@ async def _response_reader(self):
         finally:
             logger_debug("Response reader exiting gracefully")
 
-    def _ensure_reader_task(self):
+    def _ensure_reader_task(self) -> None:
         """Ensure the response reader task is running."""
         with self._reader_lock:
             if self._reader_task is None or self._reader_task.done():
@@ -254,7 +255,7 @@ def _ensure_reader_task(self):
                     # No running event loop, will be started when needed
                     pass
 
-    async def _call_async(self, method_name, *args, **kwargs):
+    async def _call_async(self, method_name: str, *args, **kwargs) -> Any:
         """Async version of RPC call.
         Args:
             method_name: Method name to call
@@ -320,7 +321,7 @@ async def _call_async(self, method_name, *args, **kwargs):
         except Exception:
             raise
 
-    def _call_sync(self, method_name, *args, **kwargs):
+    def _call_sync(self, method_name: str, *args, **kwargs) -> Any:
         """Synchronous version of RPC call."""
         logger_debug(
             f"RPC Client calling method: {method_name} with args: {args} and kwargs: {kwargs}"
@@ -331,7 +332,7 @@ def _call_sync(self, method_name, *args, **kwargs):
             asyncio.get_running_loop()
 
             # We're inside an event loop, we need to run in a thread to avoid deadlock
-            def run_in_thread():
+            def run_in_thread() -> Any:
                 return asyncio.run(
                     self._call_async(method_name, *args, **kwargs))
 
@@ -454,12 +455,12 @@ async def _call_streaming(self, name: str, *args,
             # Clean up
             self._streaming_queues.pop(request_id, None)
 
-    def get_server_attr(self, name: str):
+    def get_server_attr(self, name: str) -> Any:
         """ Get the attribute of the RPC server.
         This is mainly used for testing. """
         return self._rpc_get_attr(name).remote()
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> Callable[..., RemoteCall]:
         """
         Magically handles calls to non-existent methods.
         Returns a callable that when invoked returns a RemoteCall instance.
@@ -472,16 +473,16 @@ def __getattr__(self, name):
         """
         logger_debug(f"RPC Client getting attribute: {name}")
 
-        def method_caller(*args, **kwargs):
+        def method_caller(*args, **kwargs) -> RemoteCall:
             return RemoteCall(self, name, *args, **kwargs)
 
         return method_caller
 
-    def __enter__(self):
+    def __enter__(self) -> 'RPCClient':
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.close()
 
-    def __del__(self):
+    def __del__(self) -> None:
         self.close()

tensorrt_llm/executor/rpc/rpc_common.py

@@ -86,5 +86,5 @@ class RPCResponse(NamedTuple):
     result: Any
     error: Optional[RPCError] = None
     is_streaming: bool = False  # True if more responses coming
-    sequence_number: int = 0  # For ordering streaming responses
+    chunk_index: int = 0  # For ordering streaming responses
     stream_status: Literal['start', 'data', 'end', 'error'] = 'data'