| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from ...utils import create_new_process_for_each_test |
| 8 | |
| 9 | |
| 10 | @pytest.mark.parametrize("tp_size", [1, 2]) |
| 11 | @pytest.mark.parametrize("backend", ["mp", "ray"]) |
| 12 | @create_new_process_for_each_test() |
| 13 | def test_collective_rpc(tp_size, backend, monkeypatch, vllm_runner): |
| 14 | if torch.accelerator.device_count() < tp_size: |
| 15 | pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}") |
| 16 | if tp_size == 1 and backend == "ray": |
| 17 | pytest.skip("Skip duplicate test case") |
| 18 | if tp_size == 1: |
| 19 | backend = None |
| 20 | |
| 21 | # intentionally define the method and class in the test function, |
| 22 | # to test if they can be serialized and sent to the workers |
| 23 | def echo_rank(self): |
| 24 | return self.rank |
| 25 | |
| 26 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 27 | with vllm_runner( |
| 28 | "hmellor/tiny-random-LlamaForCausalLM", |
| 29 | enforce_eager=True, |
| 30 | load_format="dummy", |
| 31 | tensor_parallel_size=tp_size, |
| 32 | distributed_executor_backend=backend, |
| 33 | ) as runner: |
| 34 | assert runner.llm.collective_rpc(echo_rank) == list(range(tp_size)) |