scale out

Signed-off-by: Rui Qiao <[email protected]>

Author: ruisearch42

vllm/entrypoints/cli/serve.py

@@ -28,9 +28,9 @@
 from vllm.v1.executor.abstract import Executor
 from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus
 from vllm.v1.utils import (APIServerProcessManager, CoreEngine,
-                           get_engine_client_zmq_addr,
+                           CoreEngineActorManager, get_engine_client_zmq_addr,
                            wait_for_completion_or_failure,
-                           wait_for_engine_startup)
+                           wait_for_engine_startup, wait_for_ray_engine_actors)
 
 logger = init_logger(__name__)
 
@@ -212,6 +212,34 @@ def run_multi_api_server(args: argparse.Namespace):
         logger.info("Started DP Coordinator process (PID: %d)",
                     coordinator.proc.pid)
 
+    if parallel_config.data_parallel_backend == "ray":
+        logger.info("Starting ray-based data parallel backend")
+
+        engine_actor_manager = CoreEngineActorManager(
+            local_engine_count=local_engine_count,
+            start_index=args.data_parallel_start_rank,
+            local_start_index=0,
+            vllm_config=vllm_config,
+            addresses=addresses,
+            executor_class=Executor.get_class(vllm_config),
+            log_stats=not engine_args.disable_log_stats,
+        )
+        # Start API servers using the manager
+        api_server_manager = APIServerProcessManager(
+            target_server_fn=run_api_server_worker,
+            listen_address=listen_address,
+            sock=sock,
+            args=args,
+            num_servers=num_api_servers,
+            input_addresses=input_addresses,
+            output_addresses=output_addresses,
+            stats_update_address=stats_update_address)
+
+        wait_for_ray_engine_actors(api_server_manager=api_server_manager,
+                                   engine_actor_manager=engine_actor_manager,
+                                   coordinator=coordinator)
+        return
+
     handshake_address = get_engine_client_zmq_addr(
         local_only, host, parallel_config.data_parallel_rpc_port)
 

vllm/v1/engine/core_client.py

@@ -1144,17 +1144,6 @@ def __init__(
                     self.stats_update_address = \
                         coordinator.get_stats_publish_address()
 
-            # Start all engines.
-            self.resources.local_engine_manager = CoreEngineActorManager(
-                vllm_config=vllm_config,
-                executor_class=executor_class,
-                log_stats=log_stats,
-                input_address=input_address,
-                output_address=output_address,
-                local_engine_count=local_engine_count,
-                start_index=start_index,
-                local_start_index=local_start_index)
-
             self.core_engine = self.core_engines[0]
 
             self.utility_results: dict[int, AnyFuture] = {}
@@ -1180,3 +1169,36 @@ def __init__(
             self._ensure_output_queue_task()
         except RuntimeError:
             pass
+
+    def _init_engines_direct(self, vllm_config: VllmConfig, local_only: bool,
+                             local_start_index: int, input_address: str,
+                             output_address: str,
+                             executor_class: type[Executor], log_stats: bool):
+        """Self-contained client mode, launch engine and coordinator process
+        as needed."""
+
+        parallel_config = vllm_config.parallel_config
+        local_engine_count = parallel_config.data_parallel_size_local
+        start_index = parallel_config.data_parallel_rank
+
+        if len(self.core_engines) > 1:
+            self.resources.coordinator = DPCoordinator(parallel_config)
+
+        addresses: dict[str, Any] = {
+            "input_addresses": [input_address],
+            "output_addresses": [output_address],
+        }
+
+        coordinator = self.resources.coordinator
+        if coordinator is not None:
+            addresses.update(coordinator.get_engine_socket_addresses())
+
+        # Start all engines.
+        self.resources.local_engine_manager = CoreEngineActorManager(
+            vllm_config=vllm_config,
+            executor_class=executor_class,
+            log_stats=log_stats,
+            addresses=addresses,
+            local_engine_count=local_engine_count,
+            start_index=start_index,
+            local_start_index=local_start_index)

vllm/v1/utils.py

@@ -262,8 +262,7 @@ def __init__(
         start_index: int,
         local_start_index: int,
         vllm_config: VllmConfig,
-        input_address: str,
-        output_address: str,
+        addresses,
         executor_class: type[Executor],
         log_stats: bool,
     ):
@@ -284,8 +283,7 @@ def __init__(
                 vllm_config=vllm_config,
                 executor_class=executor_class,
                 log_stats=log_stats,
-                input_address=input_address,
-                output_address=output_address,
+                addresses=addresses,
                 on_head_node=True,
                 engine_index=global_index,
                 dp_rank=global_index,
@@ -301,8 +299,7 @@ def __init__(
                 vllm_config=vllm_config,
                 executor_class=executor_class,
                 log_stats=log_stats,
-                input_address=input_address,
-                output_address=output_address,
+                addresses=addresses,
                 on_head_node=False,
                 engine_index=global_index,
                 dp_rank=global_index,
@@ -490,6 +487,57 @@ def wait_for_completion_or_failure(
             local_engine_manager.close()
 
 
+def wait_for_ray_engine_actors(
+        api_server_manager: APIServerProcessManager,
+        engine_actor_manager: CoreEngineActorManager,
+        coordinator: Optional["DPCoordinator"] = None) -> None:
+    """Wait for all ray engine actors to complete or detect if any fail.
+    
+    Raises an exception if any process exits with a non-zero status.
+    """
+
+    try:
+        logger.info("Waiting for ray engine actors to complete ...")
+        # Create a mapping of sentinels to their corresponding processes
+        # for efficient lookup
+        sentinel_to_proc: dict[Any, Union[SpawnProcess, Process]] = {
+            proc.sentinel: proc
+            for proc in api_server_manager.processes
+        }
+
+        if coordinator:
+            sentinel_to_proc.update(
+                {coordinator.proc.sentinel: coordinator.proc})
+
+        # TODO(rui): check if any ray engine actor terminates
+        # Check if any process terminates
+        while sentinel_to_proc:
+            # Wait for any process to terminate
+            ready_sentinels: list[Any] = connection.wait(sentinel_to_proc)
+
+            # Process any terminated processes
+            for sentinel in ready_sentinels:
+                proc = sentinel_to_proc.pop(sentinel)
+
+                # Check if process exited with error
+                if proc.exitcode != 0:
+                    raise RuntimeError(
+                        f"Process {proc.name} (PID: {proc.pid}) "
+                        f"died with exit code {proc.exitcode}")
+    except KeyboardInterrupt:
+        logger.info("Received KeyboardInterrupt, shutting down API servers...")
+    except Exception as e:
+        logger.exception("Exception occurred while running API servers: %s",
+                         str(e))
+        raise
+    finally:
+        logger.info("Terminating remaining processes ...")
+        api_server_manager.close()
+        if coordinator:
+            coordinator.close()
+        engine_actor_manager.close()
+
+
 # Note(rob): shutdown function cannot be a bound method,
 # else the gc cannot collect the object.
 def shutdown(procs: list[Process]):