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.scale_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")