| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import io |
| 5 | import pickle |
| 6 | import random |
| 7 | import threading |
| 8 | import time |
| 9 | from types import SimpleNamespace |
| 10 | from unittest import mock |
| 11 | |
| 12 | import multiprocess as mp |
| 13 | import numpy as np |
| 14 | import pytest |
| 15 | import torch |
| 16 | import torch.distributed as dist |
| 17 | |
| 18 | from vllm.distributed.device_communicators import shm_broadcast |
| 19 | from vllm.distributed.device_communicators.shm_broadcast import ( |
| 20 | MessageQueue, |
| 21 | ShmRingBuffer, |
| 22 | _rebuild_tensor, |
| 23 | _reduce_tensor, |
| 24 | check_shm_free_space, |
| 25 | ) |
| 26 | from vllm.distributed.utils import StatelessProcessGroup |
| 27 | from vllm.utils.network_utils import get_open_port |
| 28 | from vllm.utils.system_utils import update_environment_variables |
| 29 | |
| 30 | |
| 31 | def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]: |
| 32 | np.random.seed(seed) |
| 33 | sizes = np.random.randint(1, 10_000, n) |
| 34 | # on average, each array will have 5k elements |
| 35 | # with int64, each array will have 40kb |
| 36 | return [np.random.randint(1, 100, i) for i in sizes] |
| 37 | |
| 38 | |
| 39 | def distributed_run(fn, world_size, timeout=60): |
| 40 | """Run a function in multiple processes with proper error handling. |
| 41 | |
| 42 | Args: |
| 43 | fn: Function to run in each process |
| 44 | world_size: Number of processes to spawn |
| 45 | timeout: Maximum time in seconds to wait for processes (default: 60) |
| 46 | |
| 47 | """ |
| 48 | number_of_processes = world_size |
| 49 | processes = [] |
| 50 | for i in range(number_of_processes): |
| 51 | env = {} |
| 52 | env["RANK"] = str(i) |
| 53 | env["LOCAL_RANK"] = str(i) |
| 54 | env["WORLD_SIZE"] = str(number_of_processes) |
| 55 | env["LOCAL_WORLD_SIZE"] = str(number_of_processes) |
| 56 | env["MASTER_ADDR"] = "localhost" |
| 57 | env["MASTER_PORT"] = "12345" |
| 58 | p = mp.Process(target=fn, args=(env,)) |
| 59 | processes.append(p) |
| 60 | p.start() |
| 61 | |
| 62 | # Monitor processes and fail fast if any process fails |
| 63 | start_time = time.time() |
| 64 | failed_processes = [] |
| 65 | |
| 66 | # Wait for all processes, checking for failures |
| 67 | while time.time() - start_time < timeout: |
| 68 | all_done = True |
| 69 | for i, p in enumerate(processes): |
| 70 | if p.is_alive(): |
| 71 | all_done = False |
| 72 | elif p.exitcode != 0: |
| 73 | # Process failed |
| 74 | failed_processes.append((i, p.exitcode)) |
| 75 | break |
| 76 | |
| 77 | if failed_processes or all_done: |
| 78 | break |
| 79 | time.sleep(0.1) # Check every 100ms |
| 80 | |
| 81 | # Check for timeout if no failures detected yet |
| 82 | for i, p in enumerate(processes): |
| 83 | if p.is_alive(): |
| 84 | p.kill() |
| 85 | p.join() |
| 86 | |
| 87 | # Report failures |
| 88 | if failed_processes: |
| 89 | error_msg = "Distributed test failed:\n" |
| 90 | for rank, status in failed_processes: |
| 91 | error_msg += f" Rank {rank}: Exit code {status}\n" |
| 92 | raise AssertionError(error_msg) |
| 93 | |
| 94 | |
| 95 | def worker_fn_wrapper(fn): |
| 96 | # `mp.Process` cannot accept environment variables directly |
| 97 | # so we need to pass the environment variables as arguments |
| 98 | # and update the environment variables in the function |
| 99 | def wrapped_fn(env): |
| 100 | update_environment_variables(env) |
| 101 | dist.init_process_group(backend="gloo") |
| 102 | fn() |
| 103 | |
| 104 | return wrapped_fn |
| 105 | |
| 106 | |
| 107 | @worker_fn_wrapper |
| 108 | def worker_fn(): |
| 109 | rank = dist.get_rank() |
| 110 | if rank == 0: |
| 111 | port = get_open_port() |
| 112 | ip = "127.0.0.1" |
| 113 | dist.broadcast_object_list([ip, port], src=0) |
| 114 | else: |
| 115 | recv = [None, None] |
| 116 | dist.broadcast_object_list(recv, src=0) |
| 117 | ip, port = recv # type: ignore |
| 118 | |
| 119 | stateless_pg = StatelessProcessGroup.create(ip, port, rank, dist.get_world_size()) |
| 120 | |
| 121 | for pg in [dist.group.WORLD, stateless_pg]: |
| 122 | writer_rank = 2 |
| 123 | broadcaster = MessageQueue.create_from_process_group( |
| 124 | pg, 40 * 1024, 2, writer_rank |
| 125 | ) |
| 126 | if rank == writer_rank: |
| 127 | seed = random.randint(0, 1000) |
| 128 | dist.broadcast_object_list([seed], writer_rank) |
| 129 | else: |
| 130 | recv = [None] |
| 131 | dist.broadcast_object_list(recv, writer_rank) |
| 132 | seed = recv[0] # type: ignore |
| 133 | |
| 134 | if pg == dist.group.WORLD: |
| 135 | dist.barrier() |
| 136 | else: |
| 137 | pg.barrier() |
| 138 | |
| 139 | # in case we find a race condition |
| 140 | # print the seed so that we can reproduce the error |
| 141 | print(f"Rank {rank} got seed {seed}") |
| 142 | # test broadcasting with about 400MB of data |
| 143 | N = 10_000 |
| 144 | if rank == writer_rank: |
| 145 | arrs = get_arrays(N, seed) |
| 146 | for x in arrs: |
| 147 | broadcaster.broadcast_object(x) |
| 148 | time.sleep(random.random() / 1000) |
| 149 | else: |
| 150 | arrs = get_arrays(N, seed) |
| 151 | for x in arrs: |
| 152 | y = broadcaster.broadcast_object(None) |
| 153 | assert np.array_equal(x, y) |
| 154 | time.sleep(random.random() / 1000) |
| 155 | |
| 156 | if pg == dist.group.WORLD: |
| 157 | dist.barrier() |
| 158 | print(f"torch distributed passed the test! Rank {rank}") |
| 159 | else: |
| 160 | pg.barrier() |
| 161 | print(f"StatelessProcessGroup passed the test! Rank {rank}") |
| 162 | |
| 163 | |
| 164 | def test_shm_broadcast(): |
| 165 | distributed_run(worker_fn, 4) |
| 166 | |
| 167 | |
| 168 | @worker_fn_wrapper |
| 169 | def worker_fn_test_shutdown_busy(): |
| 170 | rank = dist.get_rank() |
| 171 | writer_rank = 2 |
| 172 | message_queue = MessageQueue.create_from_process_group( |
| 173 | dist.group.WORLD, 40 * 1024, 2, writer_rank |
| 174 | ) |
| 175 | |
| 176 | if not message_queue._is_writer: |
| 177 | # Put into busy mode |
| 178 | message_queue._spin_condition.busy_loop_s = 9999 |
| 179 | |
| 180 | shutdown_event = threading.Event() |
| 181 | |
| 182 | def shutdown_thread(mq, shutdown_event): |
| 183 | shutdown_event.wait() |
| 184 | mq.shutdown() |
| 185 | |
| 186 | threading.Thread( |
| 187 | target=shutdown_thread, args=(message_queue, shutdown_event) |
| 188 | ).start() |
| 189 | |
| 190 | with pytest.raises(TimeoutError): |
| 191 | message_queue.dequeue(timeout=0.01) |
| 192 | |
| 193 | shutdown_event.set() |
| 194 | |
| 195 | with pytest.raises(RuntimeError, match="cancelled"): |
| 196 | message_queue.dequeue(timeout=1) |
| 197 | |
| 198 | assert message_queue.shutting_down |
| 199 | |
| 200 | print(f"torch distributed passed the test! Rank {rank}") |
| 201 | dist.barrier() |
| 202 | |
| 203 | |
| 204 | def test_message_queue_shutdown_busy(caplog_vllm): |
| 205 | distributed_run(worker_fn_test_shutdown_busy, 4) |
| 206 | print(caplog_vllm.text) |
| 207 | |
| 208 | |
| 209 | @worker_fn_wrapper |
| 210 | def worker_fn_test_shutdown_idle(): |
| 211 | rank = dist.get_rank() |
| 212 | writer_rank = 2 |
| 213 | message_queue = MessageQueue.create_from_process_group( |
| 214 | dist.group.WORLD, 40 * 1024, 2, writer_rank |
| 215 | ) |
| 216 | |
| 217 | if not message_queue._is_writer: |
| 218 | # Put into idle mode |
| 219 | message_queue._spin_condition.last_read = 0 |
| 220 | |
| 221 | shutdown_event = threading.Event() |
| 222 | |
| 223 | def shutdown_thread(mq, shutdown_event): |
| 224 | shutdown_event.wait() |
| 225 | mq.shutdown() |
| 226 | |
| 227 | threading.Thread( |
| 228 | target=shutdown_thread, args=(message_queue, shutdown_event) |
| 229 | ).start() |
| 230 | |
| 231 | with pytest.raises(TimeoutError): |
| 232 | message_queue.dequeue(timeout=0.01) |
| 233 | |
| 234 | shutdown_event.set() |
| 235 | |
| 236 | with pytest.raises(RuntimeError, match="cancelled"): |
| 237 | message_queue.dequeue(timeout=1) |
| 238 | |
| 239 | assert message_queue.shutting_down |
| 240 | |
| 241 | print(f"torch distributed passed the test! Rank {rank}") |
| 242 | dist.barrier() |
| 243 | |
| 244 | |
| 245 | def test_message_queue_shutdown_idle(): |
| 246 | distributed_run(worker_fn_test_shutdown_idle, 4) |
| 247 | |
| 248 | |
| 249 | @worker_fn_wrapper |
| 250 | def worker_fn_test_idle_to_busy(): |
| 251 | rank = dist.get_rank() |
| 252 | writer_rank = 2 |
| 253 | message_queue = MessageQueue.create_from_process_group( |
| 254 | dist.group.WORLD, 40 * 1024, 2, writer_rank |
| 255 | ) |
| 256 | |
| 257 | message1 = "hello world" |
| 258 | message2 = np.random.randint(1, 100, 100) |
| 259 | with mock.patch.object( |
| 260 | message_queue._spin_condition, "wait", wraps=message_queue._spin_condition.wait |
| 261 | ) as wrapped_wait: |
| 262 | if not message_queue._is_writer: |
| 263 | # Put into idle mode |
| 264 | message_queue._spin_condition.last_read = 0 |
| 265 | |
| 266 | # no messages, so expect a TimeoutError |
| 267 | with pytest.raises(TimeoutError): |
| 268 | message_queue.dequeue(timeout=0.01) |
| 269 | # wait should only be called once while idle |
| 270 | assert wrapped_wait.call_count == 1 |
| 271 | |
| 272 | # sync with the writer and wait for message1 |
| 273 | dist.barrier() |
| 274 | recv_message = message_queue.dequeue(timeout=5) |
| 275 | assert recv_message == message1 |
| 276 | # second call to wait, with a message read, this puts in a busy spin |
| 277 | assert wrapped_wait.call_count == 2 |
| 278 | |
| 279 | # sync with the writer and wait for message2 |
| 280 | dist.barrier() |
| 281 | recv_message = message_queue.dequeue(timeout=1) |
| 282 | assert np.array_equal(recv_message, message2) |
| 283 | # in busy mode, we expect wait to have been called multiple times |
| 284 | assert wrapped_wait.call_count > 3 |
| 285 | else: |
| 286 | # writer writes two messages in sync with the reader |
| 287 | dist.barrier() |
| 288 | # sleep delays the send to ensure reader enters the read loop |
| 289 | time.sleep(0.1) |
| 290 | message_queue.enqueue(message1) |
| 291 | |
| 292 | dist.barrier() |
| 293 | time.sleep(0.1) |
| 294 | message_queue.enqueue(message2) |
| 295 | |
| 296 | message_queue.shutdown() |
| 297 | assert message_queue.shutting_down |
| 298 | print(f"torch distributed passed the test! Rank {rank}") |
| 299 | |
| 300 | |
| 301 | def test_message_queue_idle_wake(): |
| 302 | distributed_run(worker_fn_test_idle_to_busy, 4) |
| 303 | |
| 304 | |
| 305 | @worker_fn_wrapper |
| 306 | def worker_fn_test_busy_to_idle(): |
| 307 | rank = dist.get_rank() |
| 308 | writer_rank = 2 |
| 309 | message_queue = MessageQueue.create_from_process_group( |
| 310 | dist.group.WORLD, 40 * 1024, 2, writer_rank |
| 311 | ) |
| 312 | |
| 313 | message1 = 12345 |
| 314 | message2 = list(range(3)) |
| 315 | with mock.patch.object( |
| 316 | message_queue._spin_condition, "wait", wraps=message_queue._spin_condition.wait |
| 317 | ) as wrapped_wait: |
| 318 | if not message_queue._is_writer: |
| 319 | # Put into busy mode |
| 320 | message_queue._spin_condition.busy_loop_s = 9999 |
| 321 | |
| 322 | # sync with the writer and wait for message1 |
| 323 | dist.barrier() |
| 324 | recv_message = message_queue.dequeue(timeout=1) |
| 325 | assert recv_message == message1 |
| 326 | # in busy mode, we expect wait to have been called many times |
| 327 | assert wrapped_wait.call_count > 1 |
| 328 | |
| 329 | # simulate busy loop ending |
| 330 | message_queue._spin_condition.busy_loop_s = 0 |
| 331 | # ensure we enter idle mode, then record call count |
| 332 | with pytest.raises(TimeoutError): |
| 333 | message_queue.dequeue(timeout=0.01) |
| 334 | call_count = wrapped_wait.call_count |
| 335 | |
| 336 | # sync with the writer and wait for message2 |
| 337 | dist.barrier() |
| 338 | recv_message = message_queue.dequeue(timeout=1) |
| 339 | assert recv_message == message2 |
| 340 | |
| 341 | # call to wait after idle should only happen once |
| 342 | assert wrapped_wait.call_count == call_count + 1 |
| 343 | else: |
| 344 | # writer writes two messages in sync with the reader |
| 345 | dist.barrier() |
| 346 | # sleep delays the send to ensure reader enters the read loop |
| 347 | time.sleep(0.1) |
| 348 | message_queue.enqueue(message1) |
| 349 | |
| 350 | dist.barrier() |
| 351 | time.sleep(0.1) |
| 352 | message_queue.enqueue(message2) |
| 353 | |
| 354 | message_queue.shutdown() |
| 355 | assert message_queue.shutting_down |
| 356 | print(f"torch distributed passed the test! Rank {rank}") |
| 357 | |
| 358 | |
| 359 | def test_message_queue_busy_to_idle(): |
| 360 | distributed_run(worker_fn_test_busy_to_idle, 4) |
| 361 | |
| 362 | |
| 363 | @worker_fn_wrapper |
| 364 | def worker_fn_tensor_broadcast(): |
| 365 | rank = dist.get_rank() |
| 366 | writer_rank = 0 |
| 367 | message_queue = MessageQueue.create_from_process_group( |
| 368 | dist.group.WORLD, 8 * 1024 * 1024, 4, writer_rank |
| 369 | ) |
| 370 | |
| 371 | # Both ranks construct the identical reference payload. |
| 372 | torch.manual_seed(42) |
| 373 | payload = { |
| 374 | # 2MiB: rides the shm ring as an out-of-band buffer (the receiving |
| 375 | # side must copy out of the reusable ring chunk). |
| 376 | "mid": torch.randn(1024, 512), |
| 377 | # 16MiB > max_chunk_bytes: overflows to the zmq socket (the |
| 378 | # receiving side aliases the zmq.Frame zero-copy). |
| 379 | "big": torch.randn(4096, 2048, dtype=torch.bfloat16), |
| 380 | "nested": ["plain", 123, {"inner": torch.arange(5)}], |
| 381 | } |
| 382 | |
| 383 | if rank == writer_rank: |
| 384 | with mock.patch( |
| 385 | "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor", |
| 386 | wraps=_reduce_tensor, |
| 387 | ) as wrapped_reduce: |
| 388 | message_queue.enqueue(payload) |
| 389 | assert wrapped_reduce.call_count == 3 |
| 390 | # Cycle the ring (max_chunks=4) several times over so that aliased |
| 391 | # ring chunks would be overwritten. |
| 392 | for i in range(16): |
| 393 | message_queue.enqueue({"junk": torch.full((1024, 512), float(i))}) |
| 394 | else: |
| 395 | received = message_queue.dequeue(timeout=30) |
| 396 | for key in ("mid", "big"): |
| 397 | assert torch.equal(received[key], payload[key]), key |
| 398 | assert received[key].dtype == payload[key].dtype, key |
| 399 | assert torch.equal(received["nested"][2]["inner"], torch.arange(5)) |
| 400 | |
| 401 | snapshot = received["mid"].clone() |
| 402 | for i in range(16): |
| 403 | junk = message_queue.dequeue(timeout=30) |
| 404 | assert torch.equal(junk["junk"], torch.full((1024, 512), float(i))) |
| 405 | # Tensors received via the shm ring must not alias chunk memory |
| 406 | # that the writer has reused for subsequent messages. |
| 407 | assert torch.equal(received["mid"], snapshot) |
| 408 | # Rebuilt tensors must be writable, like regular tensors. |
| 409 | received["mid"] += 1.0 |
| 410 | received["big"][0, 0] = 1.0 |
| 411 | |
| 412 | dist.barrier() |
| 413 | print(f"tensor broadcast passed the test! Rank {rank}") |
| 414 | |
| 415 | |
| 416 | def test_tensor_broadcast(): |
| 417 | distributed_run(worker_fn_tensor_broadcast, 2) |
| 418 | |
| 419 | |
| 420 | def _dumps_oob(obj) -> tuple[bytes, list]: |
| 421 | """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor |
| 422 | dispatch table + out-of-band buffers >= 1MiB.""" |
| 423 | buffers = [] |
| 424 | |
| 425 | def callback(buf: pickle.PickleBuffer) -> bool: |
| 426 | raw = buf.raw() |
| 427 | if raw.nbytes < 1024 * 1024: |
| 428 | return True |
| 429 | buffers.append(raw) |
| 430 | return False |
| 431 | |
| 432 | bio = io.BytesIO() |
| 433 | pickler = pickle.Pickler( |
| 434 | bio, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=callback |
| 435 | ) |
| 436 | pickler.dispatch_table = {torch.Tensor: _reduce_tensor} |
| 437 | pickler.dump(obj) |
| 438 | return bio.getvalue(), buffers |
| 439 | |
| 440 | |
| 441 | @pytest.mark.parametrize( |
| 442 | "case", |
| 443 | [ |
| 444 | "small", |
| 445 | "mid", |
| 446 | "bf16", |
| 447 | "fp8", |
| 448 | "empty", |
| 449 | "scalar", |
| 450 | "noncontig", |
| 451 | "requires_grad", |
| 452 | "conj", |
| 453 | "param", |
| 454 | ], |
| 455 | ) |
| 456 | def test_tensor_pickle_roundtrip(case: str): |
| 457 | tensor = { |
| 458 | # Inlined in-band (< 1MiB) and out-of-band (>= 1MiB) buffers. |
| 459 | "small": lambda: torch.randn(100, 10), |
| 460 | "mid": lambda: torch.randn(1024, 512), |
| 461 | # Dtypes numpy doesn't recognize. |
| 462 | "bf16": lambda: torch.randn(512, 512, dtype=torch.bfloat16), |
| 463 | "fp8": lambda: torch.randn(32, 32).to(torch.float8_e4m3fn), |
| 464 | # Shape edge cases. |
| 465 | "empty": lambda: torch.empty(0, 8), |
| 466 | "scalar": lambda: torch.tensor(3.14), |
| 467 | "noncontig": lambda: torch.randn(64, 64).t(), |
| 468 | # These fall back to torch's default reducer. |
| 469 | "requires_grad": lambda: torch.randn(8, 8, requires_grad=True), |
| 470 | "conj": lambda: torch.randn(4, dtype=torch.complex64).conj(), |
| 471 | "param": lambda: torch.nn.Parameter(torch.randn(4), requires_grad=False), |
| 472 | }[case]() |
| 473 | |
| 474 | data, buffers = _dumps_oob({"tensor": tensor, "meta": list(range(10))}) |
| 475 | received = pickle.loads(data, buffers=buffers)["tensor"] |
| 476 | |
| 477 | assert received.shape == tensor.shape |
| 478 | assert received.dtype == tensor.dtype |
| 479 | if tensor.dtype == torch.float8_e4m3fn: |
| 480 | assert torch.equal(received.view(torch.uint8), tensor.view(torch.uint8)) |
| 481 | else: |
| 482 | assert torch.equal(received, tensor) |
| 483 | assert received.requires_grad == tensor.requires_grad |
| 484 | assert isinstance(received, type(tensor)) |
| 485 | if tensor.numel() and not tensor.requires_grad: |
| 486 | # Rebuilt tensors must be writable, like regular tensors. |
| 487 | received.view(-1)[0] = 1.0 |
| 488 | |
| 489 | |
| 490 | @pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) |
| 491 | def test_reduce_tensor_fallback(case: str): |
| 492 | """Tensors the zero-copy reducer can't safely alias must fall back to |
| 493 | torch's default reduction.""" |
| 494 | if case == "cuda": |
| 495 | if not torch.cuda.is_available(): |
| 496 | pytest.skip("requires CUDA") |
| 497 | tensor = torch.randn(4, device="cuda") |
| 498 | elif case == "requires_grad": |
| 499 | tensor = torch.randn(8, requires_grad=True) |
| 500 | else: |
| 501 | tensor = torch.randn(4, dtype=torch.complex64).conj() |
| 502 | |
| 503 | reduced = _reduce_tensor(tensor) |
| 504 | assert reduced[0] is not _rebuild_tensor |
| 505 | |
| 506 | |
| 507 | @pytest.mark.parametrize("should_warn", [False, True]) |
| 508 | def test_reader_timeout_caps_indefinite_waits(should_warn): |
| 509 | with ( |
| 510 | mock.patch( |
| 511 | "vllm.distributed.device_communicators.shm_broadcast." |
| 512 | "SHM_READER_RECHECK_INTERVAL_MS", |
| 513 | new=7, |
| 514 | ), |
| 515 | mock.patch( |
| 516 | "vllm.distributed.device_communicators.shm_broadcast." |
| 517 | "VLLM_RINGBUFFER_WARNING_INTERVAL", |
| 518 | new=60, |
| 519 | ), |
| 520 | ): |
| 521 | timeout = MessageQueue.ReadTimeoutWithWarnings( |
| 522 | timeout=None, should_warn=should_warn |
| 523 | ) |
| 524 | assert timeout.timeout_ms() == 7 |
| 525 | |
| 526 | |
| 527 | def test_reader_rechecks_shm_after_idle_wait_timeout_without_notify(): |
| 528 | writer = MessageQueue( |
| 529 | n_reader=1, |
| 530 | n_local_reader=1, |
| 531 | max_chunk_bytes=1024 * 1024, |
| 532 | max_chunks=1, |
| 533 | ) |
| 534 | reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) |
| 535 | payload = 123 |
| 536 | poll_started = threading.Event() |
| 537 | allow_timeout = threading.Event() |
| 538 | result = {} |
| 539 | |
| 540 | def acquire_read_in_thread(): |
| 541 | try: |
| 542 | with reader.acquire_read(indefinite=True) as buf: |
| 543 | result["value"] = buf[0] |
| 544 | except Exception as exc: |
| 545 | result["exc"] = exc |
| 546 | |
| 547 | def poll_timeout(*, timeout: int | None = None): |
| 548 | poll_started.set() |
| 549 | assert allow_timeout.wait(timeout=5) |
| 550 | return [] |
| 551 | |
| 552 | try: |
| 553 | writer.wait_until_ready() |
| 554 | reader.wait_until_ready() |
| 555 | reader._spin_condition.last_read = 0 |
| 556 | reader._spin_condition.busy_loop_s = 0 |
| 557 | |
| 558 | with ( |
| 559 | mock.patch( |
| 560 | "vllm.distributed.device_communicators.shm_broadcast." |
| 561 | "SHM_READER_RECHECK_INTERVAL_MS", |
| 562 | new=50, |
| 563 | ), |
| 564 | mock.patch( |
| 565 | "vllm.distributed.device_communicators.shm_broadcast." |
| 566 | "VLLM_RINGBUFFER_WARNING_INTERVAL", |
| 567 | new=60, |
| 568 | ), |
| 569 | mock.patch.object( |
| 570 | reader._spin_condition.poller, |
| 571 | "poll", |
| 572 | side_effect=poll_timeout, |
| 573 | ) as poll, |
| 574 | ): |
| 575 | read_thread = threading.Thread(target=acquire_read_in_thread, daemon=True) |
| 576 | read_thread.start() |
| 577 | assert poll_started.wait(timeout=5) |
| 578 | with writer.acquire_write(timeout=0.1) as buf: |
| 579 | buf[0] = payload |
| 580 | allow_timeout.set() |
| 581 | read_thread.join(timeout=5) |
| 582 | |
| 583 | assert not read_thread.is_alive() |
| 584 | poll.assert_called_once_with(timeout=50) |
| 585 | |
| 586 | if "exc" in result: |
| 587 | raise result["exc"] |
| 588 | assert result["value"] == payload |
| 589 | with writer.buffer.get_metadata(0) as metadata_buffer: |
| 590 | assert metadata_buffer[0] == 1 |
| 591 | assert metadata_buffer[1] == 1 |
| 592 | finally: |
| 593 | writer.shutdown() |
| 594 | reader.shutdown() |
| 595 | for socket in ( |
| 596 | writer.local_socket, |
| 597 | writer._spin_condition.local_notify_socket, |
| 598 | reader.local_socket, |
| 599 | reader._spin_condition.local_notify_socket, |
| 600 | reader._spin_condition.read_cancel_socket, |
| 601 | reader._spin_condition.write_cancel_socket, |
| 602 | ): |
| 603 | socket.close(linger=0) |
| 604 | |
| 605 | |
| 606 | def test_acquire_read_releases_slot_when_reader_raises(): |
| 607 | writer = MessageQueue( |
| 608 | n_reader=1, |
| 609 | n_local_reader=1, |
| 610 | max_chunk_bytes=1024 * 1024, |
| 611 | max_chunks=1, |
| 612 | ) |
| 613 | reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) |
| 614 | try: |
| 615 | writer.wait_until_ready() |
| 616 | reader.wait_until_ready() |
| 617 | |
| 618 | writer.enqueue({"payload": "first"}) |
| 619 | |
| 620 | with ( |
| 621 | pytest.raises(RuntimeError, match="reader failed"), |
| 622 | reader.acquire_read(timeout=0.1), |
| 623 | ): |
| 624 | raise RuntimeError("reader failed") |
| 625 | |
| 626 | with writer.buffer.get_metadata(0) as metadata_buffer: |
| 627 | assert metadata_buffer[0] == 1 |
| 628 | assert metadata_buffer[1] == 1 |
| 629 | |
| 630 | with writer.acquire_write(timeout=0.1) as buf: |
| 631 | buf[0] = 0 |
| 632 | finally: |
| 633 | writer.shutdown() |
| 634 | reader.shutdown() |
| 635 | |
| 636 | |
| 637 | def test_warning_logs(caplog_vllm): |
| 638 | """Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals |
| 639 | when indefinite=False, and are not emitted when indefinite=True. |
| 640 | """ |
| 641 | # Patch the warning log interval to every 1 ms during reads |
| 642 | with mock.patch( |
| 643 | "vllm.distributed.device_communicators.shm_broadcast.VLLM_RINGBUFFER_WARNING_INTERVAL", |
| 644 | new=0.001, # 1 ms |
| 645 | ): |
| 646 | writer = MessageQueue( |
| 647 | n_reader=1, |
| 648 | n_local_reader=1, |
| 649 | max_chunk_bytes=1024 * 1024, # 1MB chunks |
| 650 | max_chunks=10, |
| 651 | ) |
| 652 | reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) |
| 653 | writer.wait_until_ready() |
| 654 | reader.wait_until_ready() |
| 655 | |
| 656 | # We should have at least one warning log here |
| 657 | # "0 seconds" expected due to rounding of 1ms test interval |
| 658 | with pytest.raises(TimeoutError): |
| 659 | reader.dequeue(timeout=0.01, indefinite=False) |
| 660 | assert any( |
| 661 | "No available shared memory broadcast block found in 0 seconds" |
| 662 | in record.message |
| 663 | for record in caplog_vllm.records |
| 664 | ) |
| 665 | caplog_vllm.clear() |
| 666 | |
| 667 | # We should have no warnings this time |
| 668 | with pytest.raises(TimeoutError): |
| 669 | reader.dequeue(timeout=0.01, indefinite=True) |
| 670 | assert all( |
| 671 | "No available shared memory broadcast block found in 0 seconds" |
| 672 | not in record.message |
| 673 | for record in caplog_vllm.records |
| 674 | ) |
| 675 | |
| 676 | # Clean up when done |
| 677 | writer.shutdown() |
| 678 | reader.shutdown() |
| 679 | |
| 680 | |
| 681 | def _fake_disk_usage(free_bytes: int): |
| 682 | return SimpleNamespace(total=free_bytes, used=0, free=free_bytes) |
| 683 | |
| 684 | |
| 685 | def test_check_shm_free_space_raises_when_insufficient(tmp_path): |
| 686 | with ( |
| 687 | mock.patch.object( |
| 688 | shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(32 << 20) |
| 689 | ), |
| 690 | pytest.raises(RuntimeError, match="Insufficient space"), |
| 691 | ): |
| 692 | check_shm_free_space(240 << 20, shm_path=str(tmp_path)) |
| 693 | |
| 694 | |
| 695 | def test_check_shm_free_space_passes_when_sufficient(tmp_path): |
| 696 | with ( |
| 697 | mock.patch.object( |
| 698 | shm_broadcast.shutil, |
| 699 | "disk_usage", |
| 700 | return_value=_fake_disk_usage(512 << 20), |
| 701 | ), |
| 702 | mock.patch.object(shm_broadcast, "check_cgroup_memory_available"), |
| 703 | ): |
| 704 | check_shm_free_space(240 << 20, shm_path=str(tmp_path)) |
| 705 | |
| 706 | |
| 707 | def test_check_shm_free_space_skipped_when_path_missing(tmp_path): |
| 708 | with mock.patch.object(shm_broadcast, "check_cgroup_memory_available"): |
| 709 | check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist")) |
| 710 | |
| 711 | |
| 712 | def test_check_shm_free_space_checks_cgroup(tmp_path): |
| 713 | with ( |
| 714 | mock.patch.object( |
| 715 | shm_broadcast.shutil, |
| 716 | "disk_usage", |
| 717 | return_value=_fake_disk_usage(512 << 20), |
| 718 | ), |
| 719 | mock.patch.object( |
| 720 | shm_broadcast, "check_cgroup_memory_available" |
| 721 | ) as check_cgroup, |
| 722 | ): |
| 723 | check_shm_free_space( |
| 724 | 240 << 20, |
| 725 | shm_path=str(tmp_path), |
| 726 | allocation_name="SHM mmap", |
| 727 | ) |
| 728 | |
| 729 | check_cgroup.assert_called_once_with( |
| 730 | 240 << 20, |
| 731 | "SHM mmap", |
| 732 | ) |
| 733 | |
| 734 | |
| 735 | def test_shm_ring_buffer_creation_checks_free_space(): |
| 736 | with ( |
| 737 | mock.patch.object( |
| 738 | shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(1 << 20) |
| 739 | ), |
| 740 | mock.patch.object(shm_broadcast.os.path, "isdir", return_value=True), |
| 741 | pytest.raises(RuntimeError, match="Insufficient space"), |
| 742 | ): |
| 743 | ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10) |
| 744 | |
| 745 | |
| 746 | def test_remote_subscribe_addr_unique_concurrent_writers( |
| 747 | monkeypatch: pytest.MonkeyPatch, |
| 748 | ): |
| 749 | """Writers bind the remote socket to port 0 (kernel-assigned), so |
| 750 | concurrent writers never race for the same probed port and the |
| 751 | announced address is connectable. |
| 752 | |
| 753 | Pre-fix, the writer probed a port with get_open_port() and bound it |
| 754 | afterwards; pinning the probe to one free port makes every writer |
| 755 | bind the same port and fail deterministically on that code path, |
| 756 | while the late-binding implementation never consults the probe.""" |
| 757 | from vllm.distributed.device_communicators import shm_broadcast |
| 758 | |
| 759 | colliding_port = get_open_port() |
| 760 | monkeypatch.setattr( |
| 761 | shm_broadcast, "get_open_port", lambda: colliding_port, raising=False |
| 762 | ) |
| 763 | |
| 764 | n_writers = 32 |
| 765 | queues: list[MessageQueue] = [] |
| 766 | lock = threading.Lock() |
| 767 | |
| 768 | def make_writer(): |
| 769 | q = MessageQueue( |
| 770 | n_reader=1, |
| 771 | n_local_reader=0, |
| 772 | max_chunk_bytes=4096, |
| 773 | max_chunks=2, |
| 774 | connect_ip="127.0.0.1", |
| 775 | ) |
| 776 | with lock: |
| 777 | queues.append(q) |
| 778 | |
| 779 | threads = [threading.Thread(target=make_writer) for _ in range(n_writers)] |
| 780 | for t in threads: |
| 781 | t.start() |
| 782 | for t in threads: |
| 783 | t.join() |
| 784 | |
| 785 | assert len(queues) == n_writers |
| 786 | addrs = [q.export_handle().remote_subscribe_addr for q in queues] |
| 787 | assert all(addr and addr.startswith("tcp://") for addr in addrs) |
| 788 | assert len(set(addrs)) == n_writers |
| 789 | |
| 790 | writer = queues[0] |
| 791 | received = [] |
| 792 | |
| 793 | def reader_main(): |
| 794 | reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) |
| 795 | reader.wait_until_ready() |
| 796 | received.append(reader.dequeue()) |
| 797 | reader.remote_socket.close(linger=0) |
| 798 | |
| 799 | reader_thread = threading.Thread(target=reader_main) |
| 800 | reader_thread.start() |
| 801 | writer.wait_until_ready() |
| 802 | writer.enqueue("ping") |
| 803 | reader_thread.join(timeout=30) |
| 804 | assert not reader_thread.is_alive() |
| 805 | assert received == ["ping"] |
| 806 | |
| 807 | for q in queues: |
| 808 | q.remote_socket.close(linger=0) |