| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | import os |
| 5 | import threading |
| 6 | import time |
| 7 | from contextlib import AsyncExitStack |
| 8 | |
| 9 | import openai # use the official client for correctness check |
| 10 | import pytest |
| 11 | import pytest_asyncio |
| 12 | import requests |
| 13 | |
| 14 | from tests.utils import RemoteOpenAIServer |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b") |
| 18 | |
| 19 | # Number of data parallel ranks for external LB testing |
| 20 | DP_SIZE = int(os.getenv("DP_SIZE", "2")) |
| 21 | # Default tensor parallel size to use |
| 22 | TP_SIZE = int(os.getenv("TP_SIZE", "1")) |
| 23 | |
| 24 | |
| 25 | class ExternalLBServerManager: |
| 26 | """Manages data parallel vLLM server instances for external |
| 27 | load balancer testing.""" |
| 28 | |
| 29 | def __init__( |
| 30 | self, |
| 31 | model_name: str, |
| 32 | dp_size: int, |
| 33 | api_server_count: int, |
| 34 | base_server_args: list, |
| 35 | tp_size: int = TP_SIZE, |
| 36 | ): |
| 37 | self.model_name = model_name |
| 38 | self.dp_size = dp_size |
| 39 | self.tp_size = tp_size |
| 40 | self.api_server_count = api_server_count |
| 41 | self.base_server_args = base_server_args |
| 42 | self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = [] |
| 43 | self.server_threads: list[threading.Thread] = [] |
| 44 | |
| 45 | def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]: |
| 46 | """Start all server instances for external LB mode.""" |
| 47 | for rank in range(self.dp_size): |
| 48 | # Create server args for this specific rank |
| 49 | server_args = self.base_server_args.copy() |
| 50 | |
| 51 | # Add external LB specific arguments |
| 52 | server_args.extend( |
| 53 | [ |
| 54 | "--data-parallel-size", |
| 55 | str(self.dp_size), |
| 56 | "--data-parallel-rank", |
| 57 | str(rank), |
| 58 | "--data-parallel-size-local", |
| 59 | "1", |
| 60 | "--tensor-parallel-size", |
| 61 | str(self.tp_size), |
| 62 | "--port", |
| 63 | str(8000 + rank), # Different port for each rank |
| 64 | "--api-server-count", |
| 65 | str(self.api_server_count), |
| 66 | ] |
| 67 | ) |
| 68 | |
| 69 | # Use a thread to start each server to allow parallel initialization |
| 70 | def start_server(r: int, sargs: list[str]): |
| 71 | try: |
| 72 | # Start the server |
| 73 | server = RemoteOpenAIServer( |
| 74 | self.model_name, |
| 75 | sargs, |
| 76 | auto_port=False, |
| 77 | env_dict={ |
| 78 | "VLLM_SERVER_DEV_MODE": "1", |
| 79 | current_platform.device_control_env_var: ",".join( |
| 80 | str(current_platform.device_id_to_physical_device_id(i)) |
| 81 | for i in range(r * TP_SIZE, (r + 1) * TP_SIZE) |
| 82 | ), |
| 83 | }, |
| 84 | ) |
| 85 | server.__enter__() |
| 86 | print( |
| 87 | f"Server rank {r} started successfully with " |
| 88 | f"{self.api_server_count} API servers" |
| 89 | ) |
| 90 | self.servers.append((server, sargs)) |
| 91 | except Exception as e: |
| 92 | print(f"Failed to start server rank {r}: {e}") |
| 93 | raise |
| 94 | |
| 95 | thread = threading.Thread(target=start_server, args=(rank, server_args)) |
| 96 | thread.start() |
| 97 | |
| 98 | self.server_threads.append(thread) |
| 99 | |
| 100 | # Wait for all servers to start |
| 101 | for thread in self.server_threads: |
| 102 | thread.join() |
| 103 | |
| 104 | # Give servers additional time to fully initialize and coordinate |
| 105 | time.sleep(2) |
| 106 | |
| 107 | if len(self.servers) != self.dp_size: |
| 108 | raise Exception("Servers failed to start") |
| 109 | |
| 110 | return self.servers |
| 111 | |
| 112 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 113 | """Stop all server instances.""" |
| 114 | servers = [s for s, _ in self.servers] |
| 115 | self.servers.clear() |
| 116 | try: |
| 117 | RemoteOpenAIServer.shutdown_many(servers) |
| 118 | except Exception as e: |
| 119 | print(f"Error stopping servers: {e}") |
| 120 | |
| 121 | |
| 122 | @pytest.fixture(scope="module") |
| 123 | def default_server_args(): |
| 124 | return [ |
| 125 | # use half precision for speed and memory savings in CI environment |
| 126 | "--dtype", |
| 127 | "bfloat16", |
| 128 | "--max-model-len", |
| 129 | "2048", |
| 130 | "--max-num-seqs", |
| 131 | "128", |
| 132 | "--enforce-eager", |
| 133 | ] |
| 134 | |
| 135 | |
| 136 | @pytest.fixture( |
| 137 | scope="module", |
| 138 | params=[(1, False), (4, False), (1, True)], |
| 139 | ids=["api-1-standard", "api-4-standard", "api-1-ep-enabled"], |
| 140 | ) |
| 141 | def server_manager(request, default_server_args): |
| 142 | api_server_count, enable_expert_parallel = request.param |
| 143 | server_args = default_server_args.copy() |
| 144 | if enable_expert_parallel: |
| 145 | server_args.append("--enable-expert-parallel") |
| 146 | |
| 147 | server_manager = ExternalLBServerManager( |
| 148 | MODEL_NAME, DP_SIZE, api_server_count, server_args |
| 149 | ) |
| 150 | |
| 151 | with server_manager: |
| 152 | yield server_manager |
| 153 | |
| 154 | |
| 155 | @pytest.fixture |
| 156 | def servers(server_manager): |
| 157 | return server_manager.servers |
| 158 | |
| 159 | |
| 160 | @pytest_asyncio.fixture |
| 161 | async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]): |
| 162 | # Create a client for each server |
| 163 | async with AsyncExitStack() as stack: |
| 164 | yield [ |
| 165 | await stack.enter_async_context(server.get_async_client()) |
| 166 | for server, _ in servers |
| 167 | ] |
| 168 | |
| 169 | |
| 170 | def _get_parallel_config(server: RemoteOpenAIServer): |
| 171 | response = requests.get(server.url_for("server_info?config_format=json")) |
| 172 | response.raise_for_status() |
| 173 | |
| 174 | vllm_config = response.json()["vllm_config"] |
| 175 | return vllm_config["parallel_config"] |
| 176 | |
| 177 | |
| 178 | def _assert_expert_parallel_config( |
| 179 | servers: list[tuple[RemoteOpenAIServer, list[str]]], |
| 180 | ): |
| 181 | for server, server_args in servers: |
| 182 | expected = "--enable-expert-parallel" in server_args |
| 183 | parallel_config = _get_parallel_config(server) |
| 184 | assert parallel_config["enable_expert_parallel"] is expected |
| 185 | |
| 186 | |
| 187 | def test_external_lb_server_info(server_manager): |
| 188 | servers = server_manager.servers |
| 189 | api_server_count = server_manager.api_server_count |
| 190 | |
| 191 | _assert_expert_parallel_config(servers) |
| 192 | |
| 193 | for i, (server, _) in enumerate(servers): |
| 194 | print(f"Testing {i=}") |
| 195 | |
| 196 | # Each request will hit one of the API servers |
| 197 | # `n_reqs` is set so that there is a good chance each server |
| 198 | # receives at least one request |
| 199 | n_reqs = 2 * api_server_count * api_server_count |
| 200 | parallel_configs = [_get_parallel_config(server) for _ in range(n_reqs)] |
| 201 | api_process_counts = [c["_api_process_count"] for c in parallel_configs] |
| 202 | api_process_ranks = [c["_api_process_rank"] for c in parallel_configs] |
| 203 | |
| 204 | assert all(c == api_server_count for c in api_process_counts), ( |
| 205 | api_process_counts |
| 206 | ) |
| 207 | assert all(0 <= r < api_server_count for r in api_process_ranks), ( |
| 208 | api_process_ranks |
| 209 | ) |
| 210 | |
| 211 | |
| 212 | @pytest.mark.asyncio |
| 213 | @pytest.mark.parametrize( |
| 214 | "model_name", |
| 215 | [MODEL_NAME], |
| 216 | ) |
| 217 | async def test_external_lb_single_completion( |
| 218 | clients: list[openai.AsyncOpenAI], |
| 219 | servers: list[tuple[RemoteOpenAIServer, list[str]]], |
| 220 | model_name: str, |
| 221 | ) -> None: |
| 222 | async def make_request(client: openai.AsyncOpenAI): |
| 223 | completion = await client.completions.create( |
| 224 | model=model_name, prompt="Hello, my name is", max_tokens=10, temperature=1.0 |
| 225 | ) |
| 226 | |
| 227 | assert completion.id is not None |
| 228 | assert completion.choices is not None and len(completion.choices) == 1 |
| 229 | |
| 230 | choice = completion.choices[0] |
| 231 | # The exact number of tokens can vary slightly with temperature=1.0, |
| 232 | # so we check for a reasonable minimum length. |
| 233 | assert len(choice.text) >= 1 |
| 234 | # Finish reason might not always be 'length' if the model finishes early |
| 235 | # or due to other reasons, especially with high temperature. |
| 236 | # So, we'll accept 'length' or 'stop'. |
| 237 | assert choice.finish_reason in ("length", "stop") |
| 238 | |
| 239 | # Token counts can also vary, so we check they are positive. |
| 240 | assert completion.usage.completion_tokens > 0 |
| 241 | assert completion.usage.prompt_tokens > 0 |
| 242 | assert completion.usage.total_tokens > 0 |
| 243 | return completion |
| 244 | |
| 245 | # Test single request to each server |
| 246 | for i, client in enumerate(clients): |
| 247 | result = await make_request(client) |
| 248 | assert result is not None |
| 249 | print(f"Server {i} handled single completion request successfully") |
| 250 | |
| 251 | await asyncio.sleep(0.5) |
| 252 | |
| 253 | # Send requests to all servers in round-robin fashion |
| 254 | num_requests_per_server = 25 # Total 50 requests across 2 servers |
| 255 | all_tasks = [] |
| 256 | |
| 257 | for i, client in enumerate(clients): |
| 258 | tasks = [make_request(client) for _ in range(num_requests_per_server)] |
| 259 | all_tasks.extend(tasks) |
| 260 | |
| 261 | results = await asyncio.gather(*all_tasks) |
| 262 | assert len(results) == num_requests_per_server * len(clients) |
| 263 | assert all(completion is not None for completion in results) |
| 264 | |
| 265 | await asyncio.sleep(0.5) |
| 266 | |
| 267 | # Second burst of requests |
| 268 | all_tasks = [] |
| 269 | for i, client in enumerate(clients): |
| 270 | tasks = [make_request(client) for _ in range(num_requests_per_server)] |
| 271 | all_tasks.extend(tasks) |
| 272 | |
| 273 | results = await asyncio.gather(*all_tasks) |
| 274 | assert len(results) == num_requests_per_server * len(clients) |
| 275 | assert all(completion is not None for completion in results) |
| 276 | |
| 277 | _, server_args = servers[0] |
| 278 | api_server_count = ( |
| 279 | server_args.count("--api-server-count") |
| 280 | and server_args[server_args.index("--api-server-count") + 1] |
| 281 | or 1 |
| 282 | ) |
| 283 | print( |
| 284 | f"Successfully completed external LB test with {len(clients)} servers " |
| 285 | f"(API server count: {api_server_count})" |
| 286 | ) |
| 287 | |
| 288 | |
| 289 | @pytest.mark.asyncio |
| 290 | @pytest.mark.parametrize( |
| 291 | "model_name", |
| 292 | [MODEL_NAME], |
| 293 | ) |
| 294 | async def test_external_lb_completion_streaming( |
| 295 | clients: list[openai.AsyncOpenAI], |
| 296 | servers: list[tuple[RemoteOpenAIServer, list[str]]], |
| 297 | model_name: str, |
| 298 | ) -> None: |
| 299 | prompt = "What is an LLM?" |
| 300 | |
| 301 | async def make_streaming_request(client: openai.AsyncOpenAI): |
| 302 | # Perform a non-streaming request to get the expected full output |
| 303 | single_completion = await client.completions.create( |
| 304 | model=model_name, |
| 305 | prompt=prompt, |
| 306 | max_tokens=5, |
| 307 | temperature=0.0, |
| 308 | ) |
| 309 | single_output = single_completion.choices[0].text |
| 310 | |
| 311 | # Perform the streaming request |
| 312 | stream = await client.completions.create( |
| 313 | model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True |
| 314 | ) |
| 315 | chunks: list[str] = [] |
| 316 | finish_reason_count = 0 |
| 317 | last_chunk = None |
| 318 | async for chunk in stream: |
| 319 | chunks.append(chunk.choices[0].text) |
| 320 | if chunk.choices[0].finish_reason is not None: |
| 321 | finish_reason_count += 1 |
| 322 | last_chunk = chunk # Keep track of the last chunk |
| 323 | |
| 324 | # finish reason should only return in the last block for OpenAI API |
| 325 | assert finish_reason_count == 1, "Finish reason should appear exactly once." |
| 326 | assert last_chunk is not None, "Stream should have yielded at least one chunk." |
| 327 | assert last_chunk.choices[0].finish_reason == "length", ( |
| 328 | "Finish reason should be 'length'." |
| 329 | ) |
| 330 | # Check that the combined text matches the non-streamed version. |
| 331 | assert "".join(chunks) == single_output, ( |
| 332 | "Streamed output should match non-streamed output." |
| 333 | ) |
| 334 | return True # Indicate success for this request |
| 335 | |
| 336 | # Test single request to each server |
| 337 | for i, client in enumerate(clients): |
| 338 | result = await make_streaming_request(client) |
| 339 | assert result is not None |
| 340 | print(f"Server {i} handled single streaming request successfully") |
| 341 | |
| 342 | await asyncio.sleep(0.5) |
| 343 | |
| 344 | # Send streaming requests to all servers in round-robin fashion |
| 345 | num_requests_per_server = 25 # Total 50 requests across 2 servers |
| 346 | all_tasks = [] |
| 347 | |
| 348 | for i, client in enumerate(clients): |
| 349 | tasks = [make_streaming_request(client) for _ in range(num_requests_per_server)] |
| 350 | all_tasks.extend(tasks) |
| 351 | |
| 352 | results = await asyncio.gather(*all_tasks) |
| 353 | assert len(results) == num_requests_per_server * len(clients) |
| 354 | assert all(results), "Not all streaming requests completed successfully." |
| 355 | |
| 356 | await asyncio.sleep(0.5) |
| 357 | |
| 358 | # Second burst of streaming requests |
| 359 | all_tasks = [] |
| 360 | for i, client in enumerate(clients): |
| 361 | tasks = [make_streaming_request(client) for _ in range(num_requests_per_server)] |
| 362 | all_tasks.extend(tasks) |
| 363 | |
| 364 | results = await asyncio.gather(*all_tasks) |
| 365 | assert len(results) == num_requests_per_server * len(clients) |
| 366 | assert all(results), "Not all streaming requests completed successfully." |
| 367 | |
| 368 | _, server_args = servers[0] |
| 369 | api_server_count = ( |
| 370 | server_args.count("--api-server-count") |
| 371 | and server_args[server_args.index("--api-server-count") + 1] |
| 372 | or 1 |
| 373 | ) |
| 374 | print( |
| 375 | f"Successfully completed external LB streaming test with " |
| 376 | f"{len(clients)} servers (API server count: {api_server_count})" |
| 377 | ) |