# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse import multiprocessing import time import weakref from collections import defaultdict from collections.abc import Sequence from dataclasses import dataclass from enum import Enum, auto from multiprocessing import Process, connection from multiprocessing.process import BaseProcess from typing import (TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar, Union, overload) import msgspec import torch import zmq from vllm.config import CacheConfig, ParallelConfig, VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.utils import extract_layer_index from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled, usage_message) from vllm.utils import (get_mp_context, get_open_port, get_open_zmq_ipc_path, get_tcp_uri, kill_process_tree) from vllm.v1.executor.abstract import Executor if TYPE_CHECKING: from ray.util.placement_group import PlacementGroup from vllm.attention.layer import Attention from vllm.v1.engine.coordinator import DPCoordinator logger = init_logger(__name__) T = TypeVar("T") STARTUP_POLL_PERIOD_MS = 10000 class ConstantList(Generic[T], Sequence): def __init__(self, x: list[T]) -> None: self._x = x def append(self, item): raise Exception("Cannot append to a constant list") def extend(self, item): raise Exception("Cannot extend a constant list") def insert(self, item): raise Exception("Cannot insert into a constant list") def pop(self, item): raise Exception("Cannot pop from a constant list") def remove(self, item): raise Exception("Cannot remove from a constant list") def clear(self): raise Exception("Cannot clear a constant list") def index(self, item: T, start: int = 0, stop: Optional[int] = None) -> int: return self._x.index(item, start, stop if stop is not None else len(self._x)) @overload def __getitem__(self, item: int) -> T: ... @overload def __getitem__(self, s: slice, /) -> list[T]: ... def __getitem__(self, item: Union[int, slice]) -> Union[T, list[T]]: return self._x[item] @overload def __setitem__(self, item: int, value: T): ... @overload def __setitem__(self, s: slice, value: T, /): ... def __setitem__(self, item: Union[int, slice], value: Union[T, list[T]]): raise Exception("Cannot set item in a constant list") def __delitem__(self, item): raise Exception("Cannot delete item from a constant list") def __iter__(self): return iter(self._x) def __contains__(self, item): return item in self._x def __len__(self): return len(self._x) def __repr__(self): return f"ConstantList({self._x})" def get_engine_client_zmq_addr(local_only: bool, host: str, port: int = 0) -> str: return get_open_zmq_ipc_path() if local_only else (get_tcp_uri( host, port or get_open_port())) class CoreEngineState(Enum): NEW = auto() CONNECTED = auto() READY = auto() class CoreEngine: """One per data parallel rank.""" def __init__(self, index: int = 0, local: bool = True): self.local = local self.index = index self.identity = index.to_bytes(2, "little") self.state = CoreEngineState.NEW @dataclass class EngineZmqAddresses: # ZMQ input socket addresses for each front-end client (requests) inputs: list[str] # ZMQ output socket addresses for each front-end client (responses) outputs: list[str] # ZMQ input socket address of DP coordinator if applicable coordinator_input: Optional[str] = None # ZMQ output socket address of DP coordinator if applicable coordinator_output: Optional[str] = None @dataclass class EngineHandshakeMetadata: """Metadata sent to each engine process during startup handshake, including addresses of the front-end ZMQ queues that they should connect to. """ addresses: EngineZmqAddresses parallel_config: dict[str, Union[int, str]] class APIServerProcessManager: """Manages a group of API server processes. Handles creation, monitoring, and termination of API server worker processes. Also monitors extra processes to check if they are healthy. """ def __init__( self, target_server_fn: Callable, listen_address: str, sock: Any, args: argparse.Namespace, num_servers: int, input_addresses: list[str], output_addresses: list[str], stats_update_address: Optional[str] = None, ): """Initialize and start API server worker processes. Args: target_server_fn: Function to call for each API server process listen_address: Address to listen for client connections sock: Socket for client connections args: Command line arguments num_servers: Number of API server processes to start input_addresses: Input addresses for each API server output_addresses: Output addresses for each API server stats_update_address: Optional stats update address """ self.listen_address = listen_address self.sock = sock self.args = args # Start API servers spawn_context = multiprocessing.get_context("spawn") self.processes: list[BaseProcess] = [] for i, in_addr, out_addr in zip(range(num_servers), input_addresses, output_addresses): client_config = { "input_address": in_addr, "output_address": out_addr, "client_index": i } if stats_update_address is not None: client_config["stats_update_address"] = stats_update_address proc = spawn_context.Process(target=target_server_fn, name=f"ApiServer_{i}", args=(listen_address, sock, args, client_config)) self.processes.append(proc) proc.start() logger.info("Started %d API server processes", len(self.processes)) # Shutdown only the API server processes on garbage collection # The extra processes are managed by their owners self._finalizer = weakref.finalize(self, shutdown, self.processes) def close(self) -> None: self._finalizer() class CoreEngineProcManager: """ Utility class to handle creation, readiness, and shutdown of background processes used by the AsyncLLM and LLMEngine. """ def __init__( self, target_fn: Callable, local_engine_count: int, start_index: int, local_start_index: int, vllm_config: VllmConfig, on_head_node: bool, handshake_address: str, executor_class: type[Executor], log_stats: bool, ): context = get_mp_context() common_kwargs = { "vllm_config": vllm_config, "on_head_node": on_head_node, "handshake_address": handshake_address, "executor_class": executor_class, "log_stats": log_stats, } self.processes: list[BaseProcess] = [] for index in range(local_engine_count): local_index = local_start_index + index global_index = start_index + index # Start EngineCore in background process. self.processes.append( context.Process(target=target_fn, name=f"EngineCore_{global_index}", kwargs=common_kwargs | { "dp_rank": global_index, "local_dp_rank": local_index, })) self._finalizer = weakref.finalize(self, shutdown, self.processes) try: for proc in self.processes: proc.start() finally: # Kill other procs if not all are running. if self.finished_procs(): self.close() def close(self): """Shutdown all procs.""" self._finalizer() def join_first(self): """Wait for any process to exit.""" connection.wait(proc.sentinel for proc in self.processes) def sentinels(self) -> list: return [proc.sentinel for proc in self.processes] def finished_procs(self) -> dict[str, int]: """Returns dict of proc name -> exit code for any finished procs.""" return { proc.name: proc.exitcode for proc in self.processes if proc.exitcode is not None } class CoreEngineActorManager: """ Utility class to handle creation, readiness, and shutdown of core engine Ray actors used by the AsyncLLM and LLMEngine. Different from CoreEngineProcManager, this class manages core engines for both local and remote nodes. """ def __init__( self, vllm_config: VllmConfig, addresses: EngineZmqAddresses, executor_class: type[Executor], log_stats: bool, placement_groups: Optional[list["PlacementGroup"]] = None, local_dp_ranks: Optional[list[int]] = None, ): import copy import ray from ray.util.scheduling_strategies import ( PlacementGroupSchedulingStrategy) from vllm.v1.engine.core import DPEngineCoreActor self.local_engine_actors: list[ray.ActorHandle] = [] self.remote_engine_actors: list[ray.ActorHandle] = [] dp_size = vllm_config.parallel_config.data_parallel_size local_engine_count = \ vllm_config.parallel_config.data_parallel_size_local world_size = vllm_config.parallel_config.world_size if ray.is_initialized(): logger.info( "Ray is already initialized. Skipping Ray initialization.") else: ray.init() if placement_groups is not None: assert local_dp_ranks is not None, ( "local_dp_ranks must be provided if " "placement_groups is provided") assert len(placement_groups) == len(local_dp_ranks), ( "placement_groups and local_dp_ranks must " "have the same length") logger.info("Using provided placement groups") # TODO(rui): validate passed-in placement groups self.created_placement_groups = [] else: placement_groups, local_dp_ranks = \ CoreEngineActorManager.create_dp_placement_groups(vllm_config) self.created_placement_groups = placement_groups assert len(placement_groups) == dp_size, ( "Number of placement groups must match data parallel size") refs = [] for index in range(dp_size): local_index = local_dp_ranks[index] dp_vllm_config = copy.deepcopy(vllm_config) pg = placement_groups[index] dp_vllm_config.parallel_config.placement_group = pg on_head_node = index < local_engine_count actor = ray.remote(DPEngineCoreActor).options( scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=world_size, )).remote(vllm_config=dp_vllm_config, executor_class=executor_class, log_stats=log_stats, on_head_node=on_head_node, addresses=addresses, dp_rank=index, local_dp_rank=local_index) if on_head_node: self.local_engine_actors.append(actor) else: self.remote_engine_actors.append(actor) refs.append(actor.wait_for_init.remote()) ray.get(refs) self.run_refs = [] for actor in self.local_engine_actors + self.remote_engine_actors: self.run_refs.append(actor.run.remote()) @staticmethod def create_dp_placement_groups( vllm_config: VllmConfig ) -> tuple[list["PlacementGroup"], list[int]]: import ray from ray._private.state import available_resources_per_node from ray.util.state import list_nodes logger.info("Creating placement groups for data parallel") dp_master_ip = \ vllm_config.parallel_config.data_parallel_master_ip dp_size = vllm_config.parallel_config.data_parallel_size local_engine_count = \ vllm_config.parallel_config.data_parallel_size_local nodes = list_nodes() nodes = sorted(list_nodes(), key=lambda node: node.node_ip != dp_master_ip) assert nodes[0].node_ip == dp_master_ip, ( "The first node must be the head node") assert len(nodes) == 1 or nodes[1].node_ip != dp_master_ip, ( "There can only be one head node") available_resources = available_resources_per_node() world_size = vllm_config.parallel_config.world_size placement_groups: list[PlacementGroup] = [] local_dp_ranks: list[int] = [] for node in nodes: node_ip = node.node_ip node_resources = available_resources[node.node_id] # For now, each DP rank can only be assigned to one node # TODO(rui): support allocating a single DP rank # to multiple nodes available_engine_count = int(node_resources["GPU"]) // world_size if node_ip == dp_master_ip: assert available_engine_count >= local_engine_count, ( "Not enough resources to allocate DP ranks " f"on DP master node {node_ip}") for i in range(local_engine_count): bundles = [{ "GPU": 1.0, "node:" + dp_master_ip: 0.001 }] * world_size + [{ "CPU": 1.0 }] pg = ray.util.placement_group( name=f"dp_rank_{len(placement_groups)}", strategy="STRICT_PACK", bundles=bundles, ) placement_groups.append(pg) local_dp_ranks.append(i) else: for i in range(available_engine_count): if len(placement_groups) == dp_size: break bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}] pg = ray.util.placement_group( name=f"dp_rank_{len(placement_groups)}", strategy="STRICT_PACK", bundles=bundles, ) placement_groups.append(pg) local_dp_ranks.append(i) return placement_groups, local_dp_ranks def get_run_refs(self): return self.run_refs def close(self): import ray for actor in self.local_engine_actors + self.remote_engine_actors: ray.kill(actor) for pg in self.created_placement_groups: ray.util.remove_placement_group(pg) def wait_for_engine_startup( handshake_socket: zmq.Socket, addresses: EngineZmqAddresses, core_engines: list[CoreEngine], parallel_config: ParallelConfig, cache_config: CacheConfig, proc_manager: Optional[CoreEngineProcManager], coord_process: Optional[Process], ): # Wait for engine core process(es) to send ready messages. local_count = parallel_config.data_parallel_size_local remote_count = len(core_engines) - local_count # [local, remote] counts conn_pending, start_pending = [local_count, remote_count], [0, 0] poller = zmq.Poller() poller.register(handshake_socket, zmq.POLLIN) if proc_manager is not None: for sentinel in proc_manager.sentinels(): poller.register(sentinel, zmq.POLLIN) if coord_process is not None: poller.register(coord_process.sentinel, zmq.POLLIN) while any(conn_pending) or any(start_pending): events = poller.poll(STARTUP_POLL_PERIOD_MS) if not events: if any(conn_pending): logger.debug( "Waiting for %d local, %d remote core engine proc(s) " "to connect.", *conn_pending) if any(start_pending): logger.debug( "Waiting for %d local, %d remote core engine proc(s) " "to start.", *start_pending) continue if len(events) > 1 or events[0][0] != handshake_socket: # One of the local core processes exited. finished = proc_manager.finished_procs() if proc_manager else {} if coord_process is not None and coord_process.exitcode is not None: finished[coord_process.name] = coord_process.exitcode raise RuntimeError("Engine core initialization failed. " "See root cause above. " f"Failed core proc(s): {finished}") # Receive HELLO and READY messages from the input socket. eng_identity, ready_msg_bytes = handshake_socket.recv_multipart() eng_index = int.from_bytes(eng_identity, "little") engine = next((e for e in core_engines if e.identity == eng_identity), None) if engine is None: raise RuntimeError(f"Message from engine with unexpected data " f"parallel rank: {eng_index}") msg = msgspec.msgpack.decode(ready_msg_bytes) status, local = msg["status"], msg["local"] if local != engine.local: raise RuntimeError(f"{status} message from " f"{'local' if local else 'remote'} " f"engine {eng_index}, expected it to be " f"{'local' if engine.local else 'remote'}") if status == "HELLO" and engine.state == CoreEngineState.NEW: # Send init message with DP config info. init_message = msgspec.msgpack.encode( EngineHandshakeMetadata( addresses=addresses, parallel_config={ "data_parallel_master_ip": parallel_config.data_parallel_master_ip, "data_parallel_master_port": parallel_config.data_parallel_master_port, "data_parallel_size": parallel_config.data_parallel_size, })) handshake_socket.send_multipart((eng_identity, init_message), copy=False) conn_pending[0 if local else 1] -= 1 start_pending[0 if local else 1] += 1 engine.state = CoreEngineState.CONNECTED elif status == "READY" and (engine.state == CoreEngineState.CONNECTED): # Setup KV cache config with initialization state from # engine core process. Sum values from all engines in DP case. num_gpu_blocks = cache_config.num_gpu_blocks or 0 num_gpu_blocks += msg["num_gpu_blocks"] cache_config.num_gpu_blocks = num_gpu_blocks start_pending[0 if local else 1] -= 1 engine.state = CoreEngineState.READY else: raise RuntimeError(f"Unexpected {status} message for " f"{'local' if local else 'remote'} engine " f"{eng_index} in {engine.state} state.") logger.debug("%s from %s core engine process %s.", status, "local" if local else "remote", eng_index) def wait_for_completion_or_failure( api_server_manager: APIServerProcessManager, engine_manager: Optional[Union[CoreEngineProcManager, CoreEngineActorManager]] = None, coordinator: Optional["DPCoordinator"] = None) -> None: """Wait for all processes to complete or detect if any fail. Raises an exception if any process exits with a non-zero status. Args: api_server_manager: The manager for API servers. engine_manager: The manager for engine processes. If CoreEngineProcManager, it manages local engines; if CoreEngineActorManager, it manages all engines. coordinator: The coordinator for data parallel. """ try: logger.info("Waiting for API servers to complete ...") # Create a mapping of sentinels to their corresponding processes # for efficient lookup sentinel_to_proc: dict[Any, BaseProcess] = { proc.sentinel: proc for proc in api_server_manager.processes } if coordinator: sentinel_to_proc[coordinator.proc.sentinel] = coordinator.proc actor_run_refs = [] if isinstance(engine_manager, CoreEngineProcManager): for proc in engine_manager.processes: sentinel_to_proc[proc.sentinel] = proc elif isinstance(engine_manager, CoreEngineActorManager): actor_run_refs = engine_manager.get_run_refs() # Check if any process terminates while sentinel_to_proc or actor_run_refs: # Wait for any process to terminate ready_sentinels: list[Any] = connection.wait(sentinel_to_proc, timeout=5) # 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}") if actor_run_refs: import ray _, actor_run_refs = ray.wait(actor_run_refs, timeout=5) 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() if engine_manager: engine_manager.close() # Note(rob): shutdown function cannot be a bound method, # else the gc cannot collect the object. def shutdown(procs: list[BaseProcess]): # Shutdown the process. for proc in procs: if proc.is_alive(): proc.terminate() # Allow 5 seconds for remaining procs to terminate. deadline = time.monotonic() + 5 for proc in procs: remaining = deadline - time.monotonic() if remaining <= 0: break if proc.is_alive(): proc.join(remaining) for proc in procs: if proc.is_alive() and (pid := proc.pid) is not None: kill_process_tree(pid) def bind_kv_cache( kv_caches: dict[str, torch.Tensor], forward_context: dict[str, "Attention"], runner_kv_caches: list[torch.Tensor], ) -> None: """ Bind the allocated KV cache to both ModelRunner and forward context so that the KV cache can be used in the forward pass. This function: 1) Fills the ModelRunner's kv cache list (`runner_kv_caches`) with kv_caches. 2) Associates each attention layer in the `forward_context` with its corresponding KV cache in kv_caches. Args: kv_caches: The allocated kv_caches with layer names as keys. forward_context: The global forward context containing all Attention layers with layer names as keys. runner_kv_caches: The kv_cache declared by ModelRunner. """ # Bind kv_caches to ModelRunner assert len(runner_kv_caches) == 0 # Convert kv_caches dict to a list of tensors in the order of layer_index. index2name = defaultdict(list) for layer_name in kv_caches: index2name[extract_layer_index(layer_name)].append(layer_name) for layer_index in sorted(index2name.keys()): layer_names = index2name[layer_index] if len(layer_names) > 1: # One typical case is encoder-decoder model, e.g., bart. # The cross attention and self attention in the same decoder layer # has different layer_name but the same layer_index. raise NotImplementedError layer_name = layer_names[0] runner_kv_caches.append(kv_caches[layer_name]) # Bind kv_caches to forward context for layer_name, kv_cache in kv_caches.items(): # NOTE: Use list because of v0 PP virtual engine. forward_context[layer_name].kv_cache = [kv_cache] def copy_slice(from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int) -> torch.Tensor: """ Copy the first length elements of a tensor into another tensor in a non-blocking manner. Used to copy pinned CPU tensor data to pre-allocated GPU tensors. Returns the sliced target tensor. """ return to_tensor[:length].copy_(from_tensor[:length], non_blocking=True) def report_usage_stats( vllm_config, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT) -> None: """Report usage statistics if enabled.""" if not is_usage_stats_enabled(): return from vllm.model_executor.model_loader import get_architecture_class_name usage_message.report_usage( get_architecture_class_name(vllm_config.model_config), usage_context, extra_kvs={ # Common configuration "dtype": str(vllm_config.model_config.dtype), "tensor_parallel_size": vllm_config.parallel_config.tensor_parallel_size, "block_size": vllm_config.cache_config.block_size, "gpu_memory_utilization": vllm_config.cache_config.gpu_memory_utilization, # Quantization "quantization": vllm_config.model_config.quantization, "kv_cache_dtype": str(vllm_config.cache_config.cache_dtype), # Feature flags "enable_lora": bool(vllm_config.lora_config), "enable_prompt_adapter": bool(vllm_config.prompt_adapter_config), "enable_prefix_caching": vllm_config.cache_config.enable_prefix_caching, "enforce_eager": vllm_config.model_config.enforce_eager, "disable_custom_all_reduce": vllm_config.parallel_config.disable_custom_all_reduce, })