| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import multiprocessing as mp |
| 5 | import os |
| 6 | import tempfile |
| 7 | from multiprocessing.queues import Queue |
| 8 | from unittest.mock import patch |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from vllm.config import set_current_vllm_config |
| 14 | from vllm.engine.arg_utils import EngineArgs |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.utils.mem_utils import MemorySnapshot |
| 17 | from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment |
| 18 | |
| 19 | # Global queue to track operation order across processes |
| 20 | _QUEUE: Queue | None = None |
| 21 | |
| 22 | |
| 23 | def track_operation(operation: str, rank: int): |
| 24 | """Track when an operation happens and its rank.""" |
| 25 | if _QUEUE is not None: |
| 26 | _QUEUE.put((operation, rank)) |
| 27 | |
| 28 | |
| 29 | def make_operation_tracker(operation_name: str, original_func): |
| 30 | """Create a mock function that tracks when an operation is called. |
| 31 | |
| 32 | Args: |
| 33 | operation_name: Name to use when tracking this operation |
| 34 | original_func: The original function to wrap |
| 35 | |
| 36 | Returns: |
| 37 | A wrapper function that tracks the operation and calls the original |
| 38 | |
| 39 | """ |
| 40 | |
| 41 | def wrapper(*args, **kwargs): |
| 42 | rank = int(os.environ.get("RANK", "-1")) |
| 43 | track_operation(operation_name, rank) |
| 44 | return original_func(*args, **kwargs) |
| 45 | |
| 46 | return wrapper |
| 47 | |
| 48 | |
| 49 | def worker_process( |
| 50 | rank: int, |
| 51 | world_size: int, |
| 52 | distributed_init_method: str, |
| 53 | queue: Queue, |
| 54 | error_queue: Queue, |
| 55 | ): |
| 56 | """Worker process that initializes a GPU worker with proper tracking.""" |
| 57 | global _QUEUE |
| 58 | _QUEUE = queue |
| 59 | |
| 60 | try: |
| 61 | # Set environment variables |
| 62 | os.environ["RANK"] = str(rank) |
| 63 | os.environ["LOCAL_RANK"] = str(rank) |
| 64 | os.environ["WORLD_SIZE"] = str(world_size) |
| 65 | |
| 66 | # Create vLLM config with small model |
| 67 | vllm_config = EngineArgs( |
| 68 | model="facebook/opt-125m", tensor_parallel_size=2, load_format="dummy" |
| 69 | ).create_engine_config() |
| 70 | |
| 71 | # Select worker class and patch target based on platform |
| 72 | if current_platform.is_xpu(): |
| 73 | from vllm.v1.worker.xpu_worker import XPUWorker |
| 74 | |
| 75 | worker_cls = XPUWorker |
| 76 | init_patch_target = ( |
| 77 | "vllm.v1.worker.xpu_worker.init_worker_distributed_environment" |
| 78 | ) |
| 79 | else: |
| 80 | worker_cls = Worker |
| 81 | init_patch_target = ( |
| 82 | "vllm.v1.worker.gpu_worker.init_worker_distributed_environment" |
| 83 | ) |
| 84 | |
| 85 | # Create worker |
| 86 | worker = worker_cls( |
| 87 | vllm_config=vllm_config, |
| 88 | local_rank=rank, |
| 89 | rank=rank, |
| 90 | distributed_init_method=distributed_init_method, |
| 91 | ) |
| 92 | |
| 93 | # Get original functions before patching |
| 94 | original_init_worker = init_worker_distributed_environment |
| 95 | original_memory_snapshot_init = MemorySnapshot.__init__ |
| 96 | original_all_reduce = torch.distributed.all_reduce |
| 97 | |
| 98 | # Apply minimal patches to track operation order |
| 99 | init_patch = patch( |
| 100 | init_patch_target, |
| 101 | side_effect=make_operation_tracker( |
| 102 | "init_distributed", original_init_worker |
| 103 | ), |
| 104 | ) |
| 105 | memory_patch = patch.object( |
| 106 | MemorySnapshot, |
| 107 | "__init__", |
| 108 | make_operation_tracker("memory_snapshot", original_memory_snapshot_init), |
| 109 | ) |
| 110 | all_reduce_patch = patch( |
| 111 | "torch.distributed.all_reduce", |
| 112 | side_effect=make_operation_tracker("dist_all_reduce", original_all_reduce), |
| 113 | ) |
| 114 | |
| 115 | with ( |
| 116 | init_patch, |
| 117 | memory_patch, |
| 118 | all_reduce_patch, |
| 119 | set_current_vllm_config(vllm_config), |
| 120 | ): |
| 121 | # Initialize device (this is where we test the order) |
| 122 | worker.init_device() |
| 123 | |
| 124 | # Load model to ensure everything works |
| 125 | worker.load_model() |
| 126 | |
| 127 | # Signal success |
| 128 | queue.put(("success", rank)) |
| 129 | |
| 130 | except Exception as e: |
| 131 | error_queue.put((rank, str(e), type(e).__name__)) |
| 132 | raise |
| 133 | |
| 134 | |
| 135 | @pytest.mark.skipif( |
| 136 | torch.accelerator.device_count() < 2, |
| 137 | reason="Need at least 2 GPUs for tensor parallelism", |
| 138 | ) |
| 139 | def test_init_distributed_is_called_before_memory_snapshot(): |
| 140 | """Test that distributed env is setup before memory snapshot. |
| 141 | |
| 142 | This test makes sure during worker initialization, the initial memory |
| 143 | snapshot is taken after distributed env is setup to include all the buffers |
| 144 | allocated by distributed env. |
| 145 | """ |
| 146 | world_size = 2 |
| 147 | |
| 148 | # Create a temporary file for distributed init |
| 149 | with tempfile.NamedTemporaryFile(delete=False) as f: |
| 150 | distributed_init_method = f"file://{f.name}" |
| 151 | |
| 152 | # Create queues for inter-process communication |
| 153 | ctx = mp.get_context("spawn") |
| 154 | operation_queue = ctx.Queue() |
| 155 | error_queue = ctx.Queue() |
| 156 | |
| 157 | # Start worker processes |
| 158 | processes = [] |
| 159 | for rank in range(world_size): |
| 160 | p = ctx.Process( |
| 161 | target=worker_process, |
| 162 | args=( |
| 163 | rank, |
| 164 | world_size, |
| 165 | distributed_init_method, |
| 166 | operation_queue, |
| 167 | error_queue, |
| 168 | ), |
| 169 | ) |
| 170 | p.start() |
| 171 | processes.append(p) |
| 172 | |
| 173 | # Wait for all processes to complete |
| 174 | for p in processes: |
| 175 | p.join(timeout=60) # 60 second timeout |
| 176 | |
| 177 | # Check for errors |
| 178 | errors = [] |
| 179 | while not error_queue.empty(): |
| 180 | rank, error_msg, error_type = error_queue.get() |
| 181 | errors.append(f"Rank {rank}: {error_type}: {error_msg}") |
| 182 | |
| 183 | if errors: |
| 184 | pytest.fail("Worker processes failed:\n" + "\n".join(errors)) |
| 185 | |
| 186 | # Collect all operations from the queue |
| 187 | operations = [] |
| 188 | while not operation_queue.empty(): |
| 189 | operations.append(operation_queue.get()) |
| 190 | |
| 191 | # Verify we got operations from both ranks |
| 192 | print(f"Collected operations: {operations}") |
| 193 | |
| 194 | # Check operations for each rank |
| 195 | for rank in range(world_size): |
| 196 | rank_ops = [op for op, r in operations if r == rank] |
| 197 | print(f"\nRank {rank} operations: {rank_ops}") |
| 198 | |
| 199 | # Raises ValueError if the operation is not found |
| 200 | init_distributed = rank_ops.index("init_distributed") |
| 201 | dist_all_reduce = rank_ops.index("dist_all_reduce") |
| 202 | memory_snapshot = rank_ops.index("memory_snapshot") |
| 203 | |
| 204 | # Verify order: init_distributed should happen before memory_snapshot |
| 205 | assert init_distributed < dist_all_reduce < memory_snapshot, ( |
| 206 | f"Rank {rank}: init_distributed (index {init_distributed}) " |
| 207 | f"must happen before dist_all_reduce (index {dist_all_reduce}) " |
| 208 | f"and memory_snapshot (index {memory_snapshot})" |
| 209 | ) |
| 210 | |
| 211 | # Clean up |
| 212 | os.unlink(distributed_init_method.replace("file://", "")) |