| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | |
| 6 | import torch.distributed as dist |
| 7 | |
| 8 | from vllm.distributed.parallel_state import _node_count |
| 9 | from vllm.distributed.utils import StatelessProcessGroup |
| 10 | from vllm.utils.network_utils import get_ip, get_open_port |
| 11 | |
| 12 | if __name__ == "__main__": |
| 13 | dist.init_process_group(backend="gloo") |
| 14 | |
| 15 | rank = dist.get_rank() |
| 16 | world_size = dist.get_world_size() |
| 17 | |
| 18 | if rank == 0: |
| 19 | port = get_open_port() |
| 20 | ip = get_ip() |
| 21 | dist.broadcast_object_list([ip, port], src=0) |
| 22 | else: |
| 23 | recv = [None, None] |
| 24 | dist.broadcast_object_list(recv, src=0) |
| 25 | ip, port = recv |
| 26 | |
| 27 | stateless_pg = StatelessProcessGroup.create(ip, port, rank, world_size) |
| 28 | |
| 29 | for pg in [dist.group.WORLD, stateless_pg]: |
| 30 | test_result = _node_count(pg) |
| 31 | |
| 32 | # Expected node count based on environment variable) |
| 33 | expected = int(os.environ.get("NUM_NODES", "1")) |
| 34 | |
| 35 | assert test_result == expected, f"Expected {expected} nodes, got {test_result}" |
| 36 | |
| 37 | if pg == dist.group.WORLD: |
| 38 | print( |
| 39 | f"Node count test passed! Got {test_result} nodes " |
| 40 | f"when using torch distributed!" |
| 41 | ) |
| 42 | else: |
| 43 | print( |
| 44 | f"Node count test passed! Got {test_result} nodes " |
| 45 | f"when using StatelessProcessGroup!" |
| 46 | ) |