| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import uuid |
| 4 | from typing import Any |
| 5 | |
| 6 | import ray |
| 7 | import torch |
| 8 | from ray.exceptions import RayChannelError |
| 9 | from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator |
| 10 | from torch.distributed import ReduceOp |
| 11 | |
| 12 | from vllm.distributed.device_communicators.base_device_communicator import ( |
| 13 | DeviceCommunicatorBase, |
| 14 | ) |
| 15 | from vllm.distributed.parallel_state import get_pp_group |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.utils.torch_utils import current_stream |
| 18 | |
| 19 | logger = init_logger(__name__) |
| 20 | |
| 21 | |
| 22 | class RayPPCommunicator(Communicator): |
| 23 | """Communicator to be used for pipeline parallelism in Ray Compiled Graph. |
| 24 | This is wraps around the vLLM _PP GroupCoordinator. |
| 25 | |
| 26 | This class is not thread-safe. |
| 27 | """ |
| 28 | |
| 29 | _comm: DeviceCommunicatorBase | None |
| 30 | |
| 31 | def __init__( |
| 32 | self, |
| 33 | world_size: int, |
| 34 | comm_id: Any, |
| 35 | rank: int | None, |
| 36 | actor_handles: list["ray.actor.ActorHandle"], |
| 37 | cuda_stream: torch.cuda.Stream | None, |
| 38 | use_communication_streams: bool = False, |
| 39 | ): |
| 40 | """Initialize a RayPPCommunicator that can be used to communicate with |
| 41 | other Ray Compiled Graph actors for pipeline parallelism. |
| 42 | |
| 43 | Args: |
| 44 | world_size: The number of participating actors. |
| 45 | comm_id: A unique communicator ID. This is just to conform with |
| 46 | the Ray Communicator API and is not used. |
| 47 | rank: The rank of this actor. If None, then the caller is not a |
| 48 | participant of the RayPPCommunicator group (e.g., the Ray |
| 49 | driver). |
| 50 | actor_handles: A list of actor handles. |
| 51 | cuda_stream: A CUDA stream to dispatch communication ops to. This |
| 52 | is not supported. |
| 53 | use_communication_streams: Whether to use communication streams. |
| 54 | This is not supported. |
| 55 | |
| 56 | """ |
| 57 | self._world_size = world_size |
| 58 | self._rank: int | None = None |
| 59 | self._actor_handles = actor_handles |
| 60 | if use_communication_streams: |
| 61 | raise NotImplementedError("use_communication_streams is not supported") |
| 62 | if cuda_stream is not None and cuda_stream != current_stream(): |
| 63 | raise ValueError( |
| 64 | "cuda_stream other than the current stream is not supported" |
| 65 | ) |
| 66 | |
| 67 | if rank is not None: |
| 68 | # Rank is not None, this is Ray worker |
| 69 | assert ray.get_gpu_ids(), "RayPPCommunicator has no GPUs assigned" |
| 70 | |
| 71 | self._comm = get_pp_group().device_communicator |
| 72 | assert self._comm is not None |
| 73 | |
| 74 | # Since we wrap around the vLLM _PP communicator, we use |
| 75 | # the rank from the vLLM communicator, and ignore the rank |
| 76 | # passed in from Ray. |
| 77 | # TODO(rui): refactor the Ray Communicator API so that |
| 78 | # it also supports no rank passed in. |
| 79 | self._rank = self._comm.rank_in_group |
| 80 | |
| 81 | self._build_actor_rank_mapping() |
| 82 | else: |
| 83 | # Rank is None, this is Ray driver |
| 84 | self._comm = None |
| 85 | |
| 86 | self._closed = False |
| 87 | |
| 88 | def _build_actor_rank_mapping(self): |
| 89 | """Use collective communication to build a mapping from actor IDs to ranks. |
| 90 | This should be called once during initialization. |
| 91 | """ |
| 92 | if self._comm is None: |
| 93 | return {} |
| 94 | |
| 95 | current_actor = ray.get_runtime_context().current_actor |
| 96 | actor_id_str = current_actor._actor_id.hex() |
| 97 | |
| 98 | # Ray actor IDs are 32-character hex strings (128 bits) |
| 99 | ACTOR_ID_LEN = 32 |
| 100 | actor_id_bytes = bytearray(actor_id_str.encode("utf-8")) |
| 101 | assert len(actor_id_bytes) == ACTOR_ID_LEN, ( |
| 102 | f"Unexpected actor ID length: {len(actor_id_bytes)}" |
| 103 | ) |
| 104 | |
| 105 | actor_id_tensor = torch.frombuffer(actor_id_bytes, dtype=torch.uint8).to( |
| 106 | self._comm.device |
| 107 | ) |
| 108 | |
| 109 | # All-gather full actor IDs from all actors |
| 110 | gathered_ids = self._comm.all_gather(actor_id_tensor, dim=0) |
| 111 | |
| 112 | # Build mapping: actor_id -> device_comm_rank |
| 113 | self._actor_id_to_rank = {} |
| 114 | for rank in range(self._world_size): |
| 115 | start_idx = rank * ACTOR_ID_LEN |
| 116 | end_idx = (rank + 1) * ACTOR_ID_LEN |
| 117 | actor_bytes = gathered_ids[start_idx:end_idx].cpu().numpy().tobytes() |
| 118 | actor_id = actor_bytes.decode("utf-8") |
| 119 | self._actor_id_to_rank[actor_id] = rank |
| 120 | |
| 121 | def initialize(self, rank: int) -> None: |
| 122 | # No additional initialization is needed. |
| 123 | pass |
| 124 | |
| 125 | def get_actor_handles(self) -> list["ray.actor.ActorHandle"]: |
| 126 | return self._actor_handles |
| 127 | |
| 128 | def get_rank(self, actor: ray.actor.ActorHandle) -> int: |
| 129 | """Return the given actor's rank using device communicator collective ops.""" |
| 130 | assert hasattr(self, "_actor_id_to_rank"), ( |
| 131 | "Actor rank mapping not built. " |
| 132 | "This should have been done during initialization." |
| 133 | ) |
| 134 | |
| 135 | actor_id_str = actor._actor_id.hex() |
| 136 | |
| 137 | if actor_id_str in self._actor_id_to_rank: |
| 138 | return self._actor_id_to_rank[actor_id_str] # type: ignore |
| 139 | else: |
| 140 | raise ValueError(f"Actor {actor} not found in communicator group") |
| 141 | |
| 142 | def get_self_rank(self) -> int | None: |
| 143 | """Return this actor's rank.""" |
| 144 | return self._rank |
| 145 | |
| 146 | def get_world_size(self) -> int: |
| 147 | """Return the number of ranks in the RayPPCommunicator group.""" |
| 148 | return self._world_size |
| 149 | |
| 150 | def send(self, buf: "torch.Tensor", peer_rank: int) -> None: |
| 151 | """Send a torch.Tensor to a peer. |
| 152 | |
| 153 | This returns when the send kernel has been queued, but the kernel may |
| 154 | not have completed. Therefore, the caller should ensure that there are |
| 155 | no concurrent writes to the sent `buf` until the send has finished. |
| 156 | That is, either all writes should be submitted on the current stream |
| 157 | (self._cuda_stream) or, if on a different stream, that stream should |
| 158 | synchronize with the current stream. |
| 159 | |
| 160 | Args: |
| 161 | buf: The torch.Tensor to send. It should already be on this |
| 162 | actor's default device. |
| 163 | peer_rank: The rank of the actor to send to. |
| 164 | |
| 165 | """ |
| 166 | if self._closed: |
| 167 | raise RayChannelError("RayPPCommunicator has been destroyed.") |
| 168 | |
| 169 | assert self._comm is not None |
| 170 | self._comm.send(buf, peer_rank) |
| 171 | |
| 172 | def recv( |
| 173 | self, |
| 174 | shape: tuple[int, ...], |
| 175 | dtype: "torch.dtype", |
| 176 | peer_rank: int, |
| 177 | allocator: TorchTensorAllocator, |
| 178 | ) -> "torch.Tensor": |
| 179 | """Receive a torch.Tensor from a peer and synchronize the current stream. |
| 180 | |
| 181 | After this call returns, the receive buffer is safe to read from |
| 182 | any stream. An RayChannelError will be raised if an error occurred |
| 183 | (e.g., remote actor died), and the buffer is not safe to read. |
| 184 | |
| 185 | Args: |
| 186 | shape: The shape of the tensor to receive. |
| 187 | dtype: The dtype of the tensor to receive. |
| 188 | peer_rank: The rank of the actor to receive from. |
| 189 | allocator: The allocator to use to create the received tensor. |
| 190 | This is ignored for this implementation. |
| 191 | |
| 192 | """ |
| 193 | if self._closed: |
| 194 | raise RayChannelError("RayPPCommunicator has been destroyed.") |
| 195 | |
| 196 | assert self._comm is not None |
| 197 | size = torch.Size(shape) |
| 198 | buf = self._comm.recv(size, dtype, src=peer_rank) |
| 199 | |
| 200 | # Buffer values are undefined if NCCL ops are aborted. Therefore, we |
| 201 | # need to synchronize here and check that the channel is still |
| 202 | # open to ensure that the receive buffer is valid. |
| 203 | # TODO(swang): Avoid CUDA synchronization. |
| 204 | current_stream().synchronize() |
| 205 | |
| 206 | if self._closed: |
| 207 | raise RayChannelError("RayPPCommunicator has been destroyed.") |
| 208 | return buf |
| 209 | |
| 210 | def allgather( |
| 211 | self, |
| 212 | send_buf: "torch.Tensor", |
| 213 | recv_buf: "torch.Tensor", |
| 214 | ): |
| 215 | raise NotImplementedError("allgather is not supported") |
| 216 | |
| 217 | def allreduce( |
| 218 | self, |
| 219 | send_buf: "torch.Tensor", |
| 220 | recv_buf: "torch.Tensor", |
| 221 | op: ReduceOp = ReduceOp.SUM, |
| 222 | ): |
| 223 | raise NotImplementedError("allreduce is not supported") |
| 224 | |
| 225 | def reducescatter( |
| 226 | self, |
| 227 | send_buf: "torch.Tensor", |
| 228 | recv_buf: "torch.Tensor", |
| 229 | op: ReduceOp = ReduceOp.SUM, |
| 230 | ): |
| 231 | raise NotImplementedError("reducescatter is not supported") |
| 232 | |
| 233 | @property |
| 234 | def recv_stream(self): |
| 235 | return torch.cuda.StreamContext(current_stream()) |
| 236 | |
| 237 | @property |
| 238 | def send_stream(self): |
| 239 | return torch.cuda.StreamContext(current_stream()) |
| 240 | |
| 241 | def destroy(self) -> None: |
| 242 | # Just sets a flag, vLLM manages the lifecycle of the underlying |
| 243 | # _PP GroupCoordinator. |
| 244 | self._closed = True |
| 245 | |
| 246 | def get_transport_name(self) -> str: |
| 247 | return "nccl" |
| 248 | |
| 249 | @classmethod |
| 250 | def generate_communicator_id(cls) -> Any: |
| 251 | return uuid.uuid4() |