| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | import subprocess |
| 6 | import threading |
| 7 | import time |
| 8 | from concurrent.futures import ThreadPoolExecutor |
| 9 | |
| 10 | import pytest |
| 11 | import requests |
| 12 | |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | from ..evals.gsm8k.gsm8k_eval import evaluate_gsm8k |
| 16 | from ..utils import RemoteOpenAIServer, multi_gpu_test |
| 17 | |
| 18 | |
| 19 | @pytest.fixture(autouse=True) |
| 20 | def cleanup_ray_between_tests(): |
| 21 | """Force-stop any lingering Ray processes between tests.""" |
| 22 | subprocess.run(["ray", "stop", "--force"], timeout=30, capture_output=True) |
| 23 | time.sleep(5) |
| 24 | yield |
| 25 | |
| 26 | |
| 27 | MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat" |
| 28 | |
| 29 | NUM_GSM8K_QUESTIONS = 256 |
| 30 | EXPECTED_ACCURACY = 0.58 |
| 31 | ACCURACY_TOL = 0.08 |
| 32 | DEFAULT_ALL2ALL_BACKEND = "allgather_reducescatter" |
| 33 | |
| 34 | |
| 35 | def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool: |
| 36 | url = server.url_for("scale_elastic_ep") |
| 37 | payload = {"new_data_parallel_size": new_dp_size} |
| 38 | headers = {"Content-Type": "application/json"} |
| 39 | |
| 40 | try: |
| 41 | response = requests.post(url, json=payload, headers=headers, timeout=300) |
| 42 | return response.status_code == 200 |
| 43 | except requests.exceptions.RequestException: |
| 44 | return False |
| 45 | |
| 46 | |
| 47 | def _traffic_loop( |
| 48 | server: RemoteOpenAIServer, |
| 49 | dp_rank: int | None, |
| 50 | ready: threading.Barrier, |
| 51 | stop: threading.Event, |
| 52 | finished: threading.Event, |
| 53 | is_probe: bool = False, |
| 54 | ) -> list[tuple[float, float, int | None]]: |
| 55 | url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions") |
| 56 | payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4} |
| 57 | headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)} |
| 58 | request_payload = None if is_probe else payload |
| 59 | responses = [] |
| 60 | is_ready = False |
| 61 | while not stop.is_set(): |
| 62 | request_start = time.perf_counter() |
| 63 | try: |
| 64 | response = requests.post( |
| 65 | url, json=request_payload, headers=headers, timeout=120 |
| 66 | ) |
| 67 | status_code = response.status_code |
| 68 | except requests.exceptions.RequestException: |
| 69 | status_code = None |
| 70 | responses.append((request_start, time.perf_counter(), status_code)) |
| 71 | if status_code == 200: |
| 72 | if not is_ready: |
| 73 | ready.wait(timeout=120) |
| 74 | is_ready = True |
| 75 | if finished.is_set(): |
| 76 | return responses |
| 77 | time.sleep(0.05) |
| 78 | return responses |
| 79 | |
| 80 | |
| 81 | def _downtime(responses: list[tuple[float, float, int | None]]) -> float: |
| 82 | rejected = [end for _, end, status in responses if status == 503] |
| 83 | if not rejected: |
| 84 | return 0 |
| 85 | recovered = next( |
| 86 | end for _, end, status in responses if status == 200 and end > rejected[-1] |
| 87 | ) |
| 88 | return recovered - rejected[0] |
| 89 | |
| 90 | |
| 91 | def _scale_with_traffic( |
| 92 | server: RemoteOpenAIServer, |
| 93 | source_dp_size: int, |
| 94 | new_dp_size: int, |
| 95 | traffic_mode: str, |
| 96 | ) -> None: |
| 97 | traffic_clients: list[int | None] = [] |
| 98 | if traffic_mode == "light": |
| 99 | traffic_clients = [0] |
| 100 | elif traffic_mode == "heavy": |
| 101 | traffic_clients = [None] * source_dp_size |
| 102 | clients = [(None, True)] + [(rank, False) for rank in traffic_clients] |
| 103 | ready = threading.Barrier(len(clients) + 1) |
| 104 | stop = threading.Event() |
| 105 | finished = threading.Event() |
| 106 | |
| 107 | with ThreadPoolExecutor(max_workers=len(clients)) as executor: |
| 108 | futures = [ |
| 109 | executor.submit( |
| 110 | _traffic_loop, server, rank, ready, stop, finished, is_probe |
| 111 | ) |
| 112 | for rank, is_probe in clients |
| 113 | ] |
| 114 | try: |
| 115 | ready.wait(timeout=120) |
| 116 | start_time = time.perf_counter() |
| 117 | assert _send_scale_command(server, new_dp_size) |
| 118 | scale_seconds = time.perf_counter() - start_time |
| 119 | finished.set() |
| 120 | probe_result, *results = [future.result(timeout=120) for future in futures] |
| 121 | finally: |
| 122 | stop.set() |
| 123 | |
| 124 | bad_statuses = { |
| 125 | status |
| 126 | for responses in [probe_result, *results] |
| 127 | for _, _, status in responses |
| 128 | if status not in (200, 503) |
| 129 | } |
| 130 | assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}" |
| 131 | probe_503 = [start for start, _, status in probe_result if status == 503] |
| 132 | assert probe_503, "Scaling probe did not observe commit" |
| 133 | assert not results or any( |
| 134 | status == 200 and start_time <= request_start and request_end < probe_503[0] |
| 135 | for responses in results |
| 136 | for request_start, request_end, status in responses |
| 137 | ), "No request completed successfully during preparation" |
| 138 | |
| 139 | print( |
| 140 | f"[Elastic EP timing][{source_dp_size}->{new_dp_size}]" |
| 141 | f"[traffic={traffic_mode}] " |
| 142 | f"scale_seconds={scale_seconds:.3f} " |
| 143 | f"downtime_seconds={_downtime(probe_result):.3f}" |
| 144 | ) |
| 145 | |
| 146 | |
| 147 | def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: |
| 148 | assert server.port is not None |
| 149 | result = evaluate_gsm8k( |
| 150 | num_questions=NUM_GSM8K_QUESTIONS, |
| 151 | host=f"http://{server.host}", |
| 152 | port=server.port, |
| 153 | ) |
| 154 | accuracy = result["accuracy"] |
| 155 | print( |
| 156 | f"[{stage}] GSM8K accuracy: {accuracy:.3f} " |
| 157 | f"({result['num_questions']} questions)" |
| 158 | ) |
| 159 | assert accuracy >= EXPECTED_ACCURACY, ( |
| 160 | f"[{stage}] GSM8K accuracy {accuracy:.3f} is below " |
| 161 | f"expected threshold {EXPECTED_ACCURACY}" |
| 162 | ) |
| 163 | return accuracy |
| 164 | |
| 165 | |
| 166 | def _base_serve_args( |
| 167 | dp_size: int = 2, |
| 168 | max_dp_size: int = 4, |
| 169 | enforce_eager: bool = False, |
| 170 | all2all_backend: str = DEFAULT_ALL2ALL_BACKEND, |
| 171 | ) -> list[str]: |
| 172 | args = [ |
| 173 | "--trust-remote-code", |
| 174 | "--tensor-parallel-size", |
| 175 | "1", |
| 176 | "--gpu-memory-utilization", |
| 177 | "0.8", |
| 178 | "--max-model-len", |
| 179 | "4096", |
| 180 | "--enable-expert-parallel", |
| 181 | "--all2all-backend", |
| 182 | all2all_backend, |
| 183 | "--enable-elastic-ep", |
| 184 | "--elastic-ep-max-dp-size", |
| 185 | str(max_dp_size), |
| 186 | "--enable-eplb", |
| 187 | "--eplb-config.num_redundant_experts", |
| 188 | "0", |
| 189 | "--eplb-config.use_async", |
| 190 | "true", |
| 191 | "--eplb-config.step_interval", |
| 192 | "300", |
| 193 | "--eplb-config.window_size", |
| 194 | "5", |
| 195 | "--data-parallel-backend", |
| 196 | "ray", |
| 197 | "--data-parallel-size", |
| 198 | str(dp_size), |
| 199 | "--api-server-count", |
| 200 | "1", |
| 201 | "--disable-access-log-for-endpoints", |
| 202 | "/is_scaling_elastic_ep", |
| 203 | ] |
| 204 | |
| 205 | leader_address = os.environ.get("LEADER_ADDRESS") |
| 206 | if leader_address: |
| 207 | args.extend(["--data-parallel-address", leader_address]) |
| 208 | if enforce_eager: |
| 209 | args.append("--enforce-eager") |
| 210 | |
| 211 | return args |
| 212 | |
| 213 | |
| 214 | @pytest.mark.parametrize( |
| 215 | ("enforce_eager", "traffic_mode", "all2all_backend"), |
| 216 | [ |
| 217 | pytest.param(True, "none", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_none"), |
| 218 | pytest.param(True, "light", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_light"), |
| 219 | pytest.param(True, "heavy", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_heavy"), |
| 220 | pytest.param( |
| 221 | False, |
| 222 | "heavy", |
| 223 | DEFAULT_ALL2ALL_BACKEND, |
| 224 | id="cuda_graphs_heavy", |
| 225 | ), |
| 226 | pytest.param(False, "heavy", "nixl_ep", id="cuda_graphs_heavy_nixl_ep"), |
| 227 | ], |
| 228 | ) |
| 229 | @multi_gpu_test(num_gpus=4) |
| 230 | def test_elastic_ep_scaling( |
| 231 | enforce_eager: bool, traffic_mode: str, all2all_backend: str |
| 232 | ): |
| 233 | from vllm.distributed.eplb.eplb_communicator import has_nixl |
| 234 | from vllm.utils.import_utils import has_nixl_ep |
| 235 | |
| 236 | if not has_nixl(): |
| 237 | pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") |
| 238 | if current_platform.is_xpu(): |
| 239 | pytest.skip("Elastic EP with NIXL is not supported on XPU") |
| 240 | if all2all_backend == "nixl_ep" and not has_nixl_ep(): |
| 241 | pytest.skip("NIXL EP is not installed") |
| 242 | |
| 243 | initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2")) |
| 244 | target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4")) |
| 245 | assert target_dp_size > initial_dp_size |
| 246 | vllm_serve_args = _base_serve_args( |
| 247 | initial_dp_size, target_dp_size, enforce_eager, all2all_backend |
| 248 | ) |
| 249 | |
| 250 | with RemoteOpenAIServer( |
| 251 | MODEL_NAME, |
| 252 | vllm_serve_args, |
| 253 | max_wait_seconds=1200, |
| 254 | ) as server: |
| 255 | initial_accuracy = _run_gsm8k_eval(server, "Initial") |
| 256 | |
| 257 | _scale_with_traffic(server, initial_dp_size, target_dp_size, traffic_mode) |
| 258 | scale_up_accuracy = _run_gsm8k_eval(server, "After scale up") |
| 259 | assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( |
| 260 | f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than " |
| 261 | f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" |
| 262 | ) |
| 263 | |
| 264 | _scale_with_traffic(server, target_dp_size, initial_dp_size, traffic_mode) |
| 265 | scale_down_accuracy = _run_gsm8k_eval(server, "After scale down") |
| 266 | assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( |
| 267 | f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than " |
| 268 | f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" |
| 269 | ) |
| 270 | |
| 271 | print("\nAccuracy Summary:") |
| 272 | print(f" Initial: {initial_accuracy:.3f}") |
| 273 | print( |
| 274 | f" Scale up: {scale_up_accuracy:.3f} " |
| 275 | f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})" |
| 276 | ) |
| 277 | print( |
| 278 | f" Scale down: {scale_down_accuracy:.3f} " |
| 279 | f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})" |
| 280 | ) |
| 281 | print(f" Tolerance: {ACCURACY_TOL:.3f}") |