tentatively add back ray queues

Signed-off-by: Erin Ho <[email protected]>

Author: hchings

tensorrt_llm/_utils.py

@@ -524,6 +524,13 @@ def mpi_disabled() -> bool:
     return os.environ.get("TLLM_DISABLE_MPI") == "1"
 
 
+def ray_use_rpc() -> bool:
+    """True if TLLM_RAY_USE_RPC is set to "1", False otherwise.
+    # TODO: deprecate this once Ray is fully moved to use RPC client/server.
+    """
+    return os.environ.get("TLLM_RAY_USE_RPC") == "1"
+
+
 def mpi_rank():
     if mpi_disabled():
         try:

tensorrt_llm/executor/ray_executor.py

@@ -13,13 +13,15 @@
                                       placement_group)
 
 from tensorrt_llm._ray_utils import unwrap_ray_errors
-from tensorrt_llm._utils import get_free_port
+from tensorrt_llm._utils import get_free_port, nvtx_range_debug, ray_use_rpc
 from tensorrt_llm.logger import logger
 
 from ..llmapi.utils import logger_debug
 from .executor import GenerationExecutor
 from .postproc_worker import PostprocWorkerConfig
 from .ray_gpu_worker import RayGPUWorker, RayWorkerWrapper
+from .request import GenerationRequest
+from .result import GenerationResult, RayAsyncQueue, RaySyncQueue
 from .rpc_proxy import RpcExecutorMixin
 
 __all__ = [
@@ -74,18 +76,40 @@ def __init__(self,
             self.tp_size = tp_size
             self.master_address = ray.util.get_node_ip_address()
             self.master_port = get_free_port()
-            self.init_rpc_executor()
+            self.use_rpc = ray_use_rpc()
 
             worker_kwargs = dict(**worker_kwargs,
                                  postproc_worker_config=postproc_worker_config,
-                                 is_llm_executor=is_llm_executor,
-                                 rpc_addr=self.rpc_addr)
-            self.create_workers(RayGPUWorker, worker_kwargs)
-            self.setup_engine_remote()
-            self.setup_mainloop(tasks=[self._fetch_responses_loop_async],
-                                thread_name="ray_executor_main_loop")
+                                 is_llm_executor=is_llm_executor)
+
+            if self.use_rpc:
+                self.init_rpc_executor()
+                worker_kwargs['rpc_addr'] = self.rpc_addr
+                self.create_workers(RayGPUWorker, worker_kwargs)
+                self.setup_engine_remote()
+                self.setup_mainloop(tasks=[self._fetch_responses_loop_async],
+                                    thread_name="ray_executor_main_loop")
+                logger.info(f"Connecting to RPC server at {self.rpc_addr}")
+            else:
+                self.response_queue = RayAsyncQueue.options(runtime_env={
+                    "env_vars": {
+                        "TLLM_DISABLE_MPI": "1"
+                    }
+                }).remote()
+                self.response_sync_queue = RaySyncQueue.options(runtime_env={
+                    "env_vars": {
+                        "TLLM_DISABLE_MPI": "1"
+                    }
+                }).remote()
+                self.async_response_queue_weakref = self.create_actor_weak_ref(
+                    self.response_queue)
+                self.sync_response_queue_weakref = self.create_actor_weak_ref(
+                    self.response_sync_queue)
+                self.response_queue.warmup.remote()
+                self.response_sync_queue.warmup.remote()
+                self.create_workers(RayGPUWorker, worker_kwargs)
+
         except Exception as e:
-            # Clean up the Ray resources early during exception
             self.shutdown()
             logger.error(f"Failed to initialize RayExecutor: {e}")
             raise e
@@ -165,6 +189,43 @@ def collective_rpc(self,
                                                         **kwargs))
         return refs if non_block else ray.get(refs)
 
+    def submit(self, request: "GenerationRequest") -> "GenerationResult":
+        """
+        Low-level API to the executor. Return a "future" GenerationResult
+        which can be waited.
+        Forwards the request to the workers through RPC or Ray queues depending on mode.
+        """
+        request.set_id(self._get_next_client_id())
+        logprob_params = self._get_logprob_params(request)
+
+        if self.use_rpc:
+            with nvtx_range_debug("rpc_submit"):
+                self.rpc_client.submit(request).remote(need_response=False)
+
+            result = GenerationResult(
+                request,
+                background_error_handler=self._handle_background_error,
+                executor=self,
+                disaggregated_params=request.disaggregated_params,
+                logprob_params=logprob_params)
+            self._results[request.id] = result
+        else:
+            result = GenerationResult(
+                request,
+                background_error_handler=self._handle_background_error,
+                executor=self,
+                disaggregated_params=request.disaggregated_params,
+                logprob_params=logprob_params)
+
+            with nvtx_range_debug("request_queue.put"):
+                self.call_all_ray_workers("enqueue_request",
+                                          leader_only=True,
+                                          request=request,
+                                          async_call=True,
+                                          result_wait_queue=result.queue)
+
+        return result
+
     def start(self):
         pass
 
@@ -177,50 +238,69 @@ def report_device_ids(self) -> list[str]:
                                             async_call=False)
         return sorted(gpu_ids)
 
+    def use_ray_queue(self) -> bool:
+        return not self.use_rpc
+
     def abort_request(self, request_id: int) -> None:
         self.call_all_ray_workers("abort_request",
                                   leader_only=True,
                                   async_call=False,
                                   request_id=request_id)
 
-    # TODO: Use Ray RPC to shutdown RPC server, and then close client
     def shutdown(self):
-        if self._shutdown_event.is_set():
+        if hasattr(self, '_shutdown_event') and self._shutdown_event.is_set():
             return
-        self._shutdown_event.set()
-        logger_debug(f"Shutting down RayExecutor (RPC mode)", color="yellow")
+        if hasattr(self, '_shutdown_event'):
+            self._shutdown_event.set()
 
-        # First, cancel the main loop to stop fetching responses
-        if hasattr(self, 'main_loop') and self.main_loop and hasattr(
-                self, 'main_loop_task_obj') and self.main_loop_task_obj:
-            logger_debug("Cancelling main loop task.", color="yellow")
-            try:
-                self.main_loop.call_soon_threadsafe(
-                    self.main_loop_task_obj.cancel)
-            except Exception as e:
-                logger_debug(f"Error cancelling main loop task: {e}",
-                             color="yellow")
+        mode_str = "RPC mode" if self.use_rpc else "Ray queue mode"
+        logger_debug(f"Shutting down RayExecutor ({mode_str})", color="yellow")
 
-            if hasattr(self, 'main_loop_thread'):
-                self.main_loop_thread.join()
+        if self.use_rpc:
+            if hasattr(self, 'main_loop') and self.main_loop and hasattr(
+                    self, 'main_loop_task_obj') and self.main_loop_task_obj:
+                logger_debug("Cancelling main loop task.", color="yellow")
+                try:
+                    self.main_loop.call_soon_threadsafe(
+                        self.main_loop_task_obj.cancel)
+                except Exception as e:
+                    logger_debug(f"Error cancelling main loop task: {e}",
+                                 color="yellow")
 
-        # Then, shutdown the workers
-        if hasattr(self, 'workers') and self.workers is not None:
-            try:
-                logger_debug("Shutting down RPC remote", color="yellow")
-                shutdown_refs = [
-                    worker.shutdown.remote() for worker in self.workers
-                ]
-                # Add timeout to prevent indefinite hanging
-                ray.get(shutdown_refs, timeout=30.0)
-            except ray.exceptions.GetTimeoutError:
-                logger.warning(
-                    "Timeout waiting for workers to shutdown after 30 seconds")
-            except Exception as e:
-                logger.warning(f"Error shutting down RPC remote: {e}")
+                if hasattr(self, 'main_loop_thread'):
+                    self.main_loop_thread.join()
 
-        if hasattr(self, 'rpc_client') and self.rpc_client is not None:
-            self.rpc_client.close()
+            # Then, shutdown the workers
+            if hasattr(self, 'workers') and self.workers is not None:
+                try:
+                    logger_debug("Shutting down RPC remote", color="yellow")
+                    shutdown_refs = [
+                        worker.shutdown.remote() for worker in self.workers
+                    ]
+                    # Add timeout to prevent indefinite hanging
+                    ray.get(shutdown_refs, timeout=30.0)
+                except ray.exceptions.GetTimeoutError:
+                    logger.warning(
+                        "Timeout waiting for workers to shutdown after 30 seconds"
+                    )
+                except Exception as e:
+                    logger.warning(f"Error shutting down RPC remote: {e}")
+
+            if hasattr(self, 'rpc_client') and self.rpc_client is not None:
+                try:
+                    self.rpc_client.close()
+                except Exception as e:
+                    # Suppress errors during RPC client shutdown
+                    # These can occur if the client is already closed or if there are
+                    # pending operations that get cancelled during cleanup
+                    logger_debug(
+                        f"Suppressed error during RPC client close: {e}")
+        else:
+            # Release actors
+            self.response_queue = None
+            self.response_sync_queue = None
+            self.async_response_queue_weakref = None
+            self.sync_response_queue_weakref = None
 
         self.workers = None
         if hasattr(self,
@@ -236,12 +316,6 @@ def shutdown(self):
             logger.debug("Shutting down Ray cluster")
             ray.shutdown()
 
-    @property
-    def enable_postprocess_parallel(self) -> bool:
-        ret = super().enable_postprocess_parallel
-        assert ret == False, "Postprocess parallel is not supported in RayExecutor"
-        return ret
-
     def _get_placement_group(self,
                              tp_size: int) -> Tuple[PlacementGroup, List[int]]:
         """
@@ -307,3 +381,15 @@ def _get_placement_group(self,
         pg = placement_group(bundles, strategy=strategy)
 
         return pg, bundle_indices
+
+    @property
+    def enable_postprocess_parallel(self) -> bool:
+        ret = super().enable_postprocess_parallel
+        assert ret == False, "Postprocess parallel is not supported in RayExecutor"
+        return ret
+
+    @staticmethod
+    def create_actor_weak_ref(actor_handle: ray.actor.ActorHandle):
+        state, _, _ = actor_handle._serialization_helper()
+        return ray.actor.ActorHandle._deserialization_helper(state,
+                                                             weak_ref=True)