| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.config |
| 9 | import vllm.distributed.parallel_state as parallel_state |
| 10 | from vllm.distributed.parallel_state import ( |
| 11 | destroy_distributed_environment, |
| 12 | init_distributed_environment, |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | def _file_store_worker(rank, init_method, result_queue): |
| 17 | try: |
| 18 | init_distributed_environment( |
| 19 | world_size=2, |
| 20 | rank=rank, |
| 21 | local_rank=rank, |
| 22 | backend="gloo", |
| 23 | distributed_init_method=init_method, |
| 24 | ) |
| 25 | result = torch.tensor(rank) |
| 26 | torch.distributed.all_reduce(result) |
| 27 | result_queue.put(result.item()) |
| 28 | finally: |
| 29 | destroy_distributed_environment() |
| 30 | |
| 31 | |
| 32 | def test_file_store_preserves_raw_path(tmp_path, monkeypatch): |
| 33 | monkeypatch.setenv("VLLM_DISTRIBUTED_USE_SPLIT_GROUP", "0") |
| 34 | store_path = tmp_path / "store#?" |
| 35 | |
| 36 | try: |
| 37 | init_distributed_environment( |
| 38 | world_size=1, |
| 39 | rank=0, |
| 40 | local_rank=0, |
| 41 | backend="gloo", |
| 42 | distributed_init_method=f"file://{store_path}", |
| 43 | ) |
| 44 | assert store_path.exists() |
| 45 | finally: |
| 46 | destroy_distributed_environment() |
| 47 | |
| 48 | |
| 49 | def test_file_store_rendezvouses_local_processes(tmp_path, monkeypatch): |
| 50 | monkeypatch.setenv("VLLM_DISTRIBUTED_USE_SPLIT_GROUP", "0") |
| 51 | context = torch.multiprocessing.get_context("spawn") |
| 52 | result_queue = context.SimpleQueue() |
| 53 | |
| 54 | torch.multiprocessing.spawn( |
| 55 | _file_store_worker, |
| 56 | args=(f"file://{tmp_path / 'store#?'}", result_queue), |
| 57 | nprocs=2, |
| 58 | join=True, |
| 59 | ) |
| 60 | |
| 61 | assert sorted(result_queue.get() for _ in range(2)) == [1, 1] |
| 62 | |
| 63 | |
| 64 | def test_multinode_overrides_file_store_rendezvous(monkeypatch): |
| 65 | monkeypatch.setenv("VLLM_DISTRIBUTED_USE_SPLIT_GROUP", "0") |
| 66 | parallel_config = SimpleNamespace( |
| 67 | distributed_executor_backend="mp", |
| 68 | nnodes=2, |
| 69 | data_parallel_size=1, |
| 70 | enable_elastic_ep=False, |
| 71 | data_parallel_rank=0, |
| 72 | world_size_across_dp=2, |
| 73 | master_addr="192.0.2.1", |
| 74 | master_port=29500, |
| 75 | nnodes_within_dp=2, |
| 76 | ) |
| 77 | config = SimpleNamespace(parallel_config=parallel_config) |
| 78 | calls = [] |
| 79 | original_world = parallel_state._WORLD |
| 80 | original_node_count = parallel_state._NODE_COUNT |
| 81 | |
| 82 | monkeypatch.setattr(vllm.config, "get_current_vllm_config_or_none", lambda: config) |
| 83 | monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) |
| 84 | monkeypatch.setattr(torch.distributed, "is_backend_available", lambda _: True) |
| 85 | monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) |
| 86 | monkeypatch.setattr( |
| 87 | torch.distributed, "init_process_group", lambda **kwargs: calls.append(kwargs) |
| 88 | ) |
| 89 | monkeypatch.setattr( |
| 90 | parallel_state, |
| 91 | "init_world_group", |
| 92 | lambda *_: SimpleNamespace(world_size=2, cpu_group=None), |
| 93 | ) |
| 94 | |
| 95 | try: |
| 96 | init_distributed_environment( |
| 97 | world_size=2, |
| 98 | rank=0, |
| 99 | local_rank=0, |
| 100 | backend="gloo", |
| 101 | distributed_init_method="file:///tmp/vllm_dist_test", |
| 102 | ) |
| 103 | finally: |
| 104 | parallel_state._WORLD = original_world |
| 105 | parallel_state._NODE_COUNT = original_node_count |
| 106 | |
| 107 | assert calls == [ |
| 108 | { |
| 109 | "backend": "gloo", |
| 110 | "init_method": "tcp://192.0.2.1:29500", |
| 111 | "store": None, |
| 112 | "world_size": 2, |
| 113 | "rank": 0, |
| 114 | "timeout": None, |
| 115 | } |
| 116 | ] |