amend test_worker_base

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

Author: Superjomn

tensorrt_llm/executor/rpc_proxy.py

@@ -12,7 +12,6 @@
 from .request import GenerationRequest
 from .result import GenerationResult
 from .rpc import RPCClient
-from .rpc_worker import rpc_worker_main
 from .utils import (ErrorResponse, create_mpi_comm_session,
                     get_spawn_proxy_process_env, is_llm_response)
 
@@ -42,7 +41,7 @@ def __init__(self,
         """
 
         GenerationExecutorRpcProxy.INSTANCE_COUNTER += 1
-        self.rpc_addr = self._gen_rpc_addr()
+        self.rpc_addr = self.gen_uniq_rpc_addr()
         self.rpc_client = RPCClient(self.rpc_addr)
 
         postproc_worker_config = postproc_worker_config or PostprocWorkerConfig(
@@ -157,6 +156,7 @@ def _create_mpi_session(self, model_world_size: int,
             print_colored_debug('using external mpi session ...\n', "yellow")
             self.mpi_session = mpi_session
 
-    def _gen_rpc_addr(self):
+    @staticmethod
+    def gen_uniq_rpc_addr() -> str:
         process_id = os.getpid()
         return f"ipc:///tmp/rpc-proxy-{process_id}-{GenerationExecutorRpcProxy.INSTANCE_COUNTER}"

tensorrt_llm/executor/rpc_worker.py

@@ -1,17 +1,30 @@
 from pathlib import Path
+from queue import Queue
 from threading import Event
 from typing import Optional, Union
 
+from .._utils import mpi_rank
 from ..bindings import executor as tllm
 from ..builder import Engine
+from ..logger import logger
 from ..lora_manager import LoraConfig
 from ..sampling_params import BatchedLogitsProcessor
 from .postproc_worker import PostprocWorkerConfig
-from .rpc import RpcService
+from .rpc import RPCServer
 from .worker_base import WorkerBase
 
 
 class RpcWorker(WorkerBase):
+    """
+    A RPC wrapper for the WorkerBase class.
+
+    Actions:
+        - `setup_engine`: Setup the engine.
+        - `fetch_responses`: Fetch the latest responses from engine.
+        - `fetch_stats`: Fetch the latest stats from engine.
+        - `fetch_kv_cache_events`: Fetch the latest kv cache events from engine.
+        - `shutdown`: Shutdown the worker.
+    """
 
     def __init__(
         self,
@@ -24,37 +37,54 @@ def __init__(
                          is_llm_executor=is_llm_executor)
         self.shutdown_event = Event()
 
+        self._response_queue = Queue()
+        self.set_result_queue(self._response_queue)
+
+    def fetch_responses(self) -> list:
+        super().await_responses()
+        qsize = self._response_queue.qsize()
+        return [self._response_queue.get() for _ in range(qsize)]
+
     def shutdown(self):
         self.shutdown_event.set()
         super().shutdown()
 
+    @staticmethod
+    def main_task(
+        engine: Union[Path, Engine],
+        rpc_addr: str,
+        *,
+        executor_config: Optional[tllm.ExecutorConfig] = None,
+        batched_logits_processor: Optional[BatchedLogitsProcessor] = None,
+        postproc_worker_config: Optional[PostprocWorkerConfig] = None,
+        is_llm_executor: Optional[bool] = None,
+        lora_config: Optional[LoraConfig] = None,
+        garbage_collection_gen0_threshold: Optional[int] = None,
+    ) -> None:
+        # Step 1: Create the worker instance
+        worker = RpcWorker(engine=engine, executor_config=executor_config)
+
+        if mpi_rank() != 0:
+            logger.debug(f"Worker {mpi_rank()} is setting up the engine")
+            # The non-leader worker will setup the engine immediately.
+            # The leader worker will wait for the RPC call to propagate the
+            # potential error.
+            worker.setup_engine(
+                engine=engine,
+                executor_config=executor_config,
+                batched_logits_processor=batched_logits_processor,
+                postproc_worker_config=postproc_worker_config,
+                is_llm_executor=is_llm_executor,
+                lora_config=lora_config,
+                garbage_collection_gen0_threshold=
+                garbage_collection_gen0_threshold)
+
+        if mpi_rank() == 0:
+            # Step 2: Create the RPC service, it will expose all the APIs of the worker as remote call to the client
+            rpc_server = RPCServer(worker)
+            rpc_server.bind(rpc_addr)
+            rpc_server.start()
 
-def rpc_worker_main(
-    engine: Union[Path, Engine],
-    rpc_addr: str,
-    executor_config: Optional[tllm.ExecutorConfig] = None,
-    batched_logits_processor: Optional[BatchedLogitsProcessor] = None,
-    postproc_worker_config: Optional[PostprocWorkerConfig] = None,
-    is_llm_executor: Optional[bool] = None,
-    lora_config: Optional[LoraConfig] = None,
-    garbage_collection_gen0_threshold: Optional[int] = None,
-) -> None:
-    # Step 1: Create the worker instance
-    worker = RpcWorker(engine=engine, executor_config=executor_config)
-    worker.create_engine(
-        engine=engine,
-        executor_config=executor_config,
-        batched_logits_processor=batched_logits_processor,
-        postproc_worker_config=postproc_worker_config,
-        is_llm_executor=is_llm_executor,
-        lora_config=lora_config,
-        garbage_collection_gen0_threshold=garbage_collection_gen0_threshold)
-
-    # Step 2: Create the RPC service, it will expose all the APIs of the worker as remote call to the client
-    rpc_service = RpcService(worker)
-    rpc_service.bind(rpc_addr)
-    rpc_service.start()
-
-    # Step 3: Wait for the worker to shutdown
-    worker.shutdown_event.wait()
-    rpc_service.shutdown()
+            # Step 3: Wait for the worker to shutdown
+            worker.shutdown_event.wait()
+            rpc_server.shutdown()

tensorrt_llm/executor/worker.py

@@ -66,7 +66,7 @@ def __init__(
         executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig(
             processor_batched=batched_logits_processor, replicate=False)
 
-        self.create_engine(
+        self.setup_engine(
             engine=engine,
             executor_config=executor_config,
             lora_config=lora_config,

tensorrt_llm/executor/worker_base.py

@@ -78,7 +78,7 @@ def __init__(
         self._prompt_adapter_manager: Optional[PromptAdapterManager] = None
         self._runtime_model_config: Optional[ModelConfig] = None
 
-    def create_engine(
+    def setup_engine(
             self,
             engine: Union[Path, Engine],
             executor_config: Optional[tllm.ExecutorConfig] = None,

tests/unittest/executor/test_rpc_worker.py

@@ -0,0 +1,52 @@
+import multiprocessing
+import os
+import sys
+import time
+from concurrent.futures import ProcessPoolExecutor
+
+from tensorrt_llm.executor.request import GenerationRequest
+from tensorrt_llm.executor.rpc import RPCClient
+from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy
+from tensorrt_llm.executor.rpc_worker import RpcWorker
+from tensorrt_llm.sampling_params import SamplingParams
+
+# isort: off
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..")
+from utils.llm_data import llm_models_root
+# isort: on
+
+model_path = llm_models_root() / "llama-models-v2/TinyLlama-1.1B-Chat-v1.0"
+
+
+class TestRpcWorker:
+
+    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)
+        return pool, addr
+
+    def create_rpc_client(self, addr: str):
+        client = RPCClient(addr)
+        return client
+
+    def test_main(self):
+        pool, addr = self.create_tp1_worker_process()
+        client = self.create_rpc_client(addr)
+        client.setup_engine(engine=model_path)
+        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()
+        pool.shutdown()
+
+
+if __name__ == '__main__':
+    worker = TestRpcWorker()
+    worker.test_main()

tests/unittest/executor/test_worker_base.py

@@ -1,7 +1,6 @@
 import os
 import sys
 import time
-from queue import Queue
 
 # isort: off
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..")
@@ -21,37 +20,37 @@
 
 class TestWorkerBase:
 
+    class FakeWorker(WorkerBase):
+
+        def __init__(self, engine: str):
+            super().__init__(engine=engine)
+            executor_config = TestWorkerBase.create_fake_executor_config(engine)
+            self.setup_engine(engine=engine, executor_config=executor_config)
+
     def test_create_engine(self):
-        with WorkerBase(engine=model_path) as worker:
-            pass
+        with self.FakeWorker(engine=model_path) as worker:
+            print(f"Created engine: {worker.engine}")
 
     def test_submit_request(self):
         sampling_params = SamplingParams(max_tokens=10)
         request = GenerationRequest(prompt_token_ids=[3, 4, 5],
                                     sampling_params=sampling_params)
-        with WorkerBase(engine=model_path) as worker:
-            worker.submit(request)
-
-    def test_await_responses(self):
-        sampling_params = SamplingParams(max_tokens=10)
-        request = GenerationRequest(prompt_token_ids=[3, 4, 5],
-                                    sampling_params=sampling_params)
-        with WorkerBase(engine=model_path) as worker:
-            result_queue = Queue()
-            worker.set_result_queue(result_queue)
-
+        with self.FakeWorker(engine=model_path) as worker:
+            print(f"Created engine: {worker.engine}")
             worker.submit(request)
             for i in range(10):
+                time.sleep(0.5)
                 worker.await_responses()
-
-            assert result_queue.qsize() > 0
+            print(f"Submitted request: {request}")
+            time.sleep(6)
 
     def test_fetch_stats(self):
         request = GenerationRequest(
             prompt_token_ids=[3, 4, 5],
             sampling_params=SamplingParams(max_tokens=10))
-        with WorkerBase(engine=model_path) as worker:
+        with self.FakeWorker(engine=model_path) as worker:
             worker.submit(request)
+            time.sleep(1)
             worker.await_responses()
             stats = worker.fetch_stats()
             assert len(stats) > 0
@@ -60,15 +59,16 @@ def test_dispatch_stats_task(self):
         request = GenerationRequest(
             prompt_token_ids=[3, 4, 5],
             sampling_params=SamplingParams(max_tokens=10))
-        with WorkerBase(engine=model_path) as worker:
+        with self.FakeWorker(engine=model_path) as worker:
             worker.submit(request)
             worker.await_responses()
             worker.dispatch_stats_task()
             time.sleep(10)
             stats = worker.fetch_stats()
             assert len(stats) == 1
 
-    def _create_executor_config(self):
+    @staticmethod
+    def create_fake_executor_config(model_path):
         llm_args = LlmArgs(model=model_path, cuda_graph_config=None)
 
         executor_config = tllm.ExecutorConfig(1)
@@ -92,4 +92,4 @@ def _create_executor_config(self):
 
 if __name__ == "__main__":
     test_worker_base = TestWorkerBase()
-    test_worker_base.test_create_engine()
+    test_worker_base.test_fetch_stats()