| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Integration tests for RayExecutorV2 at the executor level. |
| 5 | Validates executor initialization, placement group support, RPC calls, |
| 6 | and distributed execution with various TP/PP configurations. |
| 7 | """ |
| 8 | |
| 9 | import errno |
| 10 | import gc |
| 11 | import socket |
| 12 | import threading |
| 13 | from types import SimpleNamespace |
| 14 | from unittest.mock import patch |
| 15 | from urllib.parse import urlsplit |
| 16 | |
| 17 | import pytest |
| 18 | import ray |
| 19 | from torch.distributed import TCPStore |
| 20 | |
| 21 | from vllm import LLM |
| 22 | from vllm.config import VllmConfig |
| 23 | from vllm.engine.arg_utils import EngineArgs |
| 24 | from vllm.v1.executor import ray_executor_v2 |
| 25 | from vllm.v1.executor.ray_executor_v2 import RayExecutorV2 |
| 26 | |
| 27 | pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend") |
| 28 | |
| 29 | MODEL = "facebook/opt-125m" |
| 30 | |
| 31 | |
| 32 | def create_vllm_config( |
| 33 | tensor_parallel_size: int = 1, |
| 34 | pipeline_parallel_size: int = 1, |
| 35 | max_model_len: int = 256, |
| 36 | gpu_memory_utilization: float = 0.3, |
| 37 | placement_group=None, |
| 38 | ) -> VllmConfig: |
| 39 | engine_args = EngineArgs( |
| 40 | model=MODEL, |
| 41 | tensor_parallel_size=tensor_parallel_size, |
| 42 | pipeline_parallel_size=pipeline_parallel_size, |
| 43 | max_model_len=max_model_len, |
| 44 | gpu_memory_utilization=gpu_memory_utilization, |
| 45 | distributed_executor_backend="ray", |
| 46 | enforce_eager=True, |
| 47 | ) |
| 48 | vllm_config = engine_args.create_engine_config() |
| 49 | |
| 50 | if placement_group is not None: |
| 51 | vllm_config.parallel_config.placement_group = placement_group |
| 52 | |
| 53 | return vllm_config |
| 54 | |
| 55 | |
| 56 | def ensure_ray_initialized(): |
| 57 | if not ray.is_initialized(): |
| 58 | ray.init(ignore_reinit_error=True) |
| 59 | |
| 60 | |
| 61 | @pytest.fixture |
| 62 | def create_placement_group(request): |
| 63 | ensure_ray_initialized() |
| 64 | num_gpus = request.param |
| 65 | bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)] |
| 66 | pg = ray.util.placement_group(bundles, strategy="PACK") |
| 67 | ray.get(pg.ready()) |
| 68 | yield pg |
| 69 | ray.util.remove_placement_group(pg) |
| 70 | |
| 71 | |
| 72 | @pytest.fixture |
| 73 | def executor(request): |
| 74 | """Create a RayExecutorV2 and shut it down after the test.""" |
| 75 | executor = RayExecutorV2(vllm_config=request.param) |
| 76 | yield executor |
| 77 | executor.shutdown() |
| 78 | |
| 79 | |
| 80 | def assert_executor(executor, tp_size, pp_size): |
| 81 | """Common assertions for executor initialization tests.""" |
| 82 | world_size = tp_size * pp_size |
| 83 | expected_output_rank = (pp_size - 1) * tp_size |
| 84 | |
| 85 | assert executor.world_size == world_size |
| 86 | assert len(executor.ray_worker_handles) == world_size |
| 87 | assert len(executor.response_mqs) == world_size |
| 88 | assert executor._get_output_rank() == expected_output_rank |
| 89 | |
| 90 | if pp_size > 1: |
| 91 | expected_concurrent_batches = pp_size + int( |
| 92 | executor.vllm_config.scheduler_config.async_scheduling |
| 93 | and executor.vllm_config.use_v2_model_runner |
| 94 | ) |
| 95 | assert ( |
| 96 | executor.vllm_config.max_concurrent_batches == expected_concurrent_batches |
| 97 | ) |
| 98 | |
| 99 | executor.check_health() |
| 100 | assert not executor.is_failed |
| 101 | |
| 102 | ranks = sorted(h.rank for h in executor.ray_worker_handles) |
| 103 | assert ranks == list(range(world_size)) |
| 104 | |
| 105 | for handle in executor.ray_worker_handles: |
| 106 | assert handle.node_id is not None |
| 107 | |
| 108 | |
| 109 | def _detached_rank0_worker() -> ray_executor_v2.RayWorkerProc: |
| 110 | """A RayWorkerProc that can create its dist-init store without Ray.""" |
| 111 | worker = ray_executor_v2.RayWorkerProc.__new__(ray_executor_v2.RayWorkerProc) |
| 112 | worker._parallel_config = SimpleNamespace(world_size=1) |
| 113 | return worker |
| 114 | |
| 115 | |
| 116 | def _assert_port_held(port: int): |
| 117 | with socket.socket() as contender, pytest.raises(OSError) as exc_info: |
| 118 | contender.bind(("127.0.0.1", port)) |
| 119 | assert exc_info.value.errno == errno.EADDRINUSE |
| 120 | |
| 121 | |
| 122 | def test_dist_init_store_owns_port(monkeypatch): |
| 123 | """The rank-0 actor keeps the selected TCPStore port bound.""" |
| 124 | monkeypatch.setattr(ray.util, "get_node_ip_address", lambda: "127.0.0.1") |
| 125 | worker = _detached_rank0_worker() |
| 126 | |
| 127 | init_method = worker.create_dist_init_method() |
| 128 | port = urlsplit(init_method).port |
| 129 | assert port is not None and port != 0 |
| 130 | _assert_port_held(port) |
| 131 | |
| 132 | reused_store = TCPStore( |
| 133 | "127.0.0.1", |
| 134 | port, |
| 135 | world_size=1, |
| 136 | is_master=True, |
| 137 | wait_for_workers=False, |
| 138 | multi_tenant=True, |
| 139 | ) |
| 140 | worker._dist_init_store.set("key", "value") |
| 141 | assert reused_store.get("key") == b"value" |
| 142 | |
| 143 | |
| 144 | def test_colocated_dist_init_stores_hold_distinct_ports(monkeypatch): |
| 145 | """Replicas sharing a host must never publish the same TCPStore port. |
| 146 | |
| 147 | Each rank-0 store binds a kernel-assigned ephemeral port and keeps it |
| 148 | for the actor's lifetime, so concurrently live replicas cannot collide. |
| 149 | """ |
| 150 | monkeypatch.setattr(ray.util, "get_node_ip_address", lambda: "127.0.0.1") |
| 151 | workers = [_detached_rank0_worker() for _ in range(3)] |
| 152 | |
| 153 | ports = [urlsplit(w.create_dist_init_method()).port for w in workers] |
| 154 | assert len(set(ports)) == len(workers) |
| 155 | for worker, port in zip(workers, ports): |
| 156 | assert port == worker._dist_init_store.port |
| 157 | assert port >= 1024 |
| 158 | _assert_port_held(port) |
| 159 | |
| 160 | |
| 161 | @pytest.mark.parametrize("tp_size, pp_size", [(1, 1), (2, 1), (4, 1), (2, 2)]) |
| 162 | def test_ray_v2_executor(tp_size, pp_size): |
| 163 | """Validate RayExecutorV2 with various TP/PP configs.""" |
| 164 | vllm_config = create_vllm_config( |
| 165 | tensor_parallel_size=tp_size, |
| 166 | pipeline_parallel_size=pp_size, |
| 167 | ) |
| 168 | executor = RayExecutorV2(vllm_config=vllm_config) |
| 169 | try: |
| 170 | assert_executor(executor, tp_size, pp_size) |
| 171 | finally: |
| 172 | executor.shutdown() |
| 173 | |
| 174 | |
| 175 | @pytest.mark.parametrize( |
| 176 | "tp_size, pp_size, create_placement_group", |
| 177 | [(2, 1, 2), (4, 1, 4), (2, 2, 4)], |
| 178 | indirect=["create_placement_group"], |
| 179 | ) |
| 180 | def test_ray_v2_executor_pg(tp_size, pp_size, create_placement_group): |
| 181 | """Validate RayExecutorV2 with various TP/PP configs using external PG.""" |
| 182 | vllm_config = create_vllm_config( |
| 183 | tensor_parallel_size=tp_size, |
| 184 | pipeline_parallel_size=pp_size, |
| 185 | placement_group=create_placement_group, |
| 186 | ) |
| 187 | executor = RayExecutorV2(vllm_config=vllm_config) |
| 188 | try: |
| 189 | assert_executor(executor, tp_size, pp_size) |
| 190 | finally: |
| 191 | executor.shutdown() |
| 192 | |
| 193 | |
| 194 | @pytest.mark.parametrize( |
| 195 | "executor", |
| 196 | [create_vllm_config(tensor_parallel_size=2)], |
| 197 | indirect=True, |
| 198 | ) |
| 199 | def test_ray_v2_executor_failure_callback(executor): |
| 200 | """Validate failure callback registration.""" |
| 201 | callback_invoked = False |
| 202 | |
| 203 | def test_callback(): |
| 204 | nonlocal callback_invoked |
| 205 | callback_invoked = True |
| 206 | |
| 207 | executor.register_failure_callback(test_callback) |
| 208 | assert not callback_invoked |
| 209 | |
| 210 | executor.is_failed = True |
| 211 | executor.register_failure_callback(test_callback) |
| 212 | assert callback_invoked |
| 213 | |
| 214 | |
| 215 | @pytest.mark.parametrize( |
| 216 | "executor", |
| 217 | [create_vllm_config(tensor_parallel_size=2)], |
| 218 | indirect=True, |
| 219 | ) |
| 220 | def test_ray_v2_executor_collective_rpc(executor): |
| 221 | """Validate collective RPC calls through MessageQueue.""" |
| 222 | executor.check_health() |
| 223 | assert not executor.is_failed |
| 224 | assert executor.rpc_broadcast_mq is not None |
| 225 | |
| 226 | |
| 227 | @pytest.mark.parametrize( |
| 228 | "executor", |
| 229 | [create_vllm_config(tensor_parallel_size=2)], |
| 230 | indirect=True, |
| 231 | ) |
| 232 | def test_ray_v2_executor_driver_node_rank_0(executor): |
| 233 | """Validate that driver node workers get the lowest ranks.""" |
| 234 | driver_node = ray.get_runtime_context().get_node_id() |
| 235 | |
| 236 | for handle in executor.ray_worker_handles: |
| 237 | assert handle.node_id == driver_node |
| 238 | |
| 239 | rank0_handle = next(h for h in executor.ray_worker_handles if h.rank == 0) |
| 240 | assert rank0_handle.node_id == driver_node |
| 241 | |
| 242 | |
| 243 | @pytest.mark.parametrize( |
| 244 | "executor", |
| 245 | [create_vllm_config(tensor_parallel_size=2)], |
| 246 | indirect=True, |
| 247 | ) |
| 248 | def test_ray_v2_executor_worker_death(executor): |
| 249 | """Validate executor detects worker death via ray.wait().""" |
| 250 | callback_event = threading.Event() |
| 251 | |
| 252 | def on_failure(): |
| 253 | callback_event.set() |
| 254 | |
| 255 | executor.register_failure_callback(on_failure) |
| 256 | assert not executor.is_failed |
| 257 | |
| 258 | # Kill one worker actor externally |
| 259 | victim = executor.ray_worker_handles[1].actor |
| 260 | ray.kill(victim, no_restart=True) |
| 261 | |
| 262 | # Monitor thread should detect the death and invoke callback |
| 263 | assert callback_event.wait(timeout=30) |
| 264 | assert executor.is_failed |
| 265 | assert executor.shutting_down |
| 266 | |
| 267 | |
| 268 | def test_ray_v2_executor_shutdown(): |
| 269 | """Validate graceful shutdown: ray.kill() terminates all worker actors.""" |
| 270 | executor = RayExecutorV2(vllm_config=create_vllm_config(tensor_parallel_size=2)) |
| 271 | assert executor.rpc_broadcast_mq is not None |
| 272 | assert len(executor.response_mqs) == executor.world_size |
| 273 | |
| 274 | actors = [h.actor for h in executor.ray_worker_handles] |
| 275 | executor.shutdown() |
| 276 | |
| 277 | for actor in actors: |
| 278 | with pytest.raises(ray.exceptions.RayActorError): |
| 279 | ray.get(actor.wait_for_init.remote(), timeout=5) |
| 280 | |
| 281 | assert executor.rpc_broadcast_mq is None |
| 282 | assert len(executor.response_mqs) == 0 |
| 283 | |
| 284 | |
| 285 | @pytest.mark.parametrize( |
| 286 | "executor", |
| 287 | [create_vllm_config(tensor_parallel_size=2)], |
| 288 | indirect=True, |
| 289 | ) |
| 290 | def test_ray_v2_run_refs_stored_for_monitoring(executor): |
| 291 | """Validate worker handles store run_ref for monitoring.""" |
| 292 | for handle in executor.ray_worker_handles: |
| 293 | assert handle.run_ref is not None |
| 294 | ready, _ = ray.wait([handle.run_ref], timeout=0) |
| 295 | assert len(ready) == 0, "run_ref should be pending" |
| 296 | |
| 297 | |
| 298 | @pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)]) |
| 299 | def test_ray_v2_single_node_generation(tp_size, pp_size): |
| 300 | """End-to-end LLM generation with RayExecutorV2.""" |
| 301 | llm = LLM( |
| 302 | model=MODEL, |
| 303 | tensor_parallel_size=tp_size, |
| 304 | pipeline_parallel_size=pp_size, |
| 305 | distributed_executor_backend="ray", |
| 306 | enforce_eager=True, |
| 307 | max_model_len=256, |
| 308 | gpu_memory_utilization=0.3, |
| 309 | ) |
| 310 | try: |
| 311 | prompts = [ |
| 312 | "Hello, my name is", |
| 313 | "The capital of France is", |
| 314 | "The future of AI is", |
| 315 | ] |
| 316 | outputs = llm.generate(prompts) |
| 317 | |
| 318 | assert len(outputs) == len(prompts) |
| 319 | for output in outputs: |
| 320 | assert len(output.outputs) > 0 |
| 321 | assert len(output.outputs[0].text) > 0 |
| 322 | finally: |
| 323 | llm.llm_engine.model_executor.shutdown() |
| 324 | del llm |
| 325 | gc.collect() |
| 326 | |
| 327 | |
| 328 | @pytest.mark.parametrize( |
| 329 | "bundle_indices, expected_bundle_ids, create_placement_group", |
| 330 | [("2,3", [2, 3], 4), ("3,2", [3, 2], 4)], |
| 331 | indirect=["create_placement_group"], |
| 332 | ) |
| 333 | def test_ray_v2_bundle_indices_env( |
| 334 | bundle_indices, expected_bundle_ids, create_placement_group, monkeypatch |
| 335 | ): |
| 336 | """Validate explicit VLLM_RAY_BUNDLE_INDICES bundle placement.""" |
| 337 | monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", bundle_indices) |
| 338 | vllm_config = create_vllm_config( |
| 339 | tensor_parallel_size=2, |
| 340 | placement_group=create_placement_group, |
| 341 | ) |
| 342 | executor = RayExecutorV2(vllm_config=vllm_config) |
| 343 | try: |
| 344 | actual = [ |
| 345 | h.bundle_id_idx |
| 346 | for h in sorted(executor.ray_worker_handles, key=lambda h: h.rank) |
| 347 | ] |
| 348 | assert actual == expected_bundle_ids |
| 349 | assert_executor(executor, tp_size=2, pp_size=1) |
| 350 | finally: |
| 351 | executor.shutdown() |
| 352 | |
| 353 | |
| 354 | @pytest.mark.parametrize( |
| 355 | "bundle_indices, expected_error, create_placement_group", |
| 356 | [ |
| 357 | ("1,1", "cannot have duplicate values,", 4), |
| 358 | ("0,1,2", "must have the same size", 4), |
| 359 | ], |
| 360 | indirect=["create_placement_group"], |
| 361 | ) |
| 362 | def test_ray_v2_invalid_bundle_indices( |
| 363 | bundle_indices, expected_error, create_placement_group, monkeypatch |
| 364 | ): |
| 365 | """Validate invalid bundle indices are rejected.""" |
| 366 | monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", bundle_indices) |
| 367 | vllm_config = create_vllm_config( |
| 368 | tensor_parallel_size=2, placement_group=create_placement_group |
| 369 | ) |
| 370 | with pytest.raises(AssertionError, match=expected_error): |
| 371 | RayExecutorV2(vllm_config=vllm_config) |
| 372 | |
| 373 | |
| 374 | @pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)]) |
| 375 | def test_ray_v2_single_node_generation_with_pg(tp_size, pp_size): |
| 376 | """E2E LLM generation with a user-provided placement group.""" |
| 377 | ensure_ray_initialized() |
| 378 | bundles = [{"GPU": 1, "CPU": 1} for _ in range(tp_size * pp_size)] |
| 379 | pg = ray.util.placement_group(bundles, strategy="PACK") |
| 380 | ray.get(pg.ready()) |
| 381 | |
| 382 | try: |
| 383 | with patch.object(ray.util, "get_current_placement_group", return_value=pg): |
| 384 | llm = LLM( |
| 385 | model=MODEL, |
| 386 | tensor_parallel_size=tp_size, |
| 387 | pipeline_parallel_size=pp_size, |
| 388 | distributed_executor_backend="ray", |
| 389 | enforce_eager=True, |
| 390 | max_model_len=256, |
| 391 | gpu_memory_utilization=0.3, |
| 392 | ) |
| 393 | prompts = [ |
| 394 | "Hello, my name is", |
| 395 | "The capital of France is", |
| 396 | "The future of AI is", |
| 397 | ] |
| 398 | outputs = llm.generate(prompts) |
| 399 | |
| 400 | assert len(outputs) == len(prompts) |
| 401 | for output in outputs: |
| 402 | assert len(output.outputs) > 0 |
| 403 | assert len(output.outputs[0].text) > 0 |
| 404 | finally: |
| 405 | llm.llm_engine.model_executor.shutdown() |
| 406 | del llm |
| 407 | gc.collect() |