| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copy |
| 4 | import multiprocessing |
| 5 | import multiprocessing.connection |
| 6 | import time |
| 7 | import weakref |
| 8 | |
| 9 | import msgspec.msgpack |
| 10 | import zmq |
| 11 | |
| 12 | from vllm.config import LoggingConfig, ParallelConfig |
| 13 | from vllm.logger import configure_logging, init_logger |
| 14 | from vllm.utils.network_utils import make_zmq_socket |
| 15 | from vllm.utils.system_utils import get_mp_context, set_process_title |
| 16 | from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType |
| 17 | from vllm.v1.serial_utils import MsgpackDecoder |
| 18 | from vllm.v1.utils import get_engine_client_zmq_addr, shutdown |
| 19 | |
| 20 | logger = init_logger(__name__) |
| 21 | |
| 22 | |
| 23 | class DPCoordinator: |
| 24 | """Coordinator process used for data-parallel deployments (DP>1). |
| 25 | |
| 26 | Intermediates between multiple DP engine rank processes and one or more |
| 27 | front-end API server processes. |
| 28 | |
| 29 | * Collects stats from each DP engine (currently just waiting and running |
| 30 | queue lengths), and publishes these to all front-ends for use in |
| 31 | load-balancing decisions. |
| 32 | |
| 33 | * Keeps track of the current DP "request wave" number and running state |
| 34 | of the engines. This is received from the DP rank 0 engine and published |
| 35 | to the front-end processes along with the current load stats. |
| 36 | |
| 37 | The engines alternate between a global running/paused state. The global |
| 38 | "request wave" number is a count of the number of times that the workers |
| 39 | collectively move from a running state to a paused state. This transition |
| 40 | is synchronized via the all-reduce operation performed in the |
| 41 | DPEngineCoreProc._has_global_unfinished_reqs method. |
| 42 | |
| 43 | * Broadcasts the START_DP_WAVE message to engines to move them from paused |
| 44 | to running state when one engine receives a new request. This can happen |
| 45 | in two cases: |
| 46 | 1) A front-end sending a new request while the engines are paused will |
| 47 | concurrently notify the coordinator. |
| 48 | 2) An engine receiving a request for a stale request wave while in paused |
| 49 | state will notify the coordinator. |
| 50 | |
| 51 | Engines will move into running state when receiving a new request or |
| 52 | START_DP_WAVE message. |
| 53 | |
| 54 | Note that when deployed in External LB mode, no stats will be published by |
| 55 | the engines and thus updates will only be sent to front-ends when the |
| 56 | request wave / running state changes. |
| 57 | """ |
| 58 | |
| 59 | def _wait_for_zmq_addrs(self, zmq_addr_pipe) -> tuple[str, str, str]: |
| 60 | try: |
| 61 | timeout = 120 |
| 62 | ready = multiprocessing.connection.wait( |
| 63 | [zmq_addr_pipe, self.proc.sentinel], timeout=timeout |
| 64 | ) |
| 65 | if not ready: |
| 66 | raise RuntimeError( |
| 67 | "DP Coordinator process failed to report ZMQ addresses " |
| 68 | f"within timeout={timeout} seconds during startup." |
| 69 | ) |
| 70 | try: |
| 71 | return zmq_addr_pipe.recv() |
| 72 | except EOFError: |
| 73 | raise RuntimeError( |
| 74 | "DP Coordinator process failed during startup." |
| 75 | ) from None |
| 76 | finally: |
| 77 | zmq_addr_pipe.close() |
| 78 | |
| 79 | def __init__( |
| 80 | self, |
| 81 | parallel_config: ParallelConfig, |
| 82 | enable_wave_coordination: bool = True, |
| 83 | logging_config: LoggingConfig | None = None, |
| 84 | ): |
| 85 | dp_size = parallel_config.data_parallel_size |
| 86 | assert dp_size > 1, "Coordinator only used for data parallel" |
| 87 | |
| 88 | host = parallel_config.data_parallel_master_ip |
| 89 | |
| 90 | # Assume coordinator is colocated with front-end procs when not in |
| 91 | # either external or hybrid DP LB mode. |
| 92 | local_only = not parallel_config.local_engines_only |
| 93 | local_only_eng = dp_size == parallel_config.data_parallel_size_local |
| 94 | # NOTE(yongji): handling scaling from intra-node to inter-node |
| 95 | if parallel_config.enable_elastic_ep: |
| 96 | local_only_eng = False |
| 97 | |
| 98 | front_publish_address = get_engine_client_zmq_addr(local_only, host=host) |
| 99 | back_publish_address = get_engine_client_zmq_addr(local_only_eng, host=host) |
| 100 | back_output_address = get_engine_client_zmq_addr(local_only_eng, host=host) |
| 101 | |
| 102 | context = get_mp_context() |
| 103 | parent_zmq_addr_pipe, child_zmq_addr_pipe = context.Pipe(duplex=False) |
| 104 | self.proc: multiprocessing.Process = context.Process( |
| 105 | target=DPCoordinatorProc.run_coordinator, |
| 106 | name="VLLM_DP_Coordinator", |
| 107 | kwargs={ |
| 108 | "engine_count": parallel_config.data_parallel_size, |
| 109 | "front_publish_address": front_publish_address, |
| 110 | "back_output_address": back_output_address, |
| 111 | "back_publish_address": back_publish_address, |
| 112 | "zmq_addr_pipe": child_zmq_addr_pipe, |
| 113 | "enable_wave_coordination": enable_wave_coordination, |
| 114 | "logging_config": logging_config, |
| 115 | }, |
| 116 | daemon=True, |
| 117 | ) |
| 118 | self.proc.start() |
| 119 | child_zmq_addr_pipe.close() |
| 120 | ( |
| 121 | front_publish_address, |
| 122 | back_output_address, |
| 123 | back_publish_address, |
| 124 | ) = self._wait_for_zmq_addrs(parent_zmq_addr_pipe) |
| 125 | |
| 126 | self.stats_publish_address = front_publish_address |
| 127 | self.coord_in_address = back_publish_address |
| 128 | self.coord_out_address = back_output_address |
| 129 | self._finalizer = weakref.finalize(self, shutdown, [self.proc]) |
| 130 | |
| 131 | def get_stats_publish_address(self) -> str: |
| 132 | return self.stats_publish_address |
| 133 | |
| 134 | def get_engine_socket_addresses(self) -> tuple[str, str]: |
| 135 | """Returns tuple of ZMQ input address, output address.""" |
| 136 | return self.coord_in_address, self.coord_out_address |
| 137 | |
| 138 | def shutdown(self, timeout: float | None = None) -> None: |
| 139 | """Shutdown coordinator process with configurable timeout.""" |
| 140 | if self._finalizer.detach() is not None: |
| 141 | shutdown([self.proc], timeout=timeout) |
| 142 | |
| 143 | |
| 144 | class EngineState: |
| 145 | def __init__(self): |
| 146 | # [waiting, running, kv_cache_usage] |
| 147 | self.request_counts: list[int | float] = [0, 0, 0.0] |
| 148 | |
| 149 | |
| 150 | class DPCoordinatorProc: |
| 151 | def __init__( |
| 152 | self, |
| 153 | engine_count: int, |
| 154 | min_stats_update_interval_ms: int = 100, |
| 155 | enable_wave_coordination: bool = True, |
| 156 | ): |
| 157 | set_process_title("DPCoordinator") |
| 158 | self.ctx = zmq.Context() |
| 159 | |
| 160 | self.engines = [EngineState() for _ in range(engine_count)] |
| 161 | |
| 162 | self.stats_update_interval_ms = min_stats_update_interval_ms |
| 163 | self.enable_wave_coordination = enable_wave_coordination |
| 164 | |
| 165 | @staticmethod |
| 166 | def run_coordinator( |
| 167 | engine_count: int, |
| 168 | front_publish_address: str, |
| 169 | back_output_address: str, |
| 170 | back_publish_address: str, |
| 171 | zmq_addr_pipe=None, |
| 172 | min_stats_update_interval_ms: int = 100, |
| 173 | enable_wave_coordination: bool = True, |
| 174 | logging_config: LoggingConfig | None = None, |
| 175 | ): |
| 176 | if logging_config is not None: |
| 177 | configure_logging(logging_config) |
| 178 | |
| 179 | coordinator = DPCoordinatorProc( |
| 180 | engine_count=engine_count, |
| 181 | min_stats_update_interval_ms=min_stats_update_interval_ms, |
| 182 | enable_wave_coordination=enable_wave_coordination, |
| 183 | ) |
| 184 | try: |
| 185 | coordinator.process_input_socket( |
| 186 | front_publish_address, |
| 187 | back_output_address, |
| 188 | back_publish_address, |
| 189 | zmq_addr_pipe, |
| 190 | ) |
| 191 | except KeyboardInterrupt: |
| 192 | logger.info("DP Coordinator process exiting") |
| 193 | finally: |
| 194 | if zmq_addr_pipe is not None: |
| 195 | zmq_addr_pipe.close() |
| 196 | |
| 197 | def process_input_socket( |
| 198 | self, |
| 199 | front_publish_address: str, |
| 200 | back_output_address: str, |
| 201 | back_publish_address: str, |
| 202 | zmq_addr_pipe=None, |
| 203 | ): |
| 204 | decoder = MsgpackDecoder(EngineCoreOutputs) |
| 205 | |
| 206 | # For tracking request wave progression. |
| 207 | current_wave = 0 |
| 208 | engines_running = False |
| 209 | |
| 210 | # For tracking request counts for internal load-balancing. |
| 211 | stats_changed = False |
| 212 | last_stats_step = -1 |
| 213 | last_stats_wave = -1 |
| 214 | last_step_counts: list[list[int | float]] | None = None |
| 215 | |
| 216 | with ( |
| 217 | make_zmq_socket( |
| 218 | path=front_publish_address, # IPC |
| 219 | ctx=self.ctx, |
| 220 | socket_type=zmq.XPUB, |
| 221 | bind=True, |
| 222 | ) as publish_front, |
| 223 | make_zmq_socket( |
| 224 | path=back_output_address, # IPC or TCP |
| 225 | ctx=self.ctx, |
| 226 | socket_type=zmq.PULL, |
| 227 | bind=True, |
| 228 | ) as output_back, |
| 229 | make_zmq_socket( |
| 230 | path=back_publish_address, # IPC or TCP |
| 231 | ctx=self.ctx, |
| 232 | socket_type=zmq.XPUB, |
| 233 | bind=True, |
| 234 | ) as publish_back, |
| 235 | ): |
| 236 | if zmq_addr_pipe is not None: |
| 237 | try: |
| 238 | zmq_addr_pipe.send( |
| 239 | ( |
| 240 | publish_front.getsockopt(zmq.LAST_ENDPOINT).decode(), |
| 241 | output_back.getsockopt(zmq.LAST_ENDPOINT).decode(), |
| 242 | publish_back.getsockopt(zmq.LAST_ENDPOINT).decode(), |
| 243 | ) |
| 244 | ) |
| 245 | finally: |
| 246 | zmq_addr_pipe.close() |
| 247 | # Wait until all engines subscribe. |
| 248 | for _ in self.engines: |
| 249 | if publish_back.recv() != b"\x01": |
| 250 | logger.error( |
| 251 | "DP Coordinator received unexpected message while " |
| 252 | "waiting for engines to subscribe" |
| 253 | ) |
| 254 | return |
| 255 | # Send ready message to engines. |
| 256 | publish_back.send(b"READY") |
| 257 | |
| 258 | logger.info("All engine subscriptions received by DP coordinator") |
| 259 | |
| 260 | poller = zmq.Poller() |
| 261 | poller.register(publish_front, zmq.POLLIN) |
| 262 | poller.register(publish_back, zmq.POLLIN) |
| 263 | poller.register(output_back, zmq.POLLIN) |
| 264 | last_publish_time = 0 |
| 265 | while True: |
| 266 | elapsed = int(time.time() * 1000) - last_publish_time |
| 267 | # Send at stats_update_interval_ms interval if the stats have |
| 268 | # changed, or otherwise every 5 seconds. |
| 269 | wait_for = self.stats_update_interval_ms if stats_changed else 5000 |
| 270 | |
| 271 | # Wait at least 50ms to ensure we've received all stats for |
| 272 | # the current step. Only applicable to lockstep (MoE) DP; |
| 273 | # non-lockstep engines have no synchronized step boundaries. |
| 274 | if self.enable_wave_coordination and last_step_counts is None: |
| 275 | min_timeout = 50 |
| 276 | else: |
| 277 | min_timeout = 0 |
| 278 | |
| 279 | events = poller.poll(timeout=max(min_timeout, wait_for - elapsed)) |
| 280 | if not events: |
| 281 | # Poller timeout - publish current stats to front-ends. |
| 282 | if last_step_counts is not None: |
| 283 | engine_req_counts_list = last_step_counts |
| 284 | last_step_counts = None |
| 285 | else: |
| 286 | engine_req_counts_list = self._get_engine_counts() |
| 287 | stats_changed = False |
| 288 | |
| 289 | to_publish = (engine_req_counts_list, current_wave, engines_running) |
| 290 | publish_front.send(msgspec.msgpack.encode(to_publish)) |
| 291 | last_publish_time = int(time.time() * 1000) |
| 292 | continue |
| 293 | |
| 294 | events = dict(events) |
| 295 | wave_state_changed = False |
| 296 | |
| 297 | if publish_back in events: |
| 298 | buffer = publish_back.recv() |
| 299 | if buffer == b"\x01": |
| 300 | # NOTE(yongji): newly started engine subscribed |
| 301 | # We need to send READY message here instead of receiving |
| 302 | # SCALE_ELASTIC_EP notification from engine core client |
| 303 | # as SCALE_ELASTIC_EP is only sent when |
| 304 | # new engines finished initialization. |
| 305 | # Subscription message, on the other hand, is sent |
| 306 | # by each engine during initialization |
| 307 | publish_back.send(b"READY") |
| 308 | elif buffer != b"\x00": |
| 309 | logger.error( |
| 310 | "DP Coordinator received unexpected message from engines" |
| 311 | ) |
| 312 | |
| 313 | if publish_front in events: |
| 314 | buffer = publish_front.recv() |
| 315 | if buffer in (b"\x01", b"\x00"): |
| 316 | # Ignore subscription messages. |
| 317 | continue |
| 318 | |
| 319 | decoded = msgspec.msgpack.decode(buffer) |
| 320 | if ( |
| 321 | isinstance(decoded, (list, tuple)) |
| 322 | and len(decoded) == 2 |
| 323 | and decoded[0] == "SCALE_ELASTIC_EP" |
| 324 | ): |
| 325 | # Handle scale up notification |
| 326 | new_engine_count = decoded[1] |
| 327 | current_count = len(self.engines) |
| 328 | if new_engine_count > current_count: |
| 329 | for _ in range(new_engine_count - current_count): |
| 330 | self.engines.append(EngineState()) |
| 331 | # NOTE(yongji): handle the case |
| 332 | # where newly started engines have current_wave = 0 |
| 333 | # if existing engines just finished a wave |
| 334 | # and engine_running isn't updated yet at |
| 335 | # CoordinatorProc requests routed to newly started |
| 336 | # engines may not wake up existing engines, as long |
| 337 | # as 0 < request.wave < existing engines' |
| 338 | # current_wave |
| 339 | # we note that 0 is the wave number for the new |
| 340 | # engine |
| 341 | logger.info( |
| 342 | "DPCoordinator scaled up from %s to %s engines", |
| 343 | current_count, |
| 344 | new_engine_count, |
| 345 | ) |
| 346 | else: |
| 347 | self.engines = self.engines[:new_engine_count] |
| 348 | logger.info( |
| 349 | "DPCoordinator scaled down from %s to %s engines", |
| 350 | current_count, |
| 351 | new_engine_count, |
| 352 | ) |
| 353 | continue # Skip normal engine notification processing |
| 354 | |
| 355 | # Wave coordination: handle new-request messages from front-end. |
| 356 | # Only process these when wave coordination is enabled |
| 357 | if self.enable_wave_coordination: |
| 358 | # We received a message on the front-end XPUB socket, |
| 359 | # from an API server sending a new request while the |
| 360 | # engines are paused, so that we can wake the other |
| 361 | # engines. |
| 362 | engine_to_exclude, wave = decoded |
| 363 | if not engines_running: |
| 364 | if wave < current_wave: |
| 365 | # If the wave number is stale, ensure the message |
| 366 | # is handled by all the engines. |
| 367 | engine_to_exclude = None |
| 368 | |
| 369 | # engines_running is only set from the engines' |
| 370 | # own notifications; a paused engine discards |
| 371 | # START_DP_WAVE, so sending it is not evidence |
| 372 | # that the engines are running. |
| 373 | self._send_start_wave( |
| 374 | publish_back, current_wave, engine_to_exclude |
| 375 | ) |
| 376 | |
| 377 | if output_back in events: |
| 378 | # We received a message from one of the engines. |
| 379 | |
| 380 | buffer = output_back.recv() |
| 381 | outputs: EngineCoreOutputs = decoder.decode(buffer) |
| 382 | |
| 383 | assert not outputs.outputs |
| 384 | assert outputs.utility_output is None |
| 385 | |
| 386 | eng_index = outputs.engine_index |
| 387 | scheduler_stats = outputs.scheduler_stats |
| 388 | if scheduler_stats: |
| 389 | # Elastic EP stats may arrive while the engine list changes. |
| 390 | if eng_index >= len(self.engines): |
| 391 | continue |
| 392 | # 1. Updated request load stats - update our local |
| 393 | # state with these. |
| 394 | stats = self.engines[eng_index].request_counts |
| 395 | if self.enable_wave_coordination: |
| 396 | # Steps are synchronized across lockstep (MoE) DP |
| 397 | # ranks; snapshot counts at step boundaries. |
| 398 | stats_step = scheduler_stats.step_counter |
| 399 | stats_wave = scheduler_stats.current_wave |
| 400 | if ( |
| 401 | stats_wave > last_stats_wave |
| 402 | or stats_wave == last_stats_wave |
| 403 | and stats_step > last_stats_step |
| 404 | ): |
| 405 | if stats_changed: |
| 406 | last_step_counts = self._get_engine_counts( |
| 407 | do_copy=True |
| 408 | ) |
| 409 | last_stats_step = stats_step |
| 410 | last_stats_wave = stats_wave |
| 411 | elif stats_wave != last_stats_wave or ( |
| 412 | stats_step != last_stats_step |
| 413 | ): |
| 414 | logger.warning( |
| 415 | "Received stats for out-of-order " |
| 416 | "step (%d, %d) from engine %d (expected " |
| 417 | "> (%d, %d))", |
| 418 | stats_wave, |
| 419 | stats_step, |
| 420 | eng_index, |
| 421 | last_stats_wave, |
| 422 | last_stats_step, |
| 423 | ) |
| 424 | stats[0] = scheduler_stats.num_waiting_reqs |
| 425 | stats[1] = scheduler_stats.num_running_reqs |
| 426 | stats[2] = scheduler_stats.kv_cache_usage |
| 427 | stats_changed = True |
| 428 | |
| 429 | # Wave coordination: handle wave completion and start notifications |
| 430 | # Only process these when wave coordination is enabled |
| 431 | if self.enable_wave_coordination: |
| 432 | if (wave := outputs.wave_complete) is not None: |
| 433 | # 2. Notification from rank 0 engine that we've |
| 434 | # moved into the global paused state |
| 435 | # (engines_running==False). |
| 436 | if current_wave <= wave: |
| 437 | new_wave = wave + 1 |
| 438 | logger.debug( |
| 439 | "Moving DP wave from %d to %d.", |
| 440 | current_wave, |
| 441 | new_wave, |
| 442 | ) |
| 443 | current_wave = new_wave |
| 444 | engines_running = False |
| 445 | wave_state_changed = True |
| 446 | elif (wave := outputs.start_wave) is not None and ( |
| 447 | wave > current_wave |
| 448 | or (wave == current_wave and not engines_running) |
| 449 | ): |
| 450 | # 3. The engine received request for a non-current wave |
| 451 | # so we must ensure that other engines progress to the |
| 452 | # next wave (race condition handling). |
| 453 | logger.debug( |
| 454 | "Starting wave %d after notification of " |
| 455 | "stale wave request from engine.", |
| 456 | wave, |
| 457 | ) |
| 458 | current_wave = wave |
| 459 | engines_running = True |
| 460 | wave_state_changed = True |
| 461 | self._send_start_wave(publish_back, wave, eng_index) |
| 462 | |
| 463 | if wave_state_changed: |
| 464 | message = (None, current_wave, engines_running) |
| 465 | publish_front.send(msgspec.msgpack.encode(message)) |
| 466 | |
| 467 | @staticmethod |
| 468 | def _send_start_wave( |
| 469 | socket: zmq.Socket, wave: int, exclude_engine_index: int | None |
| 470 | ): |
| 471 | """Broadcast the START_DP_WAVE message to all the engines. |
| 472 | It includes the current wave number and index of engine which |
| 473 | has already received a request with this wave number and so doesn't |
| 474 | require additional notification. |
| 475 | """ |
| 476 | wave_encoded = msgspec.msgpack.encode((wave, exclude_engine_index)) |
| 477 | socket.send_multipart((EngineCoreRequestType.START_DP_WAVE.value, wave_encoded)) |
| 478 | |
| 479 | def _get_engine_counts(self, do_copy=False) -> list[list[int | float]]: |
| 480 | """Return list of [waiting, running] count lists for each engine.""" |
| 481 | if do_copy: |
| 482 | return [copy.copy(e.request_counts) for e in self.engines] |
| 483 | return [e.request_counts for e in self.engines] |