| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | from collections import defaultdict |
| 6 | from collections.abc import Callable |
| 7 | from concurrent.futures import Future |
| 8 | from dataclasses import dataclass |
| 9 | from typing import TYPE_CHECKING, Any |
| 10 | |
| 11 | import cloudpickle |
| 12 | |
| 13 | import vllm.envs as envs |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.ray.ray_env import get_env_vars_to_copy |
| 17 | from vllm.utils.network_utils import ( |
| 18 | get_distributed_init_method, |
| 19 | get_ip, |
| 20 | get_open_port, |
| 21 | ) |
| 22 | from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput |
| 23 | from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType |
| 24 | from vllm.v1.executor.abstract import Executor |
| 25 | from vllm.v1.executor.ray_env_utils import ( |
| 26 | update_runtime_env_for_worker_import, |
| 27 | ) |
| 28 | from vllm.v1.executor.ray_utils import ( |
| 29 | WORKER_SPECIFIC_ENV_VARS, |
| 30 | FutureWrapper, |
| 31 | RayWorkerWrapper, |
| 32 | detach_zero_copy_from_model_runner_output, |
| 33 | initialize_ray_cluster, |
| 34 | ray, |
| 35 | ) |
| 36 | from vllm.v1.outputs import ModelRunnerOutput |
| 37 | |
| 38 | if ray is not None: |
| 39 | from ray.actor import ActorHandle |
| 40 | from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy |
| 41 | else: |
| 42 | ActorHandle = None |
| 43 | |
| 44 | if TYPE_CHECKING: |
| 45 | from ray.util.placement_group import PlacementGroup |
| 46 | |
| 47 | logger = init_logger(__name__) |
| 48 | |
| 49 | COMPLETED_NONE_FUTURE: Future[ModelRunnerOutput | None] = Future() |
| 50 | COMPLETED_NONE_FUTURE.set_result(None) |
| 51 | |
| 52 | |
| 53 | @dataclass |
| 54 | class RayWorkerMetaData: |
| 55 | """Metadata for a Ray worker. |
| 56 | The order of ray worker creation can be random, |
| 57 | and we need to reset the rank after creating all workers. |
| 58 | """ |
| 59 | |
| 60 | worker: ActorHandle |
| 61 | created_rank: int |
| 62 | adjusted_rank: int = -1 |
| 63 | ip: str = "" |
| 64 | |
| 65 | |
| 66 | class RayDistributedExecutor(Executor): |
| 67 | """Ray-based distributed executor.""" |
| 68 | |
| 69 | uses_ray: bool = True |
| 70 | supports_pp: bool = True |
| 71 | |
| 72 | def _init_executor(self) -> None: |
| 73 | self.forward_dag: ray.dag.CompiledDAG | None = None |
| 74 | |
| 75 | # For TPU or XPU, avoid compiling NVIDIA's NCCL |
| 76 | if current_platform.is_tpu() or current_platform.is_xpu(): |
| 77 | os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm" |
| 78 | |
| 79 | assert self.uses_ray |
| 80 | initialize_ray_cluster(self.parallel_config) |
| 81 | placement_group = self.parallel_config.placement_group |
| 82 | |
| 83 | # Disable Ray usage stats collection. |
| 84 | ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0") |
| 85 | if ray_usage != "1": |
| 86 | os.environ["RAY_USAGE_STATS_ENABLED"] = "0" |
| 87 | |
| 88 | # Create the parallel GPU workers. |
| 89 | self._init_workers_ray(placement_group) |
| 90 | |
| 91 | # KV connector setup |
| 92 | self.has_connector = self.vllm_config.kv_transfer_config is not None |
| 93 | |
| 94 | if ( |
| 95 | self.vllm_config.ec_transfer_config is not None |
| 96 | and self.parallel_config.world_size > 1 |
| 97 | ): |
| 98 | raise NotImplementedError( |
| 99 | "EC connector worker metadata is not supported with the " |
| 100 | "legacy Ray executor when world_size > 1: only the output " |
| 101 | "of a single worker is fetched, silently dropping the " |
| 102 | "other workers' EC connector state. Set " |
| 103 | "VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 to use RayExecutorV2, " |
| 104 | "or use the multiprocessing executor instead." |
| 105 | ) |
| 106 | |
| 107 | self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and ( |
| 108 | self.vllm_config.ec_transfer_config is None |
| 109 | or self.vllm_config.ec_transfer_config.is_ec_consumer |
| 110 | ) |
| 111 | |
| 112 | self.scheduler_output: SchedulerOutput | None = None |
| 113 | |
| 114 | def shutdown(self) -> None: |
| 115 | if logger: |
| 116 | # Somehow logger can be None here. |
| 117 | logger.info( |
| 118 | "Shutting down Ray distributed executor. If you see error log " |
| 119 | "from logging.cc regarding SIGTERM received, please ignore " |
| 120 | "because this is the expected termination process in Ray." |
| 121 | ) |
| 122 | if hasattr(self, "forward_dag") and self.forward_dag is not None: |
| 123 | self.forward_dag.teardown() |
| 124 | import ray |
| 125 | |
| 126 | for worker in self.workers: |
| 127 | ray.kill(worker) |
| 128 | self.forward_dag = None |
| 129 | |
| 130 | def _configure_ray_workers_use_nsight(self, ray_remote_kwargs) -> dict[str, Any]: |
| 131 | # If nsight profiling is enabled, we need to set the profiling |
| 132 | # configuration for the ray workers as runtime env. |
| 133 | runtime_env = ray_remote_kwargs.setdefault("runtime_env", {}) |
| 134 | runtime_env.update( |
| 135 | { |
| 136 | "nsight": { |
| 137 | "t": "cuda,cudnn,cublas", |
| 138 | "o": "'worker_process_%p'", |
| 139 | "cuda-graph-trace": "node", |
| 140 | } |
| 141 | } |
| 142 | ) |
| 143 | |
| 144 | return ray_remote_kwargs |
| 145 | |
| 146 | def _update_noset_device_env_vars(self, ray_remote_kwargs): |
| 147 | runtime_env = ray_remote_kwargs.setdefault("runtime_env", {}) |
| 148 | env_vars = runtime_env.setdefault("env_vars", {}) |
| 149 | env_vars.update( |
| 150 | {env_var: "1" for env_var in current_platform.ray_noset_device_env_vars} |
| 151 | ) |
| 152 | update_runtime_env_for_worker_import(runtime_env) |
| 153 | return ray_remote_kwargs |
| 154 | |
| 155 | # child class could overwrite this to return actual env vars. |
| 156 | def _get_env_vars_to_be_updated(self): |
| 157 | return self._env_vars_for_all_workers |
| 158 | |
| 159 | def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs): |
| 160 | num_gpus = envs.VLLM_RAY_PER_WORKER_GPUS |
| 161 | |
| 162 | # The driver dummy worker does not actually use any resources. |
| 163 | # It holds the resource for the driver worker. |
| 164 | self.driver_dummy_worker: RayWorkerWrapper | None = None |
| 165 | # The remaining workers are the actual ray actors. |
| 166 | self.workers: list[RayWorkerWrapper] = [] |
| 167 | |
| 168 | # Used in ray compiled DAG: indexed first by PP rank, |
| 169 | # and then TP rank. In other words, the inner list is |
| 170 | # the TP group of workers for a PP rank. |
| 171 | self.pp_tp_workers: list[list[RayWorkerWrapper]] = [] |
| 172 | |
| 173 | if self.parallel_config.ray_workers_use_nsight: |
| 174 | ray_remote_kwargs = self._configure_ray_workers_use_nsight( |
| 175 | ray_remote_kwargs |
| 176 | ) |
| 177 | |
| 178 | # The way ray actors are setup in vllm is that the visible devices are |
| 179 | # not set by actors, they are left unset by ray. Internally we index |
| 180 | # the right gpu with local_rank. This is similar to how mp mode works. |
| 181 | self._update_noset_device_env_vars(ray_remote_kwargs) |
| 182 | |
| 183 | # Create the workers. |
| 184 | bundle_indices: list[int] |
| 185 | if envs.VLLM_RAY_BUNDLE_INDICES: |
| 186 | # Use the bundle indices specified by the user. |
| 187 | bundle_indices = list(map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(","))) |
| 188 | assert len(bundle_indices) == self.parallel_config.world_size, ( |
| 189 | "VLLM_RAY_BUNDLE_INDICES must have the same size" |
| 190 | f" as the world size, but got {bundle_indices=} " |
| 191 | f"and {self.parallel_config.world_size=}" |
| 192 | ) |
| 193 | assert len(set(bundle_indices)) == len(bundle_indices), ( |
| 194 | "VLLM_RAY_BUNDLE_INDICES cannot have duplicate values," |
| 195 | f" but got {bundle_indices=}" |
| 196 | ) |
| 197 | else: |
| 198 | # use the first N bundles that have GPU resources. |
| 199 | bundle_indices = [] |
| 200 | for bundle_id, bundle in enumerate(placement_group.bundle_specs): |
| 201 | if bundle.get(current_platform.ray_device_key, 0): |
| 202 | bundle_indices.append(bundle_id) |
| 203 | bundle_indices = bundle_indices[: self.parallel_config.world_size] |
| 204 | |
| 205 | worker_metadata: list[RayWorkerMetaData] = [] |
| 206 | driver_ip = get_ip() |
| 207 | for rank, bundle_id in enumerate(bundle_indices): |
| 208 | scheduling_strategy = PlacementGroupSchedulingStrategy( |
| 209 | placement_group=placement_group, |
| 210 | placement_group_capture_child_tasks=True, |
| 211 | placement_group_bundle_index=bundle_id, |
| 212 | ) |
| 213 | |
| 214 | if current_platform.ray_device_key == "GPU": |
| 215 | # NV+AMD GPUs, and Intel XPUs |
| 216 | worker = ray.remote( |
| 217 | num_cpus=0, |
| 218 | num_gpus=num_gpus, |
| 219 | scheduling_strategy=scheduling_strategy, |
| 220 | **ray_remote_kwargs, |
| 221 | )(RayWorkerWrapper).remote(rpc_rank=rank) |
| 222 | else: |
| 223 | worker = ray.remote( |
| 224 | num_cpus=0, |
| 225 | num_gpus=0, |
| 226 | resources={current_platform.ray_device_key: num_gpus}, |
| 227 | scheduling_strategy=scheduling_strategy, |
| 228 | **ray_remote_kwargs, |
| 229 | )(RayWorkerWrapper).remote(rpc_rank=rank) |
| 230 | |
| 231 | worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank)) |
| 232 | |
| 233 | worker_ips = ray.get( |
| 234 | [ |
| 235 | each.worker.get_node_ip.remote() # type: ignore[attr-defined] |
| 236 | for each in worker_metadata |
| 237 | ] |
| 238 | ) |
| 239 | |
| 240 | for each, ip in zip(worker_metadata, worker_ips): |
| 241 | each.ip = ip |
| 242 | |
| 243 | logger.debug("workers: %s", worker_metadata) |
| 244 | logger.debug("driver_dummy_worker: %s", self.driver_dummy_worker) |
| 245 | |
| 246 | ip_counts: dict[str, int] = {} |
| 247 | for ip in worker_ips: |
| 248 | ip_counts[ip] = ip_counts.get(ip, 0) + 1 |
| 249 | |
| 250 | def sort_by_driver_then_worker_ip(item: RayWorkerMetaData): |
| 251 | """Sort the workers based on 3 properties: |
| 252 | 1. If the worker is on the same node as the driver (vllm engine), |
| 253 | it should be placed first. |
| 254 | 2. Then, if the worker is on a node with fewer workers, it should |
| 255 | be placed first. |
| 256 | 3. Finally, if the work is on a node with smaller IP address, it |
| 257 | should be placed first. |
| 258 | """ |
| 259 | ip = item.ip |
| 260 | return 0 if ip == driver_ip else 1, ip_counts[ip], ip |
| 261 | |
| 262 | # After sorting, the workers on the same node will be |
| 263 | # close to each other, and the workers on the driver |
| 264 | # node will be placed first. |
| 265 | sorted_worker_metadata = sorted( |
| 266 | worker_metadata, key=sort_by_driver_then_worker_ip |
| 267 | ) |
| 268 | for i, item in enumerate(sorted_worker_metadata): |
| 269 | item.adjusted_rank = i |
| 270 | self.workers = [item.worker for item in sorted_worker_metadata] |
| 271 | rerank_mapping = { |
| 272 | item.created_rank: item.adjusted_rank for item in sorted_worker_metadata |
| 273 | } |
| 274 | self.collective_rpc("adjust_rank", args=(rerank_mapping,)) |
| 275 | |
| 276 | # Get the set of physical GPU IDs used on each node. |
| 277 | worker_node_and_physical_gpu_ids = [] |
| 278 | for worker in [self.driver_dummy_worker] + self.workers: |
| 279 | if worker is None: |
| 280 | # driver_dummy_worker can be None when using ray spmd worker. |
| 281 | continue |
| 282 | worker_node_and_physical_gpu_ids.append( |
| 283 | ray.get(worker.get_node_and_physical_gpu_ids.remote()) # type: ignore[attr-defined] |
| 284 | ) |
| 285 | |
| 286 | node_workers = defaultdict(list) # node id -> list of worker ranks |
| 287 | node_physical_gpu_ids = defaultdict(list) # node id -> physical GPU IDs |
| 288 | |
| 289 | for i, (node_id, physical_gpu_ids) in enumerate( |
| 290 | worker_node_and_physical_gpu_ids |
| 291 | ): |
| 292 | node_workers[node_id].append(i) |
| 293 | # `physical_gpu_ids` can be a list of strings or integers. |
| 294 | # convert them to integers for consistency. |
| 295 | # NOTE: physical GPU IDs can be larger than 9 (e.g. 16 GPUs), |
| 296 | # string sorting is not sufficient. |
| 297 | # see https://github.com/vllm-project/vllm/issues/5590 |
| 298 | physical_gpu_ids = [ |
| 299 | current_platform.device_control_id_to_physical_device_id(str(x)) |
| 300 | for x in physical_gpu_ids |
| 301 | ] |
| 302 | node_physical_gpu_ids[node_id].extend(physical_gpu_ids) |
| 303 | for node_id, physical_gpu_ids in node_physical_gpu_ids.items(): |
| 304 | node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids) |
| 305 | |
| 306 | all_ips = set(worker_ips + [driver_ip]) |
| 307 | n_ips = len(all_ips) |
| 308 | n_nodes = len(node_workers) |
| 309 | |
| 310 | if n_nodes != n_ips: |
| 311 | raise RuntimeError( |
| 312 | f"Every node should have a unique IP address. Got {n_nodes}" |
| 313 | f" nodes with node ids {list(node_workers.keys())} and " |
| 314 | f"{n_ips} unique IP addresses {all_ips}. Please check your" |
| 315 | " network configuration. If you set `VLLM_HOST_IP`" |
| 316 | " environment variable, make sure it is unique for" |
| 317 | " each node." |
| 318 | ) |
| 319 | |
| 320 | all_args_to_update_environment_variables: list[dict[str, str]] = [ |
| 321 | {} for _ in worker_node_and_physical_gpu_ids |
| 322 | ] |
| 323 | |
| 324 | # Environment variables to copy from driver to workers |
| 325 | env_vars_to_copy = get_env_vars_to_copy( |
| 326 | exclude_vars=WORKER_SPECIFIC_ENV_VARS, |
| 327 | additional_vars=set(current_platform.additional_env_vars), |
| 328 | destination="workers", |
| 329 | ) |
| 330 | |
| 331 | # Copy existing env vars to each worker's args |
| 332 | for args in all_args_to_update_environment_variables: |
| 333 | # TODO: refactor platform-specific env vars |
| 334 | for name in env_vars_to_copy: |
| 335 | if name in os.environ: |
| 336 | args[name] = os.environ[name] |
| 337 | |
| 338 | self._env_vars_for_all_workers = all_args_to_update_environment_variables |
| 339 | |
| 340 | self.collective_rpc( |
| 341 | "update_environment_variables", args=(self._get_env_vars_to_be_updated(),) |
| 342 | ) |
| 343 | |
| 344 | if len(node_physical_gpu_ids) == 1: |
| 345 | # in single node case, we don't need to get the IP address. |
| 346 | # the loopback address is sufficient |
| 347 | # NOTE: a node may have several IP addresses, one for each |
| 348 | # network interface. `get_ip()` might return any of them, |
| 349 | # while they might not work for communication inside the node |
| 350 | # if the network setup is complicated. Using the loopback address |
| 351 | # solves this issue, as it always works for communication inside |
| 352 | # the node. |
| 353 | driver_ip = "127.0.0.1" |
| 354 | distributed_init_method = get_distributed_init_method( |
| 355 | driver_ip, get_open_port() |
| 356 | ) |
| 357 | |
| 358 | # Initialize the actual workers inside worker wrapper. |
| 359 | all_kwargs = [] |
| 360 | for rank, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids): |
| 361 | local_rank = node_workers[node_id].index(rank) |
| 362 | kwargs = dict( |
| 363 | vllm_config=self.vllm_config, |
| 364 | assigned_physical_gpu_ids=sorted(node_physical_gpu_ids[node_id]), |
| 365 | local_rank=local_rank, |
| 366 | rank=rank, |
| 367 | distributed_init_method=distributed_init_method, |
| 368 | is_driver_worker=(not self.parallel_config) |
| 369 | or (rank % self.parallel_config.tensor_parallel_size == 0), |
| 370 | ) |
| 371 | all_kwargs.append(kwargs) |
| 372 | self.collective_rpc("init_worker", args=(all_kwargs,)) |
| 373 | |
| 374 | self.collective_rpc("init_device") |
| 375 | if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: |
| 376 | self.collective_rpc("elastic_ep_execute", args=("load_model",)) |
| 377 | else: |
| 378 | self.collective_rpc("load_model") |
| 379 | |
| 380 | def _update_block_size(worker): |
| 381 | current_platform.update_block_size_for_backend(worker.vllm_config) |
| 382 | |
| 383 | self.collective_rpc(_update_block_size) |
| 384 | |
| 385 | for pp_rank in range(self.parallel_config.pipeline_parallel_size): |
| 386 | self.pp_tp_workers.append([]) |
| 387 | for tp_rank in range(self.parallel_config.tensor_parallel_size): |
| 388 | # PP=2, TP=4 |
| 389 | # pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]] |
| 390 | rank = (pp_rank * self.parallel_config.tensor_parallel_size) + tp_rank |
| 391 | assert len(self.pp_tp_workers[pp_rank]) == tp_rank |
| 392 | assert pp_rank < len(self.pp_tp_workers) |
| 393 | self.pp_tp_workers[pp_rank].append(self.workers[rank]) |
| 394 | |
| 395 | def reinitialize_distributed( |
| 396 | self, reconfig_request: ReconfigureDistributedRequest |
| 397 | ) -> None: |
| 398 | self.collective_rpc("reinitialize_distributed", args=(reconfig_request,)) |
| 399 | if ( |
| 400 | reconfig_request.new_data_parallel_rank |
| 401 | == ReconfigureRankType.SHUTDOWN_CURRENT_RANK |
| 402 | ): |
| 403 | self.shutdown() |
| 404 | |
| 405 | def execute_model( # type: ignore[override] |
| 406 | self, |
| 407 | scheduler_output: SchedulerOutput, |
| 408 | non_block: bool = False, |
| 409 | ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]: |
| 410 | if self.scheduler_output is not None: |
| 411 | raise RuntimeError( |
| 412 | "State error: sample_tokens() must be called " |
| 413 | "after execute_model() returns None." |
| 414 | ) |
| 415 | |
| 416 | if not self.uses_sampler or not scheduler_output.total_num_scheduled_tokens: |
| 417 | # Model will not execute, call model runner immediately. |
| 418 | return self._execute_dag(scheduler_output, None, non_block) |
| 419 | |
| 420 | # Model will execute, defer to sample_tokens() call. |
| 421 | self.scheduler_output = scheduler_output |
| 422 | return COMPLETED_NONE_FUTURE if non_block else None |
| 423 | |
| 424 | def sample_tokens( # type: ignore[override] |
| 425 | self, |
| 426 | grammar_output: "GrammarOutput | None", |
| 427 | non_block: bool = False, |
| 428 | ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]: |
| 429 | """Execute the model on the Ray workers. |
| 430 | |
| 431 | The scheduler output to use should have been provided in |
| 432 | a prior call to execute_model(). |
| 433 | |
| 434 | Args: |
| 435 | grammar_output: The structured outputs grammar bitmask, if applicable. |
| 436 | non_block: If True, the method will return a Future. |
| 437 | |
| 438 | Returns: |
| 439 | The model runner output. |
| 440 | |
| 441 | """ |
| 442 | scheduler_output = self.scheduler_output |
| 443 | if scheduler_output is None: |
| 444 | return COMPLETED_NONE_FUTURE if non_block else None |
| 445 | |
| 446 | self.scheduler_output = None |
| 447 | |
| 448 | return self._execute_dag(scheduler_output, grammar_output, non_block) |
| 449 | |
| 450 | def _execute_dag( |
| 451 | self, |
| 452 | scheduler_output: SchedulerOutput, |
| 453 | grammar_output: "GrammarOutput | None", |
| 454 | non_block: bool = False, |
| 455 | ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]: |
| 456 | # Build the compiled DAG for the first time. |
| 457 | if self.forward_dag is None: # type: ignore |
| 458 | self.forward_dag = self._compiled_ray_dag(enable_asyncio=False) |
| 459 | |
| 460 | refs = self.forward_dag.execute((scheduler_output, grammar_output)) # type: ignore |
| 461 | |
| 462 | if not self.has_connector: |
| 463 | # Get output only from a single worker (output_rank) |
| 464 | # When PP is not used, we block here until the result is available. |
| 465 | if not non_block: |
| 466 | output = refs[0].get() |
| 467 | detach_zero_copy_from_model_runner_output(output) |
| 468 | return output |
| 469 | |
| 470 | # When PP is used, we return a FutureWrapper immediately so that |
| 471 | # the scheduler can yield to the next batch. |
| 472 | return FutureWrapper(refs[0]) |
| 473 | |
| 474 | # Get output from all workers when connector is present |
| 475 | assert self.kv_output_aggregator is not None |
| 476 | if not non_block: |
| 477 | # Block and get results from all workers |
| 478 | outputs = ray.get(refs) |
| 479 | for output in outputs: |
| 480 | detach_zero_copy_from_model_runner_output(output) |
| 481 | return self.kv_output_aggregator.aggregate(outputs) |
| 482 | |
| 483 | # Return a future that will aggregate outputs from all workers |
| 484 | return FutureWrapper(refs, self.kv_output_aggregator) |
| 485 | |
| 486 | def collective_rpc( # type: ignore[override] |
| 487 | self, |
| 488 | method: str | Callable, |
| 489 | timeout: float | None = None, |
| 490 | args: tuple = (), |
| 491 | kwargs: dict[str, Any] | None = None, |
| 492 | non_block: bool = False, |
| 493 | ) -> list[Any] | Future[list[Any]]: |
| 494 | """Runs the given method on all workers.""" |
| 495 | sent_method = method if isinstance(method, str) else cloudpickle.dumps(method) |
| 496 | del method |
| 497 | |
| 498 | if kwargs is None: |
| 499 | kwargs = {} |
| 500 | ray_worker_outputs = [ |
| 501 | worker.execute_method.remote( # type: ignore[attr-defined] |
| 502 | sent_method, *args, **kwargs |
| 503 | ) |
| 504 | for worker in self.workers |
| 505 | ] |
| 506 | |
| 507 | # Get the results of the ray workers. |
| 508 | if non_block: |
| 509 | return FutureWrapper(ray_worker_outputs) |
| 510 | |
| 511 | return ray.get(ray_worker_outputs, timeout=timeout) |
| 512 | |
| 513 | def _check_ray_cgraph_installation(self): |
| 514 | import importlib.metadata |
| 515 | |
| 516 | from packaging import version |
| 517 | |
| 518 | required_version = version.parse("2.43.0") |
| 519 | current_version = version.parse(importlib.metadata.version("ray")) |
| 520 | if current_version < required_version: |
| 521 | raise ValueError( |
| 522 | f"Ray version {required_version} is " |
| 523 | f"required, but found {current_version}" |
| 524 | ) |
| 525 | |
| 526 | import importlib.util |
| 527 | |
| 528 | cgraph_spec = importlib.util.find_spec("ray.experimental.compiled_dag_ref") |
| 529 | if cgraph_spec is None: |
| 530 | raise ValueError( |
| 531 | "Ray Compiled Graph is not installed. " |
| 532 | "Run `pip install ray[cgraph]` to install it." |
| 533 | ) |
| 534 | |
| 535 | cupy_spec = importlib.util.find_spec("cupy") |
| 536 | if cupy_spec is None and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl": |
| 537 | raise ValueError( |
| 538 | "cupy is not installed but required since " |
| 539 | "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE is set to 'nccl'. " |
| 540 | "Run `pip install ray[cgraph]` and check cupy installation." |
| 541 | ) |
| 542 | |
| 543 | def _compiled_ray_dag(self, enable_asyncio: bool): |
| 544 | assert self.parallel_config.use_ray |
| 545 | self._check_ray_cgraph_installation() |
| 546 | # Enlarge the default value of "RAY_CGRAPH_get_timeout" to 300 seconds |
| 547 | # (it is 10 seconds by default). This is a Ray environment variable to |
| 548 | # control the timeout of getting result from a compiled graph execution, |
| 549 | # i.e., the distributed execution that includes model forward runs and |
| 550 | # intermediate tensor communications, in the case of vllm. |
| 551 | # Note: we should set this env var before importing |
| 552 | # ray.dag, otherwise it will not take effect. |
| 553 | os.environ.setdefault("RAY_CGRAPH_get_timeout", "300") # noqa: SIM112 |
| 554 | from ray.dag import InputNode, MultiOutputNode |
| 555 | |
| 556 | logger.info( |
| 557 | "RAY_CGRAPH_get_timeout is set to %s", |
| 558 | os.environ["RAY_CGRAPH_get_timeout"], # noqa: SIM112 |
| 559 | ) |
| 560 | logger.info( |
| 561 | "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE = %s", |
| 562 | envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE, |
| 563 | ) |
| 564 | logger.info( |
| 565 | "VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM = %s", |
| 566 | envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM, |
| 567 | ) |
| 568 | |
| 569 | channel_type = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE |
| 570 | if channel_type not in ("auto", "nccl", "shm"): |
| 571 | raise ValueError( |
| 572 | "Invalid value for VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: " |
| 573 | f"{channel_type}. Valid values are: 'auto', 'nccl', or 'shm'." |
| 574 | ) |
| 575 | |
| 576 | with InputNode() as input_data: |
| 577 | # Example DAG: PP=2, TP=4 |
| 578 | # |
| 579 | # SchedulerOutput -> 0 -> (SchedulerOutput, IntermediateTensors) -> 4 -> ModelRunnerOutput # noqa: E501 |
| 580 | # SchedulerOutput -> 1 -> (SchedulerOutput, IntermediateTensors) -> 5 -> ModelRunnerOutput # noqa: E501 |
| 581 | # SchedulerOutput -> 2 -> (SchedulerOutput, IntermediateTensors) -> 6 -> ModelRunnerOutput # noqa: E501 |
| 582 | # SchedulerOutput -> 3 -> (SchedulerOutput, IntermediateTensors) -> 7 -> ModelRunnerOutput # noqa: E501 |
| 583 | |
| 584 | # All workers in the first TP group will take in the |
| 585 | # ExecuteModelRequest as input. |
| 586 | outputs = [input_data for _ in self.pp_tp_workers[0]] |
| 587 | for pp_rank, tp_group in enumerate(self.pp_tp_workers): |
| 588 | # Each PP worker takes in the output of the previous PP worker, |
| 589 | # and the TP group executes in SPMD fashion. |
| 590 | outputs = [ |
| 591 | worker.execute_model_ray.bind(outputs[i]) # type: ignore[attr-defined] |
| 592 | for i, worker in enumerate(tp_group) |
| 593 | ] |
| 594 | |
| 595 | last_pp_rank = len(self.pp_tp_workers) - 1 |
| 596 | if ( |
| 597 | pp_rank < last_pp_rank |
| 598 | and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE != "shm" |
| 599 | ): |
| 600 | # Specify how intermediate tensors should be passed |
| 601 | # between pp stages, no need to specify for the last |
| 602 | # pp stage or when using shared memory (the default). |
| 603 | transport = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE |
| 604 | outputs = [ |
| 605 | output.with_tensor_transport(transport=transport) |
| 606 | for output in outputs |
| 607 | ] |
| 608 | |
| 609 | forward_dag = MultiOutputNode(outputs) |
| 610 | |
| 611 | if envs.VLLM_USE_RAY_WRAPPED_PP_COMM: |
| 612 | from ray.experimental.channel.accelerator_context import ( |
| 613 | register_accelerator_context, |
| 614 | ) |
| 615 | |
| 616 | from vllm.distributed.device_communicators.ray_communicator import ( |
| 617 | RayPPCommunicator, |
| 618 | ) |
| 619 | |
| 620 | register_accelerator_context( |
| 621 | torch_module_name="cuda", communicator_cls=RayPPCommunicator |
| 622 | ) |
| 623 | logger.info( |
| 624 | "Using RayPPCommunicator " |
| 625 | "(which wraps vLLM _PP GroupCoordinator) " |
| 626 | "for Ray Compiled Graph communication." |
| 627 | ) |
| 628 | else: |
| 629 | logger.info( |
| 630 | "Using Ray's NCCL communicator for Ray Compiled Graph communication." |
| 631 | ) |
| 632 | |
| 633 | return forward_dag.experimental_compile( |
| 634 | enable_asyncio=enable_asyncio, |
| 635 | _overlap_gpu_communication=envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM, |
| 636 | ) |
| 637 | |
| 638 | def __del__(self): |
| 639 | self.shutdown() |
| 640 | |
| 641 | def check_health(self) -> None: |
| 642 | # Assume that the Ray workers are healthy. |
| 643 | # TODO: check the health of the Ray workers |
| 644 | return |