| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import atexit |
| 5 | import os |
| 6 | import random |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | import torch.multiprocessing as mp |
| 11 | |
| 12 | from vllm.config import VllmConfig, set_current_vllm_config |
| 13 | from vllm.distributed.parallel_state import ( |
| 14 | init_distributed_environment, |
| 15 | ) |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.utils.system_utils import update_environment_variables |
| 18 | |
| 19 | mp.set_start_method("spawn", force=True) |
| 20 | |
| 21 | |
| 22 | def _distributed_worker_wrapper(fn, env, world_size, args, rank, skip_queue): |
| 23 | try: |
| 24 | fn(env, world_size, *args) |
| 25 | except BaseException as exc: |
| 26 | if isinstance(exc, pytest.skip.Exception): |
| 27 | skip_queue.put((rank, str(exc))) |
| 28 | return |
| 29 | raise |
| 30 | |
| 31 | |
| 32 | def distributed_run(fn, world_size, *args): |
| 33 | number_of_processes = world_size |
| 34 | processes: list[mp.Process] = [] |
| 35 | skip_queue: mp.SimpleQueue = mp.SimpleQueue() |
| 36 | for i in range(number_of_processes): |
| 37 | env: dict[str, str] = {} |
| 38 | env["RANK"] = str(i) |
| 39 | env["LOCAL_RANK"] = str(i) |
| 40 | env["WORLD_SIZE"] = str(number_of_processes) |
| 41 | env["LOCAL_WORLD_SIZE"] = str(number_of_processes) |
| 42 | env["MASTER_ADDR"] = "localhost" |
| 43 | env["MASTER_PORT"] = "12345" |
| 44 | p = mp.Process( |
| 45 | target=_distributed_worker_wrapper, |
| 46 | args=(fn, env, world_size, args, i, skip_queue), |
| 47 | ) |
| 48 | processes.append(p) |
| 49 | p.start() |
| 50 | |
| 51 | for p in processes: |
| 52 | p.join() |
| 53 | |
| 54 | skipped: list[tuple[int, str]] = [] |
| 55 | while not skip_queue.empty(): |
| 56 | rank, reason = skip_queue.get() |
| 57 | skipped.append((rank, reason)) |
| 58 | |
| 59 | if len(skipped) == number_of_processes: |
| 60 | reason = skipped[0][1] |
| 61 | pytest.skip(reason) |
| 62 | if 0 < len(skipped) < number_of_processes: |
| 63 | skipped_ranks = sorted(rank for rank, _ in skipped) |
| 64 | raise AssertionError( |
| 65 | "Distributed test had partial skips; expected either all ranks " |
| 66 | f"to skip or none. Skipped ranks: {skipped_ranks}, " |
| 67 | f"total ranks: {number_of_processes}" |
| 68 | ) |
| 69 | |
| 70 | for p in processes: |
| 71 | assert p.exitcode == 0 |
| 72 | |
| 73 | |
| 74 | def set_env_vars_and_device(env: dict[str, str]) -> None: |
| 75 | update_environment_variables(env) |
| 76 | local_rank = os.environ["LOCAL_RANK"] |
| 77 | accelerator_type = torch.accelerator.current_accelerator().type |
| 78 | device = torch.device(f"{accelerator_type}:{local_rank}") |
| 79 | torch.accelerator.set_device_index(device) |
| 80 | |
| 81 | backend = "xccl" if current_platform.is_xpu() else "nccl" |
| 82 | |
| 83 | # Create a minimal vllm config for init_distributed_environment |
| 84 | vllm_config = VllmConfig() |
| 85 | with set_current_vllm_config(vllm_config): |
| 86 | init_distributed_environment(backend=backend) |
| 87 | atexit.register(_destroy_process_group_if_initialized) |
| 88 | # Ensure each worker process has the same random seed |
| 89 | random.seed(42) |
| 90 | torch.manual_seed(42) |
| 91 | |
| 92 | |
| 93 | def _destroy_process_group_if_initialized() -> None: |
| 94 | if torch.distributed.is_available() and torch.distributed.is_initialized(): |
| 95 | torch.distributed.destroy_process_group() |