test_utils.py 4.23 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
import socket

6
import pytest
7
import ray
8
import torch
9

10
import vllm.envs as envs
11
12
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
from vllm.distributed.utils import StatelessProcessGroup
13
14
from vllm.utils import update_environment_variables
from vllm.utils.network_utils import get_open_port
15
from vllm.utils.torch_utils import cuda_device_count_stateless
16

17
18
from ..utils import multi_gpu_test

19
20

@ray.remote
21
class _CUDADeviceCountStatelessTestActor:
22
23
24
25
    def get_count(self):
        return cuda_device_count_stateless()

    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
33
34


def test_cuda_device_count_stateless():
    """Test that cuda_device_count_stateless changes return value if
    CUDA_VISIBLE_DEVICES is changed."""
35
    actor = _CUDADeviceCountStatelessTestActor.options(  # type: ignore
36
37
38
        num_gpus=2
    ).remote()
    assert len(sorted(ray.get(actor.get_cuda_visible_devices.remote()).split(","))) == 2
39
40
41
42
43
    assert ray.get(actor.get_count.remote()) == 2
    ray.get(actor.set_cuda_visible_devices.remote("0"))
    assert ray.get(actor.get_count.remote()) == 1
    ray.get(actor.set_cuda_visible_devices.remote(""))
    assert ray.get(actor.get_count.remote()) == 0
44
45


46
def cpu_worker(rank, WORLD_SIZE, port1, port2):
47
48
49
    pg1 = StatelessProcessGroup.create(
        host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
    )
50
    if rank <= 2:
51
52
53
        pg2 = StatelessProcessGroup.create(
            host="127.0.0.1", port=port2, rank=rank, world_size=3
        )
54
    data = torch.tensor([rank])
55
56
    data = pg1.broadcast_obj(data, src=2)
    assert data.item() == 2
57
    if rank <= 2:
58
59
60
61
62
        data = torch.tensor([rank + 1])
        data = pg2.broadcast_obj(data, src=2)
        assert data.item() == 3
        pg2.barrier()
    pg1.barrier()
63
64


65
def gpu_worker(rank, WORLD_SIZE, port1, port2):
66
    torch.cuda.set_device(rank)
67
68
69
    pg1 = StatelessProcessGroup.create(
        host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
    )
70
    pynccl1 = PyNcclCommunicator(pg1, device=rank)
71
    if rank <= 2:
72
73
74
        pg2 = StatelessProcessGroup.create(
            host="127.0.0.1", port=port2, rank=rank, world_size=3
        )
75
        pynccl2 = PyNcclCommunicator(pg2, device=rank)
76
    data = torch.tensor([rank]).cuda()
77
78
79
    pynccl1.all_reduce(data)
    pg1.barrier()
    torch.cuda.synchronize()
80
    if rank <= 2:
81
82
83
        pynccl2.all_reduce(data)
        pg2.barrier()
        torch.cuda.synchronize()
84
85
86
87
88
89
90
91
    item = data[0].item()
    print(f"rank: {rank}, item: {item}")
    if rank == 3:
        assert item == 6
    else:
        assert item == 18


92
def broadcast_worker(rank, WORLD_SIZE, port1, port2):
93
94
95
    pg1 = StatelessProcessGroup.create(
        host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
    )
96
97
98
99
100
101
102
103
    if rank == 2:
        pg1.broadcast_obj("secret", src=2)
    else:
        obj = pg1.broadcast_obj(None, src=2)
        assert obj == "secret"
    pg1.barrier()


104
def allgather_worker(rank, WORLD_SIZE, port1, port2):
105
106
107
    pg1 = StatelessProcessGroup.create(
        host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
    )
108
109
110
111
112
    data = pg1.all_gather_obj(rank)
    assert data == list(range(WORLD_SIZE))
    pg1.barrier()


113
@pytest.mark.skip(reason="This test is flaky and prone to hang.")
114
@multi_gpu_test(num_gpus=4)
115
@pytest.mark.parametrize(
116
117
    "worker", [cpu_worker, gpu_worker, broadcast_worker, allgather_worker]
)
118
def test_stateless_process_group(worker):
119
120
121
122
    port1 = get_open_port()
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("", port1))
        port2 = get_open_port()
123
124
    WORLD_SIZE = 4
    from multiprocessing import get_context
125

126
127
128
129
    ctx = get_context("fork")
    processes = []
    for i in range(WORLD_SIZE):
        rank = i
130
        processes.append(
131
132
            ctx.Process(target=worker, args=(rank, WORLD_SIZE, port1, port2))
        )
133
134
135
136
137
138
139
    for p in processes:
        p.start()
    for p in processes:
        p.join()
    for p in processes:
        assert not p.exitcode
    print("All processes finished.")