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