| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | import time |
| 6 | from collections import defaultdict |
| 7 | from concurrent.futures import Future |
| 8 | from typing import TYPE_CHECKING, Union |
| 9 | |
| 10 | import numpy as np |
| 11 | |
| 12 | import vllm.platforms |
| 13 | from vllm.config import ParallelConfig |
| 14 | from vllm.distributed import get_pp_group |
| 15 | from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.platforms import current_platform |
| 18 | from vllm.sequence import IntermediateTensors |
| 19 | from vllm.utils.network_utils import get_ip |
| 20 | from vllm.v1.outputs import AsyncModelRunnerOutput |
| 21 | from vllm.v1.serial_utils import run_method |
| 22 | from vllm.v1.worker.worker_base import WorkerWrapperBase |
| 23 | |
| 24 | if TYPE_CHECKING: |
| 25 | from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput |
| 26 | from vllm.v1.outputs import ModelRunnerOutput |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | PG_WAIT_TIMEOUT = 1800 |
| 30 | |
| 31 | # Env vars that are worker-specific and must NOT be copied from the |
| 32 | # driver to Ray workers — they are set per-worker after GPU discovery. |
| 33 | WORKER_SPECIFIC_ENV_VARS: set[str] = { |
| 34 | "VLLM_HOST_IP", |
| 35 | "VLLM_HOST_PORT", |
| 36 | "VLLM_NIXL_SIDE_CHANNEL_HOST", |
| 37 | "LOCAL_RANK", |
| 38 | "CUDA_VISIBLE_DEVICES", |
| 39 | "HIP_VISIBLE_DEVICES", |
| 40 | "ROCR_VISIBLE_DEVICES", |
| 41 | } |
| 42 | |
| 43 | try: |
| 44 | import ray |
| 45 | from ray.util import placement_group_table |
| 46 | from ray.util.placement_group import PlacementGroup |
| 47 | |
| 48 | try: |
| 49 | from ray._private.state import available_resources_per_node |
| 50 | except ImportError: |
| 51 | # Ray 2.9.x doesn't expose `available_resources_per_node` |
| 52 | from ray._private.state import state as _state |
| 53 | |
| 54 | available_resources_per_node = _state._available_resources_per_node |
| 55 | |
| 56 | class RayWorkerWrapper(WorkerWrapperBase): |
| 57 | """Ray wrapper for vllm.worker.Worker, allowing Worker to be |
| 58 | lazily initialized after Ray sets CUDA_VISIBLE_DEVICES.""" |
| 59 | |
| 60 | def __init__(self, *args, **kwargs) -> None: |
| 61 | super().__init__(*args, **kwargs) |
| 62 | # Since the compiled DAG runs a main execution |
| 63 | # in a different thread that calls cuda.set_device. |
| 64 | # The flag indicates is set_device is called on |
| 65 | # that thread. |
| 66 | self.compiled_dag_cuda_device_set = False |
| 67 | |
| 68 | rpc_rank: int |
| 69 | |
| 70 | def adjust_rank(self, rank_mapping: dict[int, int]) -> None: |
| 71 | """Adjust the rpc_rank based on the given mapping. |
| 72 | It is only used during the initialization of the executor, |
| 73 | to adjust the rpc_rank of workers after we create all workers. |
| 74 | """ |
| 75 | if self.rpc_rank in rank_mapping: |
| 76 | self.rpc_rank = rank_mapping[self.rpc_rank] |
| 77 | |
| 78 | def execute_method(self, method: str | bytes, *args, **kwargs): |
| 79 | try: |
| 80 | return run_method(self, method, args, kwargs) |
| 81 | except Exception as e: |
| 82 | # if the driver worker also execute methods, |
| 83 | # exceptions in the rest worker may cause deadlock in rpc |
| 84 | # see https://github.com/vllm-project/vllm/issues/3455 |
| 85 | msg = ( |
| 86 | f"Error executing method {method!r}. " |
| 87 | "This might cause deadlock in distributed execution." |
| 88 | ) |
| 89 | logger.exception(msg) |
| 90 | raise e |
| 91 | |
| 92 | def get_node_ip(self) -> str: |
| 93 | return get_ip() |
| 94 | |
| 95 | def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]: |
| 96 | node_id = ray.get_runtime_context().get_node_id() |
| 97 | device_key = vllm.platforms.current_platform.ray_device_key |
| 98 | if not device_key: |
| 99 | raise RuntimeError( |
| 100 | "current platform %s does not support ray.", |
| 101 | vllm.platforms.current_platform.device_name, |
| 102 | ) |
| 103 | physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[ |
| 104 | device_key |
| 105 | ] |
| 106 | return node_id, physical_gpu_ids |
| 107 | |
| 108 | def setup_device_if_necessary(self): |
| 109 | # TODO(swang): This is needed right now because Ray CG executes |
| 110 | # on a background thread, so we need to reset torch's current |
| 111 | # device. |
| 112 | # We can remove this API after it is fixed in compiled graph. |
| 113 | assert self.worker is not None, "Worker is not initialized" |
| 114 | if not self.compiled_dag_cuda_device_set: |
| 115 | if current_platform.is_tpu(): |
| 116 | # Not needed |
| 117 | pass |
| 118 | else: |
| 119 | assert self.worker.device is not None |
| 120 | current_platform.set_device(self.worker.device) |
| 121 | |
| 122 | self.compiled_dag_cuda_device_set = True |
| 123 | |
| 124 | def execute_model_ray( |
| 125 | self, |
| 126 | execute_model_input: tuple["SchedulerOutput", "GrammarOutput"] |
| 127 | | tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"], |
| 128 | ) -> Union[ |
| 129 | "ModelRunnerOutput", |
| 130 | tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"], |
| 131 | ]: |
| 132 | # This method is used by Ray Compiled Graph to execute the model, |
| 133 | # and it needs a special logic of self.setup_device_if_necessary() |
| 134 | self.setup_device_if_necessary() |
| 135 | assert self.worker is not None, "Worker is not initialized" |
| 136 | if len(execute_model_input) == 3: |
| 137 | scheduler_output, grammar_output, intermediate_tensors = ( |
| 138 | execute_model_input |
| 139 | ) |
| 140 | else: |
| 141 | scheduler_output, grammar_output = execute_model_input |
| 142 | intermediate_tensors = None |
| 143 | assert self.worker.model_runner is not None |
| 144 | output = self.worker.model_runner.execute_model( |
| 145 | scheduler_output, intermediate_tensors |
| 146 | ) |
| 147 | if self._is_intermediate_tensors(output): |
| 148 | if ( |
| 149 | self.worker.model_runner.supports_mm_inputs |
| 150 | and get_pp_group().is_first_rank |
| 151 | ): |
| 152 | # Strip mm_features before Ray forwards it to the next PP Stage. |
| 153 | # PP Stage>0 only needs the intermediate tensors, |
| 154 | # not preprocessed multimodal data. |
| 155 | |
| 156 | # scheduled_new_reqs is a required field of SchedulerOutput, |
| 157 | # so accessing it directly will raise AttributeError if missing. |
| 158 | for req in scheduler_output.scheduled_new_reqs: |
| 159 | req.mm_features = [] |
| 160 | return scheduler_output, grammar_output, output |
| 161 | |
| 162 | if isinstance(output, AsyncModelRunnerOutput): |
| 163 | output = output.get_output() |
| 164 | if not self._is_last_rank(): |
| 165 | # Case where there are no scheduled requests |
| 166 | # but may still be finished requests. |
| 167 | assert not output or not output.req_ids |
| 168 | output = scheduler_output, grammar_output, None |
| 169 | elif output is None: |
| 170 | output = self.worker.model_runner.sample_tokens(grammar_output) |
| 171 | # Ensure outputs crossing Ray compiled DAG are serializable. |
| 172 | # AsyncModelRunnerOutput holds CUDA events and cannot be |
| 173 | # pickled. |
| 174 | if isinstance(output, AsyncModelRunnerOutput): |
| 175 | output = output.get_output() |
| 176 | return output |
| 177 | |
| 178 | def _is_intermediate_tensors(self, output) -> bool: |
| 179 | return isinstance(output, IntermediateTensors) |
| 180 | |
| 181 | def _is_last_rank(self) -> bool: |
| 182 | return get_pp_group().is_last_rank |
| 183 | |
| 184 | ray_import_err = None |
| 185 | |
| 186 | except ImportError as e: |
| 187 | ray = None # type: ignore |
| 188 | # only capture string to avoid variable references in the traceback that can |
| 189 | # prevent garbage collection in some cases |
| 190 | ray_import_err = str(e) |
| 191 | RayWorkerWrapper = None # type: ignore |
| 192 | |
| 193 | |
| 194 | def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> None: |
| 195 | """Detach Ray SHM-channel zero-copy buffers from a ModelRunnerOutput in-place. |
| 196 | |
| 197 | Ray compiled DAG SHM channels may return zero-copy objects (e.g. `np.ndarray`) |
| 198 | backed by Ray's shared-memory object store. Ray's channel docs explicitly |
| 199 | warn that subsequent reads may block if such an object is still in scope. |
| 200 | |
| 201 | vLLM can return numpy-backed logprobs and routed experts in |
| 202 | `ModelRunnerOutput`. If those arrays are backed by Ray SHM (commonly |
| 203 | read-only), retaining them in scope across scheduler iterations can stall |
| 204 | the channel and eventually hit `RAY_CGRAPH_get_timeout`. |
| 205 | |
| 206 | Copy read-only numpy arrays so the returned output no longer retains |
| 207 | references to Ray's shared-memory buffers. |
| 208 | |
| 209 | We intentionally do not touch `prompt_logprobs_dict`: those entries are |
| 210 | `LogprobsTensors` backed by PyTorch-owned CPU tensors (`to_cpu_nonblocking` |
| 211 | or `empty_cpu`), not NumPy views decoded from Ray channels. |
| 212 | """ |
| 213 | |
| 214 | def _copy_if_readonly(arr): |
| 215 | if isinstance(arr, np.ndarray) and not arr.flags.writeable: |
| 216 | return arr.copy() |
| 217 | return arr |
| 218 | |
| 219 | if output.logprobs is not None: |
| 220 | token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs |
| 221 | |
| 222 | # `cu_num_generated_tokens` is already a plain Python list (or None), |
| 223 | # so it never aliases Ray SHM buffers and can be reused as-is. |
| 224 | token_ids_c = _copy_if_readonly(token_ids) |
| 225 | logprobs_c = _copy_if_readonly(logprobs) |
| 226 | ranks_c = _copy_if_readonly(ranks) |
| 227 | if ( |
| 228 | token_ids_c is not token_ids |
| 229 | or logprobs_c is not logprobs |
| 230 | or ranks_c is not ranks |
| 231 | ): |
| 232 | output.logprobs = type(output.logprobs)( |
| 233 | token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens |
| 234 | ) |
| 235 | |
| 236 | aux_output = output.aux_output_connector_output |
| 237 | if aux_output is not None: |
| 238 | for request_output in aux_output.values(): |
| 239 | request_output.rows = _copy_if_readonly(request_output.rows) |
| 240 | |
| 241 | |
| 242 | class FutureWrapper(Future): |
| 243 | """A wrapper around Ray output reference to meet the interface |
| 244 | of .execute_model(): The top level (core busy loop) expects .result() api |
| 245 | to block and return a single output. |
| 246 | |
| 247 | If aggregator is provided, the outputs from all workers are aggregated upon |
| 248 | the result() call. If not only the first worker's output is returned. |
| 249 | """ |
| 250 | |
| 251 | def __init__(self, ref_or_refs, aggregator: KVOutputAggregator | None = None): |
| 252 | super().__init__() |
| 253 | self.ref_or_refs = ref_or_refs |
| 254 | self.aggregator = aggregator |
| 255 | |
| 256 | def result(self, timeout=None): |
| 257 | outputs = ray.get(self.ref_or_refs, timeout=timeout) |
| 258 | if self.aggregator is None: |
| 259 | detach_zero_copy_from_model_runner_output(outputs) |
| 260 | return outputs |
| 261 | |
| 262 | for output in outputs: |
| 263 | detach_zero_copy_from_model_runner_output(output) |
| 264 | return self.aggregator.aggregate(outputs, output_rank=0) |
| 265 | |
| 266 | |
| 267 | def ray_is_available() -> bool: |
| 268 | """Returns True if Ray is available.""" |
| 269 | return ray is not None |
| 270 | |
| 271 | |
| 272 | def assert_ray_available(): |
| 273 | """Raise an exception if Ray is not available.""" |
| 274 | if ray is None: |
| 275 | raise ValueError( |
| 276 | f"Failed to import Ray: {ray_import_err}." |
| 277 | "Please install Ray with `pip install ray`." |
| 278 | ) |
| 279 | |
| 280 | |
| 281 | def _verify_bundles( |
| 282 | placement_group: "PlacementGroup", |
| 283 | parallel_config: ParallelConfig, |
| 284 | device_str: str, |
| 285 | require_gpu_on_driver: bool = True, |
| 286 | ): |
| 287 | """Verify a given placement group has bundles located in the right place. |
| 288 | |
| 289 | There are 2 rules. |
| 290 | - Warn if all tensor parallel workers cannot fit in a single node. |
| 291 | - Fail if driver node is not included in a placement group |
| 292 | (only when require_gpu_on_driver is True). |
| 293 | """ |
| 294 | assert ray.is_initialized(), ( |
| 295 | "Ray is not initialized although distributed-executor-backend is ray." |
| 296 | ) |
| 297 | pg_data = placement_group_table(placement_group) |
| 298 | # bundle_idx -> node_id |
| 299 | bundle_to_node_ids = pg_data["bundles_to_node_id"] |
| 300 | # bundle_idx -> bundle (e.g., {"GPU": 1}) |
| 301 | bundles = pg_data["bundles"] |
| 302 | # node_id -> List of bundle (e.g., {"GPU": 1}) |
| 303 | node_id_to_bundle: dict[str, list[dict[str, float]]] = defaultdict(list) |
| 304 | |
| 305 | for bundle_idx, node_id in bundle_to_node_ids.items(): |
| 306 | node_id_to_bundle[node_id].append(bundles[bundle_idx]) |
| 307 | driver_node_id = ray.get_runtime_context().get_node_id() |
| 308 | |
| 309 | if require_gpu_on_driver and driver_node_id not in node_id_to_bundle: |
| 310 | raise RuntimeError( |
| 311 | f"driver node id {driver_node_id} is not included in a placement " |
| 312 | f"group {placement_group.id}. Node id -> bundles " |
| 313 | f"{node_id_to_bundle}. " |
| 314 | "You don't have enough GPUs available in a current node. Check " |
| 315 | "`ray status` and `ray list nodes` to see if you have available " |
| 316 | "GPUs in a node `{driver_node_id}` before starting an vLLM engine." |
| 317 | ) |
| 318 | |
| 319 | for node_id, bundles in node_id_to_bundle.items(): |
| 320 | if len(bundles) < parallel_config.tensor_parallel_size: |
| 321 | logger.warning( |
| 322 | "tensor_parallel_size=%d " |
| 323 | "is bigger than a reserved number of %ss (%d " |
| 324 | "%ss) in a node %s. Tensor parallel workers can be " |
| 325 | "spread out to 2+ nodes which can degrade the performance " |
| 326 | "unless you have fast interconnect across nodes, like " |
| 327 | "Infiniband. To resolve this issue, make sure you have more " |
| 328 | "than %d GPUs available at each node.", |
| 329 | parallel_config.tensor_parallel_size, |
| 330 | device_str, |
| 331 | len(bundles), |
| 332 | device_str, |
| 333 | node_id, |
| 334 | parallel_config.tensor_parallel_size, |
| 335 | ) |
| 336 | |
| 337 | |
| 338 | def build_actor_name( |
| 339 | instance_id: str, |
| 340 | rank: int, |
| 341 | tp_size: int, |
| 342 | pp_size: int, |
| 343 | pcp_size: int, |
| 344 | ) -> str: |
| 345 | """Build a descriptive Ray actor name for dashboard visibility.""" |
| 346 | name = f"vllm_Worker_{instance_id}" |
| 347 | if tp_size > 1: |
| 348 | name += f"_TP{rank % tp_size}" |
| 349 | if pp_size > 1: |
| 350 | name += f"_PP{(rank // tp_size) % pp_size}" |
| 351 | if pcp_size > 1: |
| 352 | name += f"_PCP{rank // (tp_size * pp_size)}" |
| 353 | return name |
| 354 | |
| 355 | |
| 356 | def get_bundles_for_indices( |
| 357 | placement_group: "PlacementGroup", |
| 358 | bundle_indices: list[int], |
| 359 | world_size: int, |
| 360 | ) -> list[tuple[int, str, str]]: |
| 361 | """Return GPU bundle indices paired with node IDs and node IPs for |
| 362 | explicit bundle indices specified via VLLM_RAY_BUNDLE_INDICES. |
| 363 | """ |
| 364 | assert len(bundle_indices) == world_size, ( |
| 365 | "VLLM_RAY_BUNDLE_INDICES must have the same size" |
| 366 | f" as the world size, but got {bundle_indices=} " |
| 367 | f"and {world_size=}" |
| 368 | ) |
| 369 | assert len(set(bundle_indices)) == len(bundle_indices), ( |
| 370 | "VLLM_RAY_BUNDLE_INDICES cannot have duplicate values," |
| 371 | f" but got {bundle_indices=}" |
| 372 | ) |
| 373 | |
| 374 | pg_data = placement_group_table(placement_group) |
| 375 | pg_bundle_to_node = pg_data["bundles_to_node_id"] |
| 376 | node_id_to_ip = { |
| 377 | n["NodeID"]: n["NodeManagerAddress"] for n in ray.nodes() if n["Alive"] |
| 378 | } |
| 379 | return [ |
| 380 | (bid, pg_bundle_to_node[bid], node_id_to_ip[pg_bundle_to_node[bid]]) |
| 381 | for bid in bundle_indices |
| 382 | ] |
| 383 | |
| 384 | |
| 385 | def get_bundles_sorted_by_node( |
| 386 | placement_group: "PlacementGroup", |
| 387 | ) -> list[tuple[int, str, str]]: |
| 388 | """Return GPU bundle indices paired with node IDs and node IPs, |
| 389 | sorted driver-first. |
| 390 | |
| 391 | This utility has to be invoked from the driver node. |
| 392 | |
| 393 | Example: 3-node cluster, driver on node-A, PG bundles spread |
| 394 | across nodes: |
| 395 | |
| 396 | Input: [ |
| 397 | (0, node-C), |
| 398 | (1, node-A), |
| 399 | (2, node-B), |
| 400 | (3, node-C), |
| 401 | (4, node-A), |
| 402 | (5, node-B), |
| 403 | ] |
| 404 | Output: [ |
| 405 | (1, node-A), |
| 406 | (4, node-A), |
| 407 | (2, node-B), |
| 408 | (5, node-B), |
| 409 | (0, node-C), |
| 410 | (3, node-C), |
| 411 | ] |
| 412 | """ |
| 413 | pg_data = placement_group_table(placement_group) |
| 414 | bundle_to_node = pg_data["bundles_to_node_id"] |
| 415 | |
| 416 | ray_device_key = current_platform.ray_device_key |
| 417 | if not ray_device_key: |
| 418 | raise ValueError( |
| 419 | f"current platform {current_platform.device_name} does not support ray." |
| 420 | ) |
| 421 | |
| 422 | node_id_to_ip = { |
| 423 | n["NodeID"]: n["NodeManagerAddress"] for n in ray.nodes() if n["Alive"] |
| 424 | } |
| 425 | |
| 426 | bundle_specs = placement_group.bundle_specs |
| 427 | assert bundle_specs is not None |
| 428 | bundle_to_node_id: list[tuple[int, str, str]] = [] |
| 429 | for bundle_idx, bundle in enumerate(bundle_specs): |
| 430 | if bundle.get(ray_device_key): |
| 431 | node_id = bundle_to_node.get(bundle_idx) |
| 432 | bundle_to_node_id.append((bundle_idx, node_id, node_id_to_ip[node_id])) |
| 433 | |
| 434 | driver_node = ray.get_runtime_context().get_node_id() |
| 435 | |
| 436 | def _sort_key(item): |
| 437 | _, node_id, _ = item |
| 438 | return (0 if node_id == driver_node else 1, node_id) |
| 439 | |
| 440 | bundle_to_node_id.sort(key=_sort_key) |
| 441 | |
| 442 | return bundle_to_node_id |
| 443 | |
| 444 | |
| 445 | def _wait_until_pg_ready(current_placement_group: "PlacementGroup"): |
| 446 | """Wait until a placement group is ready. |
| 447 | |
| 448 | It prints the informative log messages if the placement group is |
| 449 | not created within time. |
| 450 | |
| 451 | """ |
| 452 | # Wait until PG is ready - this will block until all |
| 453 | # requested resources are available, and will time out |
| 454 | # if they cannot be provisioned. |
| 455 | placement_group_specs = current_placement_group.bundle_specs |
| 456 | |
| 457 | s = time.time() |
| 458 | pg_ready_ref = current_placement_group.ready() |
| 459 | wait_interval = 10 |
| 460 | while time.time() - s < PG_WAIT_TIMEOUT: |
| 461 | ready, _ = ray.wait([pg_ready_ref], timeout=wait_interval) |
| 462 | if len(ready) > 0: |
| 463 | break |
| 464 | |
| 465 | # Exponential backoff for warning print. |
| 466 | wait_interval *= 2 |
| 467 | logger.info( |
| 468 | "Waiting for creating a placement group of specs for " |
| 469 | "%d seconds. specs=%s. Check `ray status` and " |
| 470 | "`ray list nodes` to see if you have enough resources," |
| 471 | " and make sure the IP addresses used by ray cluster" |
| 472 | " are the same as VLLM_HOST_IP environment variable" |
| 473 | " specified in each node if you are running on a multi-node.", |
| 474 | int(time.time() - s), |
| 475 | placement_group_specs, |
| 476 | ) |
| 477 | |
| 478 | try: |
| 479 | ray.get(pg_ready_ref, timeout=0) |
| 480 | except ray.exceptions.GetTimeoutError: |
| 481 | # Provide more helpful error message when GPU count is exceeded |
| 482 | total_gpu_required = sum(spec.get("GPU", 0) for spec in placement_group_specs) |
| 483 | # If more than one GPU is required for the placement group, provide a |
| 484 | # more specific error message. |
| 485 | # We use >1 here because multi-GPU (tensor parallel) jobs are more |
| 486 | # likely to fail due to insufficient cluster resources, and users may |
| 487 | # need to adjust tensor_parallel_size to fit available GPUs. |
| 488 | if total_gpu_required > 1: |
| 489 | raise ValueError( |
| 490 | f"Cannot provide a placement group requiring " |
| 491 | f"{total_gpu_required} GPUs " |
| 492 | f"(placement_group_specs={placement_group_specs}) within " |
| 493 | f"{PG_WAIT_TIMEOUT} seconds.\n" |
| 494 | f"Tensor parallel size may exceed available GPUs in your " |
| 495 | f"cluster. Check resources with `ray status` and " |
| 496 | f"`ray list nodes`.\n" |
| 497 | f"If running on K8s with limited GPUs, consider reducing " |
| 498 | f"--tensor-parallel-size to match available GPU resources." |
| 499 | ) from None |
| 500 | else: |
| 501 | raise ValueError( |
| 502 | "Cannot provide a placement group of " |
| 503 | f"{placement_group_specs=} within " |
| 504 | f"{PG_WAIT_TIMEOUT} seconds. See " |
| 505 | "`ray status` and `ray list nodes` to make sure the cluster " |
| 506 | "has enough resources." |
| 507 | ) from None |
| 508 | |
| 509 | |
| 510 | def initialize_ray_cluster( |
| 511 | parallel_config: ParallelConfig, |
| 512 | ray_address: str | None = None, |
| 513 | require_gpu_on_driver: bool = True, |
| 514 | ): |
| 515 | """Initialize the distributed cluster with Ray. |
| 516 | |
| 517 | it will connect to the Ray cluster and create a placement group |
| 518 | for the workers, which includes the specification of the resources |
| 519 | for each distributed worker. |
| 520 | |
| 521 | Args: |
| 522 | parallel_config: The configurations for parallel execution. |
| 523 | ray_address: The address of the Ray cluster. If None, uses |
| 524 | the default Ray cluster address. |
| 525 | require_gpu_on_driver: If True (default), require at least one GPU |
| 526 | on the current (driver) node and pin the first PG bundle to it. |
| 527 | Set to False for executors like RayExecutorV2 where all GPU work |
| 528 | is delegated to remote Ray actors. |
| 529 | |
| 530 | """ |
| 531 | assert_ray_available() |
| 532 | from vllm.platforms import current_platform |
| 533 | |
| 534 | # Disable Ray usage stats collection |
| 535 | if os.environ.get("RAY_USAGE_STATS_ENABLED", "0") != "1": |
| 536 | os.environ["RAY_USAGE_STATS_ENABLED"] = "0" |
| 537 | |
| 538 | # Prevalidate GPU requirements before Ray processing |
| 539 | if current_platform.is_cuda() and parallel_config.world_size > 1: |
| 540 | available_gpus = current_platform.device_count() |
| 541 | if parallel_config.world_size > available_gpus: |
| 542 | logger.warning( |
| 543 | "Tensor parallel size (%d) exceeds available GPUs (%d). " |
| 544 | "This may result in Ray placement group allocation failures. " |
| 545 | "Consider reducing tensor_parallel_size to %d or less, " |
| 546 | "or ensure your Ray cluster has %d GPUs available.", |
| 547 | parallel_config.world_size, |
| 548 | available_gpus, |
| 549 | available_gpus, |
| 550 | parallel_config.world_size, |
| 551 | ) |
| 552 | |
| 553 | if ray.is_initialized(): |
| 554 | logger.info("Ray is already initialized. Skipping Ray initialization.") |
| 555 | elif current_platform.is_rocm() or current_platform.is_xpu(): |
| 556 | # Try to connect existing ray instance and create a new one if not found |
| 557 | try: |
| 558 | ray.init("auto") |
| 559 | except ConnectionError: |
| 560 | logger.warning( |
| 561 | "No existing RAY instance detected. " |
| 562 | "A new instance will be launched with current node resources." |
| 563 | ) |
| 564 | ray.init( |
| 565 | address=ray_address, |
| 566 | num_gpus=current_platform.device_count(), |
| 567 | runtime_env=parallel_config.ray_runtime_env, |
| 568 | ) |
| 569 | else: |
| 570 | ray.init(address=ray_address, runtime_env=parallel_config.ray_runtime_env) |
| 571 | |
| 572 | device_str = current_platform.ray_device_key |
| 573 | if not device_str: |
| 574 | raise ValueError( |
| 575 | f"current platform {current_platform.device_name} does not support ray." |
| 576 | ) |
| 577 | |
| 578 | # Create or get the placement group for worker processes |
| 579 | if parallel_config.placement_group: |
| 580 | current_placement_group = parallel_config.placement_group |
| 581 | else: |
| 582 | current_placement_group = ray.util.get_current_placement_group() |
| 583 | |
| 584 | if current_placement_group: |
| 585 | logger.info("Using the existing placement group") |
| 586 | |
| 587 | # We are in a placement group |
| 588 | bundles = current_placement_group.bundle_specs |
| 589 | # Verify that we can use the placement group. |
| 590 | device_bundles = 0 |
| 591 | for bundle in bundles: |
| 592 | bundle_devices = bundle.get(device_str, 0) |
| 593 | if bundle_devices > 1: |
| 594 | raise ValueError( |
| 595 | f"Placement group bundle cannot have more than 1 {device_str}." |
| 596 | ) |
| 597 | if bundle_devices: |
| 598 | device_bundles += 1 |
| 599 | if parallel_config.world_size > device_bundles: |
| 600 | raise ValueError( |
| 601 | f"The number of required {device_str}s exceeds the total " |
| 602 | f"number of available {device_str}s in the placement group. " |
| 603 | f"Required number of devices: {parallel_config.world_size}. " |
| 604 | f"Total number of devices: {device_bundles}." |
| 605 | ) |
| 606 | else: |
| 607 | logger.info("No current placement group found. Creating a new placement group.") |
| 608 | num_devices_in_cluster = ray.cluster_resources().get(device_str, 0) |
| 609 | # Log a warning message and delay resource allocation failure response. |
| 610 | # Avoid immediate rejection to allow user-initiated placement group |
| 611 | # created and wait cluster to be ready |
| 612 | if parallel_config.world_size > num_devices_in_cluster: |
| 613 | logger.warning( |
| 614 | "The number of required %ss exceeds the total " |
| 615 | "number of available %ss in the placement group.", |
| 616 | device_str, |
| 617 | device_str, |
| 618 | ) |
| 619 | # Create a new placement group |
| 620 | placement_group_specs: list[dict[str, float]] = [ |
| 621 | {device_str: 1.0} for _ in range(parallel_config.world_size) |
| 622 | ] |
| 623 | |
| 624 | # vLLM engine is also a worker to execute model with an accelerator, |
| 625 | # so it requires to have the device in a current node. Check if |
| 626 | # the current node has at least one device. |
| 627 | current_ip = get_ip() |
| 628 | current_node_id = ray.get_runtime_context().get_node_id() |
| 629 | current_node_resource = available_resources_per_node()[current_node_id] |
| 630 | # TODO (jeffreywang): require_gpu_on_driver should be always False |
| 631 | # after deprecating RayDistributedExecutor. |
| 632 | if require_gpu_on_driver: |
| 633 | if current_node_resource.get(device_str, 0) < 1: |
| 634 | raise ValueError( |
| 635 | f"Current node has no {device_str} available. " |
| 636 | f"{current_node_resource=}. vLLM engine cannot start " |
| 637 | f"without {device_str}. Make sure you have at least 1 " |
| 638 | f"{device_str} available in a node " |
| 639 | f"{current_node_id=} {current_ip=}." |
| 640 | ) |
| 641 | # This way, at least bundle is required to be created in a |
| 642 | # current node. |
| 643 | placement_group_specs[0][f"node:{current_ip}"] = 0.001 |
| 644 | |
| 645 | # By default, Ray packs resources as much as possible. |
| 646 | current_placement_group = ray.util.placement_group( |
| 647 | placement_group_specs, strategy="PACK" |
| 648 | ) |
| 649 | _wait_until_pg_ready(current_placement_group) |
| 650 | |
| 651 | assert current_placement_group is not None |
| 652 | _verify_bundles( |
| 653 | current_placement_group, parallel_config, device_str, require_gpu_on_driver |
| 654 | ) |
| 655 | # Set the placement group in the parallel config |
| 656 | parallel_config.placement_group = current_placement_group |