"tests/vscode:/vscode.git/clone" did not exist on "e760fcef2265a62a6a9cfbafdb207e7c3d5c3b36"
utils.py 42.1 KB
Newer Older
1
2
3
4
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import contextlib
5
import os
6
import weakref
7
from collections.abc import Iterator
8
9
10
11
from dataclasses import dataclass
from enum import Enum, auto
from multiprocessing import Process, connection
from multiprocessing.process import BaseProcess
12
from typing import TYPE_CHECKING
13
from unittest.mock import patch
14
15
16
17

import msgspec
import zmq

18
from vllm import envs
19
20
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
from vllm.logger import init_logger
21
from vllm.platforms import current_platform
22
from vllm.ray.ray_env import get_env_vars_to_copy
23
from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
24
from vllm.utils.system_utils import get_mp_context
25
from vllm.v1.engine.coordinator import DPCoordinator
26
from vllm.v1.executor import Executor
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
from vllm.v1.utils import get_engine_client_zmq_addr, shutdown

if TYPE_CHECKING:
    from ray.util.placement_group import PlacementGroup

logger = init_logger(__name__)

STARTUP_POLL_PERIOD_MS = 10000


class CoreEngineState(Enum):
    NEW = auto()
    CONNECTED = auto()
    READY = auto()


class CoreEngine:
    """One per data parallel rank, used to track state during handshaking."""

    def __init__(self, index: int = 0, local: bool = True):
        self.local = local
        self.identity = index.to_bytes(2, "little")

        self.state = CoreEngineState.NEW


@dataclass
class EngineZmqAddresses:
    # ZMQ input socket addresses for each front-end client (requests)
    inputs: list[str]
    # ZMQ output socket addresses for each front-end client (responses)
    outputs: list[str]
    # ZMQ input socket address of DP coordinator if applicable
60
    coordinator_input: str | None = None
61
    # ZMQ output socket address of DP coordinator if applicable
62
    coordinator_output: str | None = None
63
64
65
    # ZMQ socket for front-end to connect to DP coordinator.
    # Not used by engine, just relayed to front-end in handshake response.
    # Only required for external DP LB case.
66
    frontend_stats_publish_address: str | None = None
67
68
69
70
71
72
73
74


@dataclass
class EngineHandshakeMetadata:
    """Metadata sent to each engine process during startup handshake,
    including addresses of the front-end ZMQ queues that they should
    connect to.
    """
75

76
    addresses: EngineZmqAddresses
77
    parallel_config: dict[str, int | str | list[int]]
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95


class CoreEngineProcManager:
    """
    Utility class to handle creation, readiness, and shutdown
    of background processes used by the AsyncLLM and LLMEngine.
    """

    def __init__(
        self,
        local_engine_count: int,
        start_index: int,
        local_start_index: int,
        vllm_config: VllmConfig,
        local_client: bool,
        handshake_address: str,
        executor_class: type[Executor],
        log_stats: bool,
96
        client_handshake_address: str | None = None,
97
98
99
100
101
102
103
104
105
106
107
    ):
        context = get_mp_context()
        common_kwargs = {
            "vllm_config": vllm_config,
            "local_client": local_client,
            "handshake_address": handshake_address,
            "executor_class": executor_class,
            "log_stats": log_stats,
        }

        if client_handshake_address:
108
            common_kwargs["client_handshake_address"] = client_handshake_address
109

110
111
112
113
        is_dp = vllm_config.parallel_config.data_parallel_size > 1

        from vllm.v1.engine.core import EngineCoreProc

114
        self.processes: list[BaseProcess] = []
115
        local_dp_ranks = []
116
117
118
        for index in range(local_engine_count):
            local_index = local_start_index + index
            global_index = start_index + index
119

120
            # Start EngineCore in background process.
121
            local_dp_ranks.append(local_index)
122
            self.processes.append(
123
                context.Process(
124
125
                    target=EngineCoreProc.run_engine_core,
                    name=f"EngineCore_DP{global_index}" if is_dp else "EngineCore",
126
                    kwargs=common_kwargs
127
                    | {"dp_rank": global_index, "local_dp_rank": local_index},
128
129
                )
            )
130
131

        self._finalizer = weakref.finalize(self, shutdown, self.processes)
132

133
        try:
134
            for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
135
136
137
                # Adjust device control in DP for non-CUDA platforms
                # as well as external and ray launchers
                # For CUDA platforms, we use torch.cuda.set_device()
138
139
140
                if is_dp and (
                    not current_platform.is_cuda_alike()
                    or vllm_config.parallel_config.use_ray
141
                ):
142
143
144
                    with set_device_control_env_var(vllm_config, local_dp_rank):
                        proc.start()
                else:
145
                    proc.start()
146
147
148
        finally:
            # Kill other procs if not all are running.
            if self.finished_procs():
149
                self.close()
150

151
152
153
    def close(self):
        """Shutdown all procs."""
        self._finalizer()
154
155
156
157
158
159
160
161
162
163
164
165

    def join_first(self):
        """Wait for any process to exit."""
        connection.wait(proc.sentinel for proc in self.processes)

    def sentinels(self) -> list:
        return [proc.sentinel for proc in self.processes]

    def finished_procs(self) -> dict[str, int]:
        """Returns dict of proc name -> exit code for any finished procs."""
        return {
            proc.name: proc.exitcode
166
167
            for proc in self.processes
            if proc.exitcode is not None
168
169
170
        }


171
@contextlib.contextmanager
172
173
174
def set_device_control_env_var(
    vllm_config: VllmConfig, local_dp_rank: int
) -> Iterator[None]:
175
176
177
178
179
    """
    Temporarily set CUDA_VISIBLE_DEVICES or equivalent
    for engine subprocess.
    """
    world_size = vllm_config.parallel_config.world_size
180
    local_world_size = vllm_config.parallel_config.local_world_size
181
    evar = current_platform.device_control_env_var
182

183
    value = get_device_indices(evar, local_dp_rank, world_size, local_world_size)
184
    with patch.dict(os.environ, values=((evar, value),)):
185
186
187
        yield


188
def get_device_indices(
189
190
191
192
    device_control_env_var: str,
    local_dp_rank: int,
    world_size: int,
    local_world_size: int | None = None,
193
):
194
195
196
197
198
199
200
    """
    Returns a comma-separated string of device indices for the specified
    data parallel rank.

    For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,
    this will select devices 2 and 3 for local_dp_rank=1.
    """
201
202
    if local_world_size is None:
        local_world_size = world_size
203
204
205
    try:
        value = ",".join(
            str(current_platform.device_id_to_physical_device_id(i))
206
207
208
209
            for i in range(
                local_dp_rank * world_size,
                local_dp_rank * world_size + local_world_size,
            )
210
        )
211
    except IndexError as e:
212
213
214
215
216
217
218
        raise Exception(
            f"Error setting {device_control_env_var}: "
            f"local range: [{local_dp_rank * world_size}, "
            f"{(local_dp_rank + 1) * world_size}) "
            "base value: "
            f'"{os.getenv(device_control_env_var)}"'
        ) from e
219
    return value
220
221


222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
class CoreEngineActorManager:
    """
    Utility class to handle creation, readiness, and shutdown
    of core engine Ray actors used by the AsyncLLM and LLMEngine.

    Different from CoreEngineProcManager, this class manages
    core engines for both local and remote nodes.
    """

    def __init__(
        self,
        vllm_config: VllmConfig,
        addresses: EngineZmqAddresses,
        executor_class: type[Executor],
        log_stats: bool,
237
238
        placement_groups: list["PlacementGroup"] | None = None,
        local_dp_ranks: list[int] | None = None,
239
240
241
242
    ):
        import copy

        import ray
243
        from ray.runtime_env import RuntimeEnv
244
        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
245

246
247
248
249
250
251
252
253
        from vllm.v1.engine.core import DPMoEEngineCoreActor, EngineCoreActor

        dp_size = vllm_config.parallel_config.data_parallel_size
        actor_class = (
            DPMoEEngineCoreActor
            if dp_size > 1 and vllm_config.model_config.is_moe
            else EngineCoreActor
        )
254
255
256

        self.local_engine_actors: list[ray.ActorHandle] = []
        self.remote_engine_actors: list[ray.ActorHandle] = []
257

258
        env_vars_list = get_env_vars_to_copy(destination=actor_class.__name__)
259
        self.env_vars_dict = {
260
            name: os.environ[name] for name in env_vars_list if name in os.environ
261
262
263
264
265
266
        }
        runtime_env = RuntimeEnv(env_vars=self.env_vars_dict)

        self.addresses = addresses
        self.executor_class = executor_class
        self.log_stats = log_stats
267
        local_engine_count = vllm_config.parallel_config.data_parallel_size_local
268
269
270
        world_size = vllm_config.parallel_config.world_size

        if ray.is_initialized():
271
            logger.info("Ray is already initialized. Skipping Ray initialization.")
272
273
274
        else:
            ray.init()

275
276
        vllm_config.parallel_config.allocate_elastic_ep_ports()

277
278
        if placement_groups is not None:
            assert local_dp_ranks is not None, (
279
280
                "local_dp_ranks must be provided if placement_groups is provided"
            )
281
            assert len(placement_groups) == len(local_dp_ranks), (
282
283
                "placement_groups and local_dp_ranks must have the same length"
            )
284
285
286
287
            logger.info("Using provided placement groups")
            # TODO(rui): validate passed-in placement groups
            self.created_placement_groups = []
        else:
288
            placement_groups, local_dp_ranks = (
289
                CoreEngineActorManager.create_dp_placement_groups(vllm_config)
290
            )
291
292
            self.created_placement_groups = placement_groups
        assert len(placement_groups) == dp_size, (
293
294
            "Number of placement groups must match data parallel size"
        )
295

296
        self.placement_group_is_local = []
297
        refs = []
298
299
300
        for index, local_index, pg in zip(
            range(dp_size), local_dp_ranks, placement_groups
        ):
301
302
303
            dp_vllm_config = copy.deepcopy(vllm_config)
            dp_vllm_config.parallel_config.placement_group = pg
            local_client = index < local_engine_count
304

305
306
307
308
309
310
311
            if dp_size > 1 and dp_vllm_config.kv_transfer_config is not None:
                # modify the engine_id and append the local_dp_rank to it to ensure
                # that the kv_transfer_config is unique for each DP rank.
                dp_vllm_config.kv_transfer_config.engine_id = (
                    f"{dp_vllm_config.kv_transfer_config.engine_id}_dp{local_index}"
                )

312
313
314
315
316
317
            # Ray XPU known issue: dpctl initializes the GPU runtime early, so
            # setting device env vars in Ray actor's initialization method
            # will not affect device selection. See:
            # https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501
            if current_platform.is_xpu():
                device_evar = current_platform.device_control_env_var
318
319
320
                device_indices = get_device_indices(
                    device_evar, local_index, world_size
                )
321
322
323
324
                actor_env_vars = self.env_vars_dict.copy()
                actor_env_vars[device_evar] = device_indices
                runtime_env = RuntimeEnv(env_vars=actor_env_vars)

325
            actor = (
326
                ray.remote(actor_class)
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
                .options(
                    scheduling_strategy=PlacementGroupSchedulingStrategy(
                        placement_group=pg,
                        placement_group_bundle_index=world_size,
                    ),
                    runtime_env=runtime_env,
                )
                .remote(
                    vllm_config=dp_vllm_config,
                    executor_class=executor_class,
                    log_stats=log_stats,
                    local_client=local_client,
                    addresses=addresses,
                    dp_rank=index,
                    local_dp_rank=local_index,
                )
            )
344
345
346
347
            if local_client:
                self.local_engine_actors.append(actor)
            else:
                self.remote_engine_actors.append(actor)
348
            self.placement_group_is_local.append(local_client)
349
350
351
352
353
354
355
356
357
            refs.append(actor.wait_for_init.remote())

        ray.get(refs)
        self.run_refs = []
        for actor in self.local_engine_actors + self.remote_engine_actors:
            self.run_refs.append(actor.run.remote())

    @staticmethod
    def create_dp_placement_groups(
358
        vllm_config: VllmConfig,
359
    ) -> tuple[list["PlacementGroup"], list[int]]:
360
361
362
        """
        Create placement groups for data parallel.
        """
363
364
365
366
367

        import ray
        from ray._private.state import available_resources_per_node

        logger.info("Creating placement groups for data parallel")
368
        dp_master_ip = vllm_config.parallel_config.data_parallel_master_ip
369
370
        dp_size = vllm_config.parallel_config.data_parallel_size
        dp_size_local = vllm_config.parallel_config.data_parallel_size_local
371
372
373
374
375

        available_resources = available_resources_per_node()
        world_size = vllm_config.parallel_config.world_size
        placement_groups: list[PlacementGroup] = []
        local_dp_ranks: list[int] = []
376

377
378
379
380
381
        dp_master_ip_key = f"node:{dp_master_ip}"
        nodes = sorted(
            available_resources.values(), key=lambda x: dp_master_ip_key not in x
        )
        assert len(nodes) > 0, "No nodes with resources found in Ray cluster."
382
        assert dp_master_ip_key in nodes[0], (
383
            f"The DP master node (ip: {dp_master_ip}) is missing or dead"
384
        )
385
        device_str = current_platform.ray_device_key
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
        n_node_devices: list[int] = [
            int(node_resources[device_str])
            for node_resources in nodes
            if device_str in node_resources
        ]
        assert n_node_devices, f"No {device_str} found in Ray cluster."
        max_device_per_node = max(n_node_devices)

        pack_strategy = envs.VLLM_RAY_DP_PACK_STRATEGY
        _supported_pack_strategies = ("strict", "fill", "span")
        if pack_strategy not in _supported_pack_strategies:
            raise ValueError(
                f"{envs.VLLM_RAY_DP_PACK_STRATEGY} is not supported. "
                "Make sure to set `VLLM_RAY_DP_PACK_STRATEGY` "
                f"to one of {_supported_pack_strategies}"
            )
402

403
        all2all_backend = vllm_config.parallel_config.all2all_backend
404
        if pack_strategy == "fill" and (
405
406
            all2all_backend == "deepep_high_throughput"
            or all2all_backend == "deepep_low_latency"
407
408
409
410
411
412
413
414
        ):
            raise ValueError(
                "DeepEP kernels require EP ranks [0,7] (same for [8,15], ...) "
                "to be on the same node, but VLLM_RAY_DP_PACK_STRATEGY=fill "
                "does not guarantee that. "
                "Please use VLLM_RAY_DP_PACK_STRATEGY=strict instead."
            )

415
416
417
418
419
420
421
422
423
424
425
426
        if pack_strategy in ("strict", "fill"):
            placement_strategy = "STRICT_PACK"
        else:
            placement_strategy = "PACK"
            assert world_size > max_device_per_node, (
                f"World size {world_size} is smaller than the "
                "maximum number of devices per node "
                f"{max_device_per_node}. Make sure to set "
                "`VLLM_RAY_DP_PACK_STRATEGY` to `strict` or `fill`"
            )

            # if we need multiple nodes per dp group, we require for now that
427
            # available nodes are homogeneous
428
            assert set(n_node_devices) == {max_device_per_node}, (
429
                f"Nodes are not homogeneous, {nodes}"
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
            )
            assert world_size % max_device_per_node == 0, (
                f"For multi-node data parallel groups, world_size ({world_size}) must "
                f"be a multiple of number of devices per node ({max_device_per_node})."
            )
            assert len(n_node_devices) * max_device_per_node >= world_size * dp_size, (
                f"Not enough total available nodes ({len(n_node_devices)}) "
                f"and devices per node ({max_device_per_node}) "
                f"to satisfy required world size {world_size} and data parallel size "
                f"{dp_size}"
            )
            assert dp_size_local == 1, (
                f"data-parallel-size-local {dp_size_local} should be set as the "
                "default (1) for VLLM_RAY_DP_PACK_STRATEGY=span. "
                "The actual data-parallel-size-local will be auto determined."
            )

        # bundles collected for a single DP rank from multiple nodes,
        # for "span" pack strategy
        collected_bundles = []
450
        for node_resources in nodes:
451
452
453
454
455
456
            node_ip_keys = [
                key
                for key in node_resources
                if key != "node:__internal_head__" and key.startswith("node:")
            ]
            assert len(node_ip_keys) == 1, (
457
                f"Zero or multiple node IP keys found in node resources: {node_ip_keys}"
458
459
460
461
            )
            node_ip_key = node_ip_keys[0]
            node_ip = node_ip_key.split(":")[1]

462
463
464
465
466
467
468
469
            n_device_on_node = int(node_resources.get(device_str, 0))
            if pack_strategy == "span" and n_device_on_node != 0:
                # Strictly speaking,
                # dp_size_available = n_device_on_node / world_size
                # and is a fraction, but we use 1 for easier processing
                dp_size_available = 1
            else:
                dp_size_available = n_device_on_node // world_size
470
471
472
473

            if node_ip == dp_master_ip:
                if dp_size_available < dp_size_local:
                    raise ValueError(
474
475
476
                        f"Not enough resources to allocate {dp_size_local} DP ranks "
                        f"on DP master node {dp_master_ip}, possible to fit "
                        f"{dp_size_available} DP ranks."
477
                    )
478
                dp_size_to_allocate = dp_size_local
479
            elif pack_strategy == "strict":
480
481
482
483
484
485
486
                if dp_size_available < dp_size_local:
                    logger.info(
                        "Skipping node %s as %s DP ranks could not fit, "
                        "possible to fit %s DP ranks",
                        node_ip,
                        dp_size_local,
                        dp_size_available,
487
                    )
488
489
490
                    continue
                dp_size_to_allocate = dp_size_local
            else:
491
492
                # for "pack_strategy" in "fill" and "span"
                # we always take everything that's available
493
494
495
                dp_size_to_allocate = dp_size_available

            for i in range(dp_size_to_allocate):
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
                device_bundle = [{device_str: 1.0, "node:" + node_ip: 0.001}]
                if pack_strategy == "span":
                    collected_bundles += device_bundle * n_device_on_node
                    assert len(collected_bundles) <= world_size, (
                        "collected_bundles should be <= world_size, "
                        f"but got {len(collected_bundles)=} and {world_size=}"
                    )

                    # we only create a placement group if we collected enough devices
                    if len(collected_bundles) < world_size:
                        continue

                    bundles = collected_bundles + [{"CPU": 1.0}]
                    collected_bundles = []
                else:
                    bundles = device_bundle * world_size + [{"CPU": 1.0}]

513
514
                pg = ray.util.placement_group(
                    name=f"dp_rank_{len(placement_groups)}",
515
                    strategy=placement_strategy,
516
517
518
519
                    bundles=bundles,
                )
                placement_groups.append(pg)
                local_dp_ranks.append(i)
520
521
                if len(placement_groups) == dp_size:
                    break
522
523

        if len(placement_groups) < dp_size:
524
            raise ValueError(
525
                f"Not enough resources to allocate {dp_size} "
526
527
528
                "placement groups, only created "
                f"{len(placement_groups)} placement groups. "
                "Available resources: "
529
530
                f"{available_resources}"
            )
531
532
533
534
535
536
537
        assert len(placement_groups) == dp_size, (
            f"Created {len(placement_groups)} DP placement groups, expected {dp_size}"
        )
        assert len(local_dp_ranks) == dp_size, (
            f"local_dp_ranks length {len(local_dp_ranks)} does not match "
            f"expected {dp_size}"
        )
538
539
        return placement_groups, local_dp_ranks

540
541
542
543
544
545
546
547
    @staticmethod
    def add_dp_placement_groups(
        old_vllm_config: VllmConfig, new_data_parallel_size: int
    ) -> tuple[list["PlacementGroup"], list[int]]:
        """
        Add placement groups for new data parallel size.
        """
        import ray
548
549
550
551
        from ray._private.state import (
            available_resources_per_node,
            total_resources_per_node,
        )
552
553
554
555
556
557
558
559
560
561
562
563
564
        from ray.util.state import list_nodes

        old_dp_size = old_vllm_config.parallel_config.data_parallel_size
        num_pg_to_create = new_data_parallel_size - old_dp_size

        if num_pg_to_create <= 0:
            return [], []

        dp_master_ip = old_vllm_config.parallel_config.data_parallel_master_ip
        world_size = old_vllm_config.parallel_config.world_size

        nodes = list_nodes()
        nodes = sorted(nodes, key=lambda node: node.node_ip != dp_master_ip)
565
        assert nodes[0].node_ip == dp_master_ip, "The first node must be the head node"
566
        assert len(nodes) == 1 or nodes[1].node_ip != dp_master_ip, (
567
568
            "There can only be one head node"
        )
569
570
571
572
573
574
575
576

        available_resources = available_resources_per_node()
        total_resources = total_resources_per_node()

        placement_groups = []
        local_dp_ranks = []
        num_pg_created = 0

577
        device_str = current_platform.ray_device_key
578
579
580
581
582
583
        for node in nodes:
            if num_pg_created >= num_pg_to_create:
                break

            node_ip = node.node_ip
            node_id = node.node_id
584
585
            if device_str not in available_resources[node_id]:
                continue
586
            available_gpus = int(available_resources[node_id][device_str])
587
588
589

            # Get total GPUs on this node from the node's resources
            # Ray stores node resources with node ID as key
590
            total_gpus = int(total_resources[node_id][device_str])
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607

            # Calculate used GPUs and used engines on this node
            used_gpus = max(0, total_gpus - available_gpus)
            used_engines_on_node = used_gpus // world_size

            # Calculate how many new engines this node can accommodate
            available_engine_count = available_gpus // world_size

            # Create placement groups for new engines on this node
            for i in range(available_engine_count):
                if num_pg_created >= num_pg_to_create:
                    break

                rank = old_dp_size + num_pg_created

                # Create bundles with node constraint for master node
                if node_ip == dp_master_ip:
608
609
610
                    bundles = [
                        {device_str: 1.0, "node:" + dp_master_ip: 0.001}
                    ] * world_size + [{"CPU": 1.0}]
611
                else:
612
                    bundles = [{device_str: 1.0}] * world_size + [{"CPU": 1.0}]
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628

                pg = ray.util.placement_group(
                    name=f"dp_rank_{rank}",
                    strategy="STRICT_PACK",
                    bundles=bundles,
                )
                placement_groups.append(pg)

                # Local rank starts from the number of engines already used
                # on this node
                local_rank = used_engines_on_node + i
                local_dp_ranks.append(local_rank)
                num_pg_created += 1

        return placement_groups, local_dp_ranks

629
630
631
    def scale_up_elastic_ep(
        self, cur_vllm_config: VllmConfig, new_data_parallel_size: int
    ) -> None:
632
633
634
635
        import copy

        import ray
        from ray.runtime_env import RuntimeEnv
636
        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
637

638
639
640
641
642
643
644
        from vllm.v1.engine.core import DPMoEEngineCoreActor, EngineCoreActor

        actor_class = (
            DPMoEEngineCoreActor
            if cur_vllm_config.model_config.is_moe
            else EngineCoreActor
        )
645

646
647
648
        cur_data_parallel_size = len(self.local_engine_actors) + len(
            self.remote_engine_actors
        )
649
650
651
652

        assert new_data_parallel_size > cur_data_parallel_size, (
            f"New data parallel size {new_data_parallel_size} must be greater "
            f"than current data parallel size {cur_data_parallel_size} "
653
654
            "for scale up"
        )
655

656
657
658
        placement_groups, local_dp_ranks = self.add_dp_placement_groups(
            cur_vllm_config, new_data_parallel_size
        )
659
660
661
662
663

        world_size = cur_vllm_config.parallel_config.world_size
        dp_master_ip = cur_vllm_config.parallel_config.data_parallel_master_ip
        new_local_engines = 0

664
665
666
667
        runtime_env = RuntimeEnv(
            env_vars=self.env_vars_dict | {"VLLM_ELASTIC_EP_SCALE_UP_LAUNCH": "1"}
        )
        for i, (pg, local_rank) in enumerate(zip(placement_groups, local_dp_ranks)):
668
669
            rank = cur_data_parallel_size + i
            dp_vllm_config = copy.deepcopy(cur_vllm_config)
670
            dp_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
671
672
673
674
            dp_vllm_config.parallel_config.placement_group = pg

            # Check if this placement group is on the head node
            local_client = any(
675
676
                bundle.get("node:" + dp_master_ip, 0) > 0 for bundle in pg.bundle_specs
            )
677
678
679
680
681

            if local_client:
                new_local_engines += 1
                # Update data_parallel_size_local
                dp_vllm_config.parallel_config.data_parallel_size_local = (
682
683
684
685
686
                    cur_vllm_config.parallel_config.data_parallel_size_local
                    + new_local_engines
                )

            actor = (
687
                ray.remote(actor_class)
688
689
690
691
692
693
694
695
                .options(
                    scheduling_strategy=PlacementGroupSchedulingStrategy(
                        placement_group=pg,
                        placement_group_bundle_index=world_size,
                    ),
                    runtime_env=runtime_env,
                )
                .remote(
696
697
698
699
700
701
                    vllm_config=dp_vllm_config,
                    executor_class=self.executor_class,
                    log_stats=self.log_stats,
                    local_client=local_client,
                    addresses=self.addresses,
                    dp_rank=rank,
702
703
704
                    local_dp_rank=local_rank,
                )
            )
705
706
707
708
709
710
711
712

            if local_client:
                self.local_engine_actors.append(actor)
            else:
                self.remote_engine_actors.append(actor)
            self.created_placement_groups.append(pg)
            self.placement_group_is_local.append(local_client)

713
714
715
716
717
718
719
720
721
722
723
724
725
        ray.get(
            [
                actor.wait_for_init.remote()
                for actor in (
                    self.local_engine_actors[-new_local_engines:]
                    if new_local_engines > 0
                    else []
                )
                + self.remote_engine_actors[
                    -(len(placement_groups) - new_local_engines) :
                ]
            ]
        )
726

727
728
729
730
731
        actors = (
            self.local_engine_actors[-new_local_engines:]
            if new_local_engines > 0
            else []
        ) + self.remote_engine_actors[-(len(placement_groups) - new_local_engines) :]
732
733
734
735

        for actor in actors:
            self.run_refs.append(actor.run.remote())

736
        cur_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
737
738
739
        # Update old_vllm_config with new data_parallel_size_local if any new
        # local engines were added
        if new_local_engines > 0:
740
            cur_vllm_config.parallel_config.data_parallel_size_local += (
741
                new_local_engines
742
            )
743

744
745
746
    def scale_down_elastic_ep(
        self, cur_data_parallel_size: int, new_data_parallel_size: int
    ) -> None:
747
        import ray
748

749
750
751
        assert cur_data_parallel_size > new_data_parallel_size, (
            f"cur_data_parallel_size {cur_data_parallel_size} must be greater "
            f"than new_data_parallel_size {new_data_parallel_size} "
752
753
            "for scale down"
        )
754
755
756
757
758
759
760
761
762
        for _ in range(cur_data_parallel_size - new_data_parallel_size):
            pg = self.created_placement_groups.pop()
            is_local = self.placement_group_is_local.pop()
            if is_local:
                self.local_engine_actors.pop()
            else:
                self.remote_engine_actors.pop()
            ray.util.remove_placement_group(pg)

763
764
765
    def get_run_refs(self):
        return self.run_refs

766
    def close(self):
767
        import ray
768

769
770
771
772
773
774
        for actor in self.local_engine_actors + self.remote_engine_actors:
            ray.kill(actor)
        for pg in self.created_placement_groups:
            ray.util.remove_placement_group(pg)


775
def get_engine_zmq_addresses(
776
777
    vllm_config: VllmConfig,
    num_api_servers: int = 1,
778
779
) -> EngineZmqAddresses:
    """Allocate ZMQ addresses for engine-client communication."""
780
781
782
    parallel_config = vllm_config.parallel_config
    local_engine_count = parallel_config.data_parallel_size_local
    local_start_index = parallel_config.data_parallel_rank_local
783
    dp_size = parallel_config.data_parallel_size
784
    host = parallel_config.data_parallel_master_ip
785
    local_engines_only = parallel_config.local_engines_only
786
787
788
789
790
791
792
793

    # In offline mode there is an LLM instance per DP rank and
    # one core engine per LLM, see
    # examples/offline_inference/data_parallel.py.
    offline_mode = local_start_index is not None

    # client_local_only = True for cases where this front-end
    # sends requests only to colocated engines.
794
795
796
    client_local_only = (
        offline_mode or local_engines_only or (local_engine_count == dp_size)
    )
797
798
799
    # NOTE(yongji): handling scaling from intra-node to inter-node
    if parallel_config.enable_elastic_ep:
        client_local_only = False
800

801
    return EngineZmqAddresses(
802
803
804
805
806
807
808
809
810
811
        inputs=[
            get_engine_client_zmq_addr(client_local_only, host)
            for _ in range(num_api_servers)
        ],
        outputs=[
            get_engine_client_zmq_addr(client_local_only, host)
            for _ in range(num_api_servers)
        ],
    )

812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838

@contextlib.contextmanager
def launch_core_engines(
    vllm_config: VllmConfig,
    executor_class: type[Executor],
    log_stats: bool,
    addresses: EngineZmqAddresses,
    num_api_servers: int = 1,
) -> Iterator[
    tuple[
        CoreEngineProcManager | CoreEngineActorManager | None,
        DPCoordinator | None,
        EngineZmqAddresses,
    ]
]:
    """Launch engine and DP coordinator processes as needed."""

    parallel_config = vllm_config.parallel_config
    dp_size = parallel_config.data_parallel_size
    local_engine_count = parallel_config.data_parallel_size_local
    local_start_index = parallel_config.data_parallel_rank_local
    dp_rank = parallel_config.data_parallel_rank
    host = parallel_config.data_parallel_master_ip
    local_engines_only = parallel_config.local_engines_only

    offline_mode = local_start_index is not None

839
840
841
842
843
844
845
    # Run the DP Coordinator process with rank 0 when in online DP mode.
    # The coordinator is needed for:
    # 1. Internal/hybrid LB: collecting and publishing queue stats for load balancing
    # 2. MoE models: wave coordination in addition to stats
    run_coordinator = (
        vllm_config.needs_dp_coordinator and not offline_mode and dp_rank == 0
    )
846
847

    if run_coordinator:
848
849
850
851
        coordinator = DPCoordinator(
            parallel_config,
            enable_wave_coordination=vllm_config.model_config.is_moe,
        )
852
853

        addresses.coordinator_input, addresses.coordinator_output = (
854
855
            coordinator.get_engine_socket_addresses()
        )
856
        addresses.frontend_stats_publish_address = (
857
858
            coordinator.get_stats_publish_address()
        )
859

860
        logger.info("Started DP Coordinator process (PID: %d)", coordinator.proc.pid)
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
    else:
        coordinator = None

    if parallel_config.data_parallel_backend == "ray":
        logger.info("Starting ray-based data parallel backend")

        engine_actor_manager = CoreEngineActorManager(
            vllm_config=vllm_config,
            addresses=addresses,
            executor_class=executor_class,
            log_stats=log_stats,
        )

        yield engine_actor_manager, coordinator, addresses
        return

877
    if offline_mode:
878
879
        assert local_engine_count == 1
        engines_to_handshake = [CoreEngine(index=dp_rank, local=True)]
880
881
882
883
884
    elif dp_rank == 0:
        # Rank 0 holds Coordinator, so it handshakes with all Cores
        # in both external dplb and internal dplb mode.
        # Note this also covers the case where we have zero local engines
        # and rank 0 is headless.
885
        engines_to_handshake = [
886
            CoreEngine(index=i, local=(i < local_engine_count)) for i in range(dp_size)
887
        ]
888
889
890
891
    else:
        # Rank > 0 handshakes with just the local cores it is managing.
        assert local_engines_only, (
            "Attempting to launch core_engines from dp_rank > 0, but "
892
893
            "found internal DPLB, which is incompatible."
        )
894
895
896
897
        engines_to_handshake = [
            CoreEngine(index=i, local=True)
            for i in range(dp_rank, dp_rank + local_engine_count)
        ]
898
899
900
901
902
903
904

    # Whether the started engines will handshake only with co-located
    # front-end processes. In external_dp_lb mode, ranks > 0 handshake with
    # their co-located frontend and also the rank 0 front-end, and hence this
    # will be False.
    handshake_local_only = offline_mode or local_engine_count == dp_size

905
906
907
908
    # NOTE(yongji): handling scaling from intra-node to inter-node
    if parallel_config.enable_elastic_ep:
        handshake_local_only = False

909
    handshake_address = get_engine_client_zmq_addr(
910
911
        handshake_local_only, host, parallel_config.data_parallel_rpc_port
    )
912

913
    if local_engines_only and dp_rank > 0:
914
915
916
917
918
919
920
        assert not handshake_local_only
        local_handshake_address = get_open_zmq_ipc_path()
        client_handshake_address = local_handshake_address
    else:
        local_handshake_address = handshake_address
        client_handshake_address = None

921
922
923
    with zmq_socket_ctx(
        local_handshake_address, zmq.ROUTER, bind=True
    ) as handshake_socket:
924
925
926
927
928
929
930
931
932
933
934
        # Start local engines.
        if local_engine_count:
            local_engine_manager = CoreEngineProcManager(
                vllm_config=vllm_config,
                executor_class=executor_class,
                log_stats=log_stats,
                handshake_address=handshake_address,
                client_handshake_address=client_handshake_address,
                local_client=True,
                local_engine_count=local_engine_count,
                start_index=dp_rank,
935
936
                local_start_index=local_start_index or 0,
            )
937
938
939
940
941
942
943
944
945
946
947
        else:
            local_engine_manager = None

        yield local_engine_manager, coordinator, addresses

        # Now wait for engines to start.
        wait_for_engine_startup(
            handshake_socket,
            addresses,
            engines_to_handshake,
            parallel_config,
948
            dp_size > 1 and vllm_config.model_config.is_moe,
949
950
951
952
953
954
955
956
957
958
959
            vllm_config.cache_config,
            local_engine_manager,
            coordinator.proc if coordinator else None,
        )


def wait_for_engine_startup(
    handshake_socket: zmq.Socket,
    addresses: EngineZmqAddresses,
    core_engines: list[CoreEngine],
    parallel_config: ParallelConfig,
960
    coordinated_dp: bool,
961
    cache_config: CacheConfig,
962
963
    proc_manager: CoreEngineProcManager | None,
    coord_process: Process | None,
964
965
966
967
968
969
970
971
972
):
    # Wait for engine core process(es) to send ready messages.
    local_count = parallel_config.data_parallel_size_local
    remote_count = len(core_engines) - local_count
    # [local, remote] counts
    conn_pending, start_pending = [local_count, remote_count], [0, 0]
    poller = zmq.Poller()
    poller.register(handshake_socket, zmq.POLLIN)

973
974
    remote_should_be_headless = (
        not parallel_config.data_parallel_hybrid_lb
975
        and not parallel_config.data_parallel_external_lb
976
    )
977

978
979
980
981
982
983
984
985
986
987
    if proc_manager is not None:
        for sentinel in proc_manager.sentinels():
            poller.register(sentinel, zmq.POLLIN)
    if coord_process is not None:
        poller.register(coord_process.sentinel, zmq.POLLIN)
    while any(conn_pending) or any(start_pending):
        events = poller.poll(STARTUP_POLL_PERIOD_MS)
        if not events:
            if any(conn_pending):
                logger.debug(
988
989
990
                    "Waiting for %d local, %d remote core engine proc(s) to connect.",
                    *conn_pending,
                )
991
992
            if any(start_pending):
                logger.debug(
993
994
995
                    "Waiting for %d local, %d remote core engine proc(s) to start.",
                    *start_pending,
                )
996
997
998
999
1000
1001
            continue
        if len(events) > 1 or events[0][0] != handshake_socket:
            # One of the local core processes exited.
            finished = proc_manager.finished_procs() if proc_manager else {}
            if coord_process is not None and coord_process.exitcode is not None:
                finished[coord_process.name] = coord_process.exitcode
1002
1003
1004
1005
1006
            raise RuntimeError(
                "Engine core initialization failed. "
                "See root cause above. "
                f"Failed core proc(s): {finished}"
            )
1007
1008
1009
1010

        # Receive HELLO and READY messages from the input socket.
        eng_identity, ready_msg_bytes = handshake_socket.recv_multipart()
        eng_index = int.from_bytes(eng_identity, "little")
1011
        engine = next((e for e in core_engines if e.identity == eng_identity), None)
1012
        if engine is None:
1013
1014
1015
            raise RuntimeError(
                f"Message from engine with unexpected data parallel rank: {eng_index}"
            )
1016
        msg = msgspec.msgpack.decode(ready_msg_bytes)
1017
        status, local, headless = msg["status"], msg["local"], msg["headless"]
1018
        if local != engine.local:
1019
1020
1021
1022
1023
1024
            raise RuntimeError(
                f"{status} message from "
                f"{'local' if local else 'remote'} "
                f"engine {eng_index}, expected it to be "
                f"{'local' if engine.local else 'remote'}"
            )
1025

1026
1027
1028
        # Remote engines must be headless iff we aren't in hybrid dp lb mode.
        if not local and headless != remote_should_be_headless:
            if headless:
1029
1030
1031
1032
1033
                raise RuntimeError(
                    f"Remote engine {eng_index} must not use "
                    f"--headless in external or hybrid dp lb "
                    f"mode"
                )
1034
            else:
1035
1036
1037
1038
1039
                raise RuntimeError(
                    f"Remote engine {eng_index} must use "
                    f"--headless unless in external or hybrid "
                    f"dp lb mode"
                )
1040

1041
        if status == "HELLO" and engine.state == CoreEngineState.NEW:
1042
            # Send init message with DP config info.
1043
1044
1045
1046
            init_message = msgspec.msgpack.encode(
                EngineHandshakeMetadata(
                    addresses=addresses,
                    parallel_config={
1047
1048
1049
1050
1051
1052
1053
                        k: getattr(parallel_config, k)
                        for k in (
                            "data_parallel_master_ip",
                            "data_parallel_master_port",
                            "_data_parallel_master_port_list",
                            "data_parallel_size",
                        )
1054
1055
1056
                    }
                    if coordinated_dp
                    else {},
1057
1058
1059
                )
            )
            handshake_socket.send_multipart((eng_identity, init_message), copy=False)
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
            conn_pending[0 if local else 1] -= 1
            start_pending[0 if local else 1] += 1
            engine.state = CoreEngineState.CONNECTED
        elif status == "READY" and engine.state == CoreEngineState.CONNECTED:
            # Setup KV cache config with initialization state from
            # engine core process. Sum values from all engines in DP case.
            num_gpu_blocks = cache_config.num_gpu_blocks or 0
            num_gpu_blocks += msg["num_gpu_blocks"]
            cache_config.num_gpu_blocks = num_gpu_blocks

            # In external DP LB mode, the coordinator address that the
            # front-end procs connect to is obtained from rank 0 via
            # one of the engine handshakes, and passed to the local
            # front-end process in the response from the other.
            if addresses.frontend_stats_publish_address is None:
1075
                addresses.frontend_stats_publish_address = msg.get("dp_stats_address")
1076

1077
1078
            # Validate config hash consistency across DP workers for MoE models.
            if coordinated_dp:
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
                worker_config_hash = msg.get("parallel_config_hash")
                expected_hash = parallel_config.compute_hash()
                if worker_config_hash != expected_hash:
                    raise RuntimeError(
                        f"Configuration mismatch detected for engine "
                        f"{eng_index}. All DP workers must have identical "
                        f"configurations for parameters that affect collective "
                        f"communication (e.g., enable_eplb, "
                        f"eplb_config.log_balancedness). "
                        f"Worker hash: {worker_config_hash}, "
                        f"Expected hash: {expected_hash}. "
                        f"Please ensure all workers are started with the same "
                        f"command-line arguments."
                    )

1094
1095
1096
            start_pending[0 if local else 1] -= 1
            engine.state = CoreEngineState.READY
        else:
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
            raise RuntimeError(
                f"Unexpected {status} message for "
                f"{'local' if local else 'remote'} engine "
                f"{eng_index} in {engine.state} state."
            )

        logger.debug(
            "%s from %s core engine process %s.",
            status,
            "local" if local else "remote",
            eng_index,
        )