ray_distributed_executor.py 30.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
import asyncio
import os
6
from collections import defaultdict
7
from dataclasses import dataclass
8
from typing import TYPE_CHECKING, Any, Callable, Optional, Union
9

10
import cloudpickle
11
12
import msgspec

13
import vllm.envs as envs
14
from vllm.executor.executor_base import DistributedExecutorBase
15
from vllm.executor.msgspec_utils import encode_hook
16
from vllm.executor.ray_utils import RayWorkerWrapper, initialize_ray_cluster, ray
17
from vllm.logger import init_logger
18
from vllm.platforms import current_platform
19
from vllm.ray.ray_env import get_env_vars_to_copy
20
from vllm.sequence import ExecuteModelRequest
21
22
23
24
25
26
27
from vllm.utils import (
    _run_task_with_lock,
    get_distributed_init_method,
    get_ip,
    get_open_port,
    make_async,
)
28
from vllm.v1.outputs import SamplerOutput
29
30

if ray is not None:
31
    from ray.actor import ActorHandle
32
    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
33
34
else:
    ActorHandle = None
35
36
37
38
39
40
41

if TYPE_CHECKING:
    from ray.util.placement_group import PlacementGroup

logger = init_logger(__name__)


42
43
44
45
46
47
48
@dataclass
class RayWorkerMetaData:
    """
    Metadata for a Ray worker.
    The order of ray worker creation can be random,
    and we need to reset the rank after creating all workers.
    """
49

50
51
52
53
54
55
56
    worker: ActorHandle
    created_rank: int
    adjusted_rank: int = -1
    ip: str = ""


class RayDistributedExecutor(DistributedExecutorBase):
57
58
59
60
61
    """Ray-based distributed executor"""

    # These env vars are worker-specific, therefore are NOT copied
    # from the driver to the workers
    WORKER_SPECIFIC_ENV_VARS = {
62
63
64
65
        "VLLM_HOST_IP",
        "VLLM_HOST_PORT",
        "LOCAL_RANK",
        "CUDA_VISIBLE_DEVICES",
66
67
    }

68
69
70
    # These non-vLLM env vars are copied from the driver to workers
    ADDITIONAL_ENV_VARS = {"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}

71
72
    uses_ray: bool = True

73
    def _init_executor(self) -> None:
74
        self.forward_dag: Optional[ray.dag.CompiledDAG] = None
75
        if envs.VLLM_USE_V1:
76
            # V1 uses SPMD worker and compiled DAG
77
78
            os.environ["VLLM_USE_RAY_SPMD_WORKER"] = "1"
            os.environ["VLLM_USE_RAY_COMPILED_DAG"] = "1"
79

80
81
            # For TPU or XPU, avoid compiling NVIDIA's NCCL
            if current_platform.is_tpu() or current_platform.is_xpu():
82
                os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"
83

84
85
86
87
88
89
90
91
92
93
94
95
        # If the env var is set, it uses the Ray's compiled DAG API
        # which optimizes the control plane overhead.
        # Run vLLM with VLLM_USE_RAY_COMPILED_DAG=1 to enable it.
        # Currently, this requires USE_RAY_SPMD_WORKER=True.
        self.use_ray_compiled_dag = envs.VLLM_USE_RAY_COMPILED_DAG
        # If the env var is set, then we do not distinguish between the
        # "driver worker" vs other workers. Also, the rank 0 worker will
        # be executed in a remote Ray worker. Currently this requires
        # USE_RAY_COMPILED_DAG=True.
        self.use_ray_spmd_worker = envs.VLLM_USE_RAY_SPMD_WORKER
        if self.use_ray_compiled_dag:
            assert self.use_ray_spmd_worker, (
96
97
                "VLLM_USE_RAY_COMPILED_DAG=1 requires VLLM_USE_RAY_SPMD_WORKER=1"
            )
98
99
100
        if self.use_ray_spmd_worker:
            # TODO: Support SPMD worker for non-DAG Ray executor.
            assert self.use_ray_compiled_dag, (
101
102
                "VLLM_USE_RAY_SPMD_WORKER=1 requires VLLM_USE_RAY_COMPILED_DAG=1"
            )
103

104
        assert self.uses_ray
105
        initialize_ray_cluster(self.parallel_config)
106
107
108
109
110
111
112
113
114
115
        placement_group = self.parallel_config.placement_group

        # Disable Ray usage stats collection.
        ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0")
        if ray_usage != "1":
            os.environ["RAY_USAGE_STATS_ENABLED"] = "0"

        # Create the parallel GPU workers.
        self._init_workers_ray(placement_group)

116
        self.input_encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)
117
        self.output_decoder = msgspec.msgpack.Decoder(Optional[list[SamplerOutput]])
118
119
        self.use_v1 = envs.VLLM_USE_V1

120
        self.pp_locks: Optional[list[asyncio.Lock]] = None
121
        if not self.use_ray_compiled_dag:
122
            self.driver_exec_method = make_async(self.driver_worker.execute_method)
123

124
    def shutdown(self) -> None:
125
126
127
128
129
        if logger:
            # Somehow logger can be None here.
            logger.info(
                "Shutting down Ray distributed executor. If you see error log "
                "from logging.cc regarding SIGTERM received, please ignore "
130
131
                "because this is the expected termination process in Ray."
            )
132
133
134
        if hasattr(self, "forward_dag") and self.forward_dag is not None:
            self.forward_dag.teardown()
            import ray
135

136
137
138
139
            for worker in self.workers:
                ray.kill(worker)
            self.forward_dag = None

140
    def _configure_ray_workers_use_nsight(self, ray_remote_kwargs) -> dict[str, Any]:
141
142
143
        # If nsight profiling is enabled, we need to set the profiling
        # configuration for the ray workers as runtime env.
        runtime_env = ray_remote_kwargs.setdefault("runtime_env", {})
144
145
146
147
148
149
150
        runtime_env.update(
            {
                "nsight": {
                    "t": "cuda,cudnn,cublas",
                    "o": "'worker_process_%p'",
                    "cuda-graph-trace": "node",
                }
151
            }
152
        )
153
154
155

        return ray_remote_kwargs

156
157
158
159
    # child class could overwrite this to return actual env vars.
    def _get_env_vars_to_be_updated(self):
        return self._env_vars_for_all_workers

160
    def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs):
161
        num_gpus = envs.VLLM_RAY_PER_WORKER_GPUS
162
163
164

        # The driver dummy worker does not actually use any resources.
        # It holds the resource for the driver worker.
165
        self.driver_dummy_worker: Optional[RayWorkerWrapper] = None
166
        # The remaining workers are the actual ray actors.
167
        self.workers: list[RayWorkerWrapper] = []
168

169
170
171
        # Used in ray compiled DAG: indexed first by PP rank,
        # and then TP rank. In other words, the inner list is
        # the TP group of workers for a PP rank.
172
        self.pp_tp_workers: list[list[RayWorkerWrapper]] = []
173

174
175
        if self.parallel_config.ray_workers_use_nsight:
            ray_remote_kwargs = self._configure_ray_workers_use_nsight(
176
177
                ray_remote_kwargs
            )
178

179
        logger.info("use_ray_spmd_worker: %s", self.use_ray_spmd_worker)
180

181
        # Create the workers.
182
        bundle_indices: list[int]
183
184
        if envs.VLLM_RAY_BUNDLE_INDICES:
            # Use the bundle indices specified by the user.
185
186
187
188
189
190
191
192
193
194
            bundle_indices = list(map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(",")))
            assert len(bundle_indices) == self.parallel_config.world_size, (
                "VLLM_RAY_BUNDLE_INDICES must have the same size"
                f" as the world size, but got {bundle_indices=} "
                f"and {self.parallel_config.world_size=}"
            )
            assert len(set(bundle_indices)) == len(bundle_indices), (
                "VLLM_RAY_BUNDLE_INDICES cannot have duplicate values,"
                f" but got {bundle_indices=}"
            )
195
196
197
198
199
200
        else:
            # use the first N bundles that have GPU resources.
            bundle_indices = []
            for bundle_id, bundle in enumerate(placement_group.bundle_specs):
                if bundle.get(current_platform.ray_device_key, 0):
                    bundle_indices.append(bundle_id)
201
            bundle_indices = bundle_indices[: self.parallel_config.world_size]
202

203
        worker_metadata: list[RayWorkerMetaData] = []
204
205
        driver_ip = get_ip()
        for rank, bundle_id in enumerate(bundle_indices):
206
207
208
209
210
            scheduling_strategy = PlacementGroupSchedulingStrategy(
                placement_group=placement_group,
                placement_group_capture_child_tasks=True,
                placement_group_bundle_index=bundle_id,
            )
211

212
213
214
215
216
217
218
            if current_platform.ray_device_key == "GPU":
                # NV+AMD GPUs, and Intel XPUs
                worker = ray.remote(
                    num_cpus=0,
                    num_gpus=num_gpus,
                    scheduling_strategy=scheduling_strategy,
                    **ray_remote_kwargs,
219
                )(RayWorkerWrapper).remote(vllm_config=self.vllm_config, rpc_rank=rank)
220
221
222
223
224
225
226
            else:
                worker = ray.remote(
                    num_cpus=0,
                    num_gpus=0,
                    resources={current_platform.ray_device_key: num_gpus},
                    scheduling_strategy=scheduling_strategy,
                    **ray_remote_kwargs,
227
228
                )(RayWorkerWrapper).remote(vllm_config=self.vllm_config, rpc_rank=rank)
            worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank))
229

230
231
232
233
234
235
        worker_ips = ray.get(
            [
                each.worker.get_node_ip.remote()  # type: ignore[attr-defined]
                for each in worker_metadata
            ]
        )
236
237
238

        for each, ip in zip(worker_metadata, worker_ips):
            each.ip = ip
239
240

        if not self.use_ray_spmd_worker:
241
242
243
244
            for i, each in enumerate(worker_metadata):
                # find and remove the dummy worker from the list
                worker = each.worker
                worker_ip = each.ip
245
                if self.driver_dummy_worker is None and worker_ip == driver_ip:
246
247
248
249
                    # If the worker is on the same node as the driver, we use it
                    # as the resource holder for the driver process.
                    self.driver_dummy_worker = worker
                    self.driver_worker = RayWorkerWrapper(
250
251
                        vllm_config=self.vllm_config, rpc_rank=0
                    )
252
                    worker_metadata.pop(i)
253
                    break
254

255
        logger.debug("workers: %s", worker_metadata)
256
        logger.debug("driver_dummy_worker: %s", self.driver_dummy_worker)
257
        if not self.use_ray_spmd_worker and self.driver_dummy_worker is None:
258
            raise ValueError(
259
260
261
                "Ray does not allocate any GPUs on the driver node."
                f"Driver IP: {driver_ip}, worker IPs: {worker_ips}."
                "Consider adjusting the Ray placement group or running "
262
263
                "the driver on a GPU node."
            )
264

265
        ip_counts: dict[str, int] = {}
266
267
268
        for ip in worker_ips:
            ip_counts[ip] = ip_counts.get(ip, 0) + 1

269
        def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
270
271
272
273
274
275
276
277
278
            """
            Sort the workers based on 3 properties:
            1. If the worker is on the same node as the driver (vllm engine),
                it should be placed first.
            2. Then, if the worker is on a node with fewer workers, it should
                be placed first.
            3. Finally, if the work is on a node with smaller IP address, it
                should be placed first.
            """
279
280
            ip = item.ip
            return (0 if ip == driver_ip else 1, ip_counts[ip], ip)
281
282
283
284

        # After sorting, the workers on the same node will be
        # close to each other, and the workers on the driver
        # node will be placed first.
285
286
287
        sorted_worker_metadata = sorted(
            worker_metadata, key=sort_by_driver_then_worker_ip
        )
288
289
290
291
292
        start_rank = 0 if self.use_ray_spmd_worker else 1
        for i, item in enumerate(sorted_worker_metadata):
            item.adjusted_rank = i + start_rank
        self.workers = [item.worker for item in sorted_worker_metadata]
        rerank_mapping = {
293
            item.created_rank: item.adjusted_rank for item in sorted_worker_metadata
294
295
        }
        self._run_workers("adjust_rank", rerank_mapping)
296

297
        # Get the set of GPU IDs used on each node.
298
299
300
301
302
303
        worker_node_and_gpu_ids = []
        for worker in [self.driver_dummy_worker] + self.workers:
            if worker is None:
                # driver_dummy_worker can be None when using ray spmd worker.
                continue
            worker_node_and_gpu_ids.append(
304
305
                ray.get(worker.get_node_and_gpu_ids.remote())
            )  # type: ignore
306

307
308
309
        node_workers = defaultdict(list)  # node id -> list of worker ranks
        node_gpus = defaultdict(list)  # node id -> list of gpu ids

310
311
        for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
            node_workers[node_id].append(i)
312
313
314
315
316
317
            # `gpu_ids` can be a list of strings or integers.
            # convert them to integers for consistency.
            # NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs),
            # string sorting is not sufficient.
            # see https://github.com/vllm-project/vllm/issues/5590
            gpu_ids = [int(x) for x in gpu_ids]
318
319
320
321
            node_gpus[node_id].extend(gpu_ids)
        for node_id, gpu_ids in node_gpus.items():
            node_gpus[node_id] = sorted(gpu_ids)

322
323
324
325
326
327
328
329
330
        all_ips = set(worker_ips + [driver_ip])
        n_ips = len(all_ips)
        n_nodes = len(node_workers)

        if n_nodes != n_ips:
            raise RuntimeError(
                f"Every node should have a unique IP address. Got {n_nodes}"
                f" nodes with node ids {list(node_workers.keys())} and "
                f"{n_ips} unique IP addresses {all_ips}. Please check your"
331
332
                " network configuration. If you set `VLLM_HOST_IP`"
                " environment variable, make sure it is unique for"
333
334
                " each node."
            )
335

336
        # Set environment variables for the driver and workers.
337
338
339
340
341
342
343
344
        all_args_to_update_environment_variables = [
            {
                current_platform.device_control_env_var: ",".join(
                    map(str, node_gpus[node_id])
                ),
            }
            for (node_id, _) in worker_node_and_gpu_ids
        ]
345

346
        # Environment variables to copy from driver to workers
347
348
        env_vars_to_copy = get_env_vars_to_copy(
            exclude_vars=self.WORKER_SPECIFIC_ENV_VARS,
349
            additional_vars=set(current_platform.additional_env_vars).union(
350
351
352
353
                self.ADDITIONAL_ENV_VARS
            ),
            destination="workers",
        )
354

355
        # Copy existing env vars to each worker's args
356
357
        for args in all_args_to_update_environment_variables:
            # TODO: refactor platform-specific env vars
358
            for name in env_vars_to_copy:
359
360
                if name in os.environ:
                    args[name] = os.environ[name]
361

362
        self._env_vars_for_all_workers = all_args_to_update_environment_variables
363

364
365
366
        self._run_workers(
            "update_environment_variables", self._get_env_vars_to_be_updated()
        )
367

368
369
370
371
372
373
374
375
376
377
        if len(node_gpus) == 1:
            # in single node case, we don't need to get the IP address.
            # the loopback address is sufficient
            # NOTE: a node may have several IP addresses, one for each
            # network interface. `get_ip()` might return any of them,
            # while they might not work for communication inside the node
            # if the network setup is complicated. Using the loopback address
            # solves this issue, as it always works for communication inside
            # the node.
            driver_ip = "127.0.0.1"
378
        distributed_init_method = get_distributed_init_method(
379
380
            driver_ip, get_open_port()
        )
381

382
        # Initialize the actual workers inside worker wrapper.
383
384
385
386
387
388
        all_kwargs = []
        for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids):
            local_rank = node_workers[node_id].index(rank)
            kwargs = dict(
                vllm_config=self.vllm_config,
                local_rank=local_rank,
389
390
                rank=rank,
                distributed_init_method=distributed_init_method,
391
392
393
394
395
                is_driver_worker=(not self.parallel_config)
                or (rank % self.parallel_config.tensor_parallel_size == 0),
            )
            all_kwargs.append(kwargs)
        self._run_workers("init_worker", all_kwargs)
396

397
        self._run_workers("init_device")
398
399
400
401
        self._run_workers(
            "load_model",
            max_concurrent_workers=self.parallel_config.max_parallel_loading_workers,
        )
402

403
404
405
        if self.use_ray_spmd_worker:
            for pp_rank in range(self.parallel_config.pipeline_parallel_size):
                self.pp_tp_workers.append([])
406
                for tp_rank in range(self.parallel_config.tensor_parallel_size):
407
408
                    # PP=2, TP=4
                    # pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
409
410
411
                    rank = (
                        pp_rank * self.parallel_config.tensor_parallel_size
                    ) + tp_rank
412
413
414
415
                    assert len(self.pp_tp_workers[pp_rank]) == tp_rank
                    assert pp_rank < len(self.pp_tp_workers)
                    self.pp_tp_workers[pp_rank].append(self.workers[rank])

416
417
418
        # This is the list of workers that are rank 0 of each TP group EXCEPT
        # global rank 0. These are the workers that will broadcast to the
        # rest of the workers.
419
        self.tp_driver_workers: list[RayWorkerWrapper] = []
420
421
422
        # This is the list of workers that are not drivers and not the first
        # worker in a TP group. These are the workers that will be
        # broadcasted to.
423
        self.non_driver_workers: list[RayWorkerWrapper] = []
424

425
        # Enforce rank order for correct rank to return final output.
426
427
428
        for index, worker in enumerate(self.workers):
            # The driver worker is rank 0 and not in self.workers.
            rank = index + 1
429
            if rank % self.parallel_config.tensor_parallel_size == 0:
430
                self.tp_driver_workers.append(worker)
431
            else:
432
                self.non_driver_workers.append(worker)
433

434
    def _driver_execute_model(
435
        self, execute_model_req: Optional[ExecuteModelRequest]
436
    ) -> Optional[list[SamplerOutput]]:
437
        """Run execute_model in the driver worker.
438

439
440
441
        Passing None will cause the driver to stop the model execution
        loop running in each of the remote workers.
        """
442
        assert not self.use_ray_spmd_worker, (
443
444
445
            "driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1"
        )
        return self.driver_worker.execute_method("execute_model", execute_model_req)
446

447
    def execute_model(
448
        self, execute_model_req: ExecuteModelRequest
449
    ) -> list[SamplerOutput]:
450
451
452
453
454
455
        if not self.use_ray_spmd_worker:
            return super().execute_model(execute_model_req)

        if self.forward_dag is None:
            self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)

456
457
458
459
        if self.use_v1:
            serialized_data = execute_model_req
        else:
            serialized_data = self.input_encoder.encode(execute_model_req)
460
        outputs = ray.get(self.forward_dag.execute(serialized_data))
461
        output = outputs[0] if self.use_v1 else self.output_decoder.decode(outputs[0])
462
        return output
463

464
465
    def _run_workers(
        self,
466
        method: Union[str, Callable],
467
        *args,
468
        async_run_tensor_parallel_workers_only: bool = False,
469
470
471
        max_concurrent_workers: Optional[int] = None,
        **kwargs,
    ) -> Any:
472
473
474
        """Runs the given method on all workers. Can be used in the following
        ways:

475
476
477
478
479
        Args:
        - async_run_tensor_parallel_workers_only: If True the method will be
          run only in the remote TP workers, not the driver worker.
          It will also be run asynchronously and return a list of futures
          rather than blocking on the results.
480
        - args/kwargs: All workers share the same args/kwargs
481
        """
482
        sent_method = method if isinstance(method, str) else cloudpickle.dumps(method)
483
        del method
484
485
        if self.use_ray_spmd_worker:
            assert not async_run_tensor_parallel_workers_only, (
486
487
                "async_run_tensor_parallel_workers_only is not supported for spmd mode."
            )
488
489

        if max_concurrent_workers:
490
            raise NotImplementedError("max_concurrent_workers is not supported yet.")
491

492
493
494
495
496
        # Start the ray workers first.
        ray_workers = self.workers
        if async_run_tensor_parallel_workers_only:
            ray_workers = self.non_driver_workers
        ray_worker_outputs = [
497
            worker.execute_method.remote(sent_method, *args, **kwargs)
498
            for worker in ray_workers
499
        ]
500

501
        if async_run_tensor_parallel_workers_only:
502
503
504
            # Just return futures
            return ray_worker_outputs

505
506
507
508
509
510
        driver_worker_output = []
        # In SPMD mode, the driver worker is the same as any other worker,
        # so we only explicitly execute on the driver worker if using a
        # non-SPMD worker class.
        if not self.use_ray_spmd_worker:
            # Start the driver worker after all the ray workers.
511
            driver_worker_output = [
512
                self.driver_worker.execute_method(sent_method, *args, **kwargs)
513
            ]
514

515
516
        # Get the results of the ray workers.
        if self.workers:
517
            ray_worker_outputs = ray.get(ray_worker_outputs)
518

519
        return driver_worker_output + ray_worker_outputs
520

521
522
523
524
525
    def _wait_for_tasks_completion(self, parallel_worker_tasks: Any) -> None:
        """Wait for futures returned from _run_workers() with
        async_run_remote_workers_only to complete."""
        ray.get(parallel_worker_tasks)

526
    def _check_ray_cgraph_installation(self):
527
528
        import importlib.metadata

529
530
        from packaging import version

Rui Qiao's avatar
Rui Qiao committed
531
        required_version = version.parse("2.43.0")
532
        current_version = version.parse(importlib.metadata.version("ray"))
533
        if current_version < required_version:
534
535
536
537
            raise ValueError(
                f"Ray version {required_version} is "
                f"required, but found {current_version}"
            )
538

539
        import importlib.util
540
541

        cgraph_spec = importlib.util.find_spec("ray.experimental.compiled_dag_ref")
542
        if cgraph_spec is None:
543
544
545
546
            raise ValueError(
                "Ray Compiled Graph is not installed. "
                "Run `pip install ray[cgraph]` to install it."
            )
547
548

        cupy_spec = importlib.util.find_spec("cupy")
549
        if cupy_spec is None and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl":
550
551
            raise ValueError(
                "cupy is not installed but required since "
552
                "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE is set to 'nccl'. "
553
554
                "Run `pip install ray[cgraph]` and check cupy installation."
            )
555
556

    def _compiled_ray_dag(self, enable_asyncio: bool):
557
        assert self.parallel_config.use_ray
558
        self._check_ray_cgraph_installation()
559
560
561
562
563
564
565
566
        # Enlarge the default value of "RAY_CGRAPH_get_timeout" to 300 seconds
        # (it is 10 seconds by default). This is a Ray environment variable to
        # control the timeout of getting result from a compiled graph execution,
        # i.e., the distributed execution that includes model forward runs and
        # intermediate tensor communications, in the case of vllm.
        # Note: we should set this env var before importing
        # ray.dag, otherwise it will not take effect.
        os.environ.setdefault("RAY_CGRAPH_get_timeout", "300")  # noqa: SIM112
567
        from ray.dag import InputNode, MultiOutputNode
568
569

        logger.info(
570
571
572
            "RAY_CGRAPH_get_timeout is set to %s",
            os.environ["RAY_CGRAPH_get_timeout"],  # noqa: SIM112
        )
573
574
575
576
577
578
579
580
        logger.info(
            "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE = %s",
            envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE,
        )
        logger.info(
            "VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM = %s",
            envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM,
        )
581
582
583
584
585

        channel_type = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
        if channel_type not in ("auto", "nccl", "shm"):
            raise ValueError(
                "Invalid value for VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: "
586
587
                f"{channel_type}. Valid values are: 'auto', 'nccl', or 'shm'."
            )
588

589
        with InputNode() as input_data:
590
            # Example DAG: PP=2, TP=4
591
592
593
594
595
596
597
598
599
600
601
602
            #
            # For V0:
            # ExecuteModelRequest -> 0 -> (ExecuteModelReq, IntermediateTensors) -> 4 -> SamplerOutput   # noqa: E501
            # ExecuteModelRequest -> 1 -> (ExecuteModelReq, IntermediateTensors) -> 5 -> SamplerOutput   # noqa: E501
            # ExecuteModelRequest -> 2 -> (ExecuteModelReq, IntermediateTensors) -> 6 -> SamplerOutput   # noqa: E501
            # ExecuteModelRequest -> 3 -> (ExecuteModelReq, IntermediateTensors) -> 7 -> SamplerOutput   # noqa: E501
            #
            # For V1:
            # SchedulerOutput -> 0 -> (SchedulerOutput, IntermediateTensors) -> 4 -> ModelRunnerOutput   # noqa: E501
            # SchedulerOutput -> 1 -> (SchedulerOutput, IntermediateTensors) -> 5 -> ModelRunnerOutput   # noqa: E501
            # SchedulerOutput -> 2 -> (SchedulerOutput, IntermediateTensors) -> 6 -> ModelRunnerOutput   # noqa: E501
            # SchedulerOutput -> 3 -> (SchedulerOutput, IntermediateTensors) -> 7 -> ModelRunnerOutput   # noqa: E501
603
604
605
606
607
608
609

            # All workers in the first TP group will take in the
            # ExecuteModelRequest as input.
            outputs = [input_data for _ in self.pp_tp_workers[0]]
            for pp_rank, tp_group in enumerate(self.pp_tp_workers):
                # Each PP worker takes in the output of the previous PP worker,
                # and the TP group executes in SPMD fashion.
610
611
                if self.use_v1:
                    outputs = [
612
613
614
615
                        worker.execute_model_ray.bind(  # type: ignore[attr-defined]
                            outputs[i]
                        )
                        for i, worker in enumerate(tp_group)
616
617
618
                    ]
                else:
                    outputs = [
619
620
621
622
                        worker.execute_model_spmd.bind(  # type: ignore[attr-defined]
                            outputs[i]
                        )
                        for i, worker in enumerate(tp_group)
623
                    ]
624
625

                last_pp_rank = len(self.pp_tp_workers) - 1
626
627
628
629
                if (
                    pp_rank < last_pp_rank
                    and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE != "shm"
                ):
630
631
                    # Specify how intermediate tensors should be passed
                    # between pp stages, no need to specify for the last
632
633
                    # pp stage or when using shared memory (the default).
                    transport = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
634
                    outputs = [
Rui Qiao's avatar
Rui Qiao committed
635
                        output.with_tensor_transport(transport=transport)
636
637
638
639
640
                        for output in outputs
                    ]

            forward_dag = MultiOutputNode(outputs)

641
642
        if envs.VLLM_USE_RAY_WRAPPED_PP_COMM:
            from ray.experimental.channel.accelerator_context import (
643
644
                register_accelerator_context,
            )
645
646

            from vllm.distributed.device_communicators.ray_communicator import (
647
648
649
650
651
652
653
654
655
656
657
                RayPPCommunicator,
            )

            register_accelerator_context(
                torch_module_name="cuda", communicator_cls=RayPPCommunicator
            )
            logger.info(
                "Using RayPPCommunicator "
                "(which wraps vLLM _PP GroupCoordinator) "
                "for Ray Compiled Graph communication."
            )
658
        else:
659
660
661
            logger.info(
                "Using Ray's NCCL communicator for Ray Compiled Graph communication."
            )
662

663
664
        return forward_dag.experimental_compile(
            enable_asyncio=enable_asyncio,
665
666
            _overlap_gpu_communication=envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM,
        )
667
668

    def __del__(self):
669
        self.shutdown()
670

671
    async def execute_model_async(
672
        self, execute_model_req: ExecuteModelRequest
673
    ) -> list[SamplerOutput]:
674
675
676
677
678
679
        if not self.use_ray_spmd_worker:
            return await super().execute_model_async(execute_model_req)

        if self.forward_dag is None:
            self.forward_dag = self._compiled_ray_dag(enable_asyncio=True)

680
681
        serialized_data = self.input_encoder.encode(execute_model_req)
        dag_future = await self.forward_dag.execute_async(serialized_data)
682
683
        output = await dag_future[0]
        return self.output_decoder.decode(output)
684

685
    async def _driver_execute_model_async(
686
        self, execute_model_req: Optional[ExecuteModelRequest] = None
687
    ) -> list[SamplerOutput]:
688
        assert not self.use_ray_spmd_worker, (
689
690
            "driver_worker does not exist for VLLM_USE_RAY_SPMD_WORKER=1"
        )
691
        if not self.tp_driver_workers:
692
            return await self.driver_exec_method("execute_model", execute_model_req)
693
694
695
696
697
698
699
700
701
        if self.pp_locks is None:
            # This locks each pipeline parallel stage so multiple virtual
            # engines can't execute on the same stage at the same time
            # We create the locks here to avoid creating them in the constructor
            # which uses a different asyncio loop.
            self.pp_locks = [
                asyncio.Lock()
                for _ in range(self.parallel_config.pipeline_parallel_size)
            ]
702

703
        tasks = [
704
            asyncio.create_task(
705
706
707
708
709
710
711
                _run_task_with_lock(
                    self.driver_exec_method,
                    self.pp_locks[0],
                    "execute_model",
                    execute_model_req,
                )
            )
712
        ]
713
        for pp_rank, driver_worker in enumerate(self.tp_driver_workers, start=1):
714
715
            tasks.append(
                asyncio.create_task(
716
717
718
719
720
721
722
723
                    _run_task_with_lock(
                        driver_worker.execute_method.remote,
                        self.pp_locks[pp_rank],
                        "execute_model",
                        execute_model_req,
                    )
                )
            )
724
725
726
727
728

        results = await asyncio.gather(*tasks)

        # Only the last PP stage has the final results.
        return results[-1]
729
730

    async def _start_worker_execution_loop(self):
731
        assert not self.use_ray_spmd_worker, (
732
733
            "worker loop is disabled for VLLM_USE_RAY_SPMD_WORKER=1"
        )
734
735
        coros = [
            worker.execute_method.remote("start_worker_execution_loop")
736
            for worker in self.non_driver_workers
737
738
        ]
        return await asyncio.gather(*coros)
739

740
741
742
743
    def check_health(self) -> None:
        # Assume that the Ray workers are healthy.
        # TODO: check the health of the Ray workers
        return