diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 1dc6b5d3bb..d99aaeff7d 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -51,7 +51,7 @@ def _traffic_loop( stop: threading.Event, finished: threading.Event, is_probe: bool = False, -) -> list[tuple[float, float, int | None]]: +) -> list[tuple[float, float, int | None, str | None]]: url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions") payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4} headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)} @@ -60,15 +60,18 @@ def _traffic_loop( is_ready = False while not stop.is_set(): request_start = time.perf_counter() + phase = None try: response = requests.post( url, json=request_payload, headers=headers, timeout=120 ) status_code = response.status_code + if is_probe and status_code == 200: + phase = response.json()["phase"] except requests.exceptions.RequestException: status_code = None - responses.append((request_start, time.perf_counter(), status_code)) - if status_code == 200: + responses.append((request_start, time.perf_counter(), status_code, phase)) + if status_code == 200 and (not is_probe or phase in ("idle", "completed")): if not is_ready: ready.wait(timeout=120) is_ready = True @@ -78,14 +81,16 @@ def _traffic_loop( return responses -def _downtime(responses: list[tuple[float, float, int | None]]) -> float: - rejected = [end for _, end, status in responses if status == 503] - if not rejected: +def _downtime(responses: list[tuple[float, float, int | None, str | None]]) -> float: + committing = [end for _, end, _, phase in responses if phase == "committing"] + if not committing: return 0 recovered = next( - end for _, end, status in responses if status == 200 and end > rejected[-1] + end + for _, end, status, phase in responses + if status == 200 and phase in ("idle", "completed") and end > committing[-1] ) - return recovered - rejected[0] + return recovered - committing[0] def _scale_with_traffic( @@ -124,16 +129,20 @@ def _scale_with_traffic( bad_statuses = { status for responses in [probe_result, *results] - for _, _, status in responses + for _, _, status, _ in responses if status not in (200, 503) } assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}" - probe_503 = [start for start, _, status in probe_result if status == 503] - assert probe_503, "Scaling probe did not observe commit" + probe_committing = [ + start for start, _, _, phase in probe_result if phase == "committing" + ] + assert probe_committing, "Scaling probe did not observe commit" assert not results or any( - status == 200 and start_time <= request_start and request_end < probe_503[0] + status == 200 + and start_time <= request_start + and request_end < probe_committing[0] for responses in results - for request_start, request_end, status in responses + for request_start, request_end, status, _ in responses ), "No request completed successfully during preparation" print( diff --git a/tests/v1/distributed/test_external_lb_dp.py b/tests/v1/distributed/test_external_lb_dp.py index b53d0d8af0..1afb134ef5 100644 --- a/tests/v1/distributed/test_external_lb_dp.py +++ b/tests/v1/distributed/test_external_lb_dp.py @@ -4,17 +4,28 @@ import asyncio import os import threading import time +from concurrent.futures import Future, ThreadPoolExecutor from contextlib import AsyncExitStack +from types import SimpleNamespace +from typing import TYPE_CHECKING, cast +import msgspec import openai # use the official client for correctness check import pytest import pytest_asyncio import requests from tests.utils import RemoteOpenAIServer +from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleCoordinator, +) from vllm.platforms import current_platform +if TYPE_CHECKING: + from vllm.v1.engine.core_client import DPAsyncMPClient + MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b") +ELASTIC_EP_MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat" # Number of data parallel ranks for external LB testing DP_SIZE = int(os.getenv("DP_SIZE", "2")) @@ -22,6 +33,111 @@ DP_SIZE = int(os.getenv("DP_SIZE", "2")) TP_SIZE = int(os.getenv("TP_SIZE", "1")) +class DictStore: + def __init__(self, values: dict[str, bytes]): + self.values = values + + def check(self, keys: list[str]) -> bool: + return all(key in self.values for key in keys) + + def get(self, key: str) -> bytes: + return self.values[key] + + def set(self, key: str, value: bytes) -> None: + self.values[key] = value + + +def test_external_elastic_ep_visible_epoch_has_bootstrap_and_phase(monkeypatch): + client = SimpleNamespace( + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_master_port=12345, + _data_parallel_master_port_list=[12346, 12347], + ) + ) + ) + coordinator = ExternalElasticEPScaleCoordinator(cast("DPAsyncMPClient", client)) + control_store, reconfig_store = DictStore({}), DictStore({}) + coordinator.reconfig_store_ref = reconfig_store + monkeypatch.setattr( + coordinator, "_setup_reconfig_bootstrap", lambda: ("127.0.0.1", 12348) + ) + epoch, bootstrap = coordinator._prepare_reconfig_bootstrap(control_store, 2, 3) + + for store in (control_store, reconfig_store): + # DictStore preserves insertion order: metadata must precede the epoch. + assert list(store.values)[-1] == coordinator.key("current_epoch") + assert store.get(coordinator.key(epoch, "bootstrap")) == msgspec.msgpack.encode( + bootstrap + ) + assert store.get(coordinator.key(epoch, "phase")) == b"preparing" + assert store.get(coordinator.key("current_epoch")) == epoch.encode() + monkeypatch.setattr( + coordinator, "_get_reconfig_store", lambda store=store: store + ) + assert coordinator.get_status()["requested_data_parallel_size"] == 3 + + +def test_external_elastic_ep_calculates_target_expert_redundancy(): + client = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(get_num_experts=lambda: 128), + parallel_config=SimpleNamespace( + eplb_config=SimpleNamespace(num_redundant_experts=0) + ), + ) + ) + coordinator = ExternalElasticEPScaleCoordinator(cast("DPAsyncMPClient", client)) + + assert coordinator._calculate_num_redundant_experts(2, 3) == 64 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("dp_rank", range(4)) +async def test_validate_expert_capacity_on_every_rank(dp_rank, monkeypatch): + """Validate expert capacity on every rank before waiting for an epoch.""" + client = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(get_num_experts=lambda: 128), + parallel_config=SimpleNamespace( + data_parallel_rank=dp_rank, + elastic_ep_max_dp_size=4, + data_parallel_master_ip="127.0.0.1", + _coord_store_port=12345, + eplb_config=SimpleNamespace(num_redundant_experts=0), + ), + ) + ) + coordinator = ExternalElasticEPScaleCoordinator(cast("DPAsyncMPClient", client)) + store = DictStore({}) + monkeypatch.setattr( + "vllm.distributed.utils.get_cached_tcp_store_client", lambda *args: store + ) + + with pytest.raises(ValueError, match="cannot retain all model experts"): + await asyncio.wait_for(coordinator.prepare(4, 2), timeout=1) + + assert not store.values + assert coordinator.prepared_scale is None + + +@pytest.mark.asyncio +async def test_external_elastic_ep_late_rank_observes_epoch_error(): + coordinator = ExternalElasticEPScaleCoordinator( + cast("DPAsyncMPClient", SimpleNamespace()) + ) + epoch = "failed-epoch" + store = DictStore( + { + coordinator.key("current_epoch"): epoch.encode(), + coordinator.key(epoch, "error"): b"scale failed", + } + ) + + with pytest.raises(RuntimeError, match="scale failed"): + await coordinator._wait_for_bootstrap(store, requested_new_dp_size=3) + + class ExternalLBServerManager: """Manages data parallel vLLM server instances for external load balancer testing.""" @@ -42,57 +158,69 @@ class ExternalLBServerManager: self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = [] self.server_threads: list[threading.Thread] = [] + def start_rank( + self, + rank: int, + dp_size: int | None = None, + elastic_ep_scale_up_launch: bool = False, + ) -> tuple[RemoteOpenAIServer, list[str]]: + """Start one externally managed DP rank.""" + dp_size = self.dp_size if dp_size is None else dp_size + server_args = self.base_server_args.copy() + server_args.extend( + [ + "--data-parallel-size", + str(dp_size), + "--data-parallel-rank", + str(rank), + "--data-parallel-size-local", + "1", + "--tensor-parallel-size", + str(self.tp_size), + "--port", + str(8000 + rank), + "--api-server-count", + str(self.api_server_count), + ] + ) + + env_dict = { + "VLLM_SERVER_DEV_MODE": "1", + current_platform.device_control_env_var: ",".join( + str(current_platform.device_id_to_physical_device_id(i)) + for i in range(rank * self.tp_size, (rank + 1) * self.tp_size) + ), + } + if elastic_ep_scale_up_launch: + env_dict["VLLM_ELASTIC_EP_SCALE_UP_LAUNCH"] = "1" + + server = RemoteOpenAIServer( + self.model_name, + server_args, + auto_port=False, + env_dict=env_dict, + ) + print( + f"Server rank {rank} started successfully with " + f"{self.api_server_count} API servers" + ) + server_and_args = (server, server_args) + self.servers.append(server_and_args) + return server_and_args + def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]: - """Start all server instances for external LB mode.""" + """Start all initial server instances for external LB mode.""" + start_errors: list[Exception] = [] for rank in range(self.dp_size): - # Create server args for this specific rank - server_args = self.base_server_args.copy() - - # Add external LB specific arguments - server_args.extend( - [ - "--data-parallel-size", - str(self.dp_size), - "--data-parallel-rank", - str(rank), - "--data-parallel-size-local", - "1", - "--tensor-parallel-size", - str(self.tp_size), - "--port", - str(8000 + rank), # Different port for each rank - "--api-server-count", - str(self.api_server_count), - ] - ) - # Use a thread to start each server to allow parallel initialization - def start_server(r: int, sargs: list[str]): + def start_server(r: int): try: - # Start the server - server = RemoteOpenAIServer( - self.model_name, - sargs, - auto_port=False, - env_dict={ - "VLLM_SERVER_DEV_MODE": "1", - current_platform.device_control_env_var: ",".join( - str(current_platform.device_id_to_physical_device_id(i)) - for i in range(r * TP_SIZE, (r + 1) * TP_SIZE) - ), - }, - ) - server.__enter__() - print( - f"Server rank {r} started successfully with " - f"{self.api_server_count} API servers" - ) - self.servers.append((server, sargs)) + self.start_rank(r) except Exception as e: print(f"Failed to start server rank {r}: {e}") - raise + start_errors.append(e) - thread = threading.Thread(target=start_server, args=(rank, server_args)) + thread = threading.Thread(target=start_server, args=(rank,)) thread.start() self.server_threads.append(thread) @@ -101,6 +229,10 @@ class ExternalLBServerManager: for thread in self.server_threads: thread.join() + if start_errors: + start_error = start_errors[0] + raise RuntimeError("Failed to start external LB servers") from start_error + # Give servers additional time to fully initialize and coordinate time.sleep(2) @@ -184,12 +316,104 @@ def _assert_expert_parallel_config( assert parallel_config["enable_expert_parallel"] is expected +def _send_scale_command( + server: RemoteOpenAIServer, new_dp_size: int +) -> requests.Response: + return requests.post( + server.url_for("scale_elastic_ep"), + json={"new_data_parallel_size": new_dp_size}, + timeout=600, + ) + + +def _wait_for_scale_requests_to_start( + scale_requests: list[Future[requests.Response]], + timeout: float = 30, +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for request in scale_requests: + if request.done(): + response = request.result() + pytest.fail( + "Scale request finished before the new rank was launched: " + f"{response.status_code} {response.text}" + ) + + if all(request.running() for request in scale_requests): + return + time.sleep(0.1) + + pytest.fail("Timed out waiting for all scale requests to start") + + +@pytest.mark.distributed(num_gpus=3) +@pytest.mark.skipif( + TP_SIZE != 1 or current_platform.device_count() < 3, + reason="External Elastic EP scale-up requires three GPUs with TP_SIZE=1", +) +@pytest.mark.skipif( + os.getenv("VLLM_USE_RUST_FRONTEND") == "1", + reason="Elastic EP scaling is not supported by the Rust frontend", +) +def test_external_lb_elastic_ep_scale_up(default_server_args) -> None: + server_args = [ + *default_server_args, + "--enable-expert-parallel", + "--all2all-backend", + "allgather_reducescatter", + "--attention-backend", + "TRITON_MLA", + "--enable-elastic-ep", + "--elastic-ep-max-dp-size", + "3", + "--enable-eplb", + "--eplb-config", + '{"communicator":"pynccl","num_redundant_experts":0,"use_async":false}', + ] + server_manager = ExternalLBServerManager( + ELASTIC_EP_MODEL_NAME, + dp_size=2, + api_server_count=1, + base_server_args=server_args, + ) + + with server_manager: + existing_servers = [server for server, _ in server_manager.servers] + with ThreadPoolExecutor(max_workers=2) as executor: + scale_requests = [ + executor.submit(_send_scale_command, server, 3) + for server in existing_servers + ] + _wait_for_scale_requests_to_start(scale_requests) + server_manager.start_rank( + rank=2, + dp_size=3, + elastic_ep_scale_up_launch=True, + ) + + for request in scale_requests: + response = request.result(timeout=600) + assert response.status_code == 200, response.text + + for server, _ in server_manager.servers: + parallel_config = _get_parallel_config(server) + assert parallel_config["data_parallel_size"] == 3 + assert parallel_config["eplb_config"]["num_redundant_experts"] == 32 + with server.get_client() as client: + completion = client.completions.create( + model=ELASTIC_EP_MODEL_NAME, + prompt="Hello, my name is", + max_tokens=5, + temperature=0.0, + ) + assert completion.choices[0].finish_reason in ("length", "stop") + + def test_external_lb_server_info(server_manager): servers = server_manager.servers api_server_count = server_manager.api_server_count - _assert_expert_parallel_config(servers) - for i, (server, _) in enumerate(servers): print(f"Testing {i=}") diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index 1b23356134..f8e2373d4c 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -141,7 +141,7 @@ def _make_addresses() -> EngineZmqAddresses: False, ), ({"data_parallel_backend": "mp", "data_parallel_rank_local": 0}, False), - ({"data_parallel_backend": "mp", "enable_elastic_ep": True}, False), + ({"data_parallel_backend": "mp", "enable_elastic_ep": True}, True), ({"data_parallel_backend": "mp", "data_parallel_size": 1}, False), ], ids=["mp", "ray", "non-master-node", "offline", "elastic-ep", "single-engine"], @@ -150,8 +150,9 @@ def test_coordination_store_held_only_by_the_online_dp_master( monkeypatch: pytest.MonkeyPatch, overrides: dict[str, Any], expect_store: bool ) -> None: """The DP master node holds a coordination store for both engine backends, - so engines there pick world-group ports at bind time. Every other - deployment keeps the pre-allocated ports and must not hold one.""" + so engines there pick world-group ports at bind time. Elastic EP also + creates its coordination store before startup. The remaining cases + below keep the pre-allocated ports and must not hold a store.""" class FakeManager: def __init__(self, **kwargs) -> None: diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 10e9f1e92f..b58aed8547 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -142,6 +142,9 @@ def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch): def setsockopt(self, *_args, **_kwargs): pass + def getsockopt_string(self, _option): + return "inproc://dummy" + monkeypatch.setattr(core_client_mod.zmq.Socket, "shadow", lambda *_: ShadowSocket()) monkeypatch.setattr( core_client_mod, "make_zmq_socket", lambda *_, **__: DummySocket() @@ -448,7 +451,7 @@ def test_dplb_streaming_session_abort_goes_to_all_engines(final_sent: bool): ("effective_size", "other_size"), [(None, None), (4224, 4224), (4224, 1056), (4224, None)], ) -def test_apply_ready_response_syncs_block_size(effective_size, other_size): +def test_apply_ready_response_syncs_engine_metadata(effective_size, other_size): import msgspec client = object.__new__(MPClient) @@ -456,8 +459,11 @@ def test_apply_ready_response_syncs_block_size(effective_size, other_size): client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), + parallel_config=SimpleNamespace(_coord_store_port=0), ) client.stats_update_address = None + client.coordinator_input_address = None + client.coordinator_output_address = None payload = msgspec.msgpack.encode( EngineCoreReadyResponse( @@ -478,6 +484,9 @@ def test_apply_ready_response_syncs_block_size(effective_size, other_size): instance_id="test-instance", supports_lora=False, max_loras=0, + coord_store_port=1234, + coordinator_input_address="tcp://127.0.0.1:1235", + coordinator_output_address="tcp://127.0.0.1:1236", ) ) fields = msgspec.msgpack.decode(payload) @@ -498,6 +507,9 @@ def test_apply_ready_response_syncs_block_size(effective_size, other_size): client._apply_ready_response(b"") assert cache_config.effective_attention_block_size is None + assert client.vllm_config.parallel_config._coord_store_port == 1234 + assert client.coordinator_input_address == "tcp://127.0.0.1:1235" + assert client.coordinator_output_address == "tcp://127.0.0.1:1236" def test_apply_ready_response_syncs_mamba_block_size(): @@ -508,6 +520,7 @@ def test_apply_ready_response_syncs_mamba_block_size(): client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), + parallel_config=SimpleNamespace(_coord_store_port=0), ) client.stats_update_address = None @@ -1578,6 +1591,7 @@ def test_apply_ready_response_retains_kv_event_sources(): client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), + parallel_config=SimpleNamespace(_coord_store_port=0), ) def ready(rank: int, config: KVEventsConfig | None) -> bytes: diff --git a/tests/v1/engine/test_startup_watch_processes.py b/tests/v1/engine/test_startup_watch_processes.py index 19b2d7bca8..bf325a6108 100644 --- a/tests/v1/engine/test_startup_watch_processes.py +++ b/tests/v1/engine/test_startup_watch_processes.py @@ -281,6 +281,7 @@ class _FinishedProcess: def test_wait_for_engine_startup_reports_watched_process_exit(): + """Update the startup test argument so it reaches frontend failure handling.""" ctx = zmq.Context() handshake_socket = ctx.socket(zmq.ROUTER) recv, send = connection.Pipe(duplex=False) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6ea95142c1..69677ee9d6 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -980,11 +980,11 @@ class ParallelConfig: "Elastic EP is not supported with pipeline parallelism " f"(pipeline_parallel_size={self.pipeline_parallel_size})." ) - if self.data_parallel_external_lb or self.data_parallel_hybrid_lb: + if self.data_parallel_hybrid_lb: raise NotImplementedError( - "Elastic EP is not compatible with data_parallel_external_lb " - "or data_parallel_hybrid_lb. Elastic EP relies on a single API " - "server and core client to coordinate scale up/down." + "Elastic EP is not compatible with data_parallel_hybrid_lb. " + "Elastic EP supports a single API server/core client path " + "or data_parallel_external_lb, but not hybrid load balancing." ) if self.eplb_config.use_async: from vllm.distributed.nixl_utils import is_nixl_available diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 5ba1891b7a..f9869b4da4 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -11,6 +11,9 @@ from vllm.config import ParallelConfig from vllm.distributed import ( stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleCoordinator, +) from vllm.distributed.utils import get_cached_tcp_store_client from vllm.logger import init_logger from vllm.v1.engine import ( @@ -256,10 +259,23 @@ class ElasticEPScalingState: if state == ScaleDownRemovingEngineState.PREPARE: assert self.old_dp_group.rank() > 0 self._commit_scale_down(removing=True) - self.state = ScaleDownRemovingEngineState.COMPLETE + external_lb = self.vllm_config.parallel_config.data_parallel_external_lb + if external_lb: + # Preserve completion delivery: client cleanup on exit must not + # cancel the receiver before it records the notification. + parallel_config = self.new_parallel_config + store = get_cached_tcp_store_client( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ) + key = ExternalElasticEPScaleCoordinator.key + epoch = store.get(key("current_epoch")).decode() self.engine_core._eep_send_engine_core_notification( EEPNotificationType.SHUTDOWN_COMPLETE ) + if external_lb: + store.wait([key(epoch, "shutdown_complete", self.engine_core.dp_rank)]) + self.state = ScaleDownRemovingEngineState.COMPLETE return True else: @@ -384,7 +400,10 @@ class ElasticEPScalingState: def _send_reconfigure_finished(self): assert self.new_dp_group is not None - if self.new_dp_group.rank() == 0: + if ( + self.new_dp_group.rank() == 0 + or self.vllm_config.parallel_config.data_parallel_external_lb + ): self.engine_core._eep_send_engine_core_notification( EEPNotificationType.RECONFIGURE_FINISHED ) diff --git a/vllm/distributed/elastic_ep/external_elastic_ep.py b/vllm/distributed/elastic_ep/external_elastic_ep.py new file mode 100644 index 0000000000..6db6e5f4b1 --- /dev/null +++ b/vllm/distributed/elastic_ep/external_elastic_ep.py @@ -0,0 +1,897 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import contextlib +import enum +import uuid +from dataclasses import dataclass +from threading import Event, Thread +from typing import TYPE_CHECKING, Any + +import msgspec +import msgspec.msgpack +import zmq + +from vllm.logger import init_logger +from vllm.utils.network_utils import make_zmq_socket +from vllm.v1.engine import ( + EEPNotificationType, + ReconfigureDistributedRequest, + ReconfigureRankType, +) +from vllm.v1.engine.utils import EngineHandshakeMetadata, EngineZmqAddresses +from vllm.v1.utils import get_engine_client_zmq_addr + +if TYPE_CHECKING: + from vllm.v1.engine.core_client import DPAsyncMPClient + +logger = init_logger(__name__) +POLL_BACKOFF_STEPS_S = (0.1, 0.2, 0.4, 0.8) + + +class ExternalElasticEPScalePhase(str, enum.Enum): + """Lifecycle phases shared for an epoch-scoped external EEP operation. + + These values are stored in the control and reconfiguration stores and + exposed through the status API so external orchestrators can observe the + same operation across API processes and reconnects. COMPLETED means the + topology commit and rank completion barrier have finished; API processes + may clear their process-local request gate shortly afterward. + """ + + IDLE = "idle" + PREPARING = "preparing" + COMMITTING = "committing" + FAILED = "failed" + COMPLETED = "completed" + + +@dataclass +class ExternalElasticEPScaleUpHandshakeMetadata: + engine_metadata: EngineHandshakeMetadata + num_redundant_experts: int + elastic_ep_max_dp_size: int + + +class ExternalElasticEPScaleUpHandshakeServer: + """Temporary rank-0 handshake server for external EEP scale-up. + + During normal external startup the global handshake listener only exists + while the rank starts. Scale-up needs the same handshake contract again for + newly launched ranks, so rank 0 re-opens a temporary listener for the + duration of the current scale operation. + + The handshake gives new ranks the information they cannot infer locally: + front-end and DP coordinator ZMQ addresses, the new DP master address/ports, + target DP size, target expert redundancy, and the coordination store port + used for EEP reconfiguration. + """ + + def __init__( + self, + *, + handshake_address: str, + expected_new_ranks: list[int], + addresses: EngineZmqAddresses, + bootstrap: ReconfigureDistributedRequest, + num_redundant_experts: int, + elastic_ep_max_dp_size: int, + ) -> None: + self.handshake_address = handshake_address + self.expected_new_ranks = set(expected_new_ranks) + self.addresses = addresses + self.bootstrap = bootstrap + self.num_redundant_experts = num_redundant_experts + self.elastic_ep_max_dp_size = elastic_ep_max_dp_size + self.started_event = Event() + self._stop_event = Event() + self._thread = Thread( + target=self._run, + name="ExternalElasticEPHandshakeServer", + daemon=True, + ) + self._error: Exception | None = None + + def start(self) -> None: + self._thread.start() + started = self.started_event.wait(timeout=5) + if self._error is not None: + raise self._error + if not started: + raise TimeoutError( + "Timed out waiting for external EEP handshake server to " + f"start listening at {self.handshake_address}" + ) + + def stop(self) -> None: + self._stop_event.set() + if self._thread.is_alive(): + self._thread.join(timeout=5) + if self._error is not None: + raise self._error + + def raise_if_failed(self) -> None: + if self._error is not None: + raise self._error + + def _run(self) -> None: + pending_ready_ranks = self.expected_new_ranks.copy() + try: + with ( + zmq.Context() as ctx, + make_zmq_socket( + ctx, + self.handshake_address, + zmq.ROUTER, + bind=True, + linger=0, + router_handover=True, + ) as handshake_socket, + ): + self.started_event.set() + poller = zmq.Poller() + poller.register(handshake_socket, zmq.POLLIN) + + while not self._stop_event.is_set(): + if not pending_ready_ranks: + return + + events = dict(poller.poll(timeout=1000)) + if handshake_socket not in events: + continue + + engine_identity, payload = handshake_socket.recv_multipart() + engine_rank = int.from_bytes(engine_identity, "little") + if engine_rank not in self.expected_new_ranks: + raise RuntimeError( + "Received scale-up handshake from unexpected " + f"dp rank {engine_rank}" + ) + + message = msgspec.msgpack.decode(payload) + status = message["status"] + if status == "HELLO": + b = self.bootstrap + parallel_config: dict[str, int | str | list[int]] = { + "data_parallel_master_ip": b.new_data_parallel_master_ip, + "data_parallel_master_port": ( + b.new_data_parallel_master_port + ), + "_data_parallel_master_port_list": ( + b.new_data_parallel_master_port_list + ), + "data_parallel_size": b.new_data_parallel_size, + "_coord_store_port": b.coord_store_port, + } + init_message = msgspec.msgpack.encode( + ExternalElasticEPScaleUpHandshakeMetadata( + engine_metadata=EngineHandshakeMetadata( + addresses=self.addresses, + parallel_config=parallel_config, + ), + num_redundant_experts=self.num_redundant_experts, + elastic_ep_max_dp_size=self.elastic_ep_max_dp_size, + ) + ) + handshake_socket.send_multipart( + (engine_identity, init_message), copy=False + ) + elif status == "READY": + pending_ready_ranks.discard(engine_rank) + else: + raise RuntimeError( + f"Unexpected handshake status {status} from dp rank " + f"{engine_rank}" + ) + except Exception as e: + self._error = e + self.started_event.set() + + +@dataclass +class _PreparedExternalElasticEPScale: + control_store: Any + reconfig_store: Any + epoch: str + bootstrap: ReconfigureDistributedRequest + dp_rank: int + cur_data_parallel_size: int + num_redundant_experts: int + scale_up: bool + handshake_server: ExternalElasticEPScaleUpHandshakeServer | None + + +class ExternalElasticEPScaleCoordinator: + def __init__(self, client: "DPAsyncMPClient") -> None: + self.client = client + self.active_reconfig_store: tuple[str, int] | None = None + # Keep the previous rank-0 TCPStore server alive while the next scale + # operation publishes bootstrap metadata through it. Overwriting + # client._coord_store with the new store can otherwise close the old + # server while other ranks are still polling it. + self.control_store_ref: Any | None = None + self.reconfig_store_ref: Any | None = None + self.active_epoch: str | None = None + self.prepared_scale: _PreparedExternalElasticEPScale | None = None + + @staticmethod + def key(*parts: str | int) -> str: + return "/".join(["elastic_ep/external", *[str(part) for part in parts]]) + + @staticmethod + async def _sleep_with_backoff(backoff_step: int) -> int: + max_step = len(POLL_BACKOFF_STEPS_S) - 1 + backoff_step = max(0, min(backoff_step, max_step)) + await asyncio.sleep(POLL_BACKOFF_STEPS_S[backoff_step]) + return min(backoff_step + 1, max_step) + + def _update_parallel_config( + self, + bootstrap: ReconfigureDistributedRequest, + num_redundant_experts: int, + ) -> None: + parallel_config = self.client.vllm_config.parallel_config + parallel_config.data_parallel_size = bootstrap.new_data_parallel_size + parallel_config.data_parallel_master_ip = bootstrap.new_data_parallel_master_ip + parallel_config.data_parallel_master_port = ( + bootstrap.new_data_parallel_master_port + ) + parallel_config._data_parallel_master_port_list = ( + bootstrap.new_data_parallel_master_port_list.copy() + ) + parallel_config._coord_store_port = bootstrap.coord_store_port + parallel_config.eplb_config.num_redundant_experts = num_redundant_experts + + def _calculate_num_redundant_experts( + self, + cur_data_parallel_size: int, + new_data_parallel_size: int, + ) -> int: + parallel_config = self.client.vllm_config.parallel_config + num_experts = self.client.vllm_config.model_config.get_num_experts() + num_redundant_experts = ( + num_experts + parallel_config.eplb_config.num_redundant_experts + ) * new_data_parallel_size // cur_data_parallel_size - num_experts + if num_redundant_experts < 0: + raise ValueError( + "Cannot scale external Elastic EP from DP size " + f"{cur_data_parallel_size} to {new_data_parallel_size}: " + "the target topology cannot retain all model experts." + ) + return num_redundant_experts + + def _get_reconfig_store(self): + from vllm.distributed.utils import get_cached_tcp_store_client + + if self.active_reconfig_store is not None: + store_addr = self.active_reconfig_store + else: + parallel_config = self.client.vllm_config.parallel_config + if not parallel_config._coord_store_port: + raise RuntimeError( + "External Elastic EP requires an active reconfiguration " + "coordination store." + ) + store_addr = ( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ) + + return get_cached_tcp_store_client(*store_addr) + + def _get_existing_engine_zmq_address(self) -> EngineZmqAddresses: + coordinator_input = self.client.coordinator_input_address + coordinator_output = self.client.coordinator_output_address + if coordinator_input is None or coordinator_output is None: + raise RuntimeError( + "External Elastic EP scale-up requires DP coordinator addresses." + ) + + return EngineZmqAddresses( + inputs=[self.client.input_address], + outputs=[self.client.output_address], + coordinator_input=coordinator_input, + coordinator_output=coordinator_output, + frontend_stats_publish_address=self.client.stats_update_address, + ) + + def _setup_reconfig_bootstrap(self) -> tuple[str, int]: + from vllm.distributed.utils import create_tcp_store + + self.control_store_ref = getattr(self.client, "_coord_store", None) + ip = self.client.vllm_config.parallel_config.data_parallel_master_ip + # Keep the serving configuration on the old store until commit succeeds. + self.reconfig_store_ref = create_tcp_store( + ip, 0, is_master=True, world_size=-1, wait_for_workers=False + ) + return ip, self.reconfig_store_ref.port + + def _get_error(self, store: Any, epoch: str) -> str | None: + error_key = self.key(epoch, "error") + if not store.check([error_key]): + return None + return store.get(error_key).decode() + + def get_status(self) -> dict[str, str | int | None]: + store = self._get_reconfig_store() + current_epoch_key = self.key("current_epoch") + if not store.check([current_epoch_key]): + return { + "phase": ExternalElasticEPScalePhase.IDLE.value, + "epoch": None, + "error": None, + "requested_data_parallel_size": None, + } + + epoch = store.get(current_epoch_key).decode() + error = self._get_error(store, epoch) + if error is not None: + phase = ExternalElasticEPScalePhase.FAILED + elif store.check([self.key(epoch, "completed")]): + phase = ExternalElasticEPScalePhase.COMPLETED + elif store.check([self.key(epoch, "phase")]): + phase = ExternalElasticEPScalePhase( + store.get(self.key(epoch, "phase")).decode() + ) + else: + phase = ExternalElasticEPScalePhase.IDLE + requested_dp_size = None + bootstrap_key = self.key(epoch, "bootstrap") + if store.check([bootstrap_key]): + bootstrap = msgspec.msgpack.decode( + store.get(bootstrap_key), type=ReconfigureDistributedRequest + ) + requested_dp_size = bootstrap.new_data_parallel_size + return { + "phase": phase.value, + "epoch": epoch, + "error": error, + "requested_data_parallel_size": requested_dp_size, + } + + def _set_phase( + self, store: Any, epoch: str, phase: ExternalElasticEPScalePhase + ) -> None: + store.set(self.key(epoch, "phase"), phase.value.encode()) + + async def _wait_for_bootstrap( + self, + store: Any, + requested_new_dp_size: int, + timeout_s: float = 300, + ) -> tuple[str, ReconfigureDistributedRequest]: + """Wait for rank 0 to publish matching scale bootstrap metadata. + Non-zero ranks poll the old control store until prepared or errored.""" + loop = asyncio.get_running_loop() + start = loop.time() + current_epoch_key = self.key("current_epoch") + previous_epoch = self.active_epoch + if store.check([current_epoch_key]): + current_epoch = store.get(current_epoch_key).decode() + if store.check([self.key(current_epoch, "completed")]): + previous_epoch = current_epoch + backoff_step = 0 + while True: + if store.check([current_epoch_key]): + epoch = store.get(current_epoch_key).decode() + if epoch != previous_epoch: + error = self._get_error(store, epoch) + if error is not None: + self.active_epoch = epoch + raise RuntimeError(error) + + bootstrap_key = self.key(epoch, "bootstrap") + if store.check([bootstrap_key]): + bootstrap = msgspec.msgpack.decode( + store.get(bootstrap_key), + type=ReconfigureDistributedRequest, + ) + prepared_key = self.key(epoch, "prepared") + completed = store.check([self.key(epoch, "completed")]) + if bootstrap.new_data_parallel_size != requested_new_dp_size: + if store.check([prepared_key]) and not completed: + raise RuntimeError( + "A different external Elastic EP scaling " + "operation is already in progress for target " + f"dp size {bootstrap.new_data_parallel_size}." + ) + elif store.check([prepared_key]): + return epoch, bootstrap + + now = loop.time() + if now - start > timeout_s: + raise TimeoutError( + "Timed out waiting for rank 0 to publish external Elastic EP " + "bootstrap metadata." + ) + backoff_step = await self._sleep_with_backoff(backoff_step) + + def _prepare_reconfig_bootstrap( + self, + store: Any, + cur_data_parallel_size: int, + new_data_parallel_size: int, + ) -> tuple[str, ReconfigureDistributedRequest]: + from vllm.utils.network_utils import get_open_ports_list + + current_epoch_key = self.key("current_epoch") + if store.check([current_epoch_key]): + current_epoch = store.get(current_epoch_key).decode() + current_error = self._get_error(store, current_epoch) + if current_error is None and not store.check( + [self.key(current_epoch, "completed")] + ): + raise RuntimeError( + "Another external Elastic EP scaling operation is already active." + ) + + ip, coord_store_port = self._setup_reconfig_bootstrap() + epoch = uuid.uuid4().hex + master_ports = get_open_ports_list(5) + bootstrap = ReconfigureDistributedRequest( + new_data_parallel_size=new_data_parallel_size, + new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK, + new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, + new_data_parallel_master_ip=ip, + new_data_parallel_master_port=master_ports.pop(), + new_data_parallel_master_port_list=master_ports, + coord_store_port=coord_store_port, + ) + + bootstrap_key = self.key(epoch, "bootstrap") + bootstrap_payload = msgspec.msgpack.encode(bootstrap) + stores = [store] + if self.reconfig_store_ref is not None: + stores.append(self.reconfig_store_ref) + for target_store in stores: + target_store.set(bootstrap_key, bootstrap_payload) + self._set_phase(target_store, epoch, ExternalElasticEPScalePhase.PREPARING) + # Publish the epoch only after its metadata is ready for readers. + target_store.set(current_epoch_key, epoch.encode()) + return epoch, bootstrap + + def _start_scale_up_handshake_server( + self, + bootstrap: ReconfigureDistributedRequest, + cur_data_parallel_size: int, + num_redundant_experts: int, + elastic_ep_max_dp_size: int, + ) -> ExternalElasticEPScaleUpHandshakeServer: + handshake_server = ExternalElasticEPScaleUpHandshakeServer( + handshake_address=get_engine_client_zmq_addr( + False, + bootstrap.new_data_parallel_master_ip, + self.client.vllm_config.parallel_config.data_parallel_rpc_port, + ), + expected_new_ranks=list( + range( + cur_data_parallel_size, + bootstrap.new_data_parallel_size, + ) + ), + addresses=self._get_existing_engine_zmq_address(), + bootstrap=bootstrap, + num_redundant_experts=num_redundant_experts, + elastic_ep_max_dp_size=elastic_ep_max_dp_size, + ) + handshake_server.start() + return handshake_server + + async def _wait_for_ready_ranks( + self, + control_store: Any, + reconfig_store: Any, + epoch: str, + new_data_parallel_size: int, + handshake_server: ExternalElasticEPScaleUpHandshakeServer | None = None, + timeout_s: float = 300, + ) -> None: + """Wait until every rank in the new DP group is ready to commit.""" + loop = asyncio.get_running_loop() + start = loop.time() + ready_keys = [f"eep_ready/{rank}" for rank in range(new_data_parallel_size)] + + backoff_step = 0 + while True: + if handshake_server is not None: + handshake_server.raise_if_failed() + + error = self._get_error(control_store, epoch) + if error is not None: + raise RuntimeError( + "External Elastic EP scaling failed while waiting for " + f"ranks to prepare: {error}" + ) + + if reconfig_store.check(ready_keys): + return + + now = loop.time() + if now - start > timeout_s: + raise TimeoutError( + "Timed out waiting for external Elastic EP ranks to prepare." + ) + backoff_step = await self._sleep_with_backoff(backoff_step) + + async def _wait_for_local_reconfig_finished( + self, + control_store: Any, + reconfig_store: Any, + epoch: str, + bootstrap: ReconfigureDistributedRequest, + dp_rank: int, + scale_up: bool, + timeout_s: float = 600, + ) -> None: + """Wait until this rank's EngineCore finishes the scale transition.""" + if not scale_up and dp_rank >= bootstrap.new_data_parallel_size: + wait_key = self.key(epoch, "shutdown_complete", dp_rank) + timeout_msg = ( + "Timed out waiting for local external Elastic EP scale-down " + "shutdown to finish." + ) + else: + wait_key = self.key(epoch, "old_rank_finished", dp_rank) + timeout_msg = ( + "Timed out waiting for local external Elastic EP " + "reconfiguration to finish." + ) + + loop = asyncio.get_running_loop() + start = loop.time() + backoff_step = 0 + while True: + error = self._get_error(control_store, epoch) + if error is not None: + raise RuntimeError( + error or "External Elastic EP scaling failed on another rank." + ) + + if reconfig_store.check([wait_key]): + return + + now = loop.time() + if now - start > timeout_s: + raise TimeoutError(timeout_msg) + backoff_step = await self._sleep_with_backoff(backoff_step) + + async def _wait_for_all_old_ranks( + self, + prepared: _PreparedExternalElasticEPScale, + timeout_s: float = 600, + ) -> None: + bootstrap = prepared.bootstrap + keys = [ + self.key( + prepared.epoch, + ( + "shutdown_complete" + if not prepared.scale_up + and rank >= bootstrap.new_data_parallel_size + else "old_rank_finished" + ), + rank, + ) + for rank in range(prepared.cur_data_parallel_size) + ] + loop = asyncio.get_running_loop() + start = loop.time() + backoff_step = 0 + while True: + error = self._get_error(prepared.control_store, prepared.epoch) + if error is not None: + raise RuntimeError(error) + if prepared.reconfig_store.check(keys): + return + if loop.time() - start > timeout_s: + raise TimeoutError( + "Timed out waiting for all old ranks to finish external " + "Elastic EP scaling." + ) + backoff_step = await self._sleep_with_backoff(backoff_step) + + @staticmethod + def _stop_handshake_server( + server: ExternalElasticEPScaleUpHandshakeServer | None, + suppress_errors: bool, + ) -> None: + if server is None: + return + if suppress_errors: + with contextlib.suppress(Exception): + server.stop() + else: + server.stop() + + def _publish_error( + self, + prepared: _PreparedExternalElasticEPScale, + error: BaseException, + ) -> None: + error_key = self.key(prepared.epoch, "error") + error_payload = (str(error) or type(error).__name__).encode() + for store in (prepared.control_store, prepared.reconfig_store): + with contextlib.suppress(Exception): + store.set(error_key, error_payload) + self._set_phase( + store, prepared.epoch, ExternalElasticEPScalePhase.FAILED + ) + + def abort(self, error: str) -> None: + prepared = self.prepared_scale + if prepared is not None: + self._publish_error(prepared, RuntimeError(error)) + self._stop_handshake_server(prepared.handshake_server, suppress_errors=True) + self.prepared_scale = None + + async def _wait_for_commit_ready( + self, prepared: _PreparedExternalElasticEPScale, timeout_s: float = 600 + ) -> None: + # No engine may enter the collective pause until all old API ranks + # have finished draining; a drain failure must reach waiting peers. + store = prepared.reconfig_store + await asyncio.to_thread( + store.set, self.key(prepared.epoch, "commit_ready", prepared.dp_rank), b"1" + ) + keys = [ + self.key(prepared.epoch, "commit_ready", rank) + for rank in range(prepared.cur_data_parallel_size) + ] + + def is_ready() -> bool: + error = self._get_error(prepared.control_store, prepared.epoch) + if error is not None: + raise RuntimeError(error) + return store.check(keys) + + start = asyncio.get_running_loop().time() + backoff_step = 0 + while True: + if await asyncio.to_thread(is_ready): + return + if asyncio.get_running_loop().time() - start > timeout_s: + raise TimeoutError("Timed out waiting for external Elastic EP drains.") + backoff_step = await self._sleep_with_backoff(backoff_step) + + async def prepare( + self, cur_data_parallel_size: int, new_data_parallel_size: int + ) -> None: + from vllm.distributed.utils import get_cached_tcp_store_client + + if self.prepared_scale is not None: + if ( + self.prepared_scale.bootstrap.new_data_parallel_size + == new_data_parallel_size + ): + return + raise RuntimeError("External Elastic EP scaling is already prepared.") + + parallel_config = self.client.vllm_config.parallel_config + dp_rank = parallel_config.data_parallel_rank + scale_up = new_data_parallel_size > cur_data_parallel_size + if new_data_parallel_size > parallel_config.elastic_ep_max_dp_size: + raise ValueError( + f"Cannot scale to data_parallel_size {new_data_parallel_size}; " + "--elastic-ep-max-dp-size is " + f"{parallel_config.elastic_ep_max_dp_size}." + ) + # Validate expert capacity on every rank: rank 0 publishes no epoch + # when validation fails, so peers must reject before polling for one. + num_redundant_experts = self._calculate_num_redundant_experts( + cur_data_parallel_size, + new_data_parallel_size, + ) + if not parallel_config._coord_store_port: + raise RuntimeError( + "External Elastic EP requires a runtime coordination store port." + ) + control_store = get_cached_tcp_store_client( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ) + handshake_server: ExternalElasticEPScaleUpHandshakeServer | None = None + bootstrap: ReconfigureDistributedRequest | None = None + epoch: str | None = None + reconfig_store: Any | None = None + + try: + if dp_rank == 0: + epoch, bootstrap = self._prepare_reconfig_bootstrap( + control_store, + cur_data_parallel_size, + new_data_parallel_size, + ) + control_store.set( + self.key(epoch, "num_redundant_experts"), + str(num_redundant_experts).encode(), + ) + if scale_up: + handshake_server = self._start_scale_up_handshake_server( + bootstrap, + cur_data_parallel_size, + num_redundant_experts, + parallel_config.elastic_ep_max_dp_size, + ) + control_store.set(self.key(epoch, "prepared"), b"1") + else: + epoch, bootstrap = await self._wait_for_bootstrap( + control_store, new_data_parallel_size + ) + num_redundant_experts = int( + control_store.get(self.key(epoch, "num_redundant_experts")).decode() + ) + + assert epoch is not None + assert bootstrap is not None + self.active_epoch = epoch + self.active_reconfig_store = ( + bootstrap.new_data_parallel_master_ip, + bootstrap.coord_store_port, + ) + reconfig_store = self._get_reconfig_store() + reconfig_store.set(self.key("current_epoch"), epoch.encode()) + self._set_phase( + reconfig_store, epoch, ExternalElasticEPScalePhase.PREPARING + ) + + if scale_up or dp_rank < bootstrap.new_data_parallel_size: + await self.client.call_utility_async( + "reinitialize_distributed", bootstrap + ) + + await self._wait_for_ready_ranks( + control_store, + reconfig_store, + epoch, + bootstrap.new_data_parallel_size, + handshake_server, + ) + self.prepared_scale = _PreparedExternalElasticEPScale( + control_store=control_store, + reconfig_store=reconfig_store, + epoch=epoch, + bootstrap=bootstrap, + dp_rank=dp_rank, + cur_data_parallel_size=cur_data_parallel_size, + num_redundant_experts=num_redundant_experts, + scale_up=scale_up, + handshake_server=handshake_server, + ) + except BaseException as e: + if epoch is not None: + self.active_epoch = epoch + error_key = self.key(epoch, "error") + error_payload = (str(e) or type(e).__name__).encode() + with contextlib.suppress(Exception): + control_store.set(error_key, error_payload) + self._set_phase( + control_store, epoch, ExternalElasticEPScalePhase.FAILED + ) + failed_reconfig_store = ( + reconfig_store + if reconfig_store is not None + else self.reconfig_store_ref + ) + if failed_reconfig_store is not None: + with contextlib.suppress(Exception): + failed_reconfig_store.set(error_key, error_payload) + self._set_phase( + failed_reconfig_store, + epoch, + ExternalElasticEPScalePhase.FAILED, + ) + self.active_reconfig_store = None + self._stop_handshake_server(handshake_server, suppress_errors=True) + raise + + async def commit(self) -> None: + prepared = self.prepared_scale + if prepared is None: + raise RuntimeError("External Elastic EP scaling has not been prepared.") + + bootstrap = prepared.bootstrap + remaining = prepared.scale_up or ( + prepared.dp_rank < bootstrap.new_data_parallel_size + ) + try: + await self._wait_for_commit_ready(prepared) + for store in (prepared.control_store, prepared.reconfig_store): + self._set_phase( + store, prepared.epoch, ExternalElasticEPScalePhase.COMMITTING + ) + await self.client.pause_scheduler_async( + mode=self.client._eep_commit_pause_mode() if remaining else "abort", + clear_cache=False, + ) + if remaining: + await self.client.call_utility_async("commit_prepared_elastic_ep") + else: + reconfig_request = msgspec.structs.replace( + bootstrap, + new_data_parallel_rank=ReconfigureRankType.SHUTDOWN_CURRENT_RANK, + ) + await self.client.call_utility_async( + "reinitialize_distributed", reconfig_request + ) + + await self._wait_for_local_reconfig_finished( + prepared.control_store, + prepared.reconfig_store, + prepared.epoch, + bootstrap, + prepared.dp_rank, + prepared.scale_up, + ) + if prepared.dp_rank == 0: + await self._wait_for_all_old_ranks(prepared) + completed_key = self.key(prepared.epoch, "completed") + prepared.control_store.set(completed_key, b"1") + prepared.reconfig_store.set(completed_key, b"1") + for store in (prepared.control_store, prepared.reconfig_store): + self._set_phase( + store, prepared.epoch, ExternalElasticEPScalePhase.COMPLETED + ) + + if remaining: + self._update_parallel_config(bootstrap, prepared.num_redundant_experts) + if prepared.dp_rank == 0: + self.client._coord_store = self.reconfig_store_ref + self.client._ensure_stats_update_task() + await self.client.first_req_send_socket.send( + msgspec.msgpack.encode( + ("SCALE_ELASTIC_EP", bootstrap.new_data_parallel_size) + ) + ) + await self.client.resume_scheduler_async() + self._stop_handshake_server( + prepared.handshake_server, suppress_errors=False + ) + except BaseException as e: + self._publish_error(prepared, e) + self._stop_handshake_server(prepared.handshake_server, suppress_errors=True) + raise + finally: + self.prepared_scale = None + + async def process_engine_core_notification( + self, notification_data: tuple[str, int] + ) -> None: + """Record scale notifications emitted by EngineCore processes. + The stored keys are later polled by the external scale coordinator.""" + parallel_config = self.client.vllm_config.parallel_config + if not ( + parallel_config.enable_elastic_ep + and parallel_config.data_parallel_external_lb + ): + return + + notification_type_str, dp_rank = notification_data + notification_type = EEPNotificationType(notification_type_str) + if not parallel_config._coord_store_port: + logger.warning( + "Ignoring external Elastic EP notification %s because coord " + "store metadata is not available yet.", + notification_type.value, + ) + return + + epoch = self.active_epoch + reconfig_store = self._get_reconfig_store() + if epoch is None: + current_epoch_key = self.key("current_epoch") + if not reconfig_store.check([current_epoch_key]): + logger.warning( + "Ignoring external Elastic EP notification %s because " + "active epoch metadata is not available yet.", + notification_type.value, + ) + return + epoch = reconfig_store.get(current_epoch_key).decode() + + if notification_type == EEPNotificationType.RECONFIGURE_FINISHED: + key = self.key(epoch, "old_rank_finished", dp_rank) + reconfig_store.set(key, b"1") + elif notification_type == EEPNotificationType.SHUTDOWN_COMPLETE: + key = self.key(epoch, "shutdown_complete", dp_rank) + reconfig_store.set(key, b"1") diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 15f8f8b744..492f8a2ce1 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -60,6 +60,10 @@ class EngineClient(ABC): @abstractmethod def errored(self) -> bool: ... + @property + def should_keep_api_server_alive(self) -> bool: + return False + @property @abstractmethod def dead_error(self) -> BaseException: ... @@ -262,6 +266,12 @@ class EngineClient(ABC): """Scale the engine.""" raise NotImplementedError + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + """Return shared external Elastic EP operation status, if enabled.""" + return None + async def collective_rpc( self, method: str, diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py index 44161439b6..1db294cfc6 100644 --- a/vllm/entrypoints/launchers/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -328,7 +328,10 @@ def terminate_if_errored(server: uvicorn.Server, engine: EngineClient): for this request. """ engine_errored = engine.errored and not engine.is_running - if not envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH and engine_errored: + keep_alive = ( + envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH or engine.should_keep_api_server_alive + ) + if not keep_alive and engine_errored: server.should_exit = True diff --git a/vllm/entrypoints/serve/elastic_ep/api_router.py b/vllm/entrypoints/serve/elastic_ep/api_router.py index 1dbfc96b9b..5a839948cd 100644 --- a/vllm/entrypoints/serve/elastic_ep/api_router.py +++ b/vllm/entrypoints/serve/elastic_ep/api_router.py @@ -8,6 +8,7 @@ from http import HTTPStatus from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse +from vllm.distributed.elastic_ep.external_elastic_ep import ExternalElasticEPScalePhase from vllm.engine.protocol import EngineClient from vllm.entrypoints.serve.elastic_ep.middleware import get_scaling_elastic_ep from vllm.entrypoints.serve.engine.protocol import ErrorResponse @@ -82,7 +83,35 @@ async def scale_elastic_ep(raw_request: Request): @router.post("/is_scaling_elastic_ep") async def is_scaling_elastic_ep(raw_request: Request): - return JSONResponse({"is_scaling_elastic_ep": get_scaling_elastic_ep()}) + # External operation status comes from the shared store. Middleware gating + # remains process-local, so scaling requests must reach every old-rank API. + try: + status = await engine_client(raw_request).get_external_elastic_ep_status() + except Exception as e: + logger.warning("Failed to query external Elastic EP status: %s", e) + raise HTTPException( + status_code=503, + detail="External Elastic EP status is temporarily unavailable", + ) from e + if status is None: + # Non-external EEP modes retain the process-local middleware state. + is_scaling = get_scaling_elastic_ep() + status = { + "phase": ( + ExternalElasticEPScalePhase.COMMITTING.value + if is_scaling + else ExternalElasticEPScalePhase.IDLE.value + ), + "epoch": None, + "error": None, + "requested_data_parallel_size": None, + } + else: + is_scaling = status["phase"] in ( + ExternalElasticEPScalePhase.PREPARING.value, + ExternalElasticEPScalePhase.COMMITTING.value, + ) + return JSONResponse({"is_scaling_elastic_ep": is_scaling, **status}) def attach_router(app: FastAPI): diff --git a/vllm/entrypoints/serve/elastic_ep/middleware.py b/vllm/entrypoints/serve/elastic_ep/middleware.py index dcc62aea2c..d7f81fd7ef 100644 --- a/vllm/entrypoints/serve/elastic_ep/middleware.py +++ b/vllm/entrypoints/serve/elastic_ep/middleware.py @@ -6,8 +6,12 @@ from collections.abc import Awaitable from fastapi.responses import JSONResponse from starlette.types import ASGIApp, Receive, Scope, Send -# Global variable to track scaling state +# Process-local states used to pause serving during scaling and after this rank +# has been removed from the external-LB topology. A retired API process remains +# available for health and scaling-status queries until its orchestrator exits it. _scaling_elastic_ep = False +_elastic_ep_rank_retired = False +_SCALING_OBSERVABILITY_PATHS = frozenset({"/health", "/is_scaling_elastic_ep"}) def get_scaling_elastic_ep(): @@ -19,12 +23,20 @@ def set_scaling_elastic_ep(value): _scaling_elastic_ep = value -class ScalingMiddleware: - """Middleware that checks if the model is currently scaling and - returns a 503 Service Unavailable response if it is. +def get_elastic_ep_rank_retired(): + return _elastic_ep_rank_retired - This middleware applies to all HTTP requests and prevents - processing when the model is in a scaling state. + +def set_elastic_ep_rank_retired(value): + global _elastic_ep_rank_retired + _elastic_ep_rank_retired = value + + +class ScalingMiddleware: + """Pause serving while scaling or after the local rank is retired. + + Health and scaling-status requests remain available so an external + orchestrator can observe the operation and remove a retired process. """ def __init__(self, app: ASGIApp) -> None: @@ -34,13 +46,18 @@ class ScalingMiddleware: if scope["type"] != "http": return self.app(scope, receive, send) - # Check global scaling state - if get_scaling_elastic_ep(): - # Return 503 Service Unavailable response + # Keep observability available while serving is paused or this rank is + # waiting for the external orchestrator to remove its process. + rank_retired = get_elastic_ep_rank_retired() + serving_paused = get_scaling_elastic_ep() or rank_retired + if serving_paused and scope["path"] not in _SCALING_OBSERVABILITY_PATHS: + error = ( + "This data-parallel rank has been retired and is awaiting shutdown." + if rank_retired + else "The model is currently scaling. Please try again later." + ) response = JSONResponse( - content={ - "error": "The model is currently scaling. Please try again later." - }, + content={"error": error}, status_code=503, ) return response(scope, receive, send) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 1b14b08246..38aee44720 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -104,6 +104,9 @@ class EngineCoreReadyResponse: supports_draft_weight_updates: bool = False # Full-attention block size in tokens after initialization, or unavailable. effective_attention_block_size: int | None = None + coord_store_port: int = 0 + coordinator_input_address: str | None = None + coordinator_output_address: str | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 4ef942d0b0..8b7f913b43 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -22,7 +22,11 @@ from vllm.distributed.weight_transfer.base import ( ) from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient, StreamingInput -from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep +from vllm.entrypoints.serve.elastic_ep.middleware import ( + get_elastic_ep_rank_retired, + set_elastic_ep_rank_retired, + set_scaling_elastic_ep, +) from vllm.exceptions import ( GracefulHTTPError, MaxQueuedTokensError, @@ -918,6 +922,12 @@ class AsyncLLM(EngineClient): mm_cache_stats=renderer.stat_mm_cache(), ) except Exception as e: + if get_elastic_ep_rank_retired() and isinstance(e, EngineDeadError): + logger.info( + "[Elastic EP] Retired rank stopped receiving EngineCore outputs" + ) + output_processor.propagate_error(e) + return logger.exception("AsyncLLM output_handler failed.") output_processor.propagate_error(e) @@ -1103,6 +1113,8 @@ class AsyncLLM(EngineClient): async def check_health(self) -> None: logger.debug("Called check_health.") + if get_elastic_ep_rank_retired(): + return if self.errored: raise self.dead_error @@ -1252,8 +1264,13 @@ class AsyncLLM(EngineClient): "waiting for requests to drain before scaling" ) await self.wait_for_requests_to_drain(drain_timeout) - except BaseException: - set_scaling_elastic_ep(False) + except BaseException as e: + if self.vllm_config.parallel_config.data_parallel_external_lb: + await self.engine_core.abort_external_elastic_ep( + str(e) or type(e).__name__ + ) + else: + set_scaling_elastic_ep(False) raise async def scale_elastic_ep( @@ -1266,7 +1283,8 @@ class AsyncLLM(EngineClient): async def _scale_elastic_ep( self, new_data_parallel_size: int, drain_timeout: int ) -> None: - old_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size + parallel_config = self.vllm_config.parallel_config + old_data_parallel_size = parallel_config.data_parallel_size if old_data_parallel_size == new_data_parallel_size: logger.info( "Data parallel size is already %s, skipping scale", @@ -1297,9 +1315,36 @@ class AsyncLLM(EngineClient): if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS: await self._drain_requests_for_elastic_ep(drain_timeout) - await self.engine_core.commit_elastic_ep() - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - set_scaling_elastic_ep(False) + rank_will_retire = ( + parallel_config.data_parallel_external_lb + and new_data_parallel_size < old_data_parallel_size + and parallel_config.data_parallel_rank >= new_data_parallel_size + ) + if rank_will_retire: + set_elastic_ep_rank_retired(True) + + commit_succeeded = False + try: + await self.engine_core.commit_elastic_ep() + parallel_config.data_parallel_size = new_data_parallel_size + commit_succeeded = True + if rank_will_retire: + logger.info( + "[Elastic EP] Data-parallel rank %s is retired and awaiting " + "external process shutdown", + parallel_config.data_parallel_rank, + ) + finally: + if rank_will_retire and not commit_succeeded: + set_elastic_ep_rank_retired(False) + # A failed commit may leave engines partially reconfigured or paused. + if commit_succeeded: + set_scaling_elastic_ep(False) + + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + return await self.engine_core.get_external_elastic_ep_status() async def handle_fault( self, fault_tolerance_request: FaultToleranceRequest @@ -1323,6 +1368,10 @@ class AsyncLLM(EngineClient): def errored(self) -> bool: return self.engine_core.resources.engine_dead or not self.is_running + @property + def should_keep_api_server_alive(self) -> bool: + return get_elastic_ep_rank_retired() + @property def dead_error(self) -> BaseException: return EngineDeadError() diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index cb120ba5d9..28a4c7af6d 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -126,6 +126,7 @@ class DPCoordinator: self.stats_publish_address = front_publish_address self.coord_in_address = back_publish_address self.coord_out_address = back_output_address + self._coord_store: object | None = None self._finalizer = weakref.finalize(self, shutdown, [self.proc]) def get_stats_publish_address(self) -> str: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index d8c976ff69..c9843dce96 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1346,6 +1346,16 @@ class EngineCoreProc(EngineCore): ready_msg["parallel_config_hash"] = ( vllm_config.parallel_config.compute_hash() ) + if vllm_config.parallel_config.enable_elastic_ep: + ready_msg["coord_store_port"] = ( + vllm_config.parallel_config._coord_store_port + ) + ready_msg["elastic_ep_max_dp_size"] = ( + vllm_config.parallel_config.elastic_ep_max_dp_size + ) + ready_msg["num_redundant_experts"] = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) handshake_socket.send(msgspec.msgpack.encode(ready_msg)) @@ -1376,14 +1386,37 @@ class EngineCoreProc(EngineCore): f"minutes" ) init_bytes = handshake_socket.recv() - init_message: EngineHandshakeMetadata = msgspec.msgpack.decode( - init_bytes, type=EngineHandshakeMetadata - ) + num_redundant_experts: int | None = None + if ( + envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH + and parallel_config is not None + and parallel_config.data_parallel_external_lb + ): + from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleUpHandshakeMetadata, + ) + + scale_up_message = msgspec.msgpack.decode( + init_bytes, type=ExternalElasticEPScaleUpHandshakeMetadata + ) + init_message = scale_up_message.engine_metadata + num_redundant_experts = scale_up_message.num_redundant_experts + parallel_config.elastic_ep_max_dp_size = ( + scale_up_message.elastic_ep_max_dp_size + ) + else: + init_message = msgspec.msgpack.decode( + init_bytes, type=EngineHandshakeMetadata + ) logger.debug("Received init message: %s", init_message) if parallel_config is not None: for key, value in init_message.parallel_config.items(): setattr(parallel_config, key, value) + if num_redundant_experts is not None: + parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts + ) return init_message.addresses @@ -1826,6 +1859,9 @@ class EngineCoreProc(EngineCore): if self.vllm_config.weight_transfer_config is not None else False ), + coord_store_port=parallel_config._coord_store_port, + coordinator_input_address=self.addresses.coordinator_input, + coordinator_output_address=self.addresses.coordinator_output, ) def process_input_sockets( @@ -2342,6 +2378,13 @@ class DPEngineCoreProc(EngineCoreProc): if state.is_complete(): if state.worker_type == "removing": raise SystemExit + if ( + state.worker_type == "new" + and self.vllm_config.parallel_config.data_parallel_external_lb + ): + # An independently launched rank has no external utility + # caller to join the post-commit resume barrier. + self.resume_scheduler() self.process_input_queue_block = True self.eep_scaling_state = None elif not state.commit_requested and state.is_ready_for_switch(): @@ -2512,12 +2555,10 @@ class DPEngineCoreProc(EngineCoreProc): self.output_queue.put_nowait((0, outputs)) else: encoder = MsgpackEncoder() - with ( - zmq.Context() as ctx, - make_zmq_socket( - ctx, self.addresses.outputs[0], zmq.PUSH, linger=4000 - ) as socket, - ): + ctx = zmq.Context.instance() + with make_zmq_socket( + ctx, self.addresses.outputs[0], zmq.PUSH, linger=4000 + ) as socket: socket.send_multipart(encoder.encode(outputs)) def _eep_scale_up_before_kv_init(self): diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index ef6b269618..cfe8976668 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -23,6 +23,9 @@ import zmq.asyncio from vllm import envs from vllm.config import VllmConfig from vllm.config.kv_events import KVEventsConfig +from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleCoordinator, +) from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S from vllm.logger import init_logger from vllm.lora.request import LoRARequest @@ -271,6 +274,14 @@ class EngineCoreClient(ABC): async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: raise NotImplementedError + async def abort_external_elastic_ep(self, error: str) -> None: + raise NotImplementedError + + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + return None + async def get_output_async(self) -> EngineCoreOutputs: raise NotImplementedError @@ -626,6 +637,8 @@ class MPClient(EngineCoreClient): enable_input_socket_handover = parallel_config.enable_elastic_ep self.stats_update_address: str | None = None + self.coordinator_input_address: str | None = None + self.coordinator_output_address: str | None = None tensor_queue: Queue | None = None if client_addresses: # Engines are managed externally to this client. @@ -702,6 +715,11 @@ class MPClient(EngineCoreClient): coordinator.get_stats_publish_address() ) + self.input_address = self.input_socket.getsockopt_string(zmq.LAST_ENDPOINT) + output_socket = self.resources.output_socket + assert output_socket is not None + self.output_address = output_socket.getsockopt_string(zmq.LAST_ENDPOINT) + # Serialization setup with tensor queues for multimodal tensor IPC. tensor_ipc_sender: TensorIpcSender | None = None model_config = getattr(vllm_config, "model_config", None) @@ -885,6 +903,25 @@ class MPClient(EngineCoreClient): response.kv_events_config ) + parallel_config = vllm_config.parallel_config + if response.coord_store_port: + if parallel_config._coord_store_port: + assert response.coord_store_port == parallel_config._coord_store_port + else: + parallel_config._coord_store_port = response.coord_store_port + + for attribute in ( + "coordinator_input_address", + "coordinator_output_address", + ): + address = getattr(response, attribute) + if address is not None: + existing_address = getattr(self, attribute) + if existing_address is None: + setattr(self, attribute, address) + else: + assert address == existing_address + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]: return dict(self._kv_event_sources) @@ -1431,6 +1468,12 @@ class DPAsyncMPClient(AsyncMPClient): renderer: BaseRenderer | None = None, ): self.current_wave = 0 + parallel_config = vllm_config.parallel_config + self.external_eep_coordinator: ExternalElasticEPScaleCoordinator | None = None + if parallel_config.enable_elastic_ep and ( + parallel_config.data_parallel_external_lb + ): + self.external_eep_coordinator = ExternalElasticEPScaleCoordinator(self) super().__init__( vllm_config, @@ -1514,23 +1557,19 @@ class DPAsyncMPClient(AsyncMPClient): 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] + assert not parallel_config.data_parallel_hybrid_lb + # External LB clients continue to manage only local ranks. + if not parallel_config.data_parallel_external_lb: + self.engine_ranks_managed = list(range(dp_size)) + if len(self.lb_engines) < new_engine_count: + 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) @@ -1599,6 +1638,93 @@ class DPAsyncMPClient(AsyncMPClient): def get_core_engine_for_request(self, request: EngineCoreRequest): return self.core_engine + def _eep_commit_pause_mode(self) -> PauseMode: + from vllm.distributed.elastic_ep.elastic_execute import ( + can_reuse_fused_moe_kernel, + ) + + # MRV2 re-warms through the request pool, so running requests must + # finish first while new ones stay queued in the scheduler. + parallel_config = self.vllm_config.parallel_config + if self.vllm_config.use_v2_model_runner and not can_reuse_fused_moe_kernel( + parallel_config + ): + return "wait" + return "keep" + + def _setup_elastic_ep_reconfig_bootstrap(self) -> tuple[str, int]: + 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 + return ip, store.port + + @staticmethod + async def eep_process_engine_core_notification( + self: "DPAsyncMPClient", notification_data: tuple[str, int] + ) -> None: + if self.external_eep_coordinator is not None: + await self.external_eep_coordinator.process_engine_core_notification( + notification_data + ) + + def _get_external_eep_coordinator(self) -> ExternalElasticEPScaleCoordinator: + parallel_config = self.vllm_config.parallel_config + if not parallel_config.enable_elastic_ep: + raise NotImplementedError( + "Elastic EP scaling requires enable_elastic_ep=True." + ) + if not parallel_config.data_parallel_external_lb: + raise NotImplementedError( + "DPAsyncMPClient only supports Elastic EP scaling in external " + "load-balancer mode." + ) + coordinator = self.external_eep_coordinator + assert coordinator is not None + return coordinator + + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: + cur_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size + + assert new_data_parallel_size != cur_data_parallel_size, ( + f"new_data_parallel_size {new_data_parallel_size} must be " + f"different from cur_data_parallel_size {cur_data_parallel_size}" + ) + await self._get_external_eep_coordinator().prepare( + cur_data_parallel_size, new_data_parallel_size + ) + + async def commit_elastic_ep(self) -> None: + await self._get_external_eep_coordinator().commit() + + async def abort_external_elastic_ep(self, error: str) -> None: + self._get_external_eep_coordinator().abort(error) + + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + coordinator = self.external_eep_coordinator + return ( + None + if coordinator is None + else await asyncio.to_thread(coordinator.get_status) + ) + class DPLBAsyncMPClient(DPAsyncMPClient): """Asyncio-compatible client for multi-proc, multi-engine (data parallel) @@ -1719,8 +1845,9 @@ class DPLBAsyncMPClient(DPAsyncMPClient): @staticmethod async def eep_process_engine_core_notification( - self: "DPLBAsyncMPClient", notification_data: tuple[str, int] - ): + self: "DPAsyncMPClient", notification_data: tuple[str, int] + ) -> None: + assert isinstance(self, DPLBAsyncMPClient) cache = self.eep_scaling_cache notification_type_str, dp_rank = notification_data try: @@ -1881,27 +2008,6 @@ class DPLBAsyncMPClient(DPAsyncMPClient): 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, @@ -1956,20 +2062,6 @@ class DPLBAsyncMPClient(DPAsyncMPClient): await asyncio.to_thread(self._coord_store.wait, ready_keys) logger.info("[Elastic EP] Successfully started new engines") - def _eep_commit_pause_mode(self) -> PauseMode: - from vllm.distributed.elastic_ep.elastic_execute import ( - can_reuse_fused_moe_kernel, - ) - - # MRV2 re-warms through the request pool, so running requests must - # finish first while new ones stay queued in the scheduler. - parallel_config = self.vllm_config.parallel_config - if self.vllm_config.use_v2_model_runner and not can_reuse_fused_moe_kernel( - parallel_config - ): - return "wait" - return "keep" - async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None: new_core_engines = [ rank.to_bytes(2, "little") diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 4f1be3bacd..96df5f14eb 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -290,6 +290,7 @@ class CoreEngineProcManager: self._finalizer = weakref.finalize(self, shutdown, self.processes) self.manager_stopped = threading.Event() self.failed_proc_name: str | None = None + self._coord_store: object | None = None # All ranks share this config object: capture the user-provided # --device-ids list before the per-rank shard overwrites it. Mutating @@ -1139,7 +1140,9 @@ def get_engine_zmq_bind_addresses( offline_mode or local_engines_only or (local_engine_count == dp_size) ) # NOTE(yongji): handling scaling from intra-node to inter-node - if parallel_config.enable_elastic_ep: + if parallel_config.enable_elastic_ep and ( + not local_engines_only or parallel_config.data_parallel_rank == 0 + ): client_local_only = False def _addr() -> str: @@ -1228,6 +1231,32 @@ def launch_core_engines( vllm_config.needs_dp_coordinator and not offline_mode and dp_rank == 0 ) + coord_store = None + + create_coord_store = ( + parallel_config.enable_elastic_ep + and parallel_config.data_parallel_backend != "ray" + and dp_rank == 0 + and not parallel_config._coord_store_port + ) + + if create_coord_store: + from vllm.distributed.utils import create_tcp_store + + coord_store = create_tcp_store( + host, + 0, + is_master=True, + world_size=-1, + wait_for_workers=False, + ) + parallel_config._coord_store_port = coord_store.port + logger.info( + "Created Elastic EP coordination TCPStore on %s:%d", + host, + coord_store.port, + ) + if run_coordinator: coordinator = DPCoordinator( parallel_config, @@ -1243,12 +1272,13 @@ def launch_core_engines( ) logger.info("Started DP Coordinator process (PID: %d)", coordinator.proc.pid) + if coord_store is not None: + coordinator._coord_store = coord_store else: coordinator = None # Hold a coordination TCPStore (alive for this frame) so engines pick DP # master ports at bind time; pre-allocated ports can be taken before use. - coord_store = None if ( dp_size > 1 and not offline_mode @@ -1344,6 +1374,13 @@ def launch_core_engines( else: local_engine_manager = None + if ( + coord_store is not None + and coordinator is None + and (local_engine_manager is not None) + ): + local_engine_manager._coord_store = coord_store + launch = CoreEngineLaunch( local_engine_manager, coordinator, addresses, tensor_queue ) @@ -1474,21 +1511,33 @@ def wait_for_engine_startup( if status == "HELLO" and engine.state == CoreEngineState.NEW: # Send init message with DP config info. + parallel_config_payload: dict[str, int | str | list[int]] = {} + if coordinated_dp and not ( + parallel_config.local_engines_only + and parallel_config.data_parallel_rank > 0 + ): + parallel_config_payload = { + k: getattr(parallel_config, k) + for k in ( + "data_parallel_master_ip", + "data_parallel_master_port", + "_data_parallel_master_port_list", + "_coord_store_port", + "data_parallel_size", + ) + } + if ( + parallel_config.enable_elastic_ep + and not parallel_config._coord_store_port + ): + raise RuntimeError( + "Elastic EP startup requires a coordination TCPStore " + "port before engine handshake." + ) init_message = msgspec.msgpack.encode( EngineHandshakeMetadata( addresses=launch.addresses, - parallel_config={ - k: getattr(parallel_config, k) - for k in ( - "data_parallel_master_ip", - "data_parallel_master_port", - "_data_parallel_master_port_list", - "_coord_store_port", - "data_parallel_size", - ) - } - if coordinated_dp - else {}, + parallel_config=parallel_config_payload, ) ) handshake_socket.send_multipart((eng_identity, init_message), copy=False) @@ -1496,6 +1545,18 @@ def wait_for_engine_startup( start_pending[0 if local else 1] += 1 engine.state = CoreEngineState.CONNECTED elif status == "READY" and engine.state == CoreEngineState.CONNECTED: + coord_store_port = msg.get("coord_store_port") + if coord_store_port is not None and local: + parallel_config._coord_store_port = coord_store_port + num_redundant_experts = msg.get("num_redundant_experts") + if num_redundant_experts is not None and local: + parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts + ) + elastic_ep_max_dp_size = msg.get("elastic_ep_max_dp_size") + if elastic_ep_max_dp_size is not None and local: + parallel_config.elastic_ep_max_dp_size = elastic_ep_max_dp_size + # Validate config hash consistency across DP workers for MoE models. if coordinated_dp: worker_config_hash = msg.get("parallel_config_hash")