| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import socket |
| 5 | |
| 6 | import pytest |
| 7 | import ray |
| 8 | import torch |
| 9 | |
| 10 | import vllm.envs as envs |
| 11 | from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator |
| 12 | from vllm.distributed.utils import StatelessProcessGroup |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.network_utils import get_open_port |
| 15 | from vllm.utils.system_utils import update_environment_variables |
| 16 | |
| 17 | from ..utils import multi_gpu_test |
| 18 | |
| 19 | |
| 20 | @ray.remote |
| 21 | class _CUDADeviceCountStatelessTestActor: |
| 22 | def get_count(self): |
| 23 | return current_platform.device_count() |
| 24 | |
| 25 | def set_cuda_visible_devices(self, cuda_visible_devices: str): |
| 26 | update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices}) |
| 27 | |
| 28 | def get_cuda_visible_devices(self): |
| 29 | return envs.CUDA_VISIBLE_DEVICES |
| 30 | |
| 31 | |
| 32 | def test_cuda_device_count_stateless(): |
| 33 | """Test that cuda_device_count_stateless changes return value if |
| 34 | CUDA_VISIBLE_DEVICES is changed.""" |
| 35 | if current_platform.is_rocm(): |
| 36 | pytest.skip("Skip for ROCm because Ray uses HIP_VISIBLE_DEVICES.") |
| 37 | actor = _CUDADeviceCountStatelessTestActor.options( # type: ignore |
| 38 | num_gpus=2 |
| 39 | ).remote() |
| 40 | assert len(sorted(ray.get(actor.get_cuda_visible_devices.remote()).split(","))) == 2 |
| 41 | assert ray.get(actor.get_count.remote()) == 2 |
| 42 | ray.get(actor.set_cuda_visible_devices.remote("0")) |
| 43 | assert ray.get(actor.get_count.remote()) == 1 |
| 44 | ray.get(actor.set_cuda_visible_devices.remote("")) |
| 45 | assert ray.get(actor.get_count.remote()) == 0 |
| 46 | |
| 47 | |
| 48 | def cpu_worker(rank, WORLD_SIZE, port1, port2): |
| 49 | pg1 = StatelessProcessGroup.create( |
| 50 | host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE |
| 51 | ) |
| 52 | if rank <= 2: |
| 53 | pg2 = StatelessProcessGroup.create( |
| 54 | host="127.0.0.1", port=port2, rank=rank, world_size=3 |
| 55 | ) |
| 56 | data = torch.tensor([rank]) |
| 57 | data = pg1.broadcast_obj(data, src=2) |
| 58 | assert data.item() == 2 |
| 59 | if rank <= 2: |
| 60 | data = torch.tensor([rank + 1]) |
| 61 | data = pg2.broadcast_obj(data, src=2) |
| 62 | assert data.item() == 3 |
| 63 | pg2.barrier() |
| 64 | pg1.barrier() |
| 65 | |
| 66 | |
| 67 | def gpu_worker(rank, WORLD_SIZE, port1, port2): |
| 68 | torch.accelerator.set_device_index(rank) |
| 69 | pg1 = StatelessProcessGroup.create( |
| 70 | host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE |
| 71 | ) |
| 72 | pynccl1 = PyNcclCommunicator(pg1, device=rank) |
| 73 | if rank <= 2: |
| 74 | pg2 = StatelessProcessGroup.create( |
| 75 | host="127.0.0.1", port=port2, rank=rank, world_size=3 |
| 76 | ) |
| 77 | pynccl2 = PyNcclCommunicator(pg2, device=rank) |
| 78 | data = torch.tensor([rank]).cuda() |
| 79 | pynccl1.all_reduce(data) |
| 80 | pg1.barrier() |
| 81 | torch.accelerator.synchronize() |
| 82 | if rank <= 2: |
| 83 | pynccl2.all_reduce(data) |
| 84 | pg2.barrier() |
| 85 | torch.accelerator.synchronize() |
| 86 | item = data[0].item() |
| 87 | print(f"rank: {rank}, item: {item}") |
| 88 | if rank == 3: |
| 89 | assert item == 6 |
| 90 | else: |
| 91 | assert item == 18 |
| 92 | |
| 93 | |
| 94 | def broadcast_worker(rank, WORLD_SIZE, port1, port2): |
| 95 | pg1 = StatelessProcessGroup.create( |
| 96 | host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE |
| 97 | ) |
| 98 | if rank == 2: |
| 99 | pg1.broadcast_obj("secret", src=2) |
| 100 | else: |
| 101 | obj = pg1.broadcast_obj(None, src=2) |
| 102 | assert obj == "secret" |
| 103 | pg1.barrier() |
| 104 | |
| 105 | |
| 106 | def allgather_worker(rank, WORLD_SIZE, port1, port2): |
| 107 | pg1 = StatelessProcessGroup.create( |
| 108 | host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE |
| 109 | ) |
| 110 | data = pg1.all_gather_obj(rank) |
| 111 | assert data == list(range(WORLD_SIZE)) |
| 112 | pg1.barrier() |
| 113 | |
| 114 | |
| 115 | @pytest.mark.skip(reason="This test is flaky and prone to hang.") |
| 116 | @multi_gpu_test(num_gpus=4) |
| 117 | @pytest.mark.parametrize( |
| 118 | "worker", [cpu_worker, gpu_worker, broadcast_worker, allgather_worker] |
| 119 | ) |
| 120 | def test_stateless_process_group(worker): |
| 121 | port1 = get_open_port() |
| 122 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 123 | s.bind(("", port1)) |
| 124 | port2 = get_open_port() |
| 125 | WORLD_SIZE = 4 |
| 126 | from multiprocessing import get_context |
| 127 | |
| 128 | ctx = get_context("fork") |
| 129 | processes = [] |
| 130 | for i in range(WORLD_SIZE): |
| 131 | rank = i |
| 132 | processes.append( |
| 133 | ctx.Process(target=worker, args=(rank, WORLD_SIZE, port1, port2)) |
| 134 | ) |
| 135 | for p in processes: |
| 136 | p.start() |
| 137 | for p in processes: |
| 138 | p.join() |
| 139 | for p in processes: |
| 140 | assert not p.exitcode |
| 141 | print("All processes finished.") |