Ning XieandCyrus Leung
2a336d8239
Signed-off-by: Andy Xie <andy.xning@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
1962 lines
79 KiB
Python
1962 lines
79 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import asyncio
|
|
import contextlib
|
|
import queue
|
|
import sys
|
|
import uuid
|
|
import weakref
|
|
from abc import ABC, abstractmethod
|
|
from collections import Counter, defaultdict
|
|
from collections.abc import Awaitable, Callable, Sequence
|
|
from concurrent.futures import Future
|
|
from dataclasses import dataclass
|
|
from multiprocessing.connection import Connection
|
|
from multiprocessing.queues import Queue
|
|
from threading import Thread
|
|
from typing import Any, TypeAlias, TypeVar
|
|
|
|
import msgspec
|
|
import msgspec.msgpack
|
|
import zmq
|
|
import zmq.asyncio
|
|
|
|
from vllm import envs
|
|
from vllm.config import VllmConfig
|
|
from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S
|
|
from vllm.logger import init_logger
|
|
from vllm.lora.request import LoRARequest
|
|
from vllm.renderers import BaseRenderer
|
|
from vllm.tasks import SupportedTask
|
|
from vllm.tracing import instrument
|
|
from vllm.utils.async_utils import in_loop
|
|
from vllm.utils.network_utils import (
|
|
close_sockets,
|
|
get_open_zmq_inproc_path,
|
|
make_zmq_socket,
|
|
)
|
|
from vllm.v1.engine import (
|
|
EEP_NOTIFICATION_CALL_ID,
|
|
FT_STATUS_CALL_ID,
|
|
EEPNotificationType,
|
|
EngineCoreOutputs,
|
|
EngineCoreReadyResponse,
|
|
EngineCoreRequest,
|
|
EngineCoreRequestType,
|
|
PauseMode,
|
|
ReconfigureDistributedRequest,
|
|
ReconfigureRankType,
|
|
UtilityOutput,
|
|
)
|
|
from vllm.v1.engine.coordinator import DPCoordinator
|
|
from vllm.v1.engine.core import EngineCore, EngineCoreProc
|
|
from vllm.v1.engine.exceptions import EngineDeadError
|
|
from vllm.v1.engine.tensor_ipc import TensorIpcSender
|
|
from vllm.v1.engine.utils import (
|
|
CoreEngineActorManager,
|
|
CoreEngineProcManager,
|
|
get_engine_zmq_addresses,
|
|
launch_core_engines,
|
|
)
|
|
from vllm.v1.executor import Executor
|
|
from vllm.v1.fault_tolerance.engine_core_sentinel import FT_UTILITY_METHOD
|
|
from vllm.v1.fault_tolerance.utils import (
|
|
FaultToleranceRequest,
|
|
FaultToleranceResult,
|
|
)
|
|
from vllm.v1.pool.late_interaction import get_late_interaction_engine_index
|
|
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
AnyFuture: TypeAlias = asyncio.Future[Any] | Future[Any]
|
|
|
|
_R = TypeVar("_R") # Return type for collective_rpc
|
|
|
|
EngineIdentity = bytes
|
|
|
|
|
|
class EngineCoreClient(ABC):
|
|
"""
|
|
EngineCoreClient: subclasses handle different methods for pushing
|
|
and pulling from the EngineCore for asyncio / multiprocessing.
|
|
|
|
Subclasses:
|
|
* InprocClient: In process EngineCore (for V0-style LLMEngine use)
|
|
* SyncMPClient: ZMQ + background proc EngineCore (for LLM)
|
|
* AsyncMPClient: ZMQ + background proc EngineCore w/ asyncio (for AsyncLLM)
|
|
"""
|
|
|
|
@staticmethod
|
|
def make_client(
|
|
multiprocess_mode: bool,
|
|
asyncio_mode: bool,
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
renderer: BaseRenderer | None = None,
|
|
) -> "EngineCoreClient":
|
|
# renderer is passed through to the multiprocess clients, which start
|
|
# the frontend MM warmup (renderer.start_mm_warmup_in_background) once
|
|
# the engine-core processes have been forked, so warmup overlaps
|
|
# engine-core load without a live thread at fork() time. In-process
|
|
# clients do not take a renderer: the engine core is built in this
|
|
# process, so there is no other process to overlap with and the MM
|
|
# warmup runs synchronously inside renderer.warmup() instead.
|
|
# TODO: support this for debugging purposes.
|
|
if asyncio_mode and not multiprocess_mode:
|
|
raise NotImplementedError(
|
|
"Running EngineCore in asyncio without multiprocessing "
|
|
"is not currently supported."
|
|
)
|
|
|
|
if multiprocess_mode and asyncio_mode:
|
|
return EngineCoreClient.make_async_mp_client(
|
|
vllm_config,
|
|
executor_class,
|
|
log_stats,
|
|
renderer=renderer,
|
|
)
|
|
|
|
if multiprocess_mode and not asyncio_mode:
|
|
return SyncMPClient(
|
|
vllm_config,
|
|
executor_class,
|
|
log_stats,
|
|
renderer=renderer,
|
|
)
|
|
|
|
return InprocClient(vllm_config, executor_class, log_stats)
|
|
|
|
@staticmethod
|
|
@instrument(span_name="Overall Loading")
|
|
def make_async_mp_client(
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
client_addresses: dict[str, Any] | None = None,
|
|
client_count: int = 1,
|
|
client_index: int = 0,
|
|
renderer: BaseRenderer | None = None,
|
|
) -> "AsyncMPClient":
|
|
parallel_config = vllm_config.parallel_config
|
|
client_args = (
|
|
vllm_config,
|
|
executor_class,
|
|
log_stats,
|
|
client_addresses,
|
|
client_count,
|
|
client_index,
|
|
)
|
|
if parallel_config.data_parallel_size > 1:
|
|
if parallel_config.data_parallel_external_lb:
|
|
# External load balancer - client per DP rank.
|
|
return DPAsyncMPClient(
|
|
*client_args,
|
|
renderer=renderer,
|
|
)
|
|
# Internal load balancer - client balances to all DP ranks.
|
|
return DPLBAsyncMPClient(
|
|
*client_args,
|
|
renderer=renderer,
|
|
)
|
|
return AsyncMPClient(
|
|
*client_args,
|
|
renderer=renderer,
|
|
)
|
|
|
|
@abstractmethod
|
|
def shutdown(self, timeout: float | None = None) -> None: ...
|
|
|
|
def get_output(self) -> EngineCoreOutputs:
|
|
raise NotImplementedError
|
|
|
|
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
|
|
raise NotImplementedError
|
|
|
|
def add_request(self, request: EngineCoreRequest) -> None:
|
|
raise NotImplementedError
|
|
|
|
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
|
|
raise NotImplementedError
|
|
|
|
def reset_mm_cache(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
def reset_prefix_cache(
|
|
self, reset_running_requests: bool = False, reset_connector: bool = False
|
|
) -> bool:
|
|
raise NotImplementedError
|
|
|
|
def reset_encoder_cache(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
|
|
raise NotImplementedError
|
|
|
|
def wake_up(self, tags: list[str] | None = None) -> None:
|
|
raise NotImplementedError
|
|
|
|
def is_sleeping(self) -> bool:
|
|
raise NotImplementedError
|
|
|
|
def execute_dummy_batch(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
def set_weight_version(self, weight_version: str) -> None:
|
|
raise NotImplementedError
|
|
|
|
def get_weight_version(self) -> str:
|
|
raise NotImplementedError
|
|
|
|
async def execute_dummy_batch_async(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def set_weight_version_async(self, weight_version: str) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def get_weight_version_async(self) -> str:
|
|
raise NotImplementedError
|
|
|
|
def abort_requests(self, request_ids: list[str]) -> None:
|
|
raise NotImplementedError
|
|
|
|
def add_lora(self, lora_request: LoRARequest) -> bool:
|
|
raise NotImplementedError
|
|
|
|
def remove_lora(self, lora_id: int) -> bool:
|
|
raise NotImplementedError
|
|
|
|
def list_loras(self) -> set[int]:
|
|
raise NotImplementedError
|
|
|
|
def pin_lora(self, lora_id: int) -> bool:
|
|
raise NotImplementedError
|
|
|
|
def save_sharded_state(
|
|
self, path: str, pattern: str | None = None, max_size: int | None = None
|
|
) -> None:
|
|
raise NotImplementedError
|
|
|
|
def collective_rpc(
|
|
self,
|
|
method: str | Callable[..., _R],
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> list[_R]:
|
|
raise NotImplementedError
|
|
|
|
def dp_engines_running(self) -> bool:
|
|
"""Returns True if data parallel engines are collectively in a
|
|
running state."""
|
|
raise NotImplementedError
|
|
|
|
async def commit_elastic_ep(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def get_output_async(self) -> EngineCoreOutputs:
|
|
raise NotImplementedError
|
|
|
|
async def get_supported_tasks_async(self) -> tuple[SupportedTask, ...]:
|
|
raise NotImplementedError
|
|
|
|
async def add_request_async(self, request: EngineCoreRequest) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def profile_async(
|
|
self, is_start: bool = True, profile_prefix: str | None = None
|
|
) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def reset_mm_cache_async(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def reset_prefix_cache_async(
|
|
self, reset_running_requests: bool = False, reset_connector: bool = False
|
|
) -> bool:
|
|
raise NotImplementedError
|
|
|
|
async def reset_encoder_cache_async(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
|
|
raise NotImplementedError
|
|
|
|
async def wake_up_async(self, tags: list[str] | None = None) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def is_sleeping_async(self) -> bool:
|
|
raise NotImplementedError
|
|
|
|
async def abort_requests_async(self, request_ids: list[str]) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def add_lora_async(self, lora_request: LoRARequest) -> bool:
|
|
raise NotImplementedError
|
|
|
|
async def remove_lora_async(self, lora_id: int) -> bool:
|
|
raise NotImplementedError
|
|
|
|
async def list_loras_async(self) -> set[int]:
|
|
raise NotImplementedError
|
|
|
|
async def pin_lora_async(self, lora_id: int) -> bool:
|
|
raise NotImplementedError
|
|
|
|
async def save_sharded_state_async(
|
|
self, path: str, pattern: str | None = None, max_size: int | None = None
|
|
) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def collective_rpc_async(
|
|
self,
|
|
method: str | Callable[..., _R],
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> list[_R]:
|
|
raise NotImplementedError
|
|
|
|
async def handle_fault(
|
|
self, fault_tolerance_request: FaultToleranceRequest
|
|
) -> FaultToleranceResult:
|
|
raise NotImplementedError
|
|
|
|
async def get_status(self):
|
|
raise NotImplementedError
|
|
|
|
|
|
class InprocClient(EngineCoreClient):
|
|
"""
|
|
InprocClient: client for in-process EngineCore. Intended
|
|
for use in LLMEngine for V0-style add_request() and step()
|
|
EngineCore setup in this process (no busy loop).
|
|
|
|
* pushes EngineCoreRequest directly into the EngineCore
|
|
* pulls EngineCoreOutputs by stepping the EngineCore
|
|
"""
|
|
|
|
# Takes no renderer: this class has no _start_mm_warmup (only MPClient
|
|
# does), so a renderer would never be used here. EngineCore is built in
|
|
# this process and MM warmup runs synchronously inside renderer.warmup().
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
*,
|
|
executor_fail_callback: Callable | None = None,
|
|
):
|
|
self.engine_core = EngineCore(
|
|
vllm_config,
|
|
executor_class,
|
|
log_stats,
|
|
executor_fail_callback=executor_fail_callback,
|
|
)
|
|
|
|
def get_output(self) -> EngineCoreOutputs:
|
|
outputs, model_executed = self.engine_core.step_fn()
|
|
self.engine_core.post_step(model_executed=model_executed)
|
|
return outputs and outputs.get(0) or EngineCoreOutputs()
|
|
|
|
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
|
|
return self.engine_core.get_supported_tasks()
|
|
|
|
def add_request(self, request: EngineCoreRequest) -> None:
|
|
req, request_wave = self.engine_core.preprocess_add_request(request)
|
|
self.engine_core.add_request(req, request_wave)
|
|
|
|
def abort_requests(self, request_ids: list[str]) -> None:
|
|
if len(request_ids) > 0:
|
|
self.engine_core.abort_requests(request_ids)
|
|
|
|
def shutdown(self, timeout: float | None = None) -> None:
|
|
self.engine_core.shutdown()
|
|
|
|
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
|
|
self.engine_core.profile(is_start, profile_prefix)
|
|
|
|
def reset_mm_cache(self) -> None:
|
|
self.engine_core.reset_mm_cache()
|
|
|
|
def reset_prefix_cache(
|
|
self, reset_running_requests: bool = False, reset_connector: bool = False
|
|
) -> bool:
|
|
return self.engine_core.reset_prefix_cache(
|
|
reset_running_requests, reset_connector
|
|
)
|
|
|
|
def reset_encoder_cache(self) -> None:
|
|
self.engine_core.reset_encoder_cache()
|
|
|
|
def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
|
|
if mode == "wait":
|
|
raise ValueError("'wait' pause mode is not supported in inproc-engine mode")
|
|
result = self.engine_core.sleep(level, mode)
|
|
assert result is None
|
|
|
|
def wake_up(self, tags: list[str] | None = None) -> None:
|
|
self.engine_core.wake_up(tags)
|
|
|
|
def is_sleeping(self) -> bool:
|
|
return self.engine_core.is_sleeping()
|
|
|
|
def execute_dummy_batch(self) -> None:
|
|
self.engine_core.execute_dummy_batch()
|
|
|
|
def set_weight_version(self, weight_version: str) -> None:
|
|
self.engine_core.set_weight_version(weight_version)
|
|
|
|
def get_weight_version(self) -> str:
|
|
return self.engine_core.get_weight_version()
|
|
|
|
def add_lora(self, lora_request: LoRARequest) -> bool:
|
|
return self.engine_core.add_lora(lora_request)
|
|
|
|
def remove_lora(self, lora_id: int) -> bool:
|
|
return self.engine_core.remove_lora(lora_id)
|
|
|
|
def list_loras(self) -> set[int]:
|
|
return self.engine_core.list_loras()
|
|
|
|
def pin_lora(self, lora_id: int) -> bool:
|
|
return self.engine_core.pin_lora(lora_id)
|
|
|
|
def save_sharded_state(
|
|
self, path: str, pattern: str | None = None, max_size: int | None = None
|
|
) -> None:
|
|
self.engine_core.save_sharded_state(path, pattern, max_size)
|
|
|
|
def collective_rpc(
|
|
self,
|
|
method: str | Callable[..., _R],
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> list[_R]:
|
|
return self.engine_core.collective_rpc(method, timeout, args, kwargs)
|
|
|
|
def dp_engines_running(self) -> bool:
|
|
return False
|
|
|
|
|
|
@dataclass
|
|
class BackgroundResources:
|
|
"""Used as a finalizer for clean shutdown, avoiding
|
|
circular reference back to the client object."""
|
|
|
|
ctx: zmq.Context
|
|
# If CoreEngineProcManager, it manages local engines;
|
|
# if CoreEngineActorManager, it manages all engines.
|
|
engine_manager: CoreEngineProcManager | CoreEngineActorManager | None = None
|
|
coordinator: DPCoordinator | None = None
|
|
output_socket: zmq.Socket | zmq.asyncio.Socket | None = None
|
|
input_socket: zmq.Socket | zmq.asyncio.Socket | None = None
|
|
first_req_send_socket: zmq.asyncio.Socket | None = None
|
|
first_req_rcv_socket: zmq.asyncio.Socket | None = None
|
|
stats_update_socket: zmq.asyncio.Socket | None = None
|
|
output_queue_task: asyncio.Task | None = None
|
|
stats_update_task: asyncio.Task | None = None
|
|
shutdown_path: str | None = None
|
|
|
|
# Set if any of the engines are dead. Here so that the output
|
|
# processing threads can access it without holding a ref to the client.
|
|
engine_dead: bool = False
|
|
|
|
def __call__(self):
|
|
"""Clean up background resources."""
|
|
|
|
logger.debug_once("[shutdown] MPClient: background resource cleanup start")
|
|
self.engine_dead = True
|
|
if self.engine_manager is not None:
|
|
self.engine_manager.shutdown(
|
|
timeout=envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS
|
|
)
|
|
if self.coordinator is not None:
|
|
self.coordinator.shutdown()
|
|
|
|
if isinstance(self.output_socket, zmq.asyncio.Socket):
|
|
# Async case.
|
|
loop = self.output_queue_task._loop if self.output_queue_task else None
|
|
|
|
sockets = (
|
|
self.output_socket,
|
|
self.input_socket,
|
|
self.first_req_send_socket,
|
|
self.first_req_rcv_socket,
|
|
self.stats_update_socket,
|
|
)
|
|
|
|
tasks = (self.output_queue_task, self.stats_update_task)
|
|
|
|
def close_sockets_and_tasks():
|
|
close_sockets(sockets)
|
|
for task in tasks:
|
|
if task is not None and not task.done():
|
|
with contextlib.suppress(Exception):
|
|
task.cancel()
|
|
|
|
if loop is not None:
|
|
if in_loop(loop):
|
|
close_sockets_and_tasks()
|
|
elif not loop.is_closed():
|
|
loop.call_soon_threadsafe(close_sockets_and_tasks)
|
|
else:
|
|
# Loop has been closed, try to clean up directly.
|
|
del tasks
|
|
del close_sockets_and_tasks
|
|
close_sockets(sockets)
|
|
del self.output_queue_task
|
|
del self.stats_update_task
|
|
else:
|
|
# Sync case.
|
|
|
|
# ZMQ context termination can hang if the sockets
|
|
# aren't explicitly closed first.
|
|
close_sockets((self.output_socket, self.input_socket))
|
|
|
|
if self.shutdown_path is not None:
|
|
# We must ensure that the sync output socket is
|
|
# closed cleanly in its own thread.
|
|
with self.ctx.socket(zmq.PAIR) as shutdown_sender:
|
|
shutdown_sender.connect(self.shutdown_path)
|
|
# Send shutdown signal.
|
|
shutdown_sender.send(b"")
|
|
|
|
logger.debug_once("[shutdown] MPClient: background resource cleanup complete")
|
|
|
|
def validate_alive(self, frames: Sequence[zmq.Frame]):
|
|
if len(frames) == 1 and (frames[0].buffer == EngineCoreProc.ENGINE_CORE_DEAD):
|
|
self.engine_dead = True
|
|
raise EngineDeadError()
|
|
|
|
|
|
@dataclass
|
|
class ElasticScalingCache:
|
|
existing_core_engines: list[EngineIdentity]
|
|
num_new_core_engines: int
|
|
pending_notifications: dict[EEPNotificationType, set[int]]
|
|
|
|
|
|
class MPClient(EngineCoreClient):
|
|
"""
|
|
MPClient: base client for multi-proc EngineCore.
|
|
EngineCore runs in a background process busy loop, getting
|
|
new EngineCoreRequests and returning EngineCoreOutputs
|
|
|
|
* pushes EngineCoreRequests via input_socket
|
|
* pulls EngineCoreOutputs via output_socket
|
|
|
|
* AsyncMPClient subclass for AsyncLLM usage
|
|
* SyncMPClient subclass for LLM usage
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
asyncio_mode: bool,
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
client_addresses: dict[str, Any] | None = None,
|
|
renderer: BaseRenderer | None = None,
|
|
):
|
|
self.vllm_config = vllm_config
|
|
self._renderer: BaseRenderer | None = renderer
|
|
|
|
# ZMQ setup.
|
|
sync_ctx = zmq.Context(io_threads=2)
|
|
self.ctx = zmq.asyncio.Context(sync_ctx) if asyncio_mode else sync_ctx
|
|
|
|
# This will ensure resources created so far are closed
|
|
# when the client is garbage collected, even if an
|
|
# exception is raised mid-construction.
|
|
self.resources = BackgroundResources(ctx=sync_ctx)
|
|
self._finalizer = weakref.finalize(self, self.resources)
|
|
success = False
|
|
try:
|
|
# State used for data parallel.
|
|
self.engines_running = False
|
|
parallel_config = vllm_config.parallel_config
|
|
# Elastic EP can remove a rank and later add it back with the same
|
|
# identity. The client input ROUTER needs handover to allow the new
|
|
# engine to replace the dead connection.
|
|
enable_input_socket_handover = parallel_config.enable_elastic_ep
|
|
|
|
self.stats_update_address: str | None = None
|
|
tensor_queue: Queue | None = None
|
|
if client_addresses:
|
|
# Engines are managed externally to this client.
|
|
input_address = client_addresses["input_address"]
|
|
output_address = client_addresses["output_address"]
|
|
self.stats_update_address = client_addresses.get("stats_update_address")
|
|
# Tensor queues passed via client_addresses for multi-API-server case
|
|
tensor_queue = client_addresses.get("tensor_queue")
|
|
self.input_socket = self.resources.input_socket = make_zmq_socket(
|
|
self.ctx,
|
|
input_address,
|
|
zmq.ROUTER,
|
|
bind=True,
|
|
router_handover=enable_input_socket_handover,
|
|
)
|
|
self.resources.output_socket = make_zmq_socket(
|
|
self.ctx, output_address, zmq.PULL
|
|
)
|
|
|
|
# Report bound endpoints back so the parent can forward
|
|
# them to engines (mirrors the DPCoordinator pattern).
|
|
actual_address_pipe: Connection | None = client_addresses.get(
|
|
"actual_address_pipe"
|
|
)
|
|
if actual_address_pipe is not None:
|
|
try:
|
|
actual_input = self.input_socket.getsockopt(
|
|
zmq.LAST_ENDPOINT
|
|
).decode()
|
|
actual_output = self.resources.output_socket.getsockopt(
|
|
zmq.LAST_ENDPOINT
|
|
).decode()
|
|
actual_address_pipe.send(
|
|
{
|
|
"input_address": actual_input,
|
|
"output_address": actual_output,
|
|
}
|
|
)
|
|
finally:
|
|
actual_address_pipe.close()
|
|
# Engines are managed externally: this process does not fork
|
|
# them, so there is no fork race; start the MM warmup now.
|
|
self._start_mm_warmup()
|
|
else:
|
|
# Engines are managed by this client.
|
|
addresses = get_engine_zmq_addresses(vllm_config)
|
|
self.input_socket = self.resources.input_socket = make_zmq_socket(
|
|
self.ctx,
|
|
addresses.inputs[0],
|
|
zmq.ROUTER,
|
|
bind=True,
|
|
router_handover=enable_input_socket_handover,
|
|
)
|
|
self.resources.output_socket = make_zmq_socket(
|
|
self.ctx, addresses.outputs[0], zmq.PULL
|
|
)
|
|
|
|
# Resolve ``tcp://host:0`` placeholders to bound endpoints
|
|
# before engines DEALER-connect. No-op for IPC.
|
|
addresses.inputs[0] = self.input_socket.getsockopt(
|
|
zmq.LAST_ENDPOINT
|
|
).decode()
|
|
addresses.outputs[0] = self.resources.output_socket.getsockopt(
|
|
zmq.LAST_ENDPOINT
|
|
).decode()
|
|
|
|
with launch_core_engines(
|
|
vllm_config, executor_class, log_stats, addresses
|
|
) as engine_launch:
|
|
self.resources.coordinator = engine_launch.coordinator
|
|
self.resources.engine_manager = engine_launch.engine_manager
|
|
coordinator = engine_launch.coordinator
|
|
addresses = engine_launch.addresses
|
|
tensor_queue = engine_launch.tensor_queue
|
|
# Engine-core processes have now all been forked/started
|
|
# (CoreEngineProcManager.proc.start()). It is now safe to
|
|
# launch the frontend background MM warmup: it must not
|
|
# run while fork() is in flight (a live thread holding a
|
|
# lock would deadlock the forked child), but the
|
|
# engine-core model load (minutes) that follows is exactly
|
|
# what the warmup should overlap with.
|
|
self._start_mm_warmup()
|
|
|
|
self.stats_update_address = addresses.frontend_stats_publish_address
|
|
if coordinator is not None:
|
|
assert self.stats_update_address == (
|
|
coordinator.get_stats_publish_address()
|
|
)
|
|
|
|
# Serialization setup with tensor queues for multimodal tensor IPC.
|
|
tensor_ipc_sender: TensorIpcSender | None = None
|
|
model_config = getattr(vllm_config, "model_config", None)
|
|
if model_config is not None and model_config.multimodal_config is not None:
|
|
mm_tensor_ipc = model_config.multimodal_config.mm_tensor_ipc
|
|
if mm_tensor_ipc == "torch_shm" and tensor_queue is not None:
|
|
tensor_ipc_sender = TensorIpcSender(tensor_queue)
|
|
|
|
self.encoder = MsgpackEncoder(oob_tensor_consumer=tensor_ipc_sender)
|
|
self.decoder = MsgpackDecoder(EngineCoreOutputs)
|
|
|
|
dp_size = parallel_config.data_parallel_size
|
|
dp_rank = parallel_config.data_parallel_index
|
|
dp_local_size = parallel_config.data_parallel_size_local
|
|
offline_mode = parallel_config.data_parallel_rank_local is not None
|
|
# Client manages local+remote EngineCores in pure internal LB case.
|
|
# Client manages local EngineCores in hybrid and external LB case.
|
|
num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size
|
|
self.engine_ranks_managed = (
|
|
[dp_rank] if offline_mode else list(range(dp_rank, dp_rank + num_ranks))
|
|
)
|
|
assert parallel_config.data_parallel_size_local <= len(
|
|
self.engine_ranks_managed
|
|
)
|
|
|
|
# ZMQ identity of each engine that this client will talk to.
|
|
self.core_engines: list[EngineIdentity] = [
|
|
rank.to_bytes(2, "little") for rank in self.engine_ranks_managed
|
|
]
|
|
|
|
# Wait for ready messages from each engine on the input socket.
|
|
identities = set(self.core_engines)
|
|
sync_input_socket = zmq.Socket.shadow(self.input_socket)
|
|
while identities:
|
|
if not sync_input_socket.poll(
|
|
timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000 # convert to ms
|
|
):
|
|
raise TimeoutError(
|
|
f"Timed out waiting for engine core processes to "
|
|
f"start. This is often caused by slow weight loading "
|
|
f"for large models. Waited "
|
|
f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by "
|
|
f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the "
|
|
f"timeout, set the environment variable: "
|
|
f"VLLM_ENGINE_READY_TIMEOUT_S=<seconds>"
|
|
)
|
|
identity, payload = sync_input_socket.recv_multipart()
|
|
identities.remove(identity)
|
|
self._apply_ready_response(payload)
|
|
|
|
self.core_engine: EngineIdentity = self.core_engines[0]
|
|
self.utility_results: dict[int, AnyFuture] = {}
|
|
|
|
# Start monitoring engine core processes for unexpected failures
|
|
self.start_engine_core_monitor()
|
|
|
|
success = True
|
|
finally:
|
|
if not success:
|
|
self._finalizer()
|
|
|
|
def shutdown(self, timeout: float | None = None) -> None:
|
|
"""Shutdown engine manager under timeout and clean up resources."""
|
|
if self._finalizer.detach() is not None:
|
|
timeout_str = "default" if timeout is None else f"{timeout}s"
|
|
logger.info("[shutdown] MPClient: start timeout=%s", timeout_str)
|
|
if self.resources.engine_manager is not None:
|
|
logger.info_once("[shutdown] MPClient: stopping engine manager")
|
|
self.resources.engine_manager.shutdown(timeout=timeout)
|
|
logger.info_once("[shutdown] MPClient: engine manager stopped")
|
|
logger.info_once("[shutdown] MPClient: cleaning up background resources")
|
|
self.resources()
|
|
logger.info_once("[shutdown] MPClient: complete")
|
|
|
|
def _format_exception(self, e: Exception) -> Exception:
|
|
"""If errored, use EngineDeadError so root cause is clear."""
|
|
return (
|
|
EngineDeadError(suppress_context=True) if self.resources.engine_dead else e
|
|
)
|
|
|
|
def ensure_alive(self):
|
|
if self.resources.engine_dead:
|
|
raise EngineDeadError()
|
|
|
|
def dp_engines_running(self) -> bool:
|
|
return self.engines_running
|
|
|
|
def _start_mm_warmup(self) -> None:
|
|
# Called once the engine-core process(es) have been created (forked or
|
|
# externally managed). Overlap the frontend MM warmup with the
|
|
# engine-core model load. This is a no-op when no renderer was passed
|
|
# (e.g. text-only serving or tests).
|
|
if self._renderer is not None:
|
|
self._renderer.start_mm_warmup_in_background()
|
|
|
|
def start_engine_core_monitor(self):
|
|
"""Start a monitor thread for engine core processes."""
|
|
engine_manager = self.resources.engine_manager
|
|
if engine_manager is None:
|
|
# No engine processes to monitor
|
|
return
|
|
|
|
self_ref = weakref.ref(self)
|
|
|
|
# Monitor engine core process liveness. If any die unexpectedly,
|
|
# marks the engine as dead, and shuts down the client.
|
|
def monitor_engine_cores():
|
|
engine_manager.monitor_engine_liveness()
|
|
_self = self_ref()
|
|
if not _self or not _self._finalizer.alive or _self.resources.engine_dead:
|
|
return
|
|
_self.resources.engine_dead = True
|
|
logger.warning_once(
|
|
"[shutdown] MPClient: engine core exited unexpectedly; starting cleanup"
|
|
)
|
|
_self.shutdown()
|
|
# Note: For MPClient, we don't have a failure callback mechanism
|
|
# like MultiprocExecutor, but we set engine_dead flag which will
|
|
# cause subsequent operations to raise EngineDeadError
|
|
|
|
Thread(
|
|
target=monitor_engine_cores, daemon=True, name="MPClientEngineMonitor"
|
|
).start()
|
|
|
|
def _apply_ready_response(self, payload: bytes) -> None:
|
|
"""Decode an EngineCoreReadyResponse and sync any post-initialization
|
|
config changes (e.g. auto-fitted max_model_len) back to the frontend."""
|
|
if not payload:
|
|
return
|
|
vllm_config = self.vllm_config
|
|
response = msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse)
|
|
vllm_config.model_config.max_model_len = min(
|
|
vllm_config.model_config.max_model_len, response.max_model_len
|
|
)
|
|
|
|
# Setup KV cache config with initialization state from
|
|
# engine core process. Sum num_gpu_blocks from all engines in DP case.
|
|
num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks or 0
|
|
num_gpu_blocks += response.num_gpu_blocks
|
|
vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
|
|
|
|
# Sync block_size: may be enlarged by _align_hybrid_block_size in the
|
|
# worker for hybrid Mamba models.
|
|
cache_config = vllm_config.cache_config
|
|
cache_config.block_size = response.block_size
|
|
cache_config.mamba_block_size = response.mamba_block_size
|
|
# Keep these as per-engine cache_config_info values; do not sum across DP.
|
|
cache_config.kv_cache_size_tokens = (
|
|
getattr(cache_config, "kv_cache_size_tokens", None)
|
|
if getattr(cache_config, "kv_cache_size_tokens", None) is not None
|
|
else response.kv_cache_size_tokens
|
|
)
|
|
cache_config.kv_cache_max_concurrency = (
|
|
getattr(cache_config, "kv_cache_max_concurrency", None)
|
|
if getattr(cache_config, "kv_cache_max_concurrency", None) is not None
|
|
else response.kv_cache_max_concurrency
|
|
)
|
|
|
|
# In external DP LB mode, the coordinator address that the
|
|
# front-end procs connect to is obtained by each engine via it's
|
|
# initial handshake with the rank 0 front-end.
|
|
if response.dp_stats_address is not None:
|
|
if self.stats_update_address is None:
|
|
self.stats_update_address = response.dp_stats_address
|
|
else:
|
|
assert response.dp_stats_address == self.stats_update_address
|
|
|
|
|
|
def _process_utility_output(
|
|
output: UtilityOutput, utility_results: dict[int, AnyFuture]
|
|
):
|
|
"""Set the result from a utility method in the waiting future."""
|
|
future = utility_results.pop(output.call_id)
|
|
failure_message = output.failure_message
|
|
try:
|
|
if failure_message is not None:
|
|
future.set_exception(Exception(failure_message))
|
|
else:
|
|
assert output.result is not None
|
|
future.set_result(output.result.result)
|
|
except asyncio.InvalidStateError:
|
|
# This can happen if the future is cancelled due to the
|
|
# original calling task being cancelled.
|
|
if failure_message is not None:
|
|
logger.error(
|
|
"Cancelled call to utility method failed with error: %s",
|
|
failure_message,
|
|
)
|
|
|
|
|
|
class SyncMPClient(MPClient):
|
|
"""Synchronous client for multi-proc EngineCore."""
|
|
|
|
@instrument(span_name="SyncMPClient init")
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
renderer: BaseRenderer | None = None,
|
|
):
|
|
super().__init__(
|
|
asyncio_mode=False,
|
|
vllm_config=vllm_config,
|
|
executor_class=executor_class,
|
|
log_stats=log_stats,
|
|
renderer=renderer,
|
|
)
|
|
|
|
self.is_dp = self.vllm_config.parallel_config.data_parallel_size > 1
|
|
self.outputs_queue = queue.Queue[EngineCoreOutputs | Exception]()
|
|
|
|
# Ensure that the outputs socket processing thread does not have
|
|
# a ref to the client which prevents gc.
|
|
ctx = self.ctx
|
|
out_socket = self.resources.output_socket
|
|
decoder = self.decoder
|
|
utility_results = self.utility_results
|
|
outputs_queue = self.outputs_queue
|
|
|
|
shutdown_path = get_open_zmq_inproc_path()
|
|
resources = self.resources
|
|
resources.shutdown_path = shutdown_path
|
|
|
|
def process_outputs_socket():
|
|
assert isinstance(out_socket, zmq.Socket)
|
|
shutdown_socket = ctx.socket(zmq.PAIR)
|
|
try:
|
|
shutdown_socket.bind(shutdown_path)
|
|
poller = zmq.Poller()
|
|
poller.register(shutdown_socket, zmq.POLLIN)
|
|
poller.register(out_socket, zmq.POLLIN)
|
|
while True:
|
|
socks = poller.poll()
|
|
if not socks:
|
|
continue
|
|
if len(socks) == 2 or socks[0][0] == shutdown_socket:
|
|
# shutdown signal, exit thread.
|
|
break
|
|
|
|
frames = out_socket.recv_multipart(copy=False)
|
|
resources.validate_alive(frames)
|
|
outputs: EngineCoreOutputs = decoder.decode(frames)
|
|
if outputs.utility_output:
|
|
_process_utility_output(outputs.utility_output, utility_results)
|
|
else:
|
|
outputs_queue.put_nowait(outputs)
|
|
except Exception as e:
|
|
outputs_queue.put_nowait(e)
|
|
finally:
|
|
# Close sockets.
|
|
shutdown_socket.close(linger=0)
|
|
out_socket.close(linger=0)
|
|
|
|
# Process outputs from engine in separate thread.
|
|
self.output_queue_thread = Thread(
|
|
target=process_outputs_socket,
|
|
name="EngineCoreOutputQueueThread",
|
|
daemon=True,
|
|
)
|
|
self.output_queue_thread.start()
|
|
|
|
# The thread takes on responsibility for closing the socket.
|
|
self.resources.output_socket = None
|
|
|
|
def get_output(self) -> EngineCoreOutputs:
|
|
# If an exception arises in process_outputs_socket task,
|
|
# it is forwarded to the outputs_queue so we can raise it
|
|
# from this (run_output_handler) task to shut down the server.
|
|
outputs = self.outputs_queue.get()
|
|
|
|
if isinstance(outputs, Exception):
|
|
raise self._format_exception(outputs) from None
|
|
if outputs.wave_complete is not None:
|
|
self.engines_running = False
|
|
return outputs
|
|
|
|
def _send_input(self, request_type: EngineCoreRequestType, request: Any):
|
|
self.ensure_alive()
|
|
# (Identity, RequestType, SerializedRequest)
|
|
msg = (self.core_engine, request_type.value, *self.encoder.encode(request))
|
|
# Any zero-copy tensor/ndarray frames are kept alive by zmq itself
|
|
# until it's finished sending them (there is a ref chain from the underlying
|
|
# memoryview back to the original owning tensor/ndarray).
|
|
self.input_socket.send_multipart(msg, copy=False)
|
|
|
|
def call_utility(self, method: str, *args) -> Any:
|
|
call_id = uuid.uuid1().int >> 64
|
|
future: Future[Any] = Future()
|
|
self.utility_results[call_id] = future
|
|
self._send_input(EngineCoreRequestType.UTILITY, (0, call_id, method, args))
|
|
|
|
return future.result()
|
|
|
|
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
|
|
return self.call_utility("get_supported_tasks")
|
|
|
|
def add_request(self, request: EngineCoreRequest) -> None:
|
|
if self.is_dp:
|
|
self.engines_running = True
|
|
self._send_input(EngineCoreRequestType.ADD, request)
|
|
|
|
def abort_requests(self, request_ids: list[str]) -> None:
|
|
if request_ids and not self.resources.engine_dead:
|
|
self._send_input(EngineCoreRequestType.ABORT, request_ids)
|
|
|
|
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
|
|
self.call_utility("profile", is_start, profile_prefix)
|
|
|
|
def reset_mm_cache(self) -> None:
|
|
self.call_utility("reset_mm_cache")
|
|
|
|
def reset_prefix_cache(
|
|
self, reset_running_requests: bool = False, reset_connector: bool = False
|
|
) -> bool:
|
|
return self.call_utility(
|
|
"reset_prefix_cache", reset_running_requests, reset_connector
|
|
)
|
|
|
|
def reset_encoder_cache(self) -> None:
|
|
self.call_utility("reset_encoder_cache")
|
|
|
|
def add_lora(self, lora_request: LoRARequest) -> bool:
|
|
return self.call_utility("add_lora", lora_request)
|
|
|
|
def remove_lora(self, lora_id: int) -> bool:
|
|
return self.call_utility("remove_lora", lora_id)
|
|
|
|
def list_loras(self) -> set[int]:
|
|
return self.call_utility("list_loras")
|
|
|
|
def pin_lora(self, lora_id: int) -> bool:
|
|
return self.call_utility("pin_lora", lora_id)
|
|
|
|
def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
|
|
self.call_utility("sleep", level, mode)
|
|
|
|
def wake_up(self, tags: list[str] | None = None) -> None:
|
|
self.call_utility("wake_up", tags)
|
|
|
|
def is_sleeping(self) -> bool:
|
|
return self.call_utility("is_sleeping")
|
|
|
|
def execute_dummy_batch(self) -> None:
|
|
self.call_utility("execute_dummy_batch")
|
|
|
|
def set_weight_version(self, weight_version: str) -> None:
|
|
self.call_utility("set_weight_version", weight_version)
|
|
|
|
def get_weight_version(self) -> str:
|
|
return self.call_utility("get_weight_version")
|
|
|
|
def collective_rpc(
|
|
self,
|
|
method: str | Callable[..., _R],
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> list[_R]:
|
|
return self.call_utility("collective_rpc", method, timeout, args, kwargs)
|
|
|
|
def save_sharded_state(
|
|
self, path: str, pattern: str | None = None, max_size: int | None = None
|
|
) -> None:
|
|
self.call_utility("save_sharded_state", path, pattern, max_size)
|
|
|
|
|
|
class AsyncMPClient(MPClient):
|
|
"""Asyncio-compatible client for multi-proc EngineCore."""
|
|
|
|
@instrument(span_name="AsyncMPClient init")
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
client_addresses: dict[str, Any] | None = None,
|
|
client_count: int = 1,
|
|
client_index: int = 0,
|
|
renderer: BaseRenderer | None = None,
|
|
):
|
|
super().__init__(
|
|
asyncio_mode=True,
|
|
vllm_config=vllm_config,
|
|
executor_class=executor_class,
|
|
log_stats=log_stats,
|
|
client_addresses=client_addresses,
|
|
renderer=renderer,
|
|
)
|
|
|
|
self.client_count = client_count
|
|
self.client_index = client_index
|
|
self.outputs_queue = asyncio.Queue[EngineCoreOutputs | Exception]()
|
|
|
|
# locally-cached engine status
|
|
self._engine_status: dict[int, dict] = {}
|
|
if self.vllm_config.parallel_config.enable_fault_tolerance:
|
|
self._engine_status = {
|
|
rank: {"id": rank, "status": "healthy"}
|
|
for rank in self.engine_ranks_managed
|
|
}
|
|
try:
|
|
# If we are running in an asyncio event loop, start the queue task.
|
|
# Otherwise, it will be started lazily. If it is not started here,
|
|
# we could miss EXECUTOR_FAILED messages from engine core if they
|
|
# occur prior to any requests being sent.
|
|
asyncio.get_running_loop()
|
|
self._ensure_output_queue_task()
|
|
except RuntimeError:
|
|
pass
|
|
|
|
def _ensure_output_queue_task(self):
|
|
resources = self.resources
|
|
if resources.output_queue_task is not None:
|
|
return
|
|
|
|
# Perform IO in separate task to parallelize as much as possible.
|
|
# Avoid task having direct reference back to the client.
|
|
decoder = self.decoder
|
|
utility_results = self.utility_results
|
|
outputs_queue = self.outputs_queue
|
|
output_handler: (
|
|
Callable[[AsyncMPClient, EngineCoreOutputs], Awaitable[None]] | None
|
|
) = getattr(self.__class__, "process_engine_outputs", None)
|
|
_self_ref = weakref.ref(self)
|
|
output_socket = resources.output_socket
|
|
assert output_socket is not None
|
|
|
|
notification_callback_handler: (
|
|
Callable[[AsyncMPClient, Sequence[Any]], Any] | None
|
|
) = getattr(self.__class__, "eep_process_engine_core_notification", None)
|
|
|
|
async def process_outputs_socket():
|
|
try:
|
|
while True:
|
|
frames = await output_socket.recv_multipart(copy=False)
|
|
resources.validate_alive(frames)
|
|
outputs: EngineCoreOutputs = decoder.decode(frames)
|
|
if outputs.utility_output:
|
|
if (
|
|
outputs.utility_output.call_id == EEP_NOTIFICATION_CALL_ID
|
|
and notification_callback_handler is not None
|
|
):
|
|
assert _self_ref is not None
|
|
_self = _self_ref()
|
|
if not _self:
|
|
return
|
|
if outputs.utility_output.result is None:
|
|
continue
|
|
notification_data = outputs.utility_output.result.result
|
|
assert isinstance(notification_data, Sequence)
|
|
assert len(notification_data) == 2
|
|
asyncio.create_task(
|
|
notification_callback_handler(_self, notification_data)
|
|
)
|
|
elif outputs.utility_output.call_id == FT_STATUS_CALL_ID:
|
|
_self = _self_ref()
|
|
if not _self:
|
|
return
|
|
if outputs.utility_output.result is not None:
|
|
_self._engine_status[outputs.engine_index] = (
|
|
outputs.utility_output.result.result
|
|
)
|
|
else:
|
|
_process_utility_output(
|
|
outputs.utility_output, utility_results
|
|
)
|
|
continue
|
|
|
|
if output_handler is not None:
|
|
assert _self_ref is not None
|
|
_self = _self_ref()
|
|
if not _self:
|
|
# Client has been garbage collected, abort.
|
|
return
|
|
await output_handler(_self, outputs)
|
|
|
|
if outputs.outputs or outputs.scheduler_stats:
|
|
outputs_queue.put_nowait(outputs)
|
|
except Exception as e:
|
|
outputs_queue.put_nowait(e)
|
|
except asyncio.CancelledError:
|
|
outputs_queue.put_nowait(EngineDeadError())
|
|
|
|
resources.output_queue_task = asyncio.create_task(
|
|
process_outputs_socket(), name="EngineCoreOutputQueueTask"
|
|
)
|
|
|
|
async def get_output_async(self) -> EngineCoreOutputs:
|
|
self._ensure_output_queue_task()
|
|
# If an exception arises in process_outputs_socket task,
|
|
# it is forwarded to the outputs_queue so we can raise it
|
|
# from this (run_output_handler) task to shut down the server.
|
|
assert self.outputs_queue is not None
|
|
outputs = await self.outputs_queue.get()
|
|
if isinstance(outputs, Exception):
|
|
raise self._format_exception(outputs) from None
|
|
return outputs
|
|
|
|
def _send_input(
|
|
self,
|
|
request_type: EngineCoreRequestType,
|
|
request: Any,
|
|
engine: EngineIdentity | None = None,
|
|
) -> Awaitable[Any]:
|
|
if engine is None:
|
|
engine = self.core_engine
|
|
|
|
message = (request_type.value, *self.encoder.encode(request))
|
|
return self._send_input_message(message, engine)
|
|
|
|
def _send_input_message(
|
|
self, message: tuple[bytestr, ...], engine: EngineIdentity
|
|
) -> Awaitable[Any]:
|
|
self.ensure_alive()
|
|
# Any zero-copy tensor/ndarray frames are kept alive by zmq itself
|
|
# until it's finished sending them (there is a ref chain from the underlying
|
|
# memoryview back to the original owning tensor/ndarray).
|
|
return self.input_socket.send_multipart((engine,) + message, copy=False)
|
|
|
|
async def call_utility_async(self, method: str, *args) -> Any:
|
|
return await self._call_utility_async(method, *args, engine=self.core_engine)
|
|
|
|
async def _call_utility_async(
|
|
self, method: str, *args, engine: EngineIdentity
|
|
) -> Any:
|
|
call_id = uuid.uuid1().int >> 64
|
|
future = asyncio.get_running_loop().create_future()
|
|
self.utility_results[call_id] = future
|
|
message = (
|
|
EngineCoreRequestType.UTILITY.value,
|
|
*self.encoder.encode((self.client_index, call_id, method, args)),
|
|
)
|
|
await self._send_input_message(message, engine)
|
|
self._ensure_output_queue_task()
|
|
return await future
|
|
|
|
async def get_supported_tasks_async(self) -> tuple[SupportedTask, ...]:
|
|
return await self.call_utility_async("get_supported_tasks")
|
|
|
|
async def add_request_async(self, request: EngineCoreRequest) -> None:
|
|
request.client_index = self.client_index
|
|
await self._send_input(EngineCoreRequestType.ADD, request)
|
|
self._ensure_output_queue_task()
|
|
|
|
async def abort_requests_async(self, request_ids: list[str]) -> None:
|
|
if request_ids and not self.resources.engine_dead:
|
|
await self._send_input(EngineCoreRequestType.ABORT, request_ids)
|
|
|
|
async def pause_scheduler_async(
|
|
self, mode: PauseMode = "abort", clear_cache: bool = True
|
|
) -> None:
|
|
await self.call_utility_async("pause_scheduler", mode, clear_cache)
|
|
|
|
async def resume_scheduler_async(self) -> None:
|
|
await self.call_utility_async("resume_scheduler")
|
|
|
|
async def is_scheduler_paused_async(self) -> bool:
|
|
return await self.call_utility_async("is_scheduler_paused")
|
|
|
|
async def profile_async(
|
|
self, is_start: bool = True, profile_prefix: str | None = None
|
|
) -> None:
|
|
await self.call_utility_async("profile", is_start, profile_prefix)
|
|
|
|
async def reset_mm_cache_async(self) -> None:
|
|
await self.call_utility_async("reset_mm_cache")
|
|
|
|
async def reset_prefix_cache_async(
|
|
self, reset_running_requests: bool = False, reset_connector: bool = False
|
|
) -> bool:
|
|
return await self.call_utility_async(
|
|
"reset_prefix_cache", reset_running_requests, reset_connector
|
|
)
|
|
|
|
async def reset_encoder_cache_async(self) -> None:
|
|
await self.call_utility_async("reset_encoder_cache")
|
|
|
|
async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
|
|
await self.call_utility_async("sleep", level, mode)
|
|
|
|
async def wake_up_async(self, tags: list[str] | None = None) -> None:
|
|
await self.call_utility_async("wake_up", tags)
|
|
|
|
async def is_sleeping_async(self) -> bool:
|
|
return await self.call_utility_async("is_sleeping")
|
|
|
|
async def execute_dummy_batch_async(self) -> None:
|
|
await self.call_utility_async("execute_dummy_batch")
|
|
|
|
async def set_weight_version_async(self, weight_version: str) -> None:
|
|
await self.call_utility_async("set_weight_version", weight_version)
|
|
|
|
async def get_weight_version_async(self) -> str:
|
|
return await self.call_utility_async("get_weight_version")
|
|
|
|
async def add_lora_async(self, lora_request: LoRARequest) -> bool:
|
|
return await self.call_utility_async("add_lora", lora_request)
|
|
|
|
async def remove_lora_async(self, lora_id: int) -> bool:
|
|
return await self.call_utility_async("remove_lora", lora_id)
|
|
|
|
async def list_loras_async(self) -> set[int]:
|
|
return await self.call_utility_async("list_loras")
|
|
|
|
async def pin_lora_async(self, lora_id: int) -> bool:
|
|
return await self.call_utility_async("pin_lora", lora_id)
|
|
|
|
async def save_sharded_state_async(
|
|
self, path: str, pattern: str | None = None, max_size: int | None = None
|
|
) -> None:
|
|
await self.call_utility_async("save_sharded_state", path, pattern, max_size)
|
|
|
|
async def collective_rpc_async(
|
|
self,
|
|
method: str | Callable[..., _R],
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> list[_R]:
|
|
return await self.call_utility_async(
|
|
"collective_rpc", method, timeout, args, kwargs
|
|
)
|
|
|
|
async def handle_fault(
|
|
self, ft_request: FaultToleranceRequest
|
|
) -> FaultToleranceResult:
|
|
res = await self.call_utility_async(FT_UTILITY_METHOD, ft_request)
|
|
result = msgspec.convert(res, FaultToleranceResult)
|
|
if not result.success:
|
|
status = self._engine_status.get(self.engine_ranks_managed[0])
|
|
if status is not None:
|
|
status["last_ft_request_id"] = result.request_id
|
|
status["ft_error"] = result.reason
|
|
return result
|
|
|
|
async def get_status(self):
|
|
return {
|
|
"schema_version": 1,
|
|
"total_engines": len(self.engine_ranks_managed),
|
|
"engines": list(self._engine_status.values()),
|
|
}
|
|
|
|
|
|
class DPAsyncMPClient(AsyncMPClient):
|
|
"""Asyncio-compatible client for multi-proc, multi-engine (data parallel)
|
|
EngineCore. Assumes external load-balancing by default."""
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
client_addresses: dict[str, Any] | None = None,
|
|
client_count: int = 1,
|
|
client_index: int = 0,
|
|
renderer: BaseRenderer | None = None,
|
|
):
|
|
self.current_wave = 0
|
|
|
|
super().__init__(
|
|
vllm_config,
|
|
executor_class,
|
|
log_stats,
|
|
client_addresses,
|
|
client_count,
|
|
client_index,
|
|
renderer,
|
|
)
|
|
|
|
# List of [waiting, running, kv_cache_usage] per engine.
|
|
# Used only by DPLBAsyncMPClient subclass.
|
|
self.lb_engines: list[list[int | float]] = [
|
|
[0, 0, 0.0] for _ in self.core_engines
|
|
]
|
|
|
|
self.eep_scaling_cache: ElasticScalingCache | None = None
|
|
|
|
self.first_req_sock_addr = get_open_zmq_inproc_path()
|
|
self.first_req_send_socket = self.resources.first_req_send_socket = (
|
|
make_zmq_socket(self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=True)
|
|
)
|
|
try:
|
|
# If we are running in an asyncio event loop, start the stats task.
|
|
# Otherwise, it will be started lazily.
|
|
asyncio.get_running_loop()
|
|
self._ensure_stats_update_task()
|
|
except RuntimeError:
|
|
pass
|
|
|
|
def _ensure_stats_update_task(self):
|
|
resources = self.resources
|
|
if resources.stats_update_task is not None:
|
|
return
|
|
|
|
assert self.stats_update_address is not None
|
|
stats_addr: str = self.stats_update_address
|
|
assert len(self.engine_ranks_managed) > 0
|
|
|
|
async def run_engine_stats_update_task():
|
|
with (
|
|
make_zmq_socket(self.ctx, stats_addr, zmq.XSUB, linger=0) as socket,
|
|
make_zmq_socket(
|
|
self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=False, linger=0
|
|
) as first_req_rcv_socket,
|
|
):
|
|
assert isinstance(socket, zmq.asyncio.Socket)
|
|
assert isinstance(first_req_rcv_socket, zmq.asyncio.Socket)
|
|
self.resources.stats_update_socket = socket
|
|
self.resources.first_req_rcv_socket = first_req_rcv_socket
|
|
# Send subscription message.
|
|
await socket.send(b"\x01")
|
|
|
|
poller = zmq.asyncio.Poller()
|
|
poller.register(socket, zmq.POLLIN)
|
|
poller.register(first_req_rcv_socket, zmq.POLLIN)
|
|
|
|
while True:
|
|
events = await poller.poll()
|
|
if (
|
|
not self.engines_running
|
|
and len(events) == 2
|
|
or (events[0][0] == first_req_rcv_socket)
|
|
):
|
|
# Check if this is a regular request notification or
|
|
# scale up notification
|
|
buf = first_req_rcv_socket.recv(flags=zmq.NOBLOCK).result()
|
|
|
|
decoded = msgspec.msgpack.decode(buf)
|
|
if (
|
|
isinstance(decoded, (list, tuple))
|
|
and len(decoded) == 2
|
|
and decoded[0] == "SCALE_ELASTIC_EP"
|
|
):
|
|
# Extract new engine count from the decoded message
|
|
new_engine_count = decoded[1]
|
|
# Update engine_ranks_managed and count_slice
|
|
parallel_config = self.vllm_config.parallel_config
|
|
dp_size = parallel_config.data_parallel_size
|
|
dp_rank = parallel_config.data_parallel_rank
|
|
assert dp_rank == 0
|
|
assert dp_size == new_engine_count
|
|
assert not (
|
|
parallel_config.data_parallel_hybrid_lb
|
|
or parallel_config.data_parallel_external_lb
|
|
)
|
|
num_ranks = dp_size
|
|
self.engine_ranks_managed = list(
|
|
range(dp_rank, dp_rank + num_ranks)
|
|
)
|
|
if len(self.lb_engines) < new_engine_count:
|
|
self.lb_engines = self.lb_engines + [
|
|
[0, 0, 0.0]
|
|
for _ in range(
|
|
new_engine_count - len(self.lb_engines)
|
|
)
|
|
]
|
|
else:
|
|
self.lb_engines = self.lb_engines[:new_engine_count]
|
|
# Send scale up notification to coordinator
|
|
scale_msg = msgspec.msgpack.encode(
|
|
("SCALE_ELASTIC_EP", new_engine_count)
|
|
)
|
|
await socket.send(scale_msg)
|
|
continue
|
|
|
|
# we're sending a request while the engines are
|
|
# paused, so that it can wake the others up
|
|
# (to run dummy EP loop).
|
|
assert decoded[0] == "FIRST_REQ"
|
|
target_eng_index = decoded[1]
|
|
self.engines_running = True
|
|
msg = msgspec.msgpack.encode(
|
|
(target_eng_index, self.current_wave)
|
|
)
|
|
await socket.send(msg)
|
|
|
|
buf = None
|
|
while True:
|
|
# Drain all stats events (we only care about latest).
|
|
future: asyncio.Future[bytes] = socket.recv(flags=zmq.NOBLOCK)
|
|
if isinstance(future.exception(), zmq.Again):
|
|
break
|
|
buf = future.result()
|
|
if buf is None:
|
|
continue
|
|
|
|
# Update local load-balancing state.
|
|
counts, wave, running = msgspec.msgpack.decode(buf)
|
|
self.current_wave = wave
|
|
self.engines_running = running
|
|
if counts is not None:
|
|
# Running and waiting counts are global from the
|
|
# Coordinator including all EngineCores. Slice to get
|
|
# just the cores managed by this client.
|
|
ranks = self.engine_ranks_managed
|
|
count_slice = slice(ranks[0], ranks[-1] + 1)
|
|
sliced_counts = counts[count_slice]
|
|
self.lb_engines = sliced_counts
|
|
logger.debug(
|
|
"Received counts: %s (%s)", sliced_counts, count_slice
|
|
)
|
|
|
|
resources.stats_update_task = asyncio.create_task(
|
|
run_engine_stats_update_task()
|
|
)
|
|
|
|
async def add_request_async(self, request: EngineCoreRequest) -> None:
|
|
self._ensure_stats_update_task()
|
|
|
|
request.current_wave = self.current_wave
|
|
request.client_index = self.client_index
|
|
|
|
chosen_engine = self.get_core_engine_for_request(request)
|
|
to_await = self._send_input(EngineCoreRequestType.ADD, request, chosen_engine)
|
|
if not self.engines_running:
|
|
# Notify coordinator that we're sending a request
|
|
req_msg = msgspec.msgpack.encode(("FIRST_REQ", chosen_engine))
|
|
await self.first_req_send_socket.send(req_msg)
|
|
|
|
await to_await
|
|
|
|
self._ensure_output_queue_task()
|
|
|
|
def get_core_engine_for_request(self, request: EngineCoreRequest):
|
|
return self.core_engine
|
|
|
|
|
|
class DPLBAsyncMPClient(DPAsyncMPClient):
|
|
"""Asyncio-compatible client for multi-proc, multi-engine (data parallel)
|
|
EngineCore. Load-balances between multiple engine processes."""
|
|
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
executor_class: type[Executor],
|
|
log_stats: bool,
|
|
client_addresses: dict[str, Any] | None = None,
|
|
client_count: int = 1,
|
|
client_index: int = 0,
|
|
renderer: BaseRenderer | None = None,
|
|
):
|
|
self.client_count = client_count
|
|
|
|
# To route aborts to the correct engine.
|
|
self.reqs_in_flight: dict[str, EngineIdentity] = {}
|
|
|
|
# Exact per-engine count of this client's unfinished requests.
|
|
self.engine_inflight: Counter[EngineIdentity] = Counter()
|
|
|
|
super().__init__(
|
|
vllm_config,
|
|
executor_class,
|
|
log_stats,
|
|
client_addresses,
|
|
client_count,
|
|
client_index,
|
|
renderer,
|
|
)
|
|
|
|
assert len(self.core_engines) > 1
|
|
self._prepared_elastic_ep: tuple[int, int] | None = None
|
|
|
|
self.eng_start_index = (
|
|
len(self.core_engines) * self.client_index
|
|
) // client_count
|
|
|
|
def get_core_engine_for_request(self, request: EngineCoreRequest) -> EngineIdentity:
|
|
# Engines are in rank order.
|
|
if (eng_index := request.data_parallel_rank) is None and (
|
|
eng_index := get_late_interaction_engine_index(
|
|
request.pooling_params, len(self.core_engines)
|
|
)
|
|
) is None:
|
|
current_counts = self.lb_engines
|
|
# TODO use P2C alg for larger DP sizes
|
|
num_engines = len(current_counts)
|
|
min_score: float = sys.maxsize
|
|
eng_index = 0
|
|
for i in range(num_engines):
|
|
# Start from client_index to help with balancing when engines
|
|
# are empty.
|
|
idx = (self.eng_start_index + i) % num_engines
|
|
waiting, running, kv_cache_usage = current_counts[idx]
|
|
# Estimate engine load as the greater of the coordinator's
|
|
# latest (waiting + running) snapshot and this client's own
|
|
# in-flight count (scaled by the number of clients). The
|
|
# in-flight floor is exact and can't be erased by a snapshot
|
|
# rebind, so a burst spreads round-robin even when snapshots
|
|
# race with routing decisions; the snapshot raises the score
|
|
# when other clients or stale requests load the engine.
|
|
inflight = self.engine_inflight[self.core_engines[idx]]
|
|
score: float = max(self.client_count * inflight, waiting + running)
|
|
if waiting:
|
|
# Waiting requests are penalized in proportion to KV cache
|
|
# pressure: a queue on a KV-bound engine drains slowly, so
|
|
# new requests should strongly prefer other engines. With
|
|
# low KV usage the queue is transient (e.g. mid-burst) and
|
|
# the penalty stays off, preserving exact round-robin.
|
|
# Ramps from 0 at <=50% usage to 3x waiting at 100%.
|
|
score += waiting * 6.0 * max(0.0, kv_cache_usage - 0.5)
|
|
if score < min_score:
|
|
min_score = score
|
|
eng_index = idx
|
|
# Increment local waiting count for better balancing between stats
|
|
# updates from the coordinator (which happen every 100ms).
|
|
current_counts[eng_index][0] += self.client_count
|
|
# Rotate the scan start so that ties (equal scores, e.g. right
|
|
# after a coordinator stats reset when engines look equally loaded)
|
|
# don't systematically favor the same engine. This removes the
|
|
# fixed tie-break bias without affecting load-aware decisions when
|
|
# scores actually differ.
|
|
self.eng_start_index = (self.eng_start_index + 1) % num_engines
|
|
|
|
chosen_engine = self.core_engines[eng_index]
|
|
# Record which engine is chosen for this request, to handle aborts.
|
|
self.reqs_in_flight[request.request_id] = chosen_engine
|
|
self.engine_inflight[chosen_engine] += 1
|
|
return chosen_engine
|
|
|
|
async def call_utility_async(self, method: str, *args) -> Any:
|
|
# Only the result from the first engine is returned.
|
|
return (
|
|
await asyncio.gather(
|
|
*[
|
|
self._call_utility_async(method, *args, engine=engine)
|
|
for engine in self.core_engines
|
|
]
|
|
)
|
|
)[0]
|
|
|
|
@staticmethod
|
|
async def process_engine_outputs(
|
|
self: "DPLBAsyncMPClient", outputs: EngineCoreOutputs
|
|
):
|
|
if outputs.finished_requests and self.reqs_in_flight:
|
|
for req_id in outputs.finished_requests:
|
|
if (engine := self.reqs_in_flight.pop(req_id, None)) is not None:
|
|
self.engine_inflight[engine] -= 1
|
|
|
|
@staticmethod
|
|
async def eep_process_engine_core_notification(
|
|
self: "DPLBAsyncMPClient", notification_data: tuple[str, int]
|
|
):
|
|
cache = self.eep_scaling_cache
|
|
notification_type_str, dp_rank = notification_data
|
|
try:
|
|
notification_type = EEPNotificationType(notification_type_str)
|
|
except ValueError as e:
|
|
raise ValueError(
|
|
f"Unknown EEP notification type: {notification_type_str}"
|
|
) from e
|
|
|
|
if notification_type == EEPNotificationType.RECONFIGURE_FINISHED:
|
|
from vllm.v1.engine import UtilityResult
|
|
|
|
# NOTE(yongji): process a dummy UtilityOutput to resolve the future
|
|
# awaited in _eep_wait_for_setup_switch_complete(), signaling that
|
|
# all engine cores have completed reconfiguration.
|
|
dummy_output = UtilityOutput(
|
|
call_id=EEP_NOTIFICATION_CALL_ID, result=UtilityResult(None)
|
|
)
|
|
_process_utility_output(dummy_output, self.utility_results)
|
|
return
|
|
assert cache is not None
|
|
if notification_type not in cache.pending_notifications:
|
|
cache.pending_notifications[notification_type] = set()
|
|
if dp_rank in cache.pending_notifications[notification_type]:
|
|
raise ValueError(
|
|
f"Duplicate notification {notification_type} from dp_rank {dp_rank}"
|
|
)
|
|
cache.pending_notifications[notification_type].add(dp_rank)
|
|
if len(cache.pending_notifications[notification_type]) >= abs(
|
|
cache.num_new_core_engines
|
|
):
|
|
engine_manager = self.resources.engine_manager
|
|
assert isinstance(engine_manager, CoreEngineActorManager)
|
|
assert cache.num_new_core_engines < 0
|
|
old_dp_size = len(cache.existing_core_engines)
|
|
new_dp_size = old_dp_size + cache.num_new_core_engines
|
|
engine_manager.scale_down_elastic_ep(old_dp_size, new_dp_size)
|
|
self.vllm_config.parallel_config.data_parallel_size_local = len(
|
|
engine_manager.local_engine_actors
|
|
)
|
|
self.eep_scaling_cache = None
|
|
|
|
async def abort_requests_async(self, request_ids: list[str]) -> None:
|
|
if not request_ids or self.resources.engine_dead:
|
|
return
|
|
|
|
if len(request_ids) == 1:
|
|
# Fast-path common case.
|
|
if engine := self.reqs_in_flight.get(request_ids[0]):
|
|
await self._abort_requests(request_ids, engine)
|
|
return
|
|
|
|
by_engine = defaultdict[EngineIdentity, list[str]](list)
|
|
for req_id in request_ids:
|
|
if engine := self.reqs_in_flight.get(req_id):
|
|
by_engine[engine].append(req_id)
|
|
for engine, req_ids in by_engine.items():
|
|
await self._abort_requests(req_ids, engine)
|
|
|
|
async def _abort_requests(
|
|
self, request_ids: list[str], engine: EngineIdentity
|
|
) -> None:
|
|
await self._send_input(EngineCoreRequestType.ABORT, request_ids, engine)
|
|
|
|
async def commit_elastic_ep(self) -> None:
|
|
"""Commit prepared elastic EP scaling."""
|
|
prepared = self._prepared_elastic_ep
|
|
if prepared is None:
|
|
raise RuntimeError("Elastic EP scaling has not been prepared")
|
|
new_data_parallel_size, num_redundant_experts = prepared
|
|
cur_data_parallel_size = len(self.core_engines)
|
|
if new_data_parallel_size > cur_data_parallel_size:
|
|
await self._commit_scale_up_elastic_ep(new_data_parallel_size)
|
|
else:
|
|
await self._commit_scale_down_elastic_ep(new_data_parallel_size)
|
|
self.vllm_config.parallel_config.eplb_config.num_redundant_experts = (
|
|
num_redundant_experts
|
|
)
|
|
self._prepared_elastic_ep = None
|
|
|
|
async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None:
|
|
"""Prepare elastic EP scaling without routing requests to new engines."""
|
|
if (prepared := self._prepared_elastic_ep) is not None:
|
|
if prepared[0] == new_data_parallel_size:
|
|
return
|
|
raise RuntimeError("Elastic EP scaling is already prepared")
|
|
cur_data_parallel_size = len(self.core_engines)
|
|
assert self.vllm_config.parallel_config.data_parallel_backend == "ray", (
|
|
"Only ray DP backend supports scaling elastic EP"
|
|
)
|
|
parallel_config = self.vllm_config.parallel_config
|
|
num_experts = self.vllm_config.model_config.get_num_experts()
|
|
num_physical_experts = (
|
|
num_experts + parallel_config.eplb_config.num_redundant_experts
|
|
)
|
|
num_redundant_experts = (
|
|
num_physical_experts * new_data_parallel_size // cur_data_parallel_size
|
|
- num_experts
|
|
)
|
|
if num_redundant_experts < 0:
|
|
# Scaling keeps physical experts per engine fixed, so below this
|
|
# size the logical experts no longer fit.
|
|
raise ValueError(
|
|
f"Cannot scale to data_parallel_size {new_data_parallel_size}, "
|
|
f"minimum is "
|
|
f"{-(-num_experts * cur_data_parallel_size // num_physical_experts)}"
|
|
)
|
|
if new_data_parallel_size < cur_data_parallel_size:
|
|
await self._prepare_scale_down_elastic_ep(new_data_parallel_size)
|
|
else:
|
|
await self._prepare_scale_up_elastic_ep(
|
|
new_data_parallel_size, num_redundant_experts
|
|
)
|
|
self._prepared_elastic_ep = new_data_parallel_size, num_redundant_experts
|
|
|
|
def _eep_wait_for_setup_switch_complete(self) -> asyncio.Future:
|
|
"""
|
|
Wait for core engines to switch to the new setup.
|
|
|
|
In eep_process_engine_core_notification(), a dummy UtilityOutput with
|
|
EEP_NOTIFICATION_CALL_ID will be set when RECONFIGURE_FINISHED
|
|
notification is received from engine 0. We create a future with
|
|
that call_id and wait for it to be resolved.
|
|
"""
|
|
future = asyncio.get_running_loop().create_future()
|
|
self.utility_results[EEP_NOTIFICATION_CALL_ID] = future
|
|
self._ensure_output_queue_task()
|
|
return future
|
|
|
|
def _wait_for_new_engine_ready(self, new_core_engines: list[bytes]) -> None:
|
|
new_engine_identities = set(new_core_engines)
|
|
sync_input_socket = zmq.Socket.shadow(self.input_socket)
|
|
while new_engine_identities:
|
|
if not sync_input_socket.poll(timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000):
|
|
raise TimeoutError(
|
|
f"Timed out waiting for new engine core processes to "
|
|
f"start. Waited "
|
|
f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by "
|
|
f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the "
|
|
f"timeout, set the environment variable: "
|
|
f"VLLM_ENGINE_READY_TIMEOUT_S=<seconds>"
|
|
)
|
|
identity, payload = sync_input_socket.recv_multipart()
|
|
new_engine_identities.discard(identity)
|
|
self._apply_ready_response(payload)
|
|
|
|
def _setup_elastic_ep_reconfig_bootstrap(self) -> None:
|
|
from vllm.distributed.utils import create_tcp_store
|
|
from vllm.utils.network_utils import get_open_ports_list
|
|
|
|
parallel_config = self.vllm_config.parallel_config
|
|
parallel_config._data_parallel_master_port_list = get_open_ports_list(5)
|
|
parallel_config.data_parallel_master_port = (
|
|
parallel_config._data_parallel_master_port_list.pop()
|
|
)
|
|
|
|
ip = parallel_config.data_parallel_master_ip
|
|
store = create_tcp_store(
|
|
ip,
|
|
0,
|
|
is_master=True,
|
|
world_size=-1,
|
|
wait_for_workers=False,
|
|
)
|
|
parallel_config._coord_store_port = store.port
|
|
self._coord_store = store
|
|
|
|
def _make_reconfig_request(
|
|
self,
|
|
new_data_parallel_size: int,
|
|
rank_type: ReconfigureRankType = ReconfigureRankType.KEEP_CURRENT_RANK,
|
|
) -> ReconfigureDistributedRequest:
|
|
parallel_config = self.vllm_config.parallel_config
|
|
return ReconfigureDistributedRequest(
|
|
new_data_parallel_size=new_data_parallel_size,
|
|
new_data_parallel_rank=rank_type,
|
|
new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,
|
|
new_data_parallel_master_ip=parallel_config.data_parallel_master_ip,
|
|
new_data_parallel_master_port=parallel_config.data_parallel_master_port,
|
|
new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,
|
|
coord_store_port=parallel_config._coord_store_port,
|
|
)
|
|
|
|
async def _prepare_scale_up_elastic_ep(
|
|
self,
|
|
new_data_parallel_size: int,
|
|
num_redundant_experts: int,
|
|
) -> None:
|
|
"""Prepare scale up by creating new engine cores and reconfiguring
|
|
existing ones."""
|
|
self._setup_elastic_ep_reconfig_bootstrap()
|
|
|
|
# Phase 1: Send reconfig messages to existing engines
|
|
reconfig_futures = []
|
|
for engine in self.core_engines:
|
|
reconfig_request = self._make_reconfig_request(new_data_parallel_size)
|
|
coro = self._call_utility_async(
|
|
"reinitialize_distributed", reconfig_request, engine=engine
|
|
)
|
|
reconfig_futures.append(asyncio.create_task(coro))
|
|
|
|
# Phase 2: Create new engines
|
|
assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
|
|
start_new_worker_future = asyncio.to_thread(
|
|
self.resources.engine_manager.scale_up_elastic_ep,
|
|
self.vllm_config,
|
|
new_data_parallel_size,
|
|
num_redundant_experts,
|
|
)
|
|
|
|
# Phase 3: Wait for new engines to be created
|
|
# and reconfig messages to be received
|
|
await asyncio.gather(start_new_worker_future, *reconfig_futures)
|
|
ready_keys = [future.result() for future in reconfig_futures]
|
|
ready_keys.extend(
|
|
f"eep_ready/{rank}"
|
|
for rank in range(len(self.core_engines), new_data_parallel_size)
|
|
)
|
|
await asyncio.to_thread(self._coord_store.wait, ready_keys)
|
|
logger.info("[Elastic EP] Successfully started new engines")
|
|
|
|
async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None:
|
|
new_core_engines = [
|
|
rank.to_bytes(2, "little")
|
|
for rank in range(len(self.core_engines), new_data_parallel_size)
|
|
]
|
|
|
|
await self.pause_scheduler_async(mode="keep", clear_cache=False)
|
|
wait_future = self._eep_wait_for_setup_switch_complete()
|
|
finish_futures = [
|
|
asyncio.create_task(
|
|
self._call_utility_async("commit_prepared_elastic_ep", engine=engine)
|
|
)
|
|
for engine in self.core_engines
|
|
]
|
|
try:
|
|
await asyncio.gather(*finish_futures)
|
|
await wait_future
|
|
self._wait_for_new_engine_ready(new_core_engines)
|
|
except Exception:
|
|
wait_future.cancel()
|
|
raise
|
|
|
|
self.core_engines.extend(new_core_engines)
|
|
# Update the parallel config
|
|
parallel_config = self.vllm_config.parallel_config
|
|
parallel_config.data_parallel_size = new_data_parallel_size
|
|
if isinstance(self.resources.engine_manager, CoreEngineActorManager):
|
|
parallel_config.data_parallel_size_local = len(
|
|
self.resources.engine_manager.local_engine_actors
|
|
)
|
|
# Notify coordinator about scale up through existing
|
|
# stats_update_task connection
|
|
self._ensure_stats_update_task()
|
|
scale_up_marker = msgspec.msgpack.encode(
|
|
("SCALE_ELASTIC_EP", new_data_parallel_size)
|
|
)
|
|
await self.first_req_send_socket.send(scale_up_marker)
|
|
|
|
logger.info(
|
|
"[Elastic EP] Scale up completed, new data parallel size: %s",
|
|
new_data_parallel_size,
|
|
)
|
|
await self.resume_scheduler_async()
|
|
|
|
async def _prepare_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None:
|
|
self._setup_elastic_ep_reconfig_bootstrap()
|
|
|
|
reconfig_futures = []
|
|
for engine in self.core_engines[:new_data_parallel_size]:
|
|
reconfig_request = self._make_reconfig_request(new_data_parallel_size)
|
|
coro = self._call_utility_async(
|
|
"reinitialize_distributed", reconfig_request, engine=engine
|
|
)
|
|
reconfig_futures.append(asyncio.create_task(coro))
|
|
|
|
ready_keys = await asyncio.gather(*reconfig_futures)
|
|
await asyncio.to_thread(self._coord_store.wait, ready_keys)
|
|
|
|
async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None:
|
|
"""Scale down the data parallel size by shutting down and
|
|
reconfiguring existing engine cores."""
|
|
cur_data_parallel_size = len(self.core_engines)
|
|
|
|
self.eep_scaling_cache = ElasticScalingCache(
|
|
existing_core_engines=self.core_engines.copy(),
|
|
num_new_core_engines=new_data_parallel_size - cur_data_parallel_size,
|
|
pending_notifications=dict(),
|
|
)
|
|
|
|
old_core_engines = self.core_engines
|
|
# NOTE(yongji): Immediately stop sending requests to the removing engines.
|
|
self.core_engines = old_core_engines[:new_data_parallel_size]
|
|
self.lb_engines = self.lb_engines[:new_data_parallel_size]
|
|
removed_dp_size = cur_data_parallel_size - new_data_parallel_size
|
|
pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size
|
|
pause_futures = [
|
|
self._call_utility_async("pause_scheduler", mode, False, engine=engine)
|
|
for mode, engine in zip(pause_modes, old_core_engines)
|
|
]
|
|
await asyncio.gather(*pause_futures)
|
|
assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
|
|
self.resources.engine_manager.remove_run_refs_for_scale_down(removed_dp_size)
|
|
wait_future = self._eep_wait_for_setup_switch_complete()
|
|
reconfig_futures = []
|
|
for cur_dp_rank, engine in enumerate(old_core_engines):
|
|
if cur_dp_rank < new_data_parallel_size:
|
|
coro = self._call_utility_async(
|
|
"commit_prepared_elastic_ep", engine=engine
|
|
)
|
|
else:
|
|
reconfig_request = self._make_reconfig_request(
|
|
new_data_parallel_size,
|
|
ReconfigureRankType.SHUTDOWN_CURRENT_RANK,
|
|
)
|
|
coro = self._call_utility_async(
|
|
"reinitialize_distributed", reconfig_request, engine=engine
|
|
)
|
|
reconfig_futures.append(asyncio.create_task(coro))
|
|
|
|
try:
|
|
await asyncio.gather(*reconfig_futures)
|
|
|
|
self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
|
|
self._ensure_stats_update_task()
|
|
scale_down_marker = msgspec.msgpack.encode(
|
|
("SCALE_ELASTIC_EP", new_data_parallel_size)
|
|
)
|
|
await self.first_req_send_socket.send(scale_down_marker)
|
|
await wait_future
|
|
await self.resume_scheduler_async()
|
|
except Exception:
|
|
wait_future.cancel()
|
|
raise
|
|
|
|
logger.info(
|
|
"[Elastic EP] Scale down completed, new data parallel size: %s",
|
|
new_data_parallel_size,
|
|
)
|