| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import enum |
| 4 | import weakref |
| 5 | from concurrent.futures import Future, ThreadPoolExecutor |
| 6 | from typing import TYPE_CHECKING, Any, Literal, TypeAlias |
| 7 | |
| 8 | import torch.distributed |
| 9 | |
| 10 | from vllm.config import ParallelConfig |
| 11 | from vllm.distributed import ( |
| 12 | stateless_destroy_torch_distributed_process_group, |
| 13 | ) |
| 14 | from vllm.distributed.utils import get_cached_tcp_store_client |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.v1.engine import ( |
| 17 | EEPNotificationType, |
| 18 | ReconfigureDistributedRequest, |
| 19 | ReconfigureRankType, |
| 20 | ) |
| 21 | from vllm.v1.engine.core import DPEngineCoreProc |
| 22 | |
| 23 | if TYPE_CHECKING: |
| 24 | from vllm.config import VllmConfig |
| 25 | from vllm.v1.executor.abstract import Executor |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | WorkerType = Literal["existing", "new", "removing"] |
| 30 | |
| 31 | |
| 32 | class ScaleUpExistingEngineState(enum.IntEnum): |
| 33 | PREPARE = 0 |
| 34 | SYNC_KV_CACHE_MEMORY_SIZE = 1 |
| 35 | COMMIT_SCALE_UP = 2 # Blocks forward passes. |
| 36 | COMPLETE = 3 |
| 37 | |
| 38 | |
| 39 | class ScaleUpNewEngineState(enum.IntEnum): |
| 40 | PRE_KV_INIT = 0 |
| 41 | PREPARE = 1 |
| 42 | COMPLETE = 2 |
| 43 | |
| 44 | |
| 45 | class ScaleDownRemainingEngineState(enum.IntEnum): |
| 46 | PREPARE = 0 |
| 47 | COMMIT_SCALE_DOWN = 1 # Blocks forward passes. |
| 48 | COMPLETE = 2 |
| 49 | |
| 50 | |
| 51 | class ScaleDownRemovingEngineState(enum.IntEnum): |
| 52 | PREPARE = 0 |
| 53 | COMPLETE = 1 |
| 54 | |
| 55 | |
| 56 | EngineState: TypeAlias = ( |
| 57 | ScaleUpExistingEngineState |
| 58 | | ScaleUpNewEngineState |
| 59 | | ScaleDownRemainingEngineState |
| 60 | | ScaleDownRemovingEngineState |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | class ElasticEPScalingState: |
| 65 | def __init__( |
| 66 | self, |
| 67 | model_executor: "Executor", |
| 68 | engine_core: "DPEngineCoreProc", |
| 69 | vllm_config: "VllmConfig", |
| 70 | new_parallel_config: ParallelConfig, |
| 71 | worker_type: WorkerType, |
| 72 | scale_type: Literal["scale_up", "scale_down"], |
| 73 | reconfig_request: ReconfigureDistributedRequest | None = None, |
| 74 | ): |
| 75 | self.model_executor_ref = weakref.ref(model_executor) |
| 76 | self.engine_core_ref = weakref.ref(engine_core) |
| 77 | self.vllm_config = vllm_config |
| 78 | self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None |
| 79 | self.new_parallel_config: ParallelConfig = new_parallel_config |
| 80 | self.new_dp_group = self.engine_core.dp_group if worker_type == "new" else None |
| 81 | self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None |
| 82 | self.worker_type = worker_type |
| 83 | self.scale_type = scale_type |
| 84 | self.reconfig_request = reconfig_request |
| 85 | self.commit_requested = False |
| 86 | self._prepare_executor = ThreadPoolExecutor( |
| 87 | max_workers=1, thread_name_prefix="ElasticEPPrepare" |
| 88 | ) |
| 89 | self._prepare_future: Future[Any] | None = None |
| 90 | self._new_dp_sync: tuple[object, Any] | None = None |
| 91 | self.state: EngineState |
| 92 | if scale_type == "scale_up": |
| 93 | self.state = ( |
| 94 | ScaleUpNewEngineState.PRE_KV_INIT |
| 95 | if worker_type == "new" |
| 96 | else ScaleUpExistingEngineState.PREPARE |
| 97 | ) |
| 98 | else: |
| 99 | self.state = ( |
| 100 | ScaleDownRemovingEngineState.PREPARE |
| 101 | if worker_type == "removing" |
| 102 | else ScaleDownRemainingEngineState.PREPARE |
| 103 | ) |
| 104 | |
| 105 | @property |
| 106 | def model_executor(self) -> "Executor": |
| 107 | model_executor = self.model_executor_ref() |
| 108 | if model_executor is None: |
| 109 | raise RuntimeError("Model executor has been garbage collected") |
| 110 | return model_executor |
| 111 | |
| 112 | @property |
| 113 | def engine_core(self) -> "DPEngineCoreProc": |
| 114 | engine_core = self.engine_core_ref() |
| 115 | if engine_core is None: |
| 116 | raise RuntimeError("Engine core has been garbage collected") |
| 117 | return engine_core |
| 118 | |
| 119 | def _collective_rpc(self, *args, **kwargs): |
| 120 | return self.model_executor.collective_rpc(*args, **kwargs) |
| 121 | |
| 122 | def _execute_async(self, execute_method: str, *args) -> bool: |
| 123 | if self._prepare_future is None: |
| 124 | done_keys = self._collective_rpc( |
| 125 | "elastic_ep_execute", |
| 126 | args=("start_async", execute_method, *args), |
| 127 | ) |
| 128 | assert self.reconfig_request is not None |
| 129 | coord_store = get_cached_tcp_store_client( |
| 130 | self.reconfig_request.new_data_parallel_master_ip, |
| 131 | self.reconfig_request.coord_store_port, |
| 132 | ) |
| 133 | self._prepare_future = self._prepare_executor.submit( |
| 134 | coord_store.wait, done_keys |
| 135 | ) |
| 136 | if not self._prepare_future.done(): |
| 137 | return False |
| 138 | |
| 139 | self._prepare_future.result() |
| 140 | self._collective_rpc("elastic_ep_execute", args=("clear_async",)) |
| 141 | self._prepare_future = None |
| 142 | return True |
| 143 | |
| 144 | def progress(self) -> bool: |
| 145 | if self.scale_type == "scale_up": |
| 146 | return ( |
| 147 | self._progress_new_engine() |
| 148 | if self.worker_type == "new" |
| 149 | else self._progress_existing_engine() |
| 150 | ) |
| 151 | return ( |
| 152 | self._progress_removing_engine() |
| 153 | if self.worker_type == "removing" |
| 154 | else self._progress_remaining_engine() |
| 155 | ) |
| 156 | |
| 157 | def run_pre_kv_init_states(self) -> None: |
| 158 | assert self.scale_type == "scale_up" and self.worker_type == "new" |
| 159 | assert self.state == ScaleUpNewEngineState.PRE_KV_INIT |
| 160 | assert self.progress() |
| 161 | assert self.state == ScaleUpNewEngineState.PREPARE |
| 162 | |
| 163 | def _progress_existing_engine(self) -> bool: |
| 164 | state = self.state |
| 165 | assert self.old_dp_group is not None |
| 166 | |
| 167 | if state == ScaleUpExistingEngineState.PREPARE: |
| 168 | if not self._prepare_workers(): |
| 169 | return False |
| 170 | self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE |
| 171 | return True |
| 172 | |
| 173 | elif state == ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE: |
| 174 | if not self._sync_kv_cache_memory_size(): |
| 175 | return False |
| 176 | self.state = ScaleUpExistingEngineState.COMMIT_SCALE_UP |
| 177 | self._mark_ready_for_switch() |
| 178 | return True |
| 179 | |
| 180 | elif state == ScaleUpExistingEngineState.COMMIT_SCALE_UP: |
| 181 | if not self.commit_requested: |
| 182 | return False |
| 183 | self._commit_new_dp_group() |
| 184 | self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", True)) |
| 185 | self.state = ScaleUpExistingEngineState.COMPLETE |
| 186 | self._update_parallel_config() |
| 187 | self._send_reconfigure_finished() |
| 188 | return True |
| 189 | |
| 190 | else: |
| 191 | assert self.state == ScaleUpExistingEngineState.COMPLETE |
| 192 | return True |
| 193 | |
| 194 | def _progress_new_engine(self) -> bool: |
| 195 | state = self.state |
| 196 | assert self.new_dp_group is not None and self.new_dp_store is not None |
| 197 | |
| 198 | if state == ScaleUpNewEngineState.PRE_KV_INIT: |
| 199 | self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",)) |
| 200 | self.engine_core.available_gpu_memory_for_kv_cache = ( |
| 201 | ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1) |
| 202 | ) |
| 203 | self.state = ScaleUpNewEngineState.PREPARE |
| 204 | return True |
| 205 | |
| 206 | elif state == ScaleUpNewEngineState.PREPARE: |
| 207 | self._collective_rpc("elastic_ep_execute", args=("warmup_new_worker",)) |
| 208 | self._mark_ready_for_switch() |
| 209 | tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu") |
| 210 | torch.distributed.all_reduce( |
| 211 | tensor, |
| 212 | op=torch.distributed.ReduceOp.MAX, |
| 213 | group=self.new_dp_group, |
| 214 | ) |
| 215 | data = tensor.tolist() |
| 216 | self.engine_core.engines_running = bool(data[0]) |
| 217 | self.engine_core.current_wave = int(data[1]) |
| 218 | self.engine_core.step_counter = int(data[2]) |
| 219 | self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", False)) |
| 220 | self.state = ScaleUpNewEngineState.COMPLETE |
| 221 | return True |
| 222 | |
| 223 | else: |
| 224 | assert self.state == ScaleUpNewEngineState.COMPLETE |
| 225 | return True |
| 226 | |
| 227 | def _progress_remaining_engine(self) -> bool: |
| 228 | state = self.state |
| 229 | assert self.old_dp_group is not None |
| 230 | |
| 231 | if state == ScaleDownRemainingEngineState.PREPARE: |
| 232 | if self._prepare_workers(): |
| 233 | self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN |
| 234 | self._mark_ready_for_switch() |
| 235 | return True |
| 236 | return False |
| 237 | |
| 238 | elif state == ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN: |
| 239 | if not self.commit_requested: |
| 240 | return False |
| 241 | self._commit_scale_down(removing=False) |
| 242 | self._commit_new_dp_group() |
| 243 | self._update_parallel_config() |
| 244 | self.state = ScaleDownRemainingEngineState.COMPLETE |
| 245 | self._send_reconfigure_finished() |
| 246 | return True |
| 247 | |
| 248 | else: |
| 249 | assert self.state == ScaleDownRemainingEngineState.COMPLETE |
| 250 | return True |
| 251 | |
| 252 | def _progress_removing_engine(self) -> bool: |
| 253 | state = self.state |
| 254 | assert self.old_dp_group is not None |
| 255 | |
| 256 | if state == ScaleDownRemovingEngineState.PREPARE: |
| 257 | assert self.old_dp_group.rank() > 0 |
| 258 | self._commit_scale_down(removing=True) |
| 259 | self.state = ScaleDownRemovingEngineState.COMPLETE |
| 260 | self.engine_core._eep_send_engine_core_notification( |
| 261 | EEPNotificationType.SHUTDOWN_COMPLETE |
| 262 | ) |
| 263 | return True |
| 264 | |
| 265 | else: |
| 266 | assert self.state == ScaleDownRemovingEngineState.COMPLETE |
| 267 | return True |
| 268 | |
| 269 | def is_ready_for_switch(self) -> bool: |
| 270 | return self.worker_type == "existing" and ( |
| 271 | self.state is ScaleUpExistingEngineState.COMMIT_SCALE_UP |
| 272 | or self.state is ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN |
| 273 | ) |
| 274 | |
| 275 | @property |
| 276 | def ready_key(self) -> str: |
| 277 | return f"eep_ready/{self.engine_core.dp_rank}" |
| 278 | |
| 279 | def _mark_ready_for_switch(self) -> None: |
| 280 | parallel_config = self.new_parallel_config |
| 281 | get_cached_tcp_store_client( |
| 282 | parallel_config.data_parallel_master_ip, |
| 283 | parallel_config._coord_store_port, |
| 284 | ).set(self.ready_key, b"1") |
| 285 | |
| 286 | def is_complete(self) -> bool: |
| 287 | if self.scale_type == "scale_up": |
| 288 | return ( |
| 289 | self.state == ScaleUpNewEngineState.COMPLETE |
| 290 | if self.worker_type == "new" |
| 291 | else self.state == ScaleUpExistingEngineState.COMPLETE |
| 292 | ) |
| 293 | return ( |
| 294 | self.state == ScaleDownRemovingEngineState.COMPLETE |
| 295 | if self.worker_type == "removing" |
| 296 | else self.state == ScaleDownRemainingEngineState.COMPLETE |
| 297 | ) |
| 298 | |
| 299 | def _init_new_dp_group(self) -> tuple[Any, Any]: |
| 300 | return self.new_parallel_config.stateless_init_dp_group(return_store=True) |
| 301 | |
| 302 | def _ensure_new_dp_group(self) -> bool: |
| 303 | if self.new_dp_group is not None: |
| 304 | return True |
| 305 | |
| 306 | if self._prepare_future is None: |
| 307 | self._prepare_future = self._prepare_executor.submit( |
| 308 | self._init_new_dp_group |
| 309 | ) |
| 310 | if not self._prepare_future.done(): |
| 311 | return False |
| 312 | |
| 313 | self.new_dp_group, self.new_dp_store = self._prepare_future.result() |
| 314 | self._prepare_future = None |
| 315 | return True |
| 316 | |
| 317 | def _prepare_workers(self) -> bool: |
| 318 | assert self.old_dp_group is not None |
| 319 | if not self._ensure_new_dp_group(): |
| 320 | return False |
| 321 | if not self._execute_async( |
| 322 | "prepare_reconfiguration", |
| 323 | self.reconfig_request, |
| 324 | self.new_parallel_config.use_all2all, |
| 325 | ): |
| 326 | return False |
| 327 | if self.old_dp_group.rank() == 0: |
| 328 | logger.info("[Elastic EP] Prepared reconfiguration") |
| 329 | return True |
| 330 | |
| 331 | def _sync_kv_cache_memory_size(self) -> bool: |
| 332 | assert self.engine_core.available_gpu_memory_for_kv_cache > 0 |
| 333 | assert self.new_dp_group is not None and self.old_dp_group is not None |
| 334 | |
| 335 | if self._new_dp_sync is None: |
| 336 | tensor = torch.tensor( |
| 337 | [self.engine_core.available_gpu_memory_for_kv_cache], |
| 338 | dtype=torch.int64, |
| 339 | device="cpu", |
| 340 | ) |
| 341 | work = torch.distributed.all_reduce( |
| 342 | tensor, |
| 343 | op=torch.distributed.ReduceOp.MIN, |
| 344 | group=self.new_dp_group, |
| 345 | async_op=True, |
| 346 | ) |
| 347 | self._new_dp_sync = (tensor, work) |
| 348 | return False |
| 349 | |
| 350 | _, work = self._new_dp_sync |
| 351 | if not work.is_completed(): |
| 352 | return False |
| 353 | work.wait() |
| 354 | self._new_dp_sync = None |
| 355 | if self.old_dp_group.rank() == 0: |
| 356 | logger.info("[Elastic EP] Synced KV cache memory size to new workers") |
| 357 | return True |
| 358 | |
| 359 | def _commit_new_dp_group(self): |
| 360 | old_dp_group = self.old_dp_group |
| 361 | stateless_destroy_torch_distributed_process_group(old_dp_group) |
| 362 | assert self.new_dp_group is not None |
| 363 | new_dp_group = self.new_dp_group |
| 364 | self.engine_core.dp_group = new_dp_group |
| 365 | self.engine_core.dp_rank = new_dp_group.rank() |
| 366 | self.engine_core.dp_store = self.new_dp_store |
| 367 | engines_running = int(self.engine_core.engines_running) |
| 368 | current_wave = self.engine_core.current_wave |
| 369 | step_counter = self.engine_core.step_counter |
| 370 | tensor = torch.tensor( |
| 371 | [engines_running, current_wave, step_counter], |
| 372 | dtype=torch.int32, |
| 373 | device="cpu", |
| 374 | ) |
| 375 | torch.distributed.all_reduce( |
| 376 | tensor, op=torch.distributed.ReduceOp.MAX, group=new_dp_group |
| 377 | ) |
| 378 | data = tensor.tolist() |
| 379 | self.engine_core.engines_running = bool(data[0]) |
| 380 | self.engine_core.current_wave = int(data[1]) |
| 381 | self.engine_core.step_counter = int(data[2]) |
| 382 | if new_dp_group.rank() == 0: |
| 383 | logger.info("[Elastic EP] Switched to new setup") |
| 384 | |
| 385 | def _send_reconfigure_finished(self): |
| 386 | assert self.new_dp_group is not None |
| 387 | if self.new_dp_group.rank() == 0: |
| 388 | self.engine_core._eep_send_engine_core_notification( |
| 389 | EEPNotificationType.RECONFIGURE_FINISHED |
| 390 | ) |
| 391 | |
| 392 | def _commit_scale_down(self, removing: bool): |
| 393 | assert self.reconfig_request is not None and self.old_dp_group is not None |
| 394 | self._collective_rpc( |
| 395 | "elastic_ep_execute", |
| 396 | args=( |
| 397 | "commit_scale_down", |
| 398 | self.reconfig_request.new_data_parallel_size, |
| 399 | removing, |
| 400 | ), |
| 401 | ) |
| 402 | if self.old_dp_group.rank() == 0: |
| 403 | logger.info("[Elastic EP] EPLB reshuffle completed") |
| 404 | |
| 405 | def _update_parallel_config(self): |
| 406 | assert self.reconfig_request is not None |
| 407 | reconfig_request = self.reconfig_request |
| 408 | parallel_config = self.vllm_config.parallel_config |
| 409 | parallel_config.data_parallel_size = reconfig_request.new_data_parallel_size |
| 410 | if ( |
| 411 | reconfig_request.new_data_parallel_rank |
| 412 | != ReconfigureRankType.KEEP_CURRENT_RANK |
| 413 | ): |
| 414 | parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank |
| 415 | if ( |
| 416 | reconfig_request.new_data_parallel_rank_local |
| 417 | != ReconfigureRankType.KEEP_CURRENT_RANK |
| 418 | ): |
| 419 | parallel_config.data_parallel_rank_local = ( |
| 420 | reconfig_request.new_data_parallel_rank_local |
| 421 | ) |
| 422 | parallel_config.data_parallel_master_ip = ( |
| 423 | reconfig_request.new_data_parallel_master_ip |
| 424 | ) |
| 425 | parallel_config.data_parallel_master_port = ( |
| 426 | reconfig_request.new_data_parallel_master_port |
| 427 | ) |
| 428 | parallel_config._data_parallel_master_port_list = ( |
| 429 | reconfig_request.new_data_parallel_master_port_list |
| 430 | ) |
| 431 | parallel_config._coord_store_port = reconfig_request.coord_store_port |