main
py 1,465 lines 56.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import contextlib
5 import os
6 import threading
7 import weakref
8 from collections.abc import Callable, Iterator, Sequence
9 from dataclasses import dataclass
10 from enum import Enum, auto
11 from multiprocessing import connection
12 from multiprocessing.process import BaseProcess
13 from multiprocessing.queues import Queue
14 from typing import TYPE_CHECKING, cast
15
16 import msgspec
17 import zmq
18
19 from vllm import envs
20 from vllm.config import CacheConfig, ParallelConfig, VllmConfig
21 from vllm.logger import init_logger
22 from vllm.platforms import current_platform
23 from vllm.ray.ray_env import get_env_vars_to_copy
24 from vllm.utils import numa_utils
25 from vllm.utils.network_utils import (
26 get_open_port,
27 get_open_zmq_ipc_path,
28 get_tcp_uri,
29 zmq_socket_ctx,
30 )
31 from vllm.utils.system_utils import get_mp_context, set_env_var
32 from vllm.v1.engine.coordinator import DPCoordinator
33 from vllm.v1.executor import Executor, UniProcExecutor
34 from vllm.v1.executor.ray_utils import WORKER_SPECIFIC_ENV_VARS
35 from vllm.v1.utils import _SubprocessWrapper, get_engine_client_zmq_addr, shutdown
36
37 if TYPE_CHECKING:
38 from ray.util.placement_group import PlacementGroup
39
40 logger = init_logger(__name__)
41
42 STARTUP_POLL_PERIOD_MS = 10000
43 ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S = 15.0
44
45
46 def get_engine_process_shutdown_timeout(
47 request_timeout: float | None,
48 process_timeout: float | None,
49 ) -> float | None:
50 """Return the EngineCore process-manager shutdown timeout.
51
52 ``VllmConfig.shutdown_timeout`` controls how long in-flight requests may
53 drain. A value of zero therefore tells EngineCore to abort requests as soon
54 as it receives SIGTERM. The parent process manager still needs a separate
55 window in which the EngineCore can release device resources before it is
56 force-killed. ROCm teardown can take longer than the generic best-effort
57 window, and force-killing during teardown can leave VRAM resident.
58
59 ``process_timeout`` may be a remaining budget computed by an outer process
60 manager. Keep it unchanged unless both values are zero: a zero remaining
61 budget for a positive request timeout must not receive a fresh grace period
62 because EngineCore relies on that deadline to enforce request draining.
63 """
64 if request_timeout == 0 and process_timeout == 0 and current_platform.is_rocm():
65 return ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S
66 return process_timeout
67
68
69 class CoreEngineState(Enum):
70 NEW = auto()
71 CONNECTED = auto()
72 READY = auto()
73
74
75 class CoreEngine:
76 """One per data parallel rank, used to track state during handshaking."""
77
78 def __init__(self, index: int = 0, local: bool = True):
79 self.local = local
80 self.identity = index.to_bytes(2, "little")
81
82 self.state = CoreEngineState.NEW
83
84
85 @dataclass
86 class EngineZmqAddresses:
87 # ZMQ input socket addresses for each front-end client (requests)
88 inputs: list[str]
89 # ZMQ output socket addresses for each front-end client (responses)
90 outputs: list[str]
91 # ZMQ input socket address of DP coordinator if applicable
92 coordinator_input: str | None = None
93 # ZMQ output socket address of DP coordinator if applicable
94 coordinator_output: str | None = None
95 # ZMQ socket for front-end to connect to DP coordinator.
96 # Not used by engine, just relayed to front-end in handshake response.
97 # Only required for external DP LB case.
98 frontend_stats_publish_address: str | None = None
99
100
101 @dataclass
102 class EngineHandshakeMetadata:
103 """Metadata sent to each engine process during startup handshake,
104 including addresses of the front-end ZMQ queues that they should
105 connect to.
106 """
107
108 addresses: EngineZmqAddresses
109 parallel_config: dict[str, int | str | list[int]]
110
111
112 def _make_control_bundle(node_ip: str) -> dict[str, float]:
113 # The engine actor is scheduled on the final CPU-only bundle. Keep that
114 # bundle colocated with the group's first GPU bundle so the actor does not
115 # float to an unrelated node and reorder worker ranks away from the
116 # advertised DP bootstrap host.
117 return {"CPU": 1.0, "node:" + node_ip: 0.001}
118
119
120 def _get_bundle_node_ip(bundle: dict[str, float]) -> str:
121 for key in bundle:
122 if key.startswith("node:"):
123 return key.split(":", 1)[1]
124 raise ValueError(f"Missing node affinity in placement bundle: {bundle}")
125
126
127 def _dp_nodes_master_first(
128 available_resources: dict[str, dict], dp_master_ip: str
129 ) -> list[tuple[str, dict]]:
130 """Live Ray nodes as ``(node_id, resources)`` with the DP master first.
131
132 The master is found through the ``node:<ip>`` resource Ray sets on every
133 node, so this works with the resource maps from ``ray._private.state``.
134 Do not replace it with ``ray.util.state.list_nodes()``: that API goes
135 through the dashboard HTTP server and only exists with ``ray[default]``
136 (see #23822).
137 """
138 master_key = f"node:{dp_master_ip}"
139 nodes = sorted(
140 available_resources.items(), key=lambda item: master_key not in item[1]
141 )
142 assert len(nodes) > 0, "No nodes with resources found in Ray cluster."
143 assert master_key in nodes[0][1], (
144 f"The DP master node (ip: {dp_master_ip}) is missing or dead"
145 )
146 assert len(nodes) == 1 or master_key not in nodes[1][1], (
147 "There can only be one head node"
148 )
149 return nodes
150
151
152 def _node_ip_from_resources(node_resources: dict) -> str | None:
153 """Return the node IP encoded in a Ray per-node resource dict, or None.
154
155 Ray advertises each node's IP as a ``node:<ip>`` resource key. The head node
156 also carries ``node:__internal_head__``, and placement groups add
157 ``..._group_...`` keys; both are ignored.
158 """
159 for key in node_resources:
160 if (
161 key.startswith("node:")
162 and key != "node:__internal_head__"
163 and "_group_" not in key
164 ):
165 return key.split(":", 1)[1]
166 return None
167
168
169 @contextlib.contextmanager
170 def _configure_uniproc_startup_threads(
171 executor_class: type[Executor], local_engine_count: int
172 ) -> Iterator[None]:
173 if (
174 not issubclass(executor_class, UniProcExecutor)
175 or current_platform.is_cpu()
176 or "OMP_NUM_THREADS" in os.environ
177 ):
178 yield
179 return
180
181 import torch
182
183 from vllm.utils.torch_utils import (
184 OMP_NUM_THREADS_SET_BY_VLLM,
185 set_default_torch_num_threads,
186 startup_omp_num_threads,
187 )
188
189 original_threads = torch.get_num_threads()
190 num_threads = min(original_threads, startup_omp_num_threads(local_engine_count))
191 torch_threads = (
192 set_default_torch_num_threads(num_threads)
193 if num_threads < original_threads
194 else contextlib.nullcontext()
195 )
196 # Fork inherits Torch's setting; spawn reads OMP_NUM_THREADS on import.
197 with (
198 set_env_var("OMP_NUM_THREADS", str(num_threads)),
199 set_env_var(OMP_NUM_THREADS_SET_BY_VLLM, "1"),
200 torch_threads,
201 ):
202 yield
203
204
205 class CoreEngineProcManager:
206 """Utility class to handle creation, readiness, and shutdown
207 of background processes used by the AsyncLLM and LLMEngine.
208 """
209
210 def __init__(
211 self,
212 local_engine_count: int,
213 start_index: int,
214 local_start_index: int,
215 vllm_config: VllmConfig,
216 local_client: bool,
217 handshake_address: str,
218 executor_class: type[Executor],
219 log_stats: bool,
220 client_handshake_address: str | None = None,
221 tensor_queue: Queue | None = None,
222 ):
223 self._request_shutdown_timeout = vllm_config.shutdown_timeout
224 context = get_mp_context()
225 common_kwargs = {
226 "vllm_config": vllm_config,
227 "local_client": local_client,
228 "handshake_address": handshake_address,
229 "executor_class": executor_class,
230 "log_stats": log_stats,
231 "tensor_queue": tensor_queue,
232 }
233
234 if client_handshake_address:
235 common_kwargs["client_handshake_address"] = client_handshake_address
236
237 is_dp = vllm_config.parallel_config.data_parallel_size > 1
238
239 from vllm.v1.engine.core import EngineCoreProc
240
241 self.processes: list[BaseProcess] = []
242 local_dp_ranks = []
243 for index in range(local_engine_count):
244 local_index = local_start_index + index
245 global_index = start_index + index
246
247 # Start EngineCore in background process.
248 local_dp_ranks.append(local_index)
249 self.processes.append(
250 context.Process(
251 target=EngineCoreProc.run_engine_core,
252 name=f"EngineCore_DP{global_index}" if is_dp else "EngineCore",
253 kwargs=common_kwargs
254 | {"dp_rank": global_index, "local_dp_rank": local_index},
255 )
256 )
257
258 self._finalizer = weakref.finalize(self, shutdown, self.processes)
259 self.manager_stopped = threading.Event()
260 self.failed_proc_name: str | None = None
261
262 # All ranks share this config object: capture the user-provided
263 # --device-ids list before the per-rank shard overwrites it. Mutating
264 # the config before each proc.start() works because the spawn method
265 # pickles process args at start() time, sequentially per rank.
266 user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids
267 try:
268 with _configure_uniproc_startup_threads(executor_class, local_engine_count):
269 for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
270 # Populate the logical-to-physical GPU mapping in DP for
271 # platforms that cannot rely on
272 # torch.accelerator.set_device_index(), and for Ray.
273 needs_device_env_isolation = not (
274 current_platform.is_cuda_alike() or current_platform.is_xpu()
275 )
276 if is_dp and (
277 needs_device_env_isolation
278 or vllm_config.parallel_config.use_ray
279 ):
280 set_assigned_physical_gpu_ids_for_dp_rank(
281 vllm_config, local_dp_rank, user_assigned_gpu_ids
282 )
283
284 with numa_utils.configure_subprocess(
285 # EngineCore itself does not have a TP/PP-local rank.
286 # set_assigned_physical_gpu_ids_for_dp_rank() populates
287 # the logical-to-physical mapping for this DP
288 # shard, so local_rank=0 means "the first local GPU in
289 # this shard". The actual TP/PP worker processes spawned
290 # by the executor are bound separately with their own
291 # local_rank values.
292 vllm_config,
293 local_rank=0,
294 dp_local_rank=local_dp_rank,
295 process_kind="EngineCore",
296 ):
297 proc.start()
298 finally:
299 # Kill other procs if not all are running.
300 if self.finished_procs():
301 self.shutdown()
302
303 def shutdown(self, timeout: float | None = None) -> None:
304 """Shutdown engine core processes with configurable timeout."""
305 self.manager_stopped.set()
306 if self._finalizer.detach() is not None:
307 process_timeout = get_engine_process_shutdown_timeout(
308 self._request_shutdown_timeout, timeout
309 )
310 if process_timeout != timeout:
311 logger.info(
312 "[shutdown] EngineCore process manager: using %ss ROCm "
313 "cleanup grace after immediate request abort",
314 process_timeout,
315 )
316 shutdown(self.processes, timeout=process_timeout)
317
318 def monitor_engine_liveness(self) -> None:
319 """Monitor engine core process liveness."""
320 sentinel_to_proc = {proc.sentinel: proc for proc in self.processes}
321 sentinels = set(sentinel_to_proc.keys())
322
323 while sentinels and not self.manager_stopped.is_set():
324 died_sentinels = connection.wait(sentinels, timeout=1)
325
326 for sentinel in died_sentinels:
327 proc = sentinel_to_proc.pop(cast(int, sentinel))
328 exitcode = proc.exitcode
329 if exitcode != 0 and not self.manager_stopped.is_set():
330 self.failed_proc_name = proc.name
331 if died_sentinels:
332 break
333
334 self.shutdown()
335
336 def sentinels(self) -> list:
337 return [proc.sentinel for proc in self.processes]
338
339 def finished_procs(self) -> dict[str, int]:
340 """Returns dict of proc name -> exit code for any finished procs."""
341 return {
342 proc.name: proc.exitcode
343 for proc in self.processes
344 if proc.exitcode is not None
345 }
346
347
348 class SignalCallback:
349 """Safely trigger a callback from signal handler context via a dedicated thread."""
350
351 def __init__(self, callback: Callable[[], None]):
352 self._callback = callback
353 self._event = threading.Event()
354 self._stopped = False
355 self._thread = threading.Thread(
356 target=self._run,
357 daemon=True,
358 name="signal-callback",
359 )
360 self._thread.start()
361
362 def _run(self):
363 self._event.wait()
364 if not self._stopped:
365 self._callback()
366
367 def trigger(self):
368 self._event.set()
369
370 def stop(self):
371 self._stopped = True
372 self._event.set()
373
374
375 def set_assigned_physical_gpu_ids_for_dp_rank(
376 vllm_config: VllmConfig,
377 local_dp_rank: int,
378 user_assigned_gpu_ids: list[int] | None = None,
379 ) -> None:
380 """Populate assigned_physical_gpu_ids on the config for the given DP rank.
381
382 user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the
383 user provided one; this DP rank's shard is sliced from it. It is passed
384 explicitly rather than read from the config because callers may reuse
385 one config object across DP ranks, overwriting the field each time.
386 """
387 world_size = vllm_config.parallel_config.world_size
388 local_world_size = vllm_config.parallel_config.local_world_size
389 evar = current_platform.device_control_env_var
390
391 physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
392 evar,
393 local_dp_rank,
394 world_size,
395 local_world_size,
396 user_assigned_gpu_ids=user_assigned_gpu_ids,
397 )
398 vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids
399
400
401 def get_physical_gpu_ids_for_local_dp_rank(
402 device_control_env_var: str,
403 local_dp_rank: int,
404 world_size: int,
405 local_world_size: int | None = None,
406 user_assigned_gpu_ids: list[int] | None = None,
407 ) -> list[int]:
408 """Returns list of physical GPU IDs for the specified
409 data parallel rank.
410
411 For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,
412 this will return [2, 3] for local_dp_rank=1.
413
414 If user_assigned_gpu_ids is provided (e.g. from --device-ids), this DP
415 rank's shard is sliced from it instead of being derived from the
416 device-control env var.
417 """
418 if local_world_size is None:
419 local_world_size = world_size
420 if user_assigned_gpu_ids is not None:
421 start = local_dp_rank * world_size
422 stop = start + local_world_size
423 if stop > len(user_assigned_gpu_ids):
424 raise ValueError(
425 f"--device-ids provides {len(user_assigned_gpu_ids)} devices, "
426 f"but DP rank {local_dp_rank} needs devices [{start}, {stop})"
427 )
428 return user_assigned_gpu_ids[start:stop]
429 try:
430 return [
431 current_platform.device_id_to_physical_device_id(i)
432 for i in range(
433 local_dp_rank * world_size,
434 local_dp_rank * world_size + local_world_size,
435 )
436 ]
437 except IndexError as e:
438 raise Exception(
439 f"Error computing device indices for "
440 f"{device_control_env_var}: "
441 f"local range: [{local_dp_rank * world_size}, "
442 f"{(local_dp_rank + 1) * world_size}) "
443 "base value: "
444 f'"{os.getenv(device_control_env_var)}"'
445 ) from e
446
447
448 def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None:
449 # Ray actor names (RayExecutorV2) and KV-connector engine_ids must
450 # be unique across sibling DP engines or registration collides.
451 # Use the global DP rank, not a node-local rank, since sibling DP
452 # engines can span multiple nodes.
453 dp_vllm_config.instance_id = f"{dp_vllm_config.instance_id}_dp{dp_rank}"
454 if dp_vllm_config.kv_transfer_config is not None:
455 dp_vllm_config.kv_transfer_config.engine_id = (
456 f"{dp_vllm_config.kv_transfer_config.engine_id}_dp{dp_rank}"
457 )
458
459
460 class CoreEngineActorManager:
461 """Utility class to handle creation, readiness, and shutdown
462 of core engine Ray actors used by the AsyncLLM and LLMEngine.
463
464 Different from CoreEngineProcManager, this class manages
465 core engines for both local and remote nodes.
466 """
467
468 def __init__(
469 self,
470 vllm_config: VllmConfig,
471 addresses: EngineZmqAddresses,
472 executor_class: type[Executor],
473 log_stats: bool,
474 placement_groups: list["PlacementGroup"] | None = None,
475 local_dp_ranks: list[int] | None = None,
476 ):
477 import copy
478
479 import ray
480 from ray.runtime_env import RuntimeEnv
481 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
482
483 from vllm.v1.engine.core import DPMoEEngineCoreActor, EngineCoreActor
484
485 dp_size = vllm_config.parallel_config.data_parallel_size
486 actor_class = (
487 DPMoEEngineCoreActor
488 if dp_size > 1 and vllm_config.model_config.is_moe
489 else EngineCoreActor
490 )
491
492 self.local_engine_actors: list[ray.ActorHandle] = []
493 self.remote_engine_actors: list[ray.ActorHandle] = []
494
495 env_vars_list = get_env_vars_to_copy(
496 destination=actor_class.__name__,
497 exclude_vars=WORKER_SPECIFIC_ENV_VARS,
498 )
499 self.env_vars_dict = {
500 name: os.environ[name] for name in env_vars_list if name in os.environ
501 }
502 runtime_env = RuntimeEnv(env_vars=self.env_vars_dict)
503
504 self.addresses = addresses
505 self.executor_class = executor_class
506 self.log_stats = log_stats
507 local_engine_count = vllm_config.parallel_config.data_parallel_size_local
508 world_size = vllm_config.parallel_config.world_size
509 self.manager_stopped = threading.Event()
510 self.failed_proc_name: str | None = None
511
512 if ray.is_initialized():
513 logger.info("Ray is already initialized. Skipping Ray initialization.")
514 else:
515 ray.init()
516
517 parallel_config = vllm_config.parallel_config
518 if parallel_config.enable_elastic_ep:
519 from vllm.distributed.utils import create_tcp_store
520
521 ip = parallel_config.data_parallel_master_ip
522 store = create_tcp_store(
523 ip,
524 0,
525 is_master=True,
526 world_size=-1,
527 wait_for_workers=False,
528 )
529 parallel_config._coord_store_port = store.port
530 self._coord_store = store
531
532 if placement_groups is not None:
533 assert local_dp_ranks is not None, (
534 "local_dp_ranks must be provided if placement_groups is provided"
535 )
536 assert len(placement_groups) == len(local_dp_ranks), (
537 "placement_groups and local_dp_ranks must have the same length"
538 )
539 logger.info("Using provided placement groups")
540 # TODO(rui): validate passed-in placement groups
541 self.created_placement_groups = []
542 else:
543 placement_groups, local_dp_ranks = (
544 CoreEngineActorManager.create_dp_placement_groups(vllm_config)
545 )
546 self.created_placement_groups = placement_groups
547 assert len(placement_groups) == dp_size, (
548 "Number of placement groups must match data parallel size"
549 )
550
551 self.placement_group_is_local = []
552 refs = []
553 for index, local_index, pg in zip(
554 range(dp_size), local_dp_ranks, placement_groups
555 ):
556 dp_vllm_config = copy.deepcopy(vllm_config)
557 if dp_size > 1:
558 _apply_dp_identity_suffix(dp_vllm_config, index)
559 dp_vllm_config.parallel_config.placement_group = pg
560 local_client = index < local_engine_count
561
562 # Ray XPU known issue: dpctl initializes the GPU runtime early, so
563 # setting device env vars in Ray actor's initialization method
564 # will not affect device selection. See:
565 # https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501
566 if current_platform.is_xpu():
567 device_evar = current_platform.device_control_env_var
568 physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
569 device_evar, local_index, world_size
570 )
571 actor_env_vars = self.env_vars_dict.copy()
572 actor_env_vars[device_evar] = ",".join(str(d) for d in physical_gpu_ids)
573 runtime_env = RuntimeEnv(env_vars=actor_env_vars)
574
575 actor = (
576 ray.remote(actor_class)
577 .options(
578 scheduling_strategy=PlacementGroupSchedulingStrategy(
579 placement_group=pg,
580 placement_group_bundle_index=world_size,
581 ),
582 runtime_env=runtime_env,
583 )
584 .remote(
585 vllm_config=dp_vllm_config,
586 executor_class=executor_class,
587 log_stats=log_stats,
588 local_client=local_client,
589 addresses=addresses,
590 dp_rank=index,
591 local_dp_rank=local_index,
592 )
593 )
594 if local_client:
595 self.local_engine_actors.append(actor)
596 else:
597 self.remote_engine_actors.append(actor)
598 self.placement_group_is_local.append(local_client)
599 refs.append(actor.wait_for_init.remote())
600
601 ray.get(refs)
602 self.run_refs = []
603 self.actor_run_ref_dict = dict()
604 for actor in self.local_engine_actors + self.remote_engine_actors:
605 ref = actor.run.remote()
606 self.run_refs.append(ref)
607 self.actor_run_ref_dict[actor] = ref
608
609 @staticmethod
610 def create_dp_placement_groups(
611 vllm_config: VllmConfig,
612 ) -> tuple[list["PlacementGroup"], list[int]]:
613 """Create placement groups for data parallel."""
614 import ray
615 from ray._private.state import available_resources_per_node
616
617 logger.info("Creating placement groups for data parallel")
618 dp_master_ip = vllm_config.parallel_config.data_parallel_master_ip
619 dp_size = vllm_config.parallel_config.data_parallel_size
620 dp_size_local = vllm_config.parallel_config.data_parallel_size_local
621
622 available_resources = available_resources_per_node()
623 world_size = vllm_config.parallel_config.world_size
624 placement_groups: list[PlacementGroup] = []
625 local_dp_ranks: list[int] = []
626
627 nodes = [
628 node_resources
629 for _, node_resources in _dp_nodes_master_first(
630 available_resources, dp_master_ip
631 )
632 ]
633
634 # optionally restrict DP placement to a caller-provided node set.
635 requested_node_ips = {
636 ip.strip()
637 for ip in envs.VLLM_RAY_DP_PLACEMENT_NODE_IPS.split(",")
638 if ip.strip()
639 }
640 if requested_node_ips:
641 allowed_node_ips = set(requested_node_ips)
642 # The master node must host the local ranks, so it has to be allowed.
643 if dp_master_ip not in allowed_node_ips:
644 allowed_node_ips.add(dp_master_ip)
645 filtered_nodes = [
646 node_resources
647 for node_resources in nodes
648 if _node_ip_from_resources(node_resources) in allowed_node_ips
649 ]
650 logger.info(
651 "VLLM_RAY_DP_PLACEMENT_NODE_IPS set; restricting DP placement "
652 "from %d to %d node(s): %s",
653 len(nodes),
654 len(filtered_nodes),
655 sorted(allowed_node_ips),
656 )
657 nodes = filtered_nodes
658
659 device_str = current_platform.ray_device_key
660 n_node_devices: list[int] = [
661 int(node_resources[device_str])
662 for node_resources in nodes
663 if device_str in node_resources
664 ]
665 assert n_node_devices, f"No {device_str} found in Ray cluster."
666 max_device_per_node = max(n_node_devices)
667
668 pack_strategy = envs.VLLM_RAY_DP_PACK_STRATEGY
669 _supported_pack_strategies = ("strict", "fill", "span")
670 if pack_strategy not in _supported_pack_strategies:
671 raise ValueError(
672 f"{envs.VLLM_RAY_DP_PACK_STRATEGY} is not supported. "
673 "Make sure to set `VLLM_RAY_DP_PACK_STRATEGY` "
674 f"to one of {_supported_pack_strategies}"
675 )
676
677 all2all_backend = vllm_config.parallel_config.all2all_backend
678 if pack_strategy == "fill" and (
679 all2all_backend == "deepep_high_throughput"
680 or all2all_backend == "deepep_low_latency"
681 ):
682 raise ValueError(
683 "DeepEP kernels require EP ranks [0,7] (same for [8,15], ...) "
684 "to be on the same node, but VLLM_RAY_DP_PACK_STRATEGY=fill "
685 "does not guarantee that. "
686 "Please use VLLM_RAY_DP_PACK_STRATEGY=strict instead."
687 )
688
689 if pack_strategy in ("strict", "fill"):
690 placement_strategy = "STRICT_PACK"
691 else:
692 placement_strategy = "PACK"
693 assert world_size > max_device_per_node, (
694 f"World size {world_size} is smaller than the "
695 "maximum number of devices per node "
696 f"{max_device_per_node}. Make sure to set "
697 "`VLLM_RAY_DP_PACK_STRATEGY` to `strict` or `fill`"
698 )
699
700 # if we need multiple nodes per dp group, we require for now that
701 # available nodes are homogeneous
702 assert set(n_node_devices) == {max_device_per_node}, (
703 f"Nodes are not homogeneous, {nodes}"
704 )
705 assert world_size % max_device_per_node == 0, (
706 f"For multi-node data parallel groups, world_size ({world_size}) must "
707 f"be a multiple of number of devices per node ({max_device_per_node})."
708 )
709 assert len(n_node_devices) * max_device_per_node >= world_size * dp_size, (
710 f"Not enough total available nodes ({len(n_node_devices)}) "
711 f"and devices per node ({max_device_per_node}) "
712 f"to satisfy required world size {world_size} and data parallel size "
713 f"{dp_size}"
714 )
715 assert dp_size_local == 1, (
716 f"data-parallel-size-local {dp_size_local} should be set as the "
717 "default (1) for VLLM_RAY_DP_PACK_STRATEGY=span. "
718 "The actual data-parallel-size-local will be auto determined."
719 )
720
721 # bundles collected for a single DP rank from multiple nodes,
722 # for "span" pack strategy
723 collected_bundles = []
724 for node_resources in nodes:
725 node_ip = _node_ip_from_resources(node_resources)
726 assert node_ip is not None, (
727 f"No node IP key found in node resources: {node_resources}"
728 )
729
730 n_device_on_node = int(node_resources.get(device_str, 0))
731 if pack_strategy == "span" and n_device_on_node != 0:
732 # Strictly speaking,
733 # dp_size_available = n_device_on_node / world_size
734 # and is a fraction, but we use 1 for easier processing
735 dp_size_available = 1
736 else:
737 dp_size_available = n_device_on_node // world_size
738
739 if node_ip == dp_master_ip:
740 if dp_size_available < dp_size_local:
741 raise ValueError(
742 f"Not enough resources to allocate {dp_size_local} DP ranks "
743 f"on DP master node {dp_master_ip}, possible to fit "
744 f"{dp_size_available} DP ranks."
745 )
746 dp_size_to_allocate = dp_size_local
747 elif pack_strategy == "strict":
748 if dp_size_available < dp_size_local:
749 logger.info(
750 "Skipping node %s as %s DP ranks could not fit, "
751 "possible to fit %s DP ranks",
752 node_ip,
753 dp_size_local,
754 dp_size_available,
755 )
756 continue
757 dp_size_to_allocate = dp_size_local
758 else:
759 # for "pack_strategy" in "fill" and "span"
760 # we always take everything that's available
761 dp_size_to_allocate = dp_size_available
762
763 for i in range(dp_size_to_allocate):
764 device_bundle = [{device_str: 1.0, "node:" + node_ip: 0.001}]
765 if pack_strategy == "span":
766 collected_bundles += device_bundle * n_device_on_node
767 assert len(collected_bundles) <= world_size, (
768 "collected_bundles should be <= world_size, "
769 f"but got {len(collected_bundles)=} and {world_size=}"
770 )
771
772 # we only create a placement group if we collected enough devices
773 if len(collected_bundles) < world_size:
774 continue
775
776 control_node_ip = _get_bundle_node_ip(collected_bundles[0])
777 bundles = collected_bundles + [
778 _make_control_bundle(control_node_ip)
779 ]
780 collected_bundles = []
781 else:
782 # STRICT_PACK already keeps every bundle in the placement
783 # group on one node, so the explicit node affinity on the
784 # control bundle is redundant for correctness here. Keep it
785 # anyway for consistency with the span path and to preserve
786 # intent if this scheduling strategy changes later.
787 bundles = device_bundle * world_size + [
788 _make_control_bundle(node_ip)
789 ]
790
791 pg = ray.util.placement_group(
792 name=f"dp_rank_{len(placement_groups)}",
793 strategy=placement_strategy,
794 bundles=bundles,
795 )
796 placement_groups.append(pg)
797 local_dp_ranks.append(i)
798 if len(placement_groups) == dp_size:
799 break
800
801 if len(placement_groups) == dp_size:
802 break
803
804 if len(placement_groups) < dp_size:
805 raise ValueError(
806 f"Not enough resources to allocate {dp_size} "
807 "placement groups, only created "
808 f"{len(placement_groups)} placement groups. "
809 "Available resources: "
810 f"{available_resources}"
811 )
812 assert len(placement_groups) == dp_size, (
813 f"Created {len(placement_groups)} DP placement groups, expected {dp_size}"
814 )
815 assert len(local_dp_ranks) == dp_size, (
816 f"local_dp_ranks length {len(local_dp_ranks)} does not match "
817 f"expected {dp_size}"
818 )
819 return placement_groups, local_dp_ranks
820
821 @staticmethod
822 def add_dp_placement_groups(
823 old_vllm_config: VllmConfig, new_data_parallel_size: int
824 ) -> tuple[list["PlacementGroup"], list[int]]:
825 """Add placement groups for new data parallel size."""
826 import ray
827 from ray._private.state import (
828 available_resources_per_node,
829 total_resources_per_node,
830 )
831
832 old_dp_size = old_vllm_config.parallel_config.data_parallel_size
833 num_pg_to_create = new_data_parallel_size - old_dp_size
834
835 if num_pg_to_create <= 0:
836 return [], []
837
838 dp_master_ip = old_vllm_config.parallel_config.data_parallel_master_ip
839 world_size = old_vllm_config.parallel_config.world_size
840
841 # Both maps are keyed by node id and only contain live nodes.
842 available_resources = available_resources_per_node()
843 total_resources = total_resources_per_node()
844 nodes = _dp_nodes_master_first(available_resources, dp_master_ip)
845
846 placement_groups = []
847 local_dp_ranks = []
848 num_pg_created = 0
849
850 device_str = current_platform.ray_device_key
851 for node_id, node_resources in nodes:
852 if num_pg_created >= num_pg_to_create:
853 break
854
855 node_ip = _node_ip_from_resources(node_resources)
856 assert node_ip is not None, (
857 f"No node IP key found in node resources: {node_resources}"
858 )
859 if device_str not in node_resources:
860 continue
861 available_gpus = int(node_resources[device_str])
862
863 # Get total GPUs on this node from the node's resources
864 # Ray stores node resources with node ID as key
865 total_gpus = int(total_resources[node_id][device_str])
866
867 # Calculate used GPUs and used engines on this node
868 used_gpus = max(0, total_gpus - available_gpus)
869 used_engines_on_node = used_gpus // world_size
870
871 # Calculate how many new engines this node can accommodate
872 available_engine_count = available_gpus // world_size
873
874 # Create placement groups for new engines on this node
875 for i in range(available_engine_count):
876 if num_pg_created >= num_pg_to_create:
877 break
878
879 rank = old_dp_size + num_pg_created
880
881 # Create bundles with node constraint for master node
882 if node_ip == dp_master_ip:
883 bundles = [
884 {device_str: 1.0, "node:" + dp_master_ip: 0.001}
885 ] * world_size + [{"CPU": 1.0}]
886 else:
887 bundles = [{device_str: 1.0}] * world_size + [{"CPU": 1.0}]
888
889 pg = ray.util.placement_group(
890 name=f"dp_rank_{rank}",
891 strategy="STRICT_PACK",
892 bundles=bundles,
893 )
894 placement_groups.append(pg)
895
896 # Local rank starts from the number of engines already used
897 # on this node
898 local_rank = used_engines_on_node + i
899 local_dp_ranks.append(local_rank)
900 num_pg_created += 1
901
902 return placement_groups, local_dp_ranks
903
904 def scale_up_elastic_ep(
905 self,
906 cur_vllm_config: VllmConfig,
907 new_data_parallel_size: int,
908 num_redundant_experts: int,
909 ) -> None:
910 import copy
911
912 import ray
913 from ray.runtime_env import RuntimeEnv
914 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
915
916 from vllm.v1.engine.core import DPMoEEngineCoreActor, EngineCoreActor
917
918 actor_class = (
919 DPMoEEngineCoreActor
920 if cur_vllm_config.model_config.is_moe
921 else EngineCoreActor
922 )
923
924 cur_data_parallel_size = len(self.local_engine_actors) + len(
925 self.remote_engine_actors
926 )
927
928 assert new_data_parallel_size > cur_data_parallel_size, (
929 f"New data parallel size {new_data_parallel_size} must be greater "
930 f"than current data parallel size {cur_data_parallel_size} "
931 "for scale up"
932 )
933
934 placement_groups, local_dp_ranks = self.add_dp_placement_groups(
935 cur_vllm_config, new_data_parallel_size
936 )
937
938 world_size = cur_vllm_config.parallel_config.world_size
939 dp_master_ip = cur_vllm_config.parallel_config.data_parallel_master_ip
940 new_local_engines = 0
941
942 runtime_env = RuntimeEnv(
943 env_vars=self.env_vars_dict | {"VLLM_ELASTIC_EP_SCALE_UP_LAUNCH": "1"}
944 )
945 for i, (pg, local_rank) in enumerate(zip(placement_groups, local_dp_ranks)):
946 rank = cur_data_parallel_size + i
947 dp_vllm_config = copy.deepcopy(cur_vllm_config)
948 if new_data_parallel_size > 1:
949 _apply_dp_identity_suffix(dp_vllm_config, rank)
950 dp_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
951 dp_vllm_config.parallel_config.eplb_config.num_redundant_experts = (
952 num_redundant_experts
953 )
954 dp_vllm_config.parallel_config.placement_group = pg
955
956 # Check if this placement group is on the head node
957 local_client = any(
958 bundle.get("node:" + dp_master_ip, 0) > 0 for bundle in pg.bundle_specs
959 )
960
961 if local_client:
962 new_local_engines += 1
963 # Update data_parallel_size_local
964 dp_vllm_config.parallel_config.data_parallel_size_local = (
965 cur_vllm_config.parallel_config.data_parallel_size_local
966 + new_local_engines
967 )
968
969 actor = (
970 ray.remote(actor_class)
971 .options(
972 scheduling_strategy=PlacementGroupSchedulingStrategy(
973 placement_group=pg,
974 placement_group_bundle_index=world_size,
975 ),
976 runtime_env=runtime_env,
977 )
978 .remote(
979 vllm_config=dp_vllm_config,
980 executor_class=self.executor_class,
981 log_stats=self.log_stats,
982 local_client=local_client,
983 addresses=self.addresses,
984 dp_rank=rank,
985 local_dp_rank=local_rank,
986 )
987 )
988
989 if local_client:
990 self.local_engine_actors.append(actor)
991 else:
992 self.remote_engine_actors.append(actor)
993 self.created_placement_groups.append(pg)
994 self.placement_group_is_local.append(local_client)
995
996 actors = (
997 self.local_engine_actors[-new_local_engines:]
998 if new_local_engines > 0
999 else []
1000 ) + self.remote_engine_actors[-(len(placement_groups) - new_local_engines) :]
1001
1002 ray.get([actor.wait_for_init.remote() for actor in actors])
1003 for actor in actors:
1004 ref = actor.run.remote()
1005 self.run_refs.append(ref)
1006 self.actor_run_ref_dict[actor] = ref
1007
1008 def scale_down_elastic_ep(
1009 self, cur_data_parallel_size: int, new_data_parallel_size: int
1010 ) -> None:
1011 import ray
1012
1013 assert cur_data_parallel_size > new_data_parallel_size, (
1014 f"cur_data_parallel_size {cur_data_parallel_size} must be greater "
1015 f"than new_data_parallel_size {new_data_parallel_size} "
1016 "for scale down"
1017 )
1018 for _ in range(cur_data_parallel_size - new_data_parallel_size):
1019 pg = self.created_placement_groups.pop()
1020 is_local = self.placement_group_is_local.pop()
1021 if is_local:
1022 self.local_engine_actors.pop()
1023 else:
1024 self.remote_engine_actors.pop()
1025 ray.util.remove_placement_group(pg)
1026
1027 def remove_run_refs_for_scale_down(self, removed_dp_size: int) -> None:
1028 if removed_dp_size <= 0:
1029 return
1030 flags = self.placement_group_is_local[-removed_dp_size:]
1031 li = len(self.local_engine_actors) - 1
1032 ri = len(self.remote_engine_actors) - 1
1033 for is_local in reversed(flags):
1034 if is_local:
1035 actor = self.local_engine_actors[li]
1036 li -= 1
1037 else:
1038 actor = self.remote_engine_actors[ri]
1039 ri -= 1
1040 ref = self.actor_run_ref_dict.pop(actor)
1041 self.run_refs.remove(ref)
1042
1043 def get_run_refs(self):
1044 return self.run_refs
1045
1046 def monitor_engine_liveness(self) -> None:
1047 import ray
1048
1049 while not self.manager_stopped.is_set():
1050 actor_run_refs = list(self.get_run_refs())
1051 if not actor_run_refs:
1052 logger.info(
1053 "There are no actors to monitor currently. "
1054 "The monitoring function is about to terminate."
1055 )
1056 break
1057 actor_done_refs, _ = ray.wait(actor_run_refs, timeout=5)
1058 unexpected_failure = False
1059 for actor_ref in actor_done_refs:
1060 if self.manager_stopped.is_set():
1061 break
1062 if actor_ref not in self.get_run_refs():
1063 # The run refs may have been updated by elastic scale-down.
1064 continue
1065 try:
1066 ray.get(actor_ref)
1067 except ray.exceptions.RayActorError:
1068 self.failed_proc_name = f"Actor {actor_ref}"
1069 unexpected_failure = True
1070
1071 if unexpected_failure:
1072 break
1073
1074 self.shutdown()
1075
1076 def shutdown(self, timeout: float | None = None) -> None:
1077 import ray
1078
1079 self.manager_stopped.set()
1080 for actor in self.local_engine_actors + self.remote_engine_actors:
1081 ray.kill(actor)
1082 for pg in self.created_placement_groups:
1083 ray.util.remove_placement_group(pg)
1084
1085
1086 def get_engine_zmq_addresses(
1087 vllm_config: VllmConfig,
1088 num_api_servers: int = 1,
1089 *,
1090 defer_api_server_ports: bool = True,
1091 ) -> EngineZmqAddresses:
1092 """Allocate ZMQ addresses for engine-client communication.
1093
1094 By default each TCP address is a ``tcp://host:0`` placeholder; the
1095 consumer (API-server child or single-process ``MPClient``) binds, then
1096 recovers the kernel-assigned port via ``getsockopt(zmq.LAST_ENDPOINT)``
1097 and writes it back into ``addresses`` before the engine handshake.
1098
1099 Set ``defer_api_server_ports=False`` only when the consumer cannot
1100 report a bound port back (e.g. the Rust front-end). IPC paths are
1101 unaffected."""
1102 parallel_config = vllm_config.parallel_config
1103 local_engine_count = parallel_config.data_parallel_size_local
1104 local_start_index = parallel_config.data_parallel_rank_local
1105 dp_size = parallel_config.data_parallel_size
1106 host = parallel_config.data_parallel_master_ip
1107 local_engines_only = parallel_config.local_engines_only
1108
1109 # In offline mode there is an LLM instance per DP rank and
1110 # one core engine per LLM, see
1111 # examples/features/data_parallel/data_parallel_offline.py.
1112 offline_mode = local_start_index is not None
1113
1114 # client_local_only = True for cases where this front-end
1115 # sends requests only to colocated engines.
1116 client_local_only = (
1117 offline_mode or local_engines_only or (local_engine_count == dp_size)
1118 )
1119 # NOTE(yongji): handling scaling from intra-node to inter-node
1120 if parallel_config.enable_elastic_ep:
1121 client_local_only = False
1122
1123 def _addr() -> str:
1124 if client_local_only:
1125 return get_open_zmq_ipc_path()
1126 return get_tcp_uri(host, 0 if defer_api_server_ports else get_open_port())
1127
1128 return EngineZmqAddresses(
1129 inputs=[_addr() for _ in range(num_api_servers)],
1130 outputs=[_addr() for _ in range(num_api_servers)],
1131 )
1132
1133
1134 FrontendProcess = BaseProcess | _SubprocessWrapper
1135
1136
1137 @dataclass
1138 class CoreEngineLaunch:
1139 """Resources and startup barrier for launched engine processes."""
1140
1141 engine_manager: CoreEngineProcManager | CoreEngineActorManager | None
1142 coordinator: DPCoordinator | None
1143 addresses: EngineZmqAddresses
1144 tensor_queue: Queue | None
1145 # Frontend processes to watch during engine startup; may be assigned by
1146 # the caller before the startup barrier runs on context manager exit.
1147 watched_frontend_processes: Sequence[FrontendProcess] = ()
1148
1149
1150 @contextlib.contextmanager
1151 def launch_core_engines(
1152 vllm_config: VllmConfig,
1153 executor_class: type[Executor],
1154 log_stats: bool,
1155 addresses: EngineZmqAddresses,
1156 ) -> Iterator[CoreEngineLaunch]:
1157 """Launch engine and DP coordinator processes as needed."""
1158 parallel_config = vllm_config.parallel_config
1159 dp_size = parallel_config.data_parallel_size
1160 local_engine_count = parallel_config.data_parallel_size_local
1161 local_start_index = parallel_config.data_parallel_rank_local
1162 dp_rank = parallel_config.data_parallel_rank
1163 host = parallel_config.data_parallel_master_ip
1164 local_engines_only = parallel_config.local_engines_only
1165
1166 offline_mode = local_start_index is not None
1167
1168 # Create a single tensor IPC queue for sharing multimodal tensors between
1169 # API servers and engine core. Returns a single queue since we only support
1170 # DP=1 for this data flow.
1171 tensor_queue: Queue | None = None
1172 multimodal_config = vllm_config.model_config.multimodal_config
1173 if multimodal_config is not None and multimodal_config.mm_tensor_ipc == "torch_shm":
1174 tensor_queue = get_mp_context().Queue()
1175
1176 # Run the DP Coordinator process with rank 0 when in online DP mode.
1177 # The coordinator is needed for:
1178 # 1. Internal/hybrid LB: collecting and publishing queue stats for load balancing
1179 # 2. MoE models: wave coordination in addition to stats
1180 run_coordinator = (
1181 vllm_config.needs_dp_coordinator and not offline_mode and dp_rank == 0
1182 )
1183
1184 if run_coordinator:
1185 coordinator = DPCoordinator(
1186 parallel_config,
1187 enable_wave_coordination=vllm_config.model_config.is_moe,
1188 logging_config=vllm_config.logging_config,
1189 )
1190
1191 addresses.coordinator_input, addresses.coordinator_output = (
1192 coordinator.get_engine_socket_addresses()
1193 )
1194 addresses.frontend_stats_publish_address = (
1195 coordinator.get_stats_publish_address()
1196 )
1197
1198 logger.info("Started DP Coordinator process (PID: %d)", coordinator.proc.pid)
1199 else:
1200 coordinator = None
1201
1202 if parallel_config.data_parallel_backend == "ray":
1203 logger.info("Starting ray-based data parallel backend")
1204
1205 engine_actor_manager = CoreEngineActorManager(
1206 vllm_config=vllm_config,
1207 addresses=addresses,
1208 executor_class=executor_class,
1209 log_stats=log_stats,
1210 )
1211
1212 yield CoreEngineLaunch(
1213 engine_actor_manager, coordinator, addresses, tensor_queue
1214 )
1215 return
1216
1217 if offline_mode:
1218 assert local_engine_count == 1
1219 engines_to_handshake = [CoreEngine(index=dp_rank, local=True)]
1220 elif dp_rank == 0:
1221 # Rank 0 holds Coordinator, so it handshakes with all Cores
1222 # in both external dplb and internal dplb mode.
1223 # Note this also covers the case where we have zero local engines
1224 # and rank 0 is headless.
1225 engines_to_handshake = [
1226 CoreEngine(index=i, local=(i < local_engine_count)) for i in range(dp_size)
1227 ]
1228 else:
1229 # Rank > 0 handshakes with just the local cores it is managing.
1230 assert local_engines_only, (
1231 "Attempting to launch core_engines from dp_rank > 0, but "
1232 "found internal DPLB, which is incompatible."
1233 )
1234 engines_to_handshake = [
1235 CoreEngine(index=i, local=True)
1236 for i in range(dp_rank, dp_rank + local_engine_count)
1237 ]
1238
1239 # Whether the started engines will handshake only with co-located
1240 # front-end processes. In external_dp_lb mode, ranks > 0 handshake with
1241 # their co-located frontend and also the rank 0 front-end, and hence this
1242 # will be False.
1243 handshake_local_only = offline_mode or local_engine_count == dp_size
1244
1245 # NOTE(yongji): handling scaling from intra-node to inter-node
1246 if parallel_config.enable_elastic_ep:
1247 handshake_local_only = False
1248
1249 handshake_address = get_engine_client_zmq_addr(
1250 handshake_local_only,
1251 host,
1252 parallel_config.data_parallel_rpc_port,
1253 )
1254
1255 if local_engines_only and dp_rank > 0:
1256 assert not handshake_local_only
1257 local_handshake_address = get_open_zmq_ipc_path()
1258 client_handshake_address = local_handshake_address
1259 else:
1260 local_handshake_address = handshake_address
1261 client_handshake_address = None
1262
1263 with zmq_socket_ctx(
1264 local_handshake_address, zmq.ROUTER, bind=True
1265 ) as handshake_socket:
1266 # Start local engines.
1267 if local_engine_count:
1268 local_engine_manager = CoreEngineProcManager(
1269 vllm_config=vllm_config,
1270 executor_class=executor_class,
1271 log_stats=log_stats,
1272 handshake_address=handshake_address,
1273 client_handshake_address=client_handshake_address,
1274 local_client=True,
1275 local_engine_count=local_engine_count,
1276 start_index=dp_rank,
1277 local_start_index=local_start_index or 0,
1278 tensor_queue=tensor_queue,
1279 )
1280 else:
1281 local_engine_manager = None
1282
1283 launch = CoreEngineLaunch(
1284 local_engine_manager, coordinator, addresses, tensor_queue
1285 )
1286 yield launch
1287 wait_for_engine_startup(
1288 handshake_socket,
1289 engines_to_handshake,
1290 parallel_config,
1291 dp_size > 1 and vllm_config.model_config.is_moe,
1292 vllm_config.cache_config,
1293 launch,
1294 )
1295
1296
1297 def wait_for_engine_startup(
1298 handshake_socket: zmq.Socket,
1299 core_engines: list[CoreEngine],
1300 parallel_config: ParallelConfig,
1301 coordinated_dp: bool,
1302 cache_config: CacheConfig,
1303 launch: CoreEngineLaunch,
1304 ):
1305 # Wait for engine core process(es) to send ready messages.
1306 local_count = parallel_config.data_parallel_size_local
1307 remote_count = len(core_engines) - local_count
1308 # [local, remote] counts
1309 conn_pending, start_pending = [local_count, remote_count], [0, 0]
1310 poller = zmq.Poller()
1311 poller.register(handshake_socket, zmq.POLLIN)
1312
1313 remote_should_be_headless = (
1314 not parallel_config.data_parallel_hybrid_lb
1315 and not parallel_config.data_parallel_external_lb
1316 )
1317
1318 # 1. Engine processes
1319 if isinstance(launch.engine_manager, CoreEngineProcManager):
1320 for sentinel in launch.engine_manager.sentinels():
1321 poller.register(sentinel, zmq.POLLIN)
1322 # 2. DP Coordinator process, if present
1323 coord_process = launch.coordinator.proc if launch.coordinator else None
1324 if coord_process is not None:
1325 poller.register(coord_process.sentinel, zmq.POLLIN)
1326 # 3. Watched frontend processes, if any
1327 frontend_process_by_fd: dict[int, FrontendProcess] = {}
1328 for proc in launch.watched_frontend_processes:
1329 fd = proc.sentinel if isinstance(proc.sentinel, int) else proc.sentinel.fileno()
1330 frontend_process_by_fd[fd] = proc
1331 poller.register(fd, zmq.POLLIN)
1332
1333 while any(conn_pending) or any(start_pending):
1334 events = poller.poll(STARTUP_POLL_PERIOD_MS)
1335 if not events:
1336 if any(conn_pending):
1337 logger.debug(
1338 "Waiting for %d local, %d remote core engine proc(s) to connect.",
1339 *conn_pending,
1340 )
1341 if any(start_pending):
1342 logger.debug(
1343 "Waiting for %d local, %d remote core engine proc(s) to start.",
1344 *start_pending,
1345 )
1346 continue
1347 if len(events) > 1 or events[0][0] != handshake_socket:
1348 # One of the local core, coordinator, or watched frontend processes exited.
1349 if isinstance(launch.engine_manager, CoreEngineProcManager):
1350 finished = launch.engine_manager.finished_procs()
1351 else:
1352 finished = {}
1353 if coord_process is not None and coord_process.exitcode is not None:
1354 finished[coord_process.name] = coord_process.exitcode
1355 failed_frontend_procs = {
1356 proc.name: proc.exitcode
1357 for fd, proc in frontend_process_by_fd.items()
1358 if proc.exitcode is not None
1359 or any(event_fd == fd for event_fd, _ in events)
1360 }
1361 if failed_frontend_procs and not finished:
1362 raise RuntimeError(
1363 "Frontend process failed during engine core initialization. "
1364 "See root cause above. "
1365 f"Failed frontend proc(s): {failed_frontend_procs}"
1366 )
1367 raise RuntimeError(
1368 "Engine core initialization failed. "
1369 "See root cause above. "
1370 f"Failed core proc(s): {finished}"
1371 + (
1372 f", failed frontend proc(s): {failed_frontend_procs}"
1373 if failed_frontend_procs
1374 else ""
1375 )
1376 )
1377
1378 # Receive HELLO and READY messages from the input socket.
1379 eng_identity, ready_msg_bytes = handshake_socket.recv_multipart()
1380 eng_index = int.from_bytes(eng_identity, "little")
1381 engine = next((e for e in core_engines if e.identity == eng_identity), None)
1382 if engine is None:
1383 raise RuntimeError(
1384 f"Message from engine with unexpected data parallel rank: {eng_index}"
1385 )
1386 msg = msgspec.msgpack.decode(ready_msg_bytes)
1387 status, local, headless = msg["status"], msg["local"], msg["headless"]
1388 if local != engine.local:
1389 raise RuntimeError(
1390 f"{status} message from "
1391 f"{'local' if local else 'remote'} "
1392 f"engine {eng_index}, expected it to be "
1393 f"{'local' if engine.local else 'remote'}"
1394 )
1395
1396 # Remote engines must be headless iff we aren't in hybrid dp lb mode.
1397 if not local and headless != remote_should_be_headless:
1398 if headless:
1399 raise RuntimeError(
1400 f"Remote engine {eng_index} must not use "
1401 f"--headless in external or hybrid dp lb "
1402 f"mode"
1403 )
1404 else:
1405 raise RuntimeError(
1406 f"Remote engine {eng_index} must use "
1407 f"--headless unless in external or hybrid "
1408 f"dp lb mode"
1409 )
1410
1411 if status == "HELLO" and engine.state == CoreEngineState.NEW:
1412 # Send init message with DP config info.
1413 init_message = msgspec.msgpack.encode(
1414 EngineHandshakeMetadata(
1415 addresses=launch.addresses,
1416 parallel_config={
1417 k: getattr(parallel_config, k)
1418 for k in (
1419 "data_parallel_master_ip",
1420 "data_parallel_master_port",
1421 "_data_parallel_master_port_list",
1422 "data_parallel_size",
1423 )
1424 }
1425 if coordinated_dp
1426 else {},
1427 )
1428 )
1429 handshake_socket.send_multipart((eng_identity, init_message), copy=False)
1430 conn_pending[0 if local else 1] -= 1
1431 start_pending[0 if local else 1] += 1
1432 engine.state = CoreEngineState.CONNECTED
1433 elif status == "READY" and engine.state == CoreEngineState.CONNECTED:
1434 # Validate config hash consistency across DP workers for MoE models.
1435 if coordinated_dp:
1436 worker_config_hash = msg.get("parallel_config_hash")
1437 expected_hash = parallel_config.compute_hash()
1438 if worker_config_hash != expected_hash:
1439 raise RuntimeError(
1440 f"Configuration mismatch detected for engine "
1441 f"{eng_index}. All DP workers must have identical "
1442 f"configurations for parameters that affect collective "
1443 f"communication (e.g., enable_eplb, "
1444 f"eplb_config.log_balancedness). "
1445 f"Worker hash: {worker_config_hash}, "
1446 f"Expected hash: {expected_hash}. "
1447 f"Please ensure all workers are started with the same "
1448 f"command-line arguments."
1449 )
1450
1451 start_pending[0 if local else 1] -= 1
1452 engine.state = CoreEngineState.READY
1453 else:
1454 raise RuntimeError(
1455 f"Unexpected {status} message for "
1456 f"{'local' if local else 'remote'} engine "
1457 f"{eng_index} in {engine.state} state."
1458 )
1459
1460 logger.debug(
1461 "%s from %s core engine process %s.",
1462 status,
1463 "local" if local else "remote",
1464 eng_index,
1465 )