update rpc, wip

Author: hchings

tensorrt_llm/bench/dataclasses/reporting.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import os
 from collections import defaultdict
 from typing import Any, Dict, List, NamedTuple
 
@@ -507,9 +508,12 @@ def report_statistics(self) -> None:
                 f"Max Sequence Length:\t{build_cfg['max_seq_len']}\n"
                 f"\n")
         else:
+            # Check MPI vs RAY and RPC status
+            comm_backend = "RAY" if os.environ.get("TLLM_DISABLE_MPI") == "1" else "MPI"
+            ray_status = "[RPC]" if os.environ.get("TLLM_RAY_USE_RPC") == "1" else "[original]"
             backend_info = (
                 "\n\n===========================================================\n"
-                "= PYTORCH BACKEND\n"
+                f"= PYTORCH BACKEND [{comm_backend}] {ray_status}\n"
                 "===========================================================\n"
                 f"Model:\t\t\t{engine['model']}\n"
                 f"Model Path:\t\t{engine['model_path']}\n"

tensorrt_llm/executor/ray_executor.py

@@ -89,7 +89,7 @@ def __init__(self,
             if self.use_rpc:
                 self.rpc_addr = get_unique_ipc_addr()
                 self.rpc_client = RPCClient(self.rpc_addr)
-                print(f"RPC client created at {self.rpc_addr}")
+                # print(f"RPC client created at {self.rpc_addr}")
 
                 self._results = {}
                 self._shutdown_event = threading.Event()
@@ -144,6 +144,13 @@ def use_ray_queue(self) -> bool:
     async def _generic_fetch_loop_async(self, fetch_method_name: str,
                                         handler_method, method_name: str):
         # TODO copied from GenerationExecutorRpcProxy, need refactoring.
+        """Generic method for fetching data in a loop from RPC worker.
+
+        Args:
+            fetch_method_name: Name of the RPC client method to call
+            handler_method: The handler method to call with the fetched data
+            method_name: Name of the method for logging
+        """
         try:
             fetch_method = getattr(self.rpc_client, fetch_method_name)
             async for data in fetch_method().remote_streaming():
@@ -169,6 +176,7 @@ async def main_loop_task():
             await self._fetch_responses_loop_async()
 
         def _run_main_loop_task():
+            """Local method to run the main loop task."""
             self.main_loop = asyncio.new_event_loop()
             asyncio.set_event_loop(self.main_loop)
 
@@ -177,18 +185,21 @@ def _run_main_loop_task():
             try:
                 self.main_loop.run_until_complete(self.main_loop_task_obj)
             except asyncio.CancelledError:
-                pass
+                pass  # Task cancellation is expected during shutdown
             finally:
                 self.main_loop.close()
 
         self.main_loop_thread = threading.Thread(target=_run_main_loop_task,
-                                                 daemon=True)
+                                                 daemon=True,
+                                                 name="ray_executor_main_loop")
         self.main_loop_thread.start()
         atexit.register(self.shutdown)
 
     def setup_engine_remote(self):
         return self.collective_rpc("setup_engine", non_block=False)
 
+    # TODO: use Ray RPC to shutdown RPC server, and then close client. 
+
     def handle_responses(self, responses: list[GenerationResult]) -> bool:
         # TODO copied from GenerationExecutorRpcProxy, need refactoring.
         async_queues = []
@@ -209,6 +220,7 @@ def process_res(res: list):
                 if isinstance(queue, _SyncQueue):
                     queue.put_nowait(r)
                     async_queues.append(queue)
+                    # all the loops are identical
                     event_loop = event_loop or queue.loop
                 else:
                     queue.put(r)
@@ -217,7 +229,9 @@ def process_res(res: list):
                         r, ErrorResponse):
                     self._results.pop(client_id)
 
+        # Handle the case where responses might not be a list of lists
         if responses and not isinstance(responses[0], list):
+            # If responses is a flat list, wrap it
             responses = [responses]
 
         for res in responses:
@@ -314,10 +328,9 @@ def submit(self, request: GenerationRequest) -> GenerationResult:
         if self.use_rpc:
             with nvtx_range_debug("rpc_submit"):
                 self.rpc_client.submit(request).remote(need_response=False)
-                print(
-                    f"[RPC] RayExecutor  submit done for request {request.id}")
+                # print(
+                #     f"[RPC] RayExecutor  submit done for request {request.id}")
 
-            # TODO. use the future return by BaseWorker submit
             result = GenerationResult(
                 request,
                 background_error_handler=self._handle_background_error,
@@ -355,6 +368,13 @@ def abort_request(self, request_id: int) -> None:
                                   request_id=request_id)
 
     def shutdown(self):
+        try:
+            self.shutdown_impl()
+        except Exception as e:
+            print(f"Error shutting down RayExecutor: {e}")
+            raise e
+
+    def shutdown_impl(self):
         if self.use_rpc:
             if self._shutdown_event.is_set():
                 return

tensorrt_llm/executor/ray_gpu_worker.py

@@ -4,12 +4,12 @@
 from pathlib import Path
 from queue import Queue
 from threading import Event
-from typing import Any, Optional, Type, Union
+from typing import Any, AsyncGenerator, Optional, Type, Union
 
 import ray
 import torch
 
-from .._utils import mpi_rank, ray_use_rpc
+from .._utils import nvtx_range_debug, ray_use_rpc
 from ..bindings import executor as tllm
 from ..builder import Engine
 from ..llmapi.llm_args import BaseLlmArgs
@@ -236,53 +236,59 @@ def enqueue_request(self,
                         request: GenerationRequest,
                         result_wait_queue: Queue | None = None) -> int:
         # TODO. remove this. originally we didn't have to handle all the req id dict
+        # raise ValueError("enqueue_request should not be called.")
         return self._enqueue_request(request, result_wait_queue)
 
     def submit(self, request: GenerationRequest):
-        print(f"RayGPUWorker {self.rank} submitted request {request.id}")
         return super().submit(request)
 
-    async def fetch_responses_async(self,
-                                    timeout: Optional[float] = None) -> list:
+    def fetch_responses(self, timeout: Optional[float] = None) -> list:
         # TODO copied from RpcWorker, need refactoring.
-        logger_debug(f"RayGPUWorker {mpi_rank()} is fetching responses async",
+        logger_debug(f"RayGPUWorker {self.rank} is fetching responses",
                      color="yellow")
-
-        responses = await asyncio.to_thread(self.await_responses,
-                                            timeout=timeout)
-        if self._await_response_helper:
+        with nvtx_range_debug("RayGPUWorker.fetch_responses",
+                              color="orange",
+                              category="Worker"):
+            # NOTE: This is a blocking call, it will wait for the responses to be available.
+            responses = super().await_responses(timeout)
             self._await_response_helper.responses_handler(responses)
 
-        if hasattr(self,
-                   '_response_queue') and self._response_queue is not None:
-            qsize = self._response_queue.qsize()
-            logger_debug(f"RayGPUWorker returning {qsize} responses",
-                         color="yellow")
+        qsize = self._response_queue.qsize()
+        logger_debug(f"RayGPUWorker returning {qsize} responses", color="yellow")
+
+        all_responses = []
+        for _ in range(qsize):
+            # The queue contains batches of responses, so extend the list
+            all_responses.extend(self._response_queue.get())
+        return all_responses
 
-            all_responses = []
-            for _ in range(qsize):
-                all_responses.extend(self._response_queue.get())
-            return all_responses
+    async def fetch_responses_async(self,
+                                    timeout: Optional[float] = None) -> list:
+        # TODO copied from RpcWorker, need refactoring.
+        # A really async version of fetch_responses
+        logger_debug(f"RayGPUWorker {self.rank} is fetching responses async",
+                     color="yellow")
 
-        return responses if responses else []
+        # First, await any pending responses without blocking the event loop
+        responses = await asyncio.to_thread(self.fetch_responses,
+                                            timeout=timeout)
+        return responses
 
     # for streaming performance
-    async def fetch_responses_loop_async(self):
+    async def fetch_responses_loop_async(self) -> AsyncGenerator[list, None]:
         # TODO copied from RpcWorker, need refactoring.
-        shutdown_event = getattr(self, 'shutdown_event', Event())
-
-        while not shutdown_event.is_set():
+        while not self.shutdown_event.is_set():
             responses = await self.fetch_responses_async()
-            if responses:
+            if responses:  # Only yield if there are actual responses
                 logger_debug(
-                    f"RayGPUWorker {mpi_rank()} yielding responses: {responses}",
+                    f"RayGPUWorker {self.rank} is yielding responses: {responses}",
                     color="yellow")
-                yield responses
+                yield responses  # batching the responses to opt IPC performance
             else:
+                # Small delay to prevent busy waiting when no responses
                 await asyncio.sleep(0)
-
         logger_debug(
-            f"RayGPUWorker {mpi_rank()} quitting fetch_responses_loop_async",
+            f"RayGPUWorker {self.rank} quitting fetch_responses_loop_async",
             color="yellow")
 
     def shutdown(self):