| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import threading |
| 5 | import time |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream |
| 11 | from vllm.platforms import current_platform |
| 12 | |
| 13 | |
| 14 | def test_wait_blocks_until_record(): |
| 15 | event = CpuGpuEvent() |
| 16 | record_stream = torch.Stream() |
| 17 | wait_stream = torch.Stream() |
| 18 | wait_returned = threading.Event() |
| 19 | |
| 20 | def waiter(): |
| 21 | event.wait(stream=wait_stream) |
| 22 | wait_returned.set() |
| 23 | |
| 24 | t = threading.Thread(target=waiter) |
| 25 | t.start() |
| 26 | |
| 27 | time.sleep(0.05) |
| 28 | assert not wait_returned.is_set(), "wait() returned before record() was called" |
| 29 | |
| 30 | event.record(stream=record_stream) |
| 31 | t.join(timeout=5.0) |
| 32 | |
| 33 | assert not event._recorded.is_set() |
| 34 | |
| 35 | |
| 36 | def test_reuse_across_multiple_cycles(): |
| 37 | wrapper = CpuGpuEvent() |
| 38 | record_stream = torch.Stream() |
| 39 | wait_stream = torch.Stream() |
| 40 | NUM_CYCLES = 8 |
| 41 | completed_cycles = [] |
| 42 | barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)] |
| 43 | |
| 44 | def waiter(): |
| 45 | for i in range(NUM_CYCLES): |
| 46 | wrapper.wait(stream=wait_stream) |
| 47 | completed_cycles.append(True) |
| 48 | barriers[i].wait() |
| 49 | |
| 50 | t = threading.Thread(target=waiter) |
| 51 | t.start() |
| 52 | |
| 53 | for i in range(NUM_CYCLES): |
| 54 | wrapper.record(stream=record_stream) |
| 55 | barriers[i].wait() |
| 56 | |
| 57 | t.join(timeout=10.0) |
| 58 | assert len(completed_cycles) == NUM_CYCLES |
| 59 | |
| 60 | |
| 61 | @pytest.mark.parametrize("device", ["cuda", "xpu"]) |
| 62 | def test_producer_consumer(device): |
| 63 | """This test uses the CpuGpuEvent to synchronize reads and writes to/from |
| 64 | a shared GPU tensor on multiple CPU threads. |
| 65 | """ |
| 66 | if device == "cuda" and not torch.cuda.is_available(): |
| 67 | pytest.skip("CUDA not available") |
| 68 | if device == "xpu" and not current_platform.is_xpu(): |
| 69 | pytest.skip("XPU not available") |
| 70 | worker_stream = torch.Stream() |
| 71 | # Create a single element counter that will be shared between two threads |
| 72 | buf = torch.zeros(1, device=device) |
| 73 | NUM_ROUNDS = 5 |
| 74 | |
| 75 | ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)] |
| 76 | events = [CpuGpuEvent() for _ in range(NUM_ROUNDS)] |
| 77 | errors: list[str] = [] |
| 78 | |
| 79 | # For each round, the worker thread (writer) sets the counter in buf and waits for |
| 80 | # the main thread to read it. |
| 81 | def worker(): |
| 82 | for i in range(NUM_ROUNDS): |
| 83 | if i > 0: |
| 84 | events[i - 1].wait(stream=worker_stream) |
| 85 | |
| 86 | with device_stream(worker_stream): |
| 87 | buf.fill_(float(i + 1)) |
| 88 | |
| 89 | worker_stream.synchronize() |
| 90 | ready_cpu[i].set() |
| 91 | |
| 92 | t = threading.Thread(target=worker) |
| 93 | t.start() |
| 94 | |
| 95 | for i in range(NUM_ROUNDS): |
| 96 | ready_cpu[i].wait() |
| 97 | snapshot = buf.clone() |
| 98 | events[i].record() |
| 99 | val = snapshot.item() |
| 100 | if val != float(i + 1): |
| 101 | errors.append(f"round {i}: expected {i + 1:.1f}, got {val:.1f}") |
| 102 | |
| 103 | t.join(timeout=10.0) |
| 104 | assert not errors, f"Buffer ordering errors: {errors}" |