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