iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Core] Remove the legacy Ray compiled-graph executor (#61712)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Woosuk Kwon
2026-10-14 18:13:27 -07:00
committed by GitHub
co-authored by Claude Opus 5.5
parent bfdbfa1632
commit e57bb4ace7
16 changed files with 14 additions and 1257 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/hardware_tests/ascend_npu.yaml
-1
View File
@@ -56,7 +56,6 @@ steps:
- "!vllm/distributed/device_communicators/shm_object_storage.py"
- "!vllm/distributed/device_communicators/symm_mem.py"
- "!vllm/distributed/device_communicators/mnnvl_compat.py"
- "!vllm/distributed/device_communicators/ray_communicator.py"
- "!vllm/third_party/flashmla/"
- "!vllm/third_party/pynvml.py"
- "!vllm/_custom_ops.py"
.buildkite/test-amd.yaml
-1
View File
@@ -940,7 +940,6 @@ steps:
- tests/basic_correctness/models/test_basic_correctness.py
- vllm/platforms/rocm.py
commands:
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
- pytest -v -s distributed/test_ray_v2_executor.py
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
.buildkite/test_areas/distributed.yaml
-2
View File
@@ -697,7 +697,6 @@ steps:
- tests/distributed/test_pipeline_parallel.py
- tests/basic_correctness/models/test_basic_correctness.py
commands:
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
- export NCCL_CUMEM_HOST_ENABLE=0
- pytest -v -s distributed/test_ray_v2_executor.py
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
@@ -722,7 +721,6 @@ steps:
- tests/basic_correctness/models/test_basic_correctness.py
- vllm/platforms/rocm.py
commands:
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
- pytest -v -s distributed/test_ray_v2_executor.py
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
docs/serving/parallelism_scaling.md
+1 -1
View File
@@ -71,7 +71,7 @@ Ray also offers high-level APIs for large-scale [offline batch inference](https:
Ray is an optional dependency. Install it explicitly before using Ray-based execution, for example:
```bash
pip install "ray[cgraph]"
pip install ray
```
For details, see the [Ray documentation](https://docs.ray.io/en/latest/index.html).
examples/rl/rdt_vllm_serve.py
-1
View File
@@ -76,7 +76,6 @@ def launch_vllm_serve(
env = dict(
os.environ,
VLLM_SERVER_DEV_MODE="1",
VLLM_USE_RAY_V2_EXECUTOR_BACKEND="1",
)
# Pin the child to THIS Ray cluster. vLLM's ray DP backend reaches Ray
# through a bare ``ray.init()``, which STARTS a cluster rather than
tests/distributed/conftest.py
+1 -5
View File
@@ -182,19 +182,15 @@ class MockSubscriber:
@pytest.fixture
def enable_ray_v2_backend():
"""Set env vars for the Ray V2 executor backend and shut down Ray
"""Run the engine in-process for Ray executor tests and shut down Ray
between tests."""
import ray
saved = {
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND": os.environ.get(
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND"
),
"VLLM_ENABLE_V1_MULTIPROCESSING": os.environ.get(
"VLLM_ENABLE_V1_MULTIPROCESSING"
),
}
os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1"
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
if ray.is_initialized():
ray.shutdown()
tests/distributed/test_pipeline_parallel.py
+1 -10
View File
@@ -294,16 +294,7 @@ def _compare_tp(
if max_num_seqs:
common_args.extend(["--max-num-seqs", f"{max_num_seqs}"])
if distributed_backend == "ray":
# Test Ray Compiled Graph for all the tests
pp_env = {
"VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL": "1",
}
elif distributed_backend == "mp":
pp_env = None
else:
pp_env = None
pp_env = None
tp_env = None
pp_args = [
tests/distributed/test_ray_v2_executor_e2e.py
-1
View File
@@ -38,7 +38,6 @@ def ray_init():
class _AsyncLLMActor:
def start(self, pg, bundle_indices=None, ray_runtime_env=None):
os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1"
# Needed so collective_rpc can pickle _get_env_var over the
# AsyncLLM -> EngineCore ZMQ boundary.
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
tests/v1/executor/test_ray_utils.py
-79
View File
@@ -1,79 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import numpy as np
from vllm.distributed.aux_output_connector.connector import (
AuxRequestOutput,
)
from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output
from vllm.v1.outputs import (
LogprobsLists,
LogprobsTensors,
ModelRunnerOutput,
)
def _make_readonly(arr: np.ndarray) -> np.ndarray:
arr.setflags(write=False)
return arr
def test_detach_zero_copy_from_model_runner_output_copies_only_numpy_views():
cu_num_generated_tokens = [0, 2]
prompt_logprobs = LogprobsTensors.empty_cpu(1, 2)
output = ModelRunnerOutput(
req_ids=["req-0"],
req_id_to_index={"req-0": 0},
logprobs=LogprobsLists(
logprob_token_ids=_make_readonly(
np.array([[1, 2], [3, 4]], dtype=np.int32)
),
logprobs=_make_readonly(
np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
),
sampled_token_ranks=_make_readonly(np.array([1, 2], dtype=np.int32)),
cu_num_generated_tokens=cu_num_generated_tokens,
),
prompt_logprobs_dict={"req-0": prompt_logprobs},
)
original_logprobs = output.logprobs
assert original_logprobs is not None
detach_zero_copy_from_model_runner_output(output)
detached_logprobs = output.logprobs
assert detached_logprobs is not None
assert detached_logprobs is not original_logprobs
assert (
detached_logprobs.logprob_token_ids is not original_logprobs.logprob_token_ids
)
assert detached_logprobs.logprobs is not original_logprobs.logprobs
assert (
detached_logprobs.sampled_token_ranks
is not original_logprobs.sampled_token_ranks
)
assert detached_logprobs.logprob_token_ids.flags.writeable
assert detached_logprobs.logprobs.flags.writeable
assert detached_logprobs.sampled_token_ranks.flags.writeable
assert detached_logprobs.cu_num_generated_tokens is cu_num_generated_tokens
assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs
def test_detach_zero_copy_aux_output_without_logprobs():
rows = _make_readonly(np.arange(12, dtype=np.uint8).reshape(2, 3, 2))
output = ModelRunnerOutput(
req_ids=["req-0"],
req_id_to_index={"req-0": 0},
aux_output_connector_output={"req-0": AuxRequestOutput(0, rows)},
)
detach_zero_copy_from_model_runner_output(output)
aux_output = output.aux_output_connector_output
assert aux_output is not None
detached = aux_output["req-0"].rows
assert detached is not rows
assert detached.flags.writeable
np.testing.assert_array_equal(detached, rows)
tools/pre_commit/check_forbidden_imports.py
-1
View File
@@ -73,7 +73,6 @@ CHECK_IMPORTS = {
"benchmarks/fused_kernels/layernorm_rms_benchmarks.py",
# cloudpickle
"vllm/v1/executor/multiproc_executor.py",
"vllm/v1/executor/ray_executor.py",
"vllm/entrypoints/llm.py",
"tests/utils.py",
# pickle and cloudpickle
vllm/distributed/device_communicators/ray_communicator.py
-251
View File
@@ -1,251 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import uuid
from typing import Any
import ray
import torch
from ray.exceptions import RayChannelError
from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
from torch.distributed import ReduceOp
from vllm.distributed.device_communicators.base_device_communicator import (
DeviceCommunicatorBase,
)
from vllm.distributed.parallel_state import get_pp_group
from vllm.logger import init_logger
from vllm.utils.torch_utils import current_stream
logger = init_logger(__name__)
class RayPPCommunicator(Communicator):
"""Communicator to be used for pipeline parallelism in Ray Compiled Graph.
This is wraps around the vLLM _PP GroupCoordinator.
This class is not thread-safe.
"""
_comm: DeviceCommunicatorBase | None
def __init__(
self,
world_size: int,
comm_id: Any,
rank: int | None,
actor_handles: list["ray.actor.ActorHandle"],
cuda_stream: torch.cuda.Stream | None,
use_communication_streams: bool = False,
):
"""Initialize a RayPPCommunicator that can be used to communicate with
other Ray Compiled Graph actors for pipeline parallelism.
Args:
world_size: The number of participating actors.
comm_id: A unique communicator ID. This is just to conform with
the Ray Communicator API and is not used.
rank: The rank of this actor. If None, then the caller is not a
participant of the RayPPCommunicator group (e.g., the Ray
driver).
actor_handles: A list of actor handles.
cuda_stream: A CUDA stream to dispatch communication ops to. This
is not supported.
use_communication_streams: Whether to use communication streams.
This is not supported.
"""
self._world_size = world_size
self._rank: int | None = None
self._actor_handles = actor_handles
if use_communication_streams:
raise NotImplementedError("use_communication_streams is not supported")
if cuda_stream is not None and cuda_stream != current_stream():
raise ValueError(
"cuda_stream other than the current stream is not supported"
)
if rank is not None:
# Rank is not None, this is Ray worker
assert ray.get_gpu_ids(), "RayPPCommunicator has no GPUs assigned"
self._comm = get_pp_group().device_communicator
assert self._comm is not None
# Since we wrap around the vLLM _PP communicator, we use
# the rank from the vLLM communicator, and ignore the rank
# passed in from Ray.
# TODO(rui): refactor the Ray Communicator API so that
# it also supports no rank passed in.
self._rank = self._comm.rank_in_group
self._build_actor_rank_mapping()
else:
# Rank is None, this is Ray driver
self._comm = None
self._closed = False
def _build_actor_rank_mapping(self):
"""Use collective communication to build a mapping from actor IDs to ranks.
This should be called once during initialization.
"""
if self._comm is None:
return {}
current_actor = ray.get_runtime_context().current_actor
actor_id_str = current_actor._actor_id.hex()
# Ray actor IDs are 32-character hex strings (128 bits)
ACTOR_ID_LEN = 32
actor_id_bytes = bytearray(actor_id_str.encode("utf-8"))
assert len(actor_id_bytes) == ACTOR_ID_LEN, (
f"Unexpected actor ID length: {len(actor_id_bytes)}"
)
actor_id_tensor = torch.frombuffer(actor_id_bytes, dtype=torch.uint8).to(
self._comm.device
)
# All-gather full actor IDs from all actors
gathered_ids = self._comm.all_gather(actor_id_tensor, dim=0)
# Build mapping: actor_id -> device_comm_rank
self._actor_id_to_rank = {}
for rank in range(self._world_size):
start_idx = rank * ACTOR_ID_LEN
end_idx = (rank + 1) * ACTOR_ID_LEN
actor_bytes = gathered_ids[start_idx:end_idx].cpu().numpy().tobytes()
actor_id = actor_bytes.decode("utf-8")
self._actor_id_to_rank[actor_id] = rank
def initialize(self, rank: int) -> None:
# No additional initialization is needed.
pass
def get_actor_handles(self) -> list["ray.actor.ActorHandle"]:
return self._actor_handles
def get_rank(self, actor: ray.actor.ActorHandle) -> int:
"""Return the given actor's rank using device communicator collective ops."""
assert hasattr(self, "_actor_id_to_rank"), (
"Actor rank mapping not built. "
"This should have been done during initialization."
)
actor_id_str = actor._actor_id.hex()
if actor_id_str in self._actor_id_to_rank:
return self._actor_id_to_rank[actor_id_str] # type: ignore
else:
raise ValueError(f"Actor {actor} not found in communicator group")
def get_self_rank(self) -> int | None:
"""Return this actor's rank."""
return self._rank
def get_world_size(self) -> int:
"""Return the number of ranks in the RayPPCommunicator group."""
return self._world_size
def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
"""Send a torch.Tensor to a peer.
This returns when the send kernel has been queued, but the kernel may
not have completed. Therefore, the caller should ensure that there are
no concurrent writes to the sent `buf` until the send has finished.
That is, either all writes should be submitted on the current stream
(self._cuda_stream) or, if on a different stream, that stream should
synchronize with the current stream.
Args:
buf: The torch.Tensor to send. It should already be on this
actor's default device.
peer_rank: The rank of the actor to send to.
"""
if self._closed:
raise RayChannelError("RayPPCommunicator has been destroyed.")
assert self._comm is not None
self._comm.send(buf, peer_rank)
def recv(
self,
shape: tuple[int, ...],
dtype: "torch.dtype",
peer_rank: int,
allocator: TorchTensorAllocator,
) -> "torch.Tensor":
"""Receive a torch.Tensor from a peer and synchronize the current stream.
After this call returns, the receive buffer is safe to read from
any stream. An RayChannelError will be raised if an error occurred
(e.g., remote actor died), and the buffer is not safe to read.
Args:
shape: The shape of the tensor to receive.
dtype: The dtype of the tensor to receive.
peer_rank: The rank of the actor to receive from.
allocator: The allocator to use to create the received tensor.
This is ignored for this implementation.
"""
if self._closed:
raise RayChannelError("RayPPCommunicator has been destroyed.")
assert self._comm is not None
size = torch.Size(shape)
buf = self._comm.recv(size, dtype, src=peer_rank)
# Buffer values are undefined if NCCL ops are aborted. Therefore, we
# need to synchronize here and check that the channel is still
# open to ensure that the receive buffer is valid.
# TODO(swang): Avoid CUDA synchronization.
current_stream().synchronize()
if self._closed:
raise RayChannelError("RayPPCommunicator has been destroyed.")
return buf
def allgather(
self,
send_buf: "torch.Tensor",
recv_buf: "torch.Tensor",
):
raise NotImplementedError("allgather is not supported")
def allreduce(
self,
send_buf: "torch.Tensor",
recv_buf: "torch.Tensor",
op: ReduceOp = ReduceOp.SUM,
):
raise NotImplementedError("allreduce is not supported")
def reducescatter(
self,
send_buf: "torch.Tensor",
recv_buf: "torch.Tensor",
op: ReduceOp = ReduceOp.SUM,
):
raise NotImplementedError("reducescatter is not supported")
@property
def recv_stream(self):
return torch.cuda.StreamContext(current_stream())
@property
def send_stream(self):
return torch.cuda.StreamContext(current_stream())
def destroy(self) -> None:
# Just sets a flag, vLLM manages the lifecycle of the underlying
# _PP GroupCoordinator.
self._closed = True
def get_transport_name(self) -> str:
return "nccl"
@classmethod
def generate_communicator_id(cls) -> Any:
return uuid.uuid4()
vllm/distributed/weight_transfer/sharded_rdt_engine.py
+2 -3
View File
@@ -316,9 +316,8 @@ class ShardedRDTWeightTransferEngine(
raise ValueError(
f"The 'sharded_rdt' weight transfer backend requires "
f"distributed_executor_backend='ray', got {backend!r}. Pass "
f"--distributed-executor-backend ray (and keep "
f"VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1, the default, so the "
f"workers are RayExecutorV2 actors with tensor transport)."
f"--distributed-executor-backend ray so the workers are "
f"RayExecutorV2 actors with tensor transport."
)
# Ownership table, routing rules, and (once bound) every producer's
# handle — pulls go to one owner per chunk, but free_group fans out to
vllm/envs.py
+7 -30
View File
@@ -60,11 +60,7 @@ if TYPE_CHECKING:
VLLM_XLA_CHECK_RECOMPILATION: bool = False
VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: int = 512
VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN: int = 8192
VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
VLLM_XPU_PP_MICROBATCH: bool = True
VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = False
VLLM_DISTRIBUTED_USE_SPLIT_GROUP: bool = False
VLLM_XLA_USE_SPMD: bool = False
VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork"
@@ -174,6 +170,7 @@ if TYPE_CHECKING:
VLLM_SERVER_DEV_MODE: bool = False
VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128
VLLM_MLA_DISABLE: bool = False
VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = True
VLLM_RAY_PER_WORKER_GPUS: float = 1.0
VLLM_RAY_BUNDLE_INDICES: str = ""
VLLM_CUDART_SO_PATH: str | None = None
@@ -914,37 +911,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
# (CPU backend only) whether to use SGLang INT4 W4A8 kernels for AWQ, and
# on Zen CPUs whether to serve int4 checkpoints as DA8W4 rather than W4A16.
"VLLM_CPU_INT4_W4A8": lambda: bool(int(os.getenv("VLLM_CPU_INT4_W4A8", "1"))),
# If the env var is set, Ray Compiled Graph uses the specified
# channel type to communicate between workers belonging to
# different pipeline-parallel stages.
# Available options:
# - "auto": use the default channel type
# - "nccl": use NCCL for communication
# - "shm": use shared memory and gRPC for communication
"VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE": env_with_choices(
"VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "auto", ["auto", "nccl", "shm"]
),
# If the env var is set, it enables GPU communication overlap
# (experimental feature) in Ray's Compiled Graph.
"VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM": lambda: bool(
int(os.getenv("VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM", "0"))
),
# If the env var is set, it uses a Ray Communicator wrapping
# vLLM's pipeline parallelism communicator to interact with Ray's
# Compiled Graph. Otherwise, it uses Ray's NCCL communicator.
"VLLM_USE_RAY_WRAPPED_PP_COMM": lambda: bool(
int(os.getenv("VLLM_USE_RAY_WRAPPED_PP_COMM", "1"))
),
# Using a flag to control microbatch on XPU device, users on XPU device can
# decide to disable it when they needed.
"VLLM_XPU_PP_MICROBATCH": lambda: bool(
int(os.getenv("VLLM_XPU_PP_MICROBATCH", "1"))
),
# When True and distributed_executor_backend="ray", use RayExecutorV2
# (MQ-based) instead of RayDistributedExecutor (compiled-graph backend).
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool(
int(os.getenv("VLLM_USE_RAY_V2_EXECUTOR_BACKEND", "1"))
),
# When True, GroupCoordinator constructs its CPU/device subgroups via
# ``torch.distributed.split_group(backend=...)``
# and ``init_distributed_environment`` initializes the default PG with
@@ -1513,6 +1484,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
# which indices are used for the Ray bundle, for every worker.
# Format: comma-separated list of integers, e.g. "0,1,2,3"
"VLLM_RAY_BUNDLE_INDICES": lambda: os.getenv("VLLM_RAY_BUNDLE_INDICES", ""),
# Deprecated and ignored: RayExecutorV2 is the only Ray executor. Kept so
# out-of-tree platforms that still read it keep working; will be removed
# in v0.34.
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool(
int(os.getenv("VLLM_USE_RAY_V2_EXECUTOR_BACKEND", "1"))
),
# In some system, find_loaded_library() may not work. So we allow users to
# specify the path through environment variable VLLM_CUDART_SO_PATH.
"VLLM_CUDART_SO_PATH": lambda: os.getenv("VLLM_CUDART_SO_PATH", None),
vllm/v1/executor/abstract.py
+2 -8
View File
@@ -8,7 +8,6 @@ from concurrent.futures import Future
from functools import cached_property
from typing import TYPE_CHECKING, Literal, TypeVar, overload
import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
@@ -75,14 +74,9 @@ class Executor(ABC):
)
executor_class = distributed_executor_backend
elif distributed_executor_backend == "ray":
if envs.VLLM_USE_RAY_V2_EXECUTOR_BACKEND:
from vllm.v1.executor.ray_executor_v2 import RayExecutorV2
from vllm.v1.executor.ray_executor_v2 import RayExecutorV2
executor_class = RayExecutorV2
else:
from vllm.v1.executor.ray_executor import RayDistributedExecutor
executor_class = RayDistributedExecutor
executor_class = RayExecutorV2
elif distributed_executor_backend == "mp":
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
vllm/v1/executor/ray_executor.py
-644
View File
@@ -1,644 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from collections import defaultdict
from collections.abc import Callable
from concurrent.futures import Future
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import cloudpickle
import vllm.envs as envs
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.ray.ray_env import get_env_vars_to_copy
from vllm.utils.network_utils import (
get_distributed_init_method,
get_ip,
get_open_port,
)
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
from vllm.v1.executor.abstract import Executor
from vllm.v1.executor.ray_env_utils import (
update_runtime_env_for_worker_import,
)
from vllm.v1.executor.ray_utils import (
WORKER_SPECIFIC_ENV_VARS,
FutureWrapper,
RayWorkerWrapper,
detach_zero_copy_from_model_runner_output,
initialize_ray_cluster,
ray,
)
from vllm.v1.outputs import ModelRunnerOutput
if ray is not None:
from ray.actor import ActorHandle
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
else:
ActorHandle = None
if TYPE_CHECKING:
from ray.util.placement_group import PlacementGroup
logger = init_logger(__name__)
COMPLETED_NONE_FUTURE: Future[ModelRunnerOutput | None] = Future()
COMPLETED_NONE_FUTURE.set_result(None)
@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.
"""
worker: ActorHandle
created_rank: int
adjusted_rank: int = -1
ip: str = ""
class RayDistributedExecutor(Executor):
"""Ray-based distributed executor."""
uses_ray: bool = True
supports_pp: bool = True
def _init_executor(self) -> None:
self.forward_dag: ray.dag.CompiledDAG | None = None
# For TPU or XPU, avoid compiling NVIDIA's NCCL
if current_platform.is_tpu() or current_platform.is_xpu():
os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"
assert self.uses_ray
initialize_ray_cluster(self.parallel_config)
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)
# KV connector setup
self.has_connector = self.vllm_config.kv_transfer_config is not None
if (
self.vllm_config.ec_transfer_config is not None
and self.parallel_config.world_size > 1
):
raise NotImplementedError(
"EC connector worker metadata is not supported with the "
"legacy Ray executor when world_size > 1: only the output "
"of a single worker is fetched, silently dropping the "
"other workers' EC connector state. Set "
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 to use RayExecutorV2, "
"or use the multiprocessing executor instead."
)
self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
self.vllm_config.ec_transfer_config is None
or self.vllm_config.ec_transfer_config.is_ec_consumer
)
self.scheduler_output: SchedulerOutput | None = None
def shutdown(self) -> None:
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 "
"because this is the expected termination process in Ray."
)
if hasattr(self, "forward_dag") and self.forward_dag is not None:
self.forward_dag.teardown()
import ray
for worker in self.workers:
ray.kill(worker)
self.forward_dag = None
def _configure_ray_workers_use_nsight(self, ray_remote_kwargs) -> dict[str, Any]:
# 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", {})
runtime_env.update(
{
"nsight": {
"t": "cuda,cudnn,cublas",
"o": "'worker_process_%p'",
"cuda-graph-trace": "node",
}
}
)
return ray_remote_kwargs
def _update_noset_device_env_vars(self, ray_remote_kwargs):
runtime_env = ray_remote_kwargs.setdefault("runtime_env", {})
env_vars = runtime_env.setdefault("env_vars", {})
env_vars.update(
{env_var: "1" for env_var in current_platform.ray_noset_device_env_vars}
)
update_runtime_env_for_worker_import(runtime_env)
return ray_remote_kwargs
# 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
def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs):
num_gpus = envs.VLLM_RAY_PER_WORKER_GPUS
# The driver dummy worker does not actually use any resources.
# It holds the resource for the driver worker.
self.driver_dummy_worker: RayWorkerWrapper | None = None
# The remaining workers are the actual ray actors.
self.workers: list[RayWorkerWrapper] = []
# 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.
self.pp_tp_workers: list[list[RayWorkerWrapper]] = []
if self.parallel_config.ray_workers_use_nsight:
ray_remote_kwargs = self._configure_ray_workers_use_nsight(
ray_remote_kwargs
)
# The way ray actors are setup in vllm is that the visible devices are
# not set by actors, they are left unset by ray. Internally we index
# the right gpu with local_rank. This is similar to how mp mode works.
self._update_noset_device_env_vars(ray_remote_kwargs)
# Create the workers.
bundle_indices: list[int]
if envs.VLLM_RAY_BUNDLE_INDICES:
# Use the bundle indices specified by the user.
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=}"
)
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)
bundle_indices = bundle_indices[: self.parallel_config.world_size]
worker_metadata: list[RayWorkerMetaData] = []
driver_ip = get_ip()
for rank, bundle_id in enumerate(bundle_indices):
scheduling_strategy = PlacementGroupSchedulingStrategy(
placement_group=placement_group,
placement_group_capture_child_tasks=True,
placement_group_bundle_index=bundle_id,
)
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,
)(RayWorkerWrapper).remote(rpc_rank=rank)
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,
)(RayWorkerWrapper).remote(rpc_rank=rank)
worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank))
worker_ips = ray.get(
[
each.worker.get_node_ip.remote() # type: ignore[attr-defined]
for each in worker_metadata
]
)
for each, ip in zip(worker_metadata, worker_ips):
each.ip = ip
logger.debug("workers: %s", worker_metadata)
logger.debug("driver_dummy_worker: %s", self.driver_dummy_worker)
ip_counts: dict[str, int] = {}
for ip in worker_ips:
ip_counts[ip] = ip_counts.get(ip, 0) + 1
def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
"""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.
"""
ip = item.ip
return 0 if ip == driver_ip else 1, ip_counts[ip], ip
# 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.
sorted_worker_metadata = sorted(
worker_metadata, key=sort_by_driver_then_worker_ip
)
for i, item in enumerate(sorted_worker_metadata):
item.adjusted_rank = i
self.workers = [item.worker for item in sorted_worker_metadata]
rerank_mapping = {
item.created_rank: item.adjusted_rank for item in sorted_worker_metadata
}
self.collective_rpc("adjust_rank", args=(rerank_mapping,))
# Get the set of physical GPU IDs used on each node.
worker_node_and_physical_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_physical_gpu_ids.append(
ray.get(worker.get_node_and_physical_gpu_ids.remote()) # type: ignore[attr-defined]
)
node_workers = defaultdict(list) # node id -> list of worker ranks
node_physical_gpu_ids = defaultdict(list) # node id -> physical GPU IDs
for i, (node_id, physical_gpu_ids) in enumerate(
worker_node_and_physical_gpu_ids
):
node_workers[node_id].append(i)
# `physical_gpu_ids` can be a list of strings or integers.
# convert them to integers for consistency.
# NOTE: physical 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
physical_gpu_ids = [
current_platform.device_control_id_to_physical_device_id(str(x))
for x in physical_gpu_ids
]
node_physical_gpu_ids[node_id].extend(physical_gpu_ids)
for node_id, physical_gpu_ids in node_physical_gpu_ids.items():
node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)
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"
" network configuration. If you set `VLLM_HOST_IP`"
" environment variable, make sure it is unique for"
" each node."
)
all_args_to_update_environment_variables: list[dict[str, str]] = [
{} for _ in worker_node_and_physical_gpu_ids
]
# Environment variables to copy from driver to workers
env_vars_to_copy = get_env_vars_to_copy(
exclude_vars=WORKER_SPECIFIC_ENV_VARS,
additional_vars=set(current_platform.additional_env_vars),
destination="workers",
)
# Copy existing env vars to each worker's args
for args in all_args_to_update_environment_variables:
# TODO: refactor platform-specific env vars
for name in env_vars_to_copy:
if name in os.environ:
args[name] = os.environ[name]
self._env_vars_for_all_workers = all_args_to_update_environment_variables
self.collective_rpc(
"update_environment_variables", args=(self._get_env_vars_to_be_updated(),)
)
if len(node_physical_gpu_ids) == 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"
distributed_init_method = get_distributed_init_method(
driver_ip, get_open_port()
)
# Initialize the actual workers inside worker wrapper.
all_kwargs = []
for rank, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):
local_rank = node_workers[node_id].index(rank)
kwargs = dict(
vllm_config=self.vllm_config,
assigned_physical_gpu_ids=sorted(node_physical_gpu_ids[node_id]),
local_rank=local_rank,
rank=rank,
distributed_init_method=distributed_init_method,
is_driver_worker=(not self.parallel_config)
or (rank % self.parallel_config.tensor_parallel_size == 0),
)
all_kwargs.append(kwargs)
self.collective_rpc("init_worker", args=(all_kwargs,))
self.collective_rpc("init_device")
if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
self.collective_rpc("elastic_ep_execute", args=("load_model",))
else:
self.collective_rpc("load_model")
def _update_block_size(worker):
current_platform.update_block_size_for_backend(worker.vllm_config)
self.collective_rpc(_update_block_size)
for pp_rank in range(self.parallel_config.pipeline_parallel_size):
self.pp_tp_workers.append([])
for tp_rank in range(self.parallel_config.tensor_parallel_size):
# PP=2, TP=4
# pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
rank = (pp_rank * self.parallel_config.tensor_parallel_size) + tp_rank
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])
def reinitialize_distributed(
self, reconfig_request: ReconfigureDistributedRequest
) -> None:
self.collective_rpc("reinitialize_distributed", args=(reconfig_request,))
if (
reconfig_request.new_data_parallel_rank
== ReconfigureRankType.SHUTDOWN_CURRENT_RANK
):
self.shutdown()
def execute_model( # type: ignore[override]
self,
scheduler_output: SchedulerOutput,
non_block: bool = False,
) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
if self.scheduler_output is not None:
raise RuntimeError(
"State error: sample_tokens() must be called "
"after execute_model() returns None."
)
if not self.uses_sampler or not scheduler_output.total_num_scheduled_tokens:
# Model will not execute, call model runner immediately.
return self._execute_dag(scheduler_output, None, non_block)
# Model will execute, defer to sample_tokens() call.
self.scheduler_output = scheduler_output
return COMPLETED_NONE_FUTURE if non_block else None
def sample_tokens( # type: ignore[override]
self,
grammar_output: "GrammarOutput | None",
non_block: bool = False,
) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
"""Execute the model on the Ray workers.
The scheduler output to use should have been provided in
a prior call to execute_model().
Args:
grammar_output: The structured outputs grammar bitmask, if applicable.
non_block: If True, the method will return a Future.
Returns:
The model runner output.
"""
scheduler_output = self.scheduler_output
if scheduler_output is None:
return COMPLETED_NONE_FUTURE if non_block else None
self.scheduler_output = None
return self._execute_dag(scheduler_output, grammar_output, non_block)
def _execute_dag(
self,
scheduler_output: SchedulerOutput,
grammar_output: "GrammarOutput | None",
non_block: bool = False,
) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
# Build the compiled DAG for the first time.
if self.forward_dag is None: # type: ignore
self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)
refs = self.forward_dag.execute((scheduler_output, grammar_output)) # type: ignore
if not self.has_connector:
# Get output only from a single worker (output_rank)
# When PP is not used, we block here until the result is available.
if not non_block:
output = refs[0].get()
detach_zero_copy_from_model_runner_output(output)
return output
# When PP is used, we return a FutureWrapper immediately so that
# the scheduler can yield to the next batch.
return FutureWrapper(refs[0])
# Get output from all workers when connector is present
assert self.kv_output_aggregator is not None
if not non_block:
# Block and get results from all workers
outputs = ray.get(refs)
for output in outputs:
detach_zero_copy_from_model_runner_output(output)
return self.kv_output_aggregator.aggregate(outputs)
# Return a future that will aggregate outputs from all workers
return FutureWrapper(refs, self.kv_output_aggregator)
def collective_rpc( # type: ignore[override]
self,
method: str | Callable,
timeout: float | None = None,
args: tuple = (),
kwargs: dict[str, Any] | None = None,
non_block: bool = False,
) -> list[Any] | Future[list[Any]]:
"""Runs the given method on all workers."""
sent_method = method if isinstance(method, str) else cloudpickle.dumps(method)
del method
if kwargs is None:
kwargs = {}
ray_worker_outputs = [
worker.execute_method.remote( # type: ignore[attr-defined]
sent_method, *args, **kwargs
)
for worker in self.workers
]
# Get the results of the ray workers.
if non_block:
return FutureWrapper(ray_worker_outputs)
return ray.get(ray_worker_outputs, timeout=timeout)
def _check_ray_cgraph_installation(self):
import importlib.metadata
from packaging import version
required_version = version.parse("2.43.0")
current_version = version.parse(importlib.metadata.version("ray"))
if current_version < required_version:
raise ValueError(
f"Ray version {required_version} is "
f"required, but found {current_version}"
)
import importlib.util
cgraph_spec = importlib.util.find_spec("ray.experimental.compiled_dag_ref")
if cgraph_spec is None:
raise ValueError(
"Ray Compiled Graph is not installed. "
"Run `pip install ray[cgraph]` to install it."
)
cupy_spec = importlib.util.find_spec("cupy")
if cupy_spec is None and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl":
raise ValueError(
"cupy is not installed but required since "
"VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE is set to 'nccl'. "
"Run `pip install ray[cgraph]` and check cupy installation."
)
def _compiled_ray_dag(self, enable_asyncio: bool):
assert self.parallel_config.use_ray
self._check_ray_cgraph_installation()
# 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
from ray.dag import InputNode, MultiOutputNode
logger.info(
"RAY_CGRAPH_get_timeout is set to %s",
os.environ["RAY_CGRAPH_get_timeout"], # noqa: SIM112
)
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,
)
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: "
f"{channel_type}. Valid values are: 'auto', 'nccl', or 'shm'."
)
with InputNode() as input_data:
# Example DAG: PP=2, TP=4
#
# 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
# 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.
outputs = [
worker.execute_model_ray.bind(outputs[i]) # type: ignore[attr-defined]
for i, worker in enumerate(tp_group)
]
last_pp_rank = len(self.pp_tp_workers) - 1
if (
pp_rank < last_pp_rank
and envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE != "shm"
):
# Specify how intermediate tensors should be passed
# between pp stages, no need to specify for the last
# pp stage or when using shared memory (the default).
transport = envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
outputs = [
output.with_tensor_transport(transport=transport)
for output in outputs
]
forward_dag = MultiOutputNode(outputs)
if envs.VLLM_USE_RAY_WRAPPED_PP_COMM:
from ray.experimental.channel.accelerator_context import (
register_accelerator_context,
)
from vllm.distributed.device_communicators.ray_communicator import (
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."
)
else:
logger.info(
"Using Ray's NCCL communicator for Ray Compiled Graph communication."
)
return forward_dag.experimental_compile(
enable_asyncio=enable_asyncio,
_overlap_gpu_communication=envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM,
)
def __del__(self):
self.shutdown()
def check_health(self) -> None:
# Assume that the Ray workers are healthy.
# TODO: check the health of the Ray workers
return
vllm/v1/executor/ray_utils.py
-219
View File
@@ -4,26 +4,11 @@
import os
import time
from collections import defaultdict
from concurrent.futures import Future
from typing import TYPE_CHECKING, Union
import numpy as np
import vllm.platforms
from vllm.config import ParallelConfig
from vllm.distributed import get_pp_group
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.utils.network_utils import get_ip
from vllm.v1.outputs import AsyncModelRunnerOutput
from vllm.v1.serial_utils import run_method
from vllm.v1.worker.worker_base import WorkerWrapperBase
if TYPE_CHECKING:
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
logger = init_logger(__name__)
PG_WAIT_TIMEOUT = 1800
@@ -53,134 +38,6 @@ try:
available_resources_per_node = _state._available_resources_per_node
class RayWorkerWrapper(WorkerWrapperBase):
"""Ray wrapper for vllm.worker.Worker, allowing Worker to be
lazily initialized after Ray sets CUDA_VISIBLE_DEVICES."""
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
# Since the compiled DAG runs a main execution
# in a different thread that calls cuda.set_device.
# The flag indicates is set_device is called on
# that thread.
self.compiled_dag_cuda_device_set = False
rpc_rank: int
def adjust_rank(self, rank_mapping: dict[int, int]) -> None:
"""Adjust the rpc_rank based on the given mapping.
It is only used during the initialization of the executor,
to adjust the rpc_rank of workers after we create all workers.
"""
if self.rpc_rank in rank_mapping:
self.rpc_rank = rank_mapping[self.rpc_rank]
def execute_method(self, method: str | bytes, *args, **kwargs):
try:
return run_method(self, method, args, kwargs)
except Exception as e:
# if the driver worker also execute methods,
# exceptions in the rest worker may cause deadlock in rpc
# see https://github.com/vllm-project/vllm/issues/3455
msg = (
f"Error executing method {method!r}. "
"This might cause deadlock in distributed execution."
)
logger.exception(msg)
raise e
def get_node_ip(self) -> str:
return get_ip()
def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:
node_id = ray.get_runtime_context().get_node_id()
device_key = vllm.platforms.current_platform.ray_device_key
if not device_key:
raise RuntimeError(
"current platform %s does not support ray.",
vllm.platforms.current_platform.device_name,
)
physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[
device_key
]
return node_id, physical_gpu_ids
def setup_device_if_necessary(self):
# TODO(swang): This is needed right now because Ray CG executes
# on a background thread, so we need to reset torch's current
# device.
# We can remove this API after it is fixed in compiled graph.
assert self.worker is not None, "Worker is not initialized"
if not self.compiled_dag_cuda_device_set:
if current_platform.is_tpu():
# Not needed
pass
else:
assert self.worker.device is not None
current_platform.set_device(self.worker.device)
self.compiled_dag_cuda_device_set = True
def execute_model_ray(
self,
execute_model_input: tuple["SchedulerOutput", "GrammarOutput"]
| tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"],
) -> Union[
"ModelRunnerOutput",
tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"],
]:
# This method is used by Ray Compiled Graph to execute the model,
# and it needs a special logic of self.setup_device_if_necessary()
self.setup_device_if_necessary()
assert self.worker is not None, "Worker is not initialized"
if len(execute_model_input) == 3:
scheduler_output, grammar_output, intermediate_tensors = (
execute_model_input
)
else:
scheduler_output, grammar_output = execute_model_input
intermediate_tensors = None
assert self.worker.model_runner is not None
output = self.worker.model_runner.execute_model(
scheduler_output, intermediate_tensors
)
if self._is_intermediate_tensors(output):
if (
self.worker.model_runner.supports_mm_inputs
and get_pp_group().is_first_rank
):
# Strip mm_features before Ray forwards it to the next PP Stage.
# PP Stage>0 only needs the intermediate tensors,
# not preprocessed multimodal data.
# scheduled_new_reqs is a required field of SchedulerOutput,
# so accessing it directly will raise AttributeError if missing.
for req in scheduler_output.scheduled_new_reqs:
req.mm_features = []
return scheduler_output, grammar_output, output
if isinstance(output, AsyncModelRunnerOutput):
output = output.get_output()
if not self._is_last_rank():
# Case where there are no scheduled requests
# but may still be finished requests.
assert not output or not output.req_ids
output = scheduler_output, grammar_output, None
elif output is None:
output = self.worker.model_runner.sample_tokens(grammar_output)
# Ensure outputs crossing Ray compiled DAG are serializable.
# AsyncModelRunnerOutput holds CUDA events and cannot be
# pickled.
if isinstance(output, AsyncModelRunnerOutput):
output = output.get_output()
return output
def _is_intermediate_tensors(self, output) -> bool:
return isinstance(output, IntermediateTensors)
def _is_last_rank(self) -> bool:
return get_pp_group().is_last_rank
ray_import_err = None
except ImportError as e:
@@ -188,80 +45,6 @@ except ImportError as e:
# only capture string to avoid variable references in the traceback that can
# prevent garbage collection in some cases
ray_import_err = str(e)
RayWorkerWrapper = None # type: ignore
def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> None:
"""Detach Ray SHM-channel zero-copy buffers from a ModelRunnerOutput in-place.
Ray compiled DAG SHM channels may return zero-copy objects (e.g. `np.ndarray`)
backed by Ray's shared-memory object store. Ray's channel docs explicitly
warn that subsequent reads may block if such an object is still in scope.
vLLM can return numpy-backed logprobs and routed experts in
`ModelRunnerOutput`. If those arrays are backed by Ray SHM (commonly
read-only), retaining them in scope across scheduler iterations can stall
the channel and eventually hit `RAY_CGRAPH_get_timeout`.
Copy read-only numpy arrays so the returned output no longer retains
references to Ray's shared-memory buffers.
We intentionally do not touch `prompt_logprobs_dict`: those entries are
`LogprobsTensors` backed by PyTorch-owned CPU tensors (`to_cpu_nonblocking`
or `empty_cpu`), not NumPy views decoded from Ray channels.
"""
def _copy_if_readonly(arr):
if isinstance(arr, np.ndarray) and not arr.flags.writeable:
return arr.copy()
return arr
if output.logprobs is not None:
token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs
# `cu_num_generated_tokens` is already a plain Python list (or None),
# so it never aliases Ray SHM buffers and can be reused as-is.
token_ids_c = _copy_if_readonly(token_ids)
logprobs_c = _copy_if_readonly(logprobs)
ranks_c = _copy_if_readonly(ranks)
if (
token_ids_c is not token_ids
or logprobs_c is not logprobs
or ranks_c is not ranks
):
output.logprobs = type(output.logprobs)(
token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
)
aux_output = output.aux_output_connector_output
if aux_output is not None:
for request_output in aux_output.values():
request_output.rows = _copy_if_readonly(request_output.rows)
class FutureWrapper(Future):
"""A wrapper around Ray output reference to meet the interface
of .execute_model(): The top level (core busy loop) expects .result() api
to block and return a single output.
If aggregator is provided, the outputs from all workers are aggregated upon
the result() call. If not only the first worker's output is returned.
"""
def __init__(self, ref_or_refs, aggregator: KVOutputAggregator | None = None):
super().__init__()
self.ref_or_refs = ref_or_refs
self.aggregator = aggregator
def result(self, timeout=None):
outputs = ray.get(self.ref_or_refs, timeout=timeout)
if self.aggregator is None:
detach_zero_copy_from_model_runner_output(outputs)
return outputs
for output in outputs:
detach_zero_copy_from_model_runner_output(output)
return self.aggregator.aggregate(outputs, output_rank=0)
def ray_is_available() -> bool:
@@ -627,8 +410,6 @@ def initialize_ray_cluster(
current_ip = get_ip()
current_node_id = ray.get_runtime_context().get_node_id()
current_node_resource = available_resources_per_node()[current_node_id]
# TODO (jeffreywang): require_gpu_on_driver should be always False
# after deprecating RayDistributedExecutor.
if require_gpu_on_driver:
if current_node_resource.get(device_str, 0) < 1:
raise ValueError(