enhance rpc_worker and test

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

Author: Superjomn

tensorrt_llm/executor/rpc.py

@@ -340,16 +340,21 @@ async def _call_async(self, name, *args, **kwargs):
             *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.
 
         Returns:
             The result of the remote method call
         """
+        logger.debug(
+            f"RPC client calling method: {name} with args: {args} and kwargs: {kwargs}"
+        )
         await self._start_reader_if_needed()
-        need_response = kwargs.pop("need_response", True)
+        need_response = kwargs.pop("__rpc_need_response", True)
+        timeout = kwargs.pop("__rpc_timeout", self._timeout)
 
         request_id = uuid.uuid4().hex
         logger.debug(f"RPC client sending request: {request_id}")
-        timeout = kwargs.pop("__rpc_timeout", self._timeout)
         request = RPCRequest(request_id,
                              name,
                              args,
@@ -395,7 +400,7 @@ 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, need_response=True)
+        return self._call_async(name, *args, **kwargs, __rpc_need_response=True)
 
     def call_future(self, name: str, *args,
                     **kwargs) -> concurrent.futures.Future:
@@ -457,9 +462,7 @@ def __call__(self, *args, **kwargs):
 
             def call_async(self, *args, **kwargs):
                 """Async call - returns coroutine"""
-                return self.client._call_async(self.method_name,
-                                               *args,
-                                               need_response=True,
+                return self.client._call_async(self.method_name, *args,
                                                **kwargs)
 
             def call_future(self, *args, **kwargs) -> concurrent.futures.Future:

tensorrt_llm/executor/rpc_worker.py

@@ -41,11 +41,13 @@ def __init__(
         self.set_result_queue(self._response_queue)
 
     def fetch_responses(self) -> list:
+        logger.debug(f"RPC worker {mpi_rank()} is fetching responses")
         super().await_responses()
         qsize = self._response_queue.qsize()
         return [self._response_queue.get() for _ in range(qsize)]
 
     def shutdown(self):
+        logger.debug(f"RPC worker {mpi_rank()} is shutting down")
         self.shutdown_event.set()
         super().shutdown()
 

tensorrt_llm/executor/worker_base.py

@@ -18,6 +18,7 @@
 from ..llmapi.llm_args import PybindMirror
 from ..llmapi.tracer import global_tracer
 from ..llmapi.utils import _SyncQueue, print_colored_debug
+from ..logger import logger
 from ..lora_manager import LoraConfig, LoraManager
 from ..prompt_adapter_manager import PromptAdapterManager
 from ..runtime import ModelConfig
@@ -29,7 +30,11 @@
 from .request import GenerationRequest, LoRARequest, PromptAdapterRequest
 from .result import (GenerationResult, LogProbsResult, ResponseWrapper,
                      compute_logprobs)
-from .utils import ErrorResponse, RequestError, is_llm_response
+from .utils import (ErrorResponse, RequestError, enable_llm_debug,
+                    is_llm_response)
+
+if enable_llm_debug():
+    logger.set_level("debug")
 
 __all__ = [
     "WorkerBase",
@@ -405,6 +410,7 @@ def __exit__(self, exc_type, exc_value, traceback):
 
     def await_responses(self) -> None:
         self._await_response_helper()
+        logger.debug(f"worker done await_responses")
 
     def fetch_kv_cache_events(self) -> list:
         if isinstance(self.engine, tllm.Executor):
@@ -472,6 +478,11 @@ def shutdown(self):
         # Check if there are any errors from the threads before shutdown.
         self._handle_background_error()
 
+    def _has_background_error(self) -> bool:
+        # TODO[Superjomn]: The worker background error should be deprecated once
+        # RPC approach is supported.
+        return not self._error_queue.empty()
+
 
 class AwaitResponseHelper:
     ''' Multiple-implementations for await_response for performance. '''
@@ -518,8 +529,11 @@ def responses_handler(self, responses: List[tllm.Response]):
 
     def __call__(self) -> bool:
         ''' This method should be called by a ManagedThread. '''
+        logger.debug(f"await_response: {self.worker.engine}")
         responses = self.worker.engine.await_responses(
             timeout=datetime.timedelta(milliseconds=100))
+        logger.debug(f"PyExecutor returned {len(responses)} responses")
+
         # filter since The _engine_response_callback may return None
         responses = list(
             filter(

tests/unittest/executor/test_rpc.py

@@ -123,7 +123,7 @@ def get_task_submitted(self) -> bool:
         server.start()
         time.sleep(0.1)
         client = RPCClient("ipc:///tmp/rpc_test_no_wait")
-        client.send_task(need_response=False)
+        client.send_task(__rpc_need_response=False)
         time.sleep(0.1)  # wait for some time to make sure the task is submitted
         assert client.get_task_submitted()
 
@@ -149,14 +149,14 @@ def send_task(self) -> None:
 
         time_start = time.time()
         for i in range(100):
-            client.send_task(need_response=False)
+            client.send_task(__rpc_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(need_response=True)
+            client.send_task(__rpc_need_response=True)
         time_end = time.time()
         wait_time = time_end - time_start
 

tests/unittest/executor/test_rpc_worker.py

@@ -4,6 +4,8 @@
 import time
 from concurrent.futures import ProcessPoolExecutor
 
+from test_worker_base import TestWorkerBase
+
 from tensorrt_llm.executor.request import GenerationRequest
 from tensorrt_llm.executor.rpc import RPCClient
 from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy
@@ -20,12 +22,19 @@
 
 class TestRpcWorker:
 
+    def __init__(self):
+        self.executor_config = TestWorkerBase.create_fake_executor_config(
+            model_path)
+
     def create_tp1_worker_process(self):
         addr = GenerationExecutorRpcProxy.gen_uniq_rpc_addr()
         # Use spawn method instead of fork
         mp_context = multiprocessing.get_context('spawn')
         pool = ProcessPoolExecutor(max_workers=1, mp_context=mp_context)
-        pool.submit(RpcWorker.main_task, engine=model_path, rpc_addr=addr)
+        pool.submit(RpcWorker.main_task,
+                    engine=model_path,
+                    rpc_addr=addr,
+                    executor_config=self.executor_config)
         return pool, addr
 
     def create_rpc_client(self, addr: str):
@@ -35,15 +44,53 @@ def create_rpc_client(self, addr: str):
     def test_main(self):
         pool, addr = self.create_tp1_worker_process()
         client = self.create_rpc_client(addr)
-        client.setup_engine(engine=model_path)
+        print("call setup_engine")
+        client.setup_engine(engine=model_path,
+                            executor_config=self.executor_config,
+                            __rpc_timeout=120)
+        print("call submit")
         time.sleep(1)
-        client.submit(
-            GenerationRequest(prompt_token_ids=[3, 4, 5],
-                              sampling_params=SamplingParams(max_tokens=10)))
-        responses = client.fetch_responses()
-        assert responses
 
-        client.shutdown()
+        def process_request():
+            ret = client.submit(GenerationRequest(
+                prompt_token_ids=[3, 4, 5],
+                sampling_params=SamplingParams(max_tokens=10)),
+                                __rpc_need_response=False)
+            assert ret is None
+
+            print(f"submit result: {ret}")
+            print("call fetch_responses")
+            # NOTE: known issue, the responses should be fetched before shutdown,
+            # or the shutdown will hang.
+            results = []
+            for i in range(3):
+                time.sleep(3)
+                results.extend(client.fetch_responses())
+                print(f"fetch_responses result: {results}")
+            assert len(results) == 1
+
+        def process_request_streaming():
+            ret = client.submit(prompt_token_ids=[3, 4, 5],
+                                sampling_params=SamplingParams(max_tokens=10),
+                                streaming=True,
+                                __rpc_need_response=False)
+            assert ret is None
+
+            print("call fetch_responses")
+            # NOTE: known issue, the responses should be fetched before shutdown,
+            # or the shutdown will hang.
+            results = []
+            for i in range(3):
+                time.sleep(3)
+                results.extend(client.fetch_responses())
+                print(f"fetch_responses result: {results}")
+            print(f"generate_async result: {results}")
+
+        process_request()
+        process_request_streaming()
+
+        print("call shutdown")
+        client.shutdown(__rpc_timeout=10)
         pool.shutdown()