mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-10 06:55:01 +08:00
678 lines
25 KiB
Python
678 lines
25 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import asyncio
|
|
import os
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from typing import Optional, cast
|
|
|
|
import openai # use the official client for correctness check
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from tests.utils import RemoteOpenAIServer
|
|
from tests.v1.test_utils import check_request_balancing
|
|
from vllm.platforms import current_platform
|
|
|
|
MODEL_NAME = "ibm-research/PowerMoE-3b"
|
|
|
|
# Number of data parallel ranks for multi-node internal LB testing
|
|
DP_SIZE = int(os.getenv("DP_SIZE", "2"))
|
|
# Default tensor parallel size to use
|
|
TP_SIZE = int(os.getenv("TP_SIZE", "1"))
|
|
|
|
# Number of nodes to simulate
|
|
NUM_NODES = 2
|
|
|
|
|
|
class MultinodeInternalLBServerManager:
|
|
"""Manages multi-node data parallel vLLM server instances for internal
|
|
load balancer testing using --headless mode."""
|
|
|
|
def __init__(self,
|
|
model_name: str,
|
|
dp_size: int,
|
|
api_server_count: int,
|
|
base_server_args: list,
|
|
dp_per_node: int = 1,
|
|
tp_size: int = TP_SIZE):
|
|
self.model_name = model_name
|
|
self.dp_size = dp_size
|
|
self.dp_per_node = dp_per_node
|
|
self.tp_size = tp_size
|
|
self.api_server_count = api_server_count
|
|
self.base_server_args = base_server_args
|
|
self.servers: list[Optional[tuple[RemoteOpenAIServer,
|
|
list[str]]]] = [None] * (dp_size //
|
|
dp_per_node)
|
|
self.server_threads: list[threading.Thread] = []
|
|
|
|
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
|
|
"""Start all server instances for multi-node internal LB mode."""
|
|
for server_idx, rank in enumerate(
|
|
range(0, self.dp_size, self.dp_per_node)):
|
|
# Create server args for this specific rank
|
|
server_args = self.base_server_args.copy()
|
|
|
|
if rank == 0:
|
|
# Head node - runs API server and first DP rank
|
|
server_args.extend([
|
|
"--data-parallel-size",
|
|
str(self.dp_size),
|
|
"--data-parallel-size-local",
|
|
str(self.dp_per_node),
|
|
"--tensor-parallel-size",
|
|
str(self.tp_size),
|
|
"--port",
|
|
"8000", # Single endpoint for all requests
|
|
"--api-server-count",
|
|
str(self.api_server_count),
|
|
"--data-parallel-address",
|
|
"127.0.0.1",
|
|
"--data-parallel-rpc-port",
|
|
"13345",
|
|
])
|
|
else:
|
|
# Secondary nodes - run in headless mode
|
|
server_args.extend([
|
|
"--headless",
|
|
"--data-parallel-size",
|
|
str(self.dp_size),
|
|
"--data-parallel-size-local",
|
|
str(self.dp_per_node),
|
|
"--data-parallel-start-rank",
|
|
str(rank),
|
|
"--tensor-parallel-size",
|
|
str(self.tp_size),
|
|
"--data-parallel-address",
|
|
"127.0.0.1",
|
|
"--data-parallel-rpc-port",
|
|
"13345",
|
|
])
|
|
|
|
# Use a thread to start each server to allow parallel initialization
|
|
def start_server(sidx: int, r: int, sargs: list[str]):
|
|
gpus_per_node = self.tp_size * self.dp_per_node
|
|
try:
|
|
# Start the server
|
|
server = RemoteOpenAIServer(
|
|
self.model_name,
|
|
sargs,
|
|
auto_port=False,
|
|
env_dict={
|
|
current_platform.device_control_env_var:
|
|
",".join(
|
|
str(
|
|
current_platform.
|
|
device_id_to_physical_device_id(i))
|
|
for i in range(r, r + gpus_per_node))
|
|
})
|
|
server.__enter__()
|
|
if r == 0:
|
|
print(
|
|
f"Head node (rank {r}) started successfully with "
|
|
f"{self.api_server_count} API servers")
|
|
else:
|
|
print(f"Headless node (rank {r}) started successfully")
|
|
self.servers[sidx] = (server, sargs)
|
|
except Exception as e:
|
|
print(f"Failed to start server rank {r}: {e}")
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
thread = threading.Thread(target=start_server,
|
|
args=(server_idx, rank, server_args))
|
|
thread.start()
|
|
|
|
self.server_threads.append(thread)
|
|
|
|
# Wait for all servers to start
|
|
for thread in self.server_threads:
|
|
thread.join()
|
|
|
|
# Give servers additional time to fully initialize and coordinate
|
|
time.sleep(3)
|
|
|
|
if not all(self.servers):
|
|
raise Exception("Servers failed to start")
|
|
|
|
return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Stop all server instances."""
|
|
while self.servers:
|
|
if server := self.servers.pop():
|
|
try:
|
|
server[0].__exit__(exc_type, exc_val, exc_tb)
|
|
except Exception as e:
|
|
print(f"Error stopping server: {e}")
|
|
traceback.print_exc()
|
|
|
|
|
|
class APIOnlyServerManager:
|
|
"""Manages API-only server (Node 0) and headless engines server (Node 1)
|
|
for testing separated API server and engine configuration."""
|
|
|
|
def __init__(self,
|
|
model_name: str,
|
|
dp_size: int,
|
|
api_server_count: int,
|
|
base_server_args: list,
|
|
tp_size: int = TP_SIZE):
|
|
self.model_name = model_name
|
|
self.dp_size = dp_size
|
|
self.tp_size = tp_size
|
|
self.api_server_count = api_server_count
|
|
self.base_server_args = base_server_args
|
|
self.servers: list[Optional[tuple[RemoteOpenAIServer,
|
|
list[str]]]] = [None] * 2
|
|
self.server_threads: list[threading.Thread] = []
|
|
|
|
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
|
|
"""Start API-only server and headless engines server."""
|
|
|
|
# Start API-only server (Node 0) - no engines, only API server
|
|
api_server_args = self.base_server_args.copy()
|
|
api_server_args.extend([
|
|
"--data-parallel-size",
|
|
str(self.dp_size),
|
|
"--data-parallel-size-local",
|
|
"0", # No engines on this node
|
|
"--tensor-parallel-size",
|
|
str(self.tp_size),
|
|
"--port",
|
|
"8000",
|
|
"--api-server-count",
|
|
str(self.api_server_count),
|
|
"--data-parallel-address",
|
|
"127.0.0.1",
|
|
"--data-parallel-rpc-port",
|
|
"13345",
|
|
])
|
|
|
|
# Start headless engines server (Node 1) - all engines, no API server
|
|
engines_server_args = self.base_server_args.copy()
|
|
engines_server_args.extend([
|
|
"--headless",
|
|
"--data-parallel-size",
|
|
str(self.dp_size),
|
|
"--data-parallel-size-local",
|
|
str(self.dp_size), # All engines on this node
|
|
"--tensor-parallel-size",
|
|
str(self.tp_size),
|
|
"--data-parallel-address",
|
|
"127.0.0.1",
|
|
"--data-parallel-rpc-port",
|
|
"13345",
|
|
])
|
|
|
|
# Use threads to start both servers in parallel
|
|
def start_api_server():
|
|
try:
|
|
server = RemoteOpenAIServer(
|
|
self.model_name,
|
|
api_server_args,
|
|
auto_port=False,
|
|
env_dict={}) # No GPUs needed for API-only server
|
|
server.__enter__()
|
|
print(f"API-only server started successfully with "
|
|
f"{self.api_server_count} API servers")
|
|
self.servers[0] = (server, api_server_args)
|
|
except Exception as e:
|
|
print(f"Failed to start API-only server: {e}")
|
|
raise
|
|
|
|
def start_engines_server():
|
|
try:
|
|
server = RemoteOpenAIServer(
|
|
self.model_name,
|
|
engines_server_args,
|
|
auto_port=False,
|
|
env_dict={
|
|
current_platform.device_control_env_var:
|
|
",".join(
|
|
str(
|
|
current_platform.
|
|
device_id_to_physical_device_id(i))
|
|
for i in range(self.dp_size * self.tp_size))
|
|
})
|
|
server.__enter__()
|
|
print(f"Headless engines server started successfully with "
|
|
f"{self.dp_size} engines")
|
|
self.servers[1] = (server, engines_server_args)
|
|
except Exception as e:
|
|
print(f"Failed to start headless engines server: {e}")
|
|
raise
|
|
|
|
# Start API server first
|
|
api_thread = threading.Thread(target=start_api_server)
|
|
api_thread.start()
|
|
self.server_threads.append(api_thread)
|
|
|
|
# Start engines server second
|
|
engines_thread = threading.Thread(target=start_engines_server)
|
|
engines_thread.start()
|
|
self.server_threads.append(engines_thread)
|
|
|
|
# Wait for both servers to start
|
|
for thread in self.server_threads:
|
|
thread.join()
|
|
|
|
# Give servers additional time to fully initialize and coordinate
|
|
time.sleep(3)
|
|
|
|
if not all(self.servers):
|
|
raise Exception("Both servers failed to start")
|
|
|
|
return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Stop both server instances."""
|
|
while self.servers:
|
|
if server := self.servers.pop():
|
|
try:
|
|
server[0].__exit__(exc_type, exc_val, exc_tb)
|
|
except Exception as e:
|
|
print(f"Error stopping server: {e}")
|
|
traceback.print_exc()
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def default_server_args():
|
|
return [
|
|
# use half precision for speed and memory savings in CI environment
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"2048",
|
|
"--max-num-seqs",
|
|
"128",
|
|
"--enforce-eager",
|
|
]
|
|
|
|
|
|
@pytest.fixture(scope="module", params=[1, 4])
|
|
def servers(request, default_server_args):
|
|
api_server_count = request.param
|
|
with MultinodeInternalLBServerManager(MODEL_NAME, DP_SIZE,
|
|
api_server_count,
|
|
default_server_args,
|
|
DP_SIZE // NUM_NODES,
|
|
TP_SIZE) as server_list:
|
|
yield server_list
|
|
|
|
|
|
@pytest.fixture(scope="module", params=[1, 4])
|
|
def api_only_servers(request, default_server_args):
|
|
"""Fixture for API-only server + headless engines configuration."""
|
|
api_server_count = request.param
|
|
with APIOnlyServerManager(MODEL_NAME, DP_SIZE, api_server_count,
|
|
default_server_args, TP_SIZE) as server_list:
|
|
yield server_list
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def client(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
|
|
# For internal LB, we only connect to the head node (rank 0)
|
|
# which provides the single API endpoint
|
|
head_server = servers[0][0]
|
|
async with head_server.get_async_client() as client:
|
|
yield client
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def api_only_client(api_only_servers: list[tuple[RemoteOpenAIServer,
|
|
list[str]]]):
|
|
"""Client fixture for API-only server configuration."""
|
|
# Connect to the API-only server (first server in the list)
|
|
api_server = api_only_servers[0][0]
|
|
async with api_server.get_async_client() as client:
|
|
yield client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"model_name",
|
|
[MODEL_NAME],
|
|
)
|
|
async def test_multinode_dp_completion(client: openai.AsyncOpenAI,
|
|
servers: list[tuple[RemoteOpenAIServer,
|
|
list[str]]],
|
|
model_name: str) -> None:
|
|
|
|
async def make_request():
|
|
completion = await client.completions.create(
|
|
model=model_name,
|
|
prompt="Hello, my name is",
|
|
max_tokens=5,
|
|
temperature=1.0)
|
|
|
|
assert completion.id is not None
|
|
assert completion.choices is not None and len(completion.choices) == 1
|
|
|
|
choice = completion.choices[0]
|
|
# The exact number of tokens can vary slightly with temperature=1.0,
|
|
# so we check for a reasonable minimum length.
|
|
assert len(choice.text) >= 1
|
|
# Finish reason might not always be 'length' if the model finishes early
|
|
# or due to other reasons, especially with high temperature.
|
|
# So, we'll accept 'length' or 'stop'.
|
|
assert choice.finish_reason in ("length", "stop")
|
|
|
|
# Token counts can also vary, so we check they are positive.
|
|
assert completion.usage.completion_tokens > 0
|
|
assert completion.usage.prompt_tokens > 0
|
|
assert completion.usage.total_tokens > 0
|
|
return completion
|
|
|
|
# Test single request
|
|
result = await make_request()
|
|
assert result is not None
|
|
print(
|
|
"Multi-node internal LB handled single completion request successfully"
|
|
)
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Send multiple requests - internal LB should distribute across DP ranks
|
|
num_requests = 200
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(completion is not None for completion in results)
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Second burst of requests
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(completion is not None for completion in results)
|
|
|
|
_, server_args = servers[0]
|
|
api_server_count = (
|
|
server_args.count('--api-server-count')
|
|
and server_args[server_args.index('--api-server-count') + 1] or 1)
|
|
print(f"Successfully completed multi-node internal LB test with "
|
|
f"{len(servers)} DP ranks (API server count: {api_server_count})")
|
|
|
|
# Check request balancing via Prometheus metrics
|
|
head_server = servers[0][0]
|
|
check_request_balancing(head_server, DP_SIZE)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"model_name",
|
|
[MODEL_NAME],
|
|
)
|
|
async def test_multinode_dp_completion_streaming(client: openai.AsyncOpenAI,
|
|
servers: list[
|
|
tuple[RemoteOpenAIServer,
|
|
list[str]]],
|
|
model_name: str) -> None:
|
|
prompt = "What is an LLM?"
|
|
|
|
async def make_streaming_request():
|
|
# Perform a non-streaming request to get the expected full output
|
|
single_completion = await client.completions.create(
|
|
model=model_name,
|
|
prompt=prompt,
|
|
max_tokens=5,
|
|
temperature=0.0,
|
|
)
|
|
single_output = single_completion.choices[0].text
|
|
|
|
# Perform the streaming request
|
|
stream = await client.completions.create(model=model_name,
|
|
prompt=prompt,
|
|
max_tokens=5,
|
|
temperature=0.0,
|
|
stream=True)
|
|
chunks: list[str] = []
|
|
finish_reason_count = 0
|
|
last_chunk = None
|
|
async for chunk in stream:
|
|
chunks.append(chunk.choices[0].text)
|
|
if chunk.choices[0].finish_reason is not None:
|
|
finish_reason_count += 1
|
|
last_chunk = chunk # Keep track of the last chunk
|
|
|
|
# finish reason should only return in the last block for OpenAI API
|
|
assert finish_reason_count == 1, (
|
|
"Finish reason should appear exactly once.")
|
|
assert last_chunk is not None, (
|
|
"Stream should have yielded at least one chunk.")
|
|
assert last_chunk.choices[
|
|
0].finish_reason == "length", "Finish reason should be 'length'."
|
|
# Check that the combined text matches the non-streamed version.
|
|
assert "".join(
|
|
chunks
|
|
) == single_output, "Streamed output should match non-streamed output."
|
|
return True # Indicate success for this request
|
|
|
|
# Test single streaming request
|
|
result = await make_streaming_request()
|
|
assert result is not None
|
|
print(
|
|
"Multi-node internal LB handled single streaming request successfully")
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Send multiple streaming requests - internal LB should distribute across
|
|
# DP ranks
|
|
num_requests = 200
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(results), "Not all streaming requests completed successfully."
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Second burst of streaming requests
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(results), "Not all streaming requests completed successfully."
|
|
|
|
_, server_args = servers[0]
|
|
api_server_count = (
|
|
server_args.count('--api-server-count')
|
|
and server_args[server_args.index('--api-server-count') + 1] or 1)
|
|
print(f"Successfully completed multi-node internal LB streaming test with "
|
|
f"{len(servers)} DP ranks (API server count: {api_server_count})")
|
|
|
|
# Check request balancing via Prometheus metrics
|
|
head_server = servers[0][0]
|
|
check_request_balancing(head_server, DP_SIZE)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"model_name",
|
|
[MODEL_NAME],
|
|
)
|
|
async def test_api_only_multinode_dp_completion(
|
|
api_only_client: openai.AsyncOpenAI,
|
|
api_only_servers: list[tuple[RemoteOpenAIServer,
|
|
list[str]]], model_name: str) -> None:
|
|
"""Test API-only server with all engines on separate headless server."""
|
|
|
|
async def make_request():
|
|
completion = await api_only_client.completions.create(
|
|
model=model_name,
|
|
prompt="Hello, my name is",
|
|
max_tokens=5,
|
|
temperature=1.0)
|
|
|
|
assert completion.id is not None
|
|
assert completion.choices is not None and len(completion.choices) == 1
|
|
|
|
choice = completion.choices[0]
|
|
# The exact number of tokens can vary slightly with temperature=1.0,
|
|
# so we check for a reasonable minimum length.
|
|
assert len(choice.text) >= 1
|
|
# Finish reason might not always be 'length' if the model finishes
|
|
# early or due to other reasons, especially with high temperature.
|
|
# So, we'll accept 'length' or 'stop'.
|
|
assert choice.finish_reason in ("length", "stop")
|
|
|
|
# Token counts can also vary, so we check they are positive.
|
|
assert completion.usage.completion_tokens > 0
|
|
assert completion.usage.prompt_tokens > 0
|
|
assert completion.usage.total_tokens > 0
|
|
return completion
|
|
|
|
# Test single request
|
|
result = await make_request()
|
|
assert result is not None
|
|
print("API-only server handled single completion request successfully")
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Send multiple requests - should be distributed across engines on
|
|
# headless server
|
|
num_requests = 200
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(completion is not None for completion in results)
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Second burst of requests
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(completion is not None for completion in results)
|
|
|
|
api_server, api_server_args = api_only_servers[0]
|
|
api_server_count = (
|
|
api_server_args.count('--api-server-count')
|
|
and api_server_args[api_server_args.index('--api-server-count') + 1]
|
|
or 1)
|
|
print(f"Successfully completed API-only multi-node test with {DP_SIZE} "
|
|
f"engines on headless server (API server count: {api_server_count})")
|
|
|
|
# Check request balancing via Prometheus metrics
|
|
check_request_balancing(api_server, DP_SIZE)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"model_name",
|
|
[MODEL_NAME],
|
|
)
|
|
async def test_api_only_multinode_dp_completion_streaming(
|
|
api_only_client: openai.AsyncOpenAI,
|
|
api_only_servers: list[tuple[RemoteOpenAIServer,
|
|
list[str]]], model_name: str) -> None:
|
|
"""Test API-only server streaming with all engines on separate
|
|
headless server."""
|
|
prompt = "What is an LLM?"
|
|
|
|
async def make_streaming_request():
|
|
# Perform a non-streaming request to get the expected full output
|
|
single_completion = await api_only_client.completions.create(
|
|
model=model_name,
|
|
prompt=prompt,
|
|
max_tokens=5,
|
|
temperature=0.0,
|
|
)
|
|
single_output = single_completion.choices[0].text
|
|
|
|
# Perform the streaming request
|
|
stream = await api_only_client.completions.create(model=model_name,
|
|
prompt=prompt,
|
|
max_tokens=5,
|
|
temperature=0.0,
|
|
stream=True)
|
|
chunks: list[str] = []
|
|
finish_reason_count = 0
|
|
last_chunk = None
|
|
async for chunk in stream:
|
|
chunks.append(chunk.choices[0].text)
|
|
if chunk.choices[0].finish_reason is not None:
|
|
finish_reason_count += 1
|
|
last_chunk = chunk # Keep track of the last chunk
|
|
|
|
# finish reason should only return in the last block for OpenAI API
|
|
assert finish_reason_count == 1, (
|
|
"Finish reason should appear exactly once.")
|
|
assert last_chunk is not None, (
|
|
"Stream should have yielded at least one chunk.")
|
|
assert last_chunk.choices[
|
|
0].finish_reason == "length", "Finish reason should be 'length'."
|
|
# Check that the combined text matches the non-streamed version.
|
|
assert "".join(
|
|
chunks
|
|
) == single_output, "Streamed output should match non-streamed output."
|
|
return True # Indicate success for this request
|
|
|
|
# Test single streaming request
|
|
result = await make_streaming_request()
|
|
assert result is not None
|
|
print("API-only server handled single streaming request successfully")
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Send multiple streaming requests - should be distributed across engines
|
|
num_requests = 200
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(results), "Not all streaming requests completed successfully."
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
# Second burst of streaming requests
|
|
all_tasks = []
|
|
for _ in range(num_requests):
|
|
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
|
await asyncio.sleep(0.01)
|
|
|
|
results = await asyncio.gather(*all_tasks)
|
|
assert len(results) == num_requests
|
|
assert all(results), "Not all streaming requests completed successfully."
|
|
|
|
_, api_server_args = api_only_servers[0]
|
|
api_server_count = (
|
|
api_server_args.count('--api-server-count')
|
|
and api_server_args[api_server_args.index('--api-server-count') + 1]
|
|
or 1)
|
|
print(f"Successfully completed API-only streaming test with {DP_SIZE} "
|
|
f"engines on headless server (API server count: {api_server_count})")
|
|
|
|
# Check request balancing via Prometheus metrics
|
|
api_server = api_only_servers[0][0]
|
|
check_request_balancing(api_server, DP_SIZE)
|