[model-gateway] rename py_test to e2e_test (#16454)

This commit is contained in:
Simo Lin
2026-01-04 22:24:34 -08:00
committed by GitHub
parent f02d82211a
commit e6fe092dcc
38 changed files with 9 additions and 9 deletions
+1
View File
@@ -0,0 +1 @@
"""Test package root for router Python tests."""
+15
View File
@@ -0,0 +1,15 @@
import sys
from importlib.util import find_spec
from pathlib import Path
# Only add bindings/python to path if the wheel is not installed (for local development)
# This ensures CI tests use the installed wheel which contains the Rust extension
_ROOT = Path(__file__).resolve().parents[1]
_SRC = _ROOT / "bindings" / "python"
# Check if sglang_router is already installed with the Rust extension
_wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None
# Only add bindings/python if wheel is not installed (development mode)
if not _wheel_installed and str(_SRC) not in sys.path:
sys.path.insert(0, str(_SRC))
@@ -0,0 +1,264 @@
"""
gRPC Router E2E Test - Embedding Correctness
Test that embeddings from the gRPC router match HuggingFace reference embeddings.
Validates numerical correctness including tokenization (BOS/EOS handling) and inference.
"""
import logging
import sys
import unittest
from pathlib import Path
from typing import Dict, List, Union
import numpy as np
import torch
import torch.nn.functional as F
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_EMBEDDING_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
kill_process_tree,
)
# Configure logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
# Test data for semantic similarity checks
SEMANTIC_TEST_SETS: List[List[str]] = [
[
"The cat sat on the mat.",
"A feline was resting on a rug.",
"Bright stars illuminate the night sky.", # Unrelated sentence
],
[
"The quick brown fox jumps over the lazy dog.",
"A fast, dark-colored fox leaps above a sluggish canine.",
"Ocean waves gently lap against the shore.", # Unrelated sentence
],
[
"An apple a day keeps the doctor away.",
"Eating a daily apple can prevent medical visits.",
"Mountains are vast and often snow-capped.", # Unrelated sentence
],
]
# Test data for relevance scoring
RELEVANCE_TEST_DATA: Dict = {
"sample_query": "Why is Oracle launching Cloud Lift Services?",
"sample_reference": [
{
"docid": 466,
"body": "What are some extended benefits of using Oracle Cloud Infrastructure? \nWhen customers migrate their on-premises Oracle applications to Oracle Cloud Infrastructure, they realize the benefits \nof the cloud without needing to rearchitect those applications. Customers can lower total cost of ownership, improve \nagility and increase workload performance. Additional benefits include: \nConsistently low global pricing and lack of hidden charges \nAutomated migration support, leveraging cloud managers and tools for key applications \nFlexible universal credits applied towards any IaaS or PaaS service \nBring Your Own License (BYOL) capabilities \nIs Oracle Cloud Lift available for PAYGO customers? \nOracle Cloud Lift Services are designed for customers who use the UCM credits (Monthly Flex). PAYGO customers can \ncontact their sales representative or cloud engineer to evaluate their eligibility. \nAre any countries excluded from Oracle Cloud Lift Services? \nAmong the countries that Oracle operates in, only China is excluded from the Oracle Cloud Lift Services program. \nHow does Oracle Cloud Lift Services impact Oracle partners? \nThe Cloud Lift Services program has been socialized with select partners to both complement and grow their existing \nbusinesses. Oracle is committed to working with different partner business models, from services partners to resellers, \nISVs, CSPs, etc. For additional details, partners should contact their Oracle representative. \nHow will a partner benefit from Oracle Cloud Lift Services model? \nThe purpose of Oracle Cloud Lift Services is to serve our collective customer base and accelerate growth in our \necosystem. Oracle Cloud Engineering will now provide guidance on planning, architecting, prototyping, and managing \ncloud migrations. When partners are leading an opportunity, Oracle will work with and through our partners to offer \nCloud Lift Services as needed to make our joint customers more successful. Public Sector accounts and partner \nengagements are not currently eligible to participate in this program.",
},
{
"docid": 636,
"body": "Cloud Lift Services as needed to make our joint customers more successful. Public Sector accounts and partner \nengagements are not currently eligible to participate in this program. \n How can I get started with Oracle Cloud? \nYou can use the Oracle Cloud Free Tier for a free trial and Contact Us for more information. \n \n \n blogs.oracle.com \n facebook.com/OracleCloud/ \n twitter.com/OracleCloud/ \n linkedin.com/showcase/oracle-cloud/\u2028 \n \n \n2 \nFrequently Asked Questions / Oracle Cloud Lift Services / Version 1.2 \n \n \nCopyright \u00a9 2021, Oracle and/or its affiliates / Public",
},
{
"docid": 545,
"body": "Frequently Asked Questions (FAQs) for \nOracle Cloud Lift Services \n \nWhy is Oracle launching Cloud Lift Services? \n \n \n \nThis program underscores Oracle\u2019s intent to better serve its customer base. Cloud Lift Services provide new and \nexisting customers expanded access to cloud engineering tools and resources to quickly migrate workloads at no \nadditional cost. \nHow are Oracle Cloud Lift Services different from pre-sales activities such as Proof-of-Concepts (POCs)? \nWhile POCs and other presales help are available from Oracle, Cloud Lift Services are post-sales and part of the \nenterprise contract. Migration and go-live support for eligible workloads mean that our experts can engage during and \nafter the sales process to help bring workloads into production faster. \nWhat\u2019s included \u2013 and excluded \u2013 from Oracle Cloud Lift Services? \nThe Oracle Cloud Lift Services web page provides details on included vs. excluded services. We encourage you to work \nwith your Oracle sales representative to talk through the details of your plans. In general terms, migrating up to ten \nOracle Databases, Oracle applications, Cloud Native or HPC applications can be included, while more than 10 \nmigrations, complex migrations involving new business logic, platform upgrades, or custom development are not. \nHow do new customers get access to Oracle Cloud Lift Services? \nNew customers work with cloud engineering during the contract process to create and agree on a documented work \nplan which lays out the specific eligible workloads, timelines, and other details. \nHow do existing customers get access to Oracle Cloud Lift Services? \nExisting customers work with cloud engineering or sales to have an addendum for Oracle Cloud Lift Services included \nas part of their existing contract. \nWhat happens if I already have a paid services engagement?",
},
{
"docid": 716,
"body": "as part of their existing contract. \nWhat happens if I already have a paid services engagement? \nPlease keep proceeding with your existing engagement. Oracle will work with you to identify expansion opportunities \nto leverage Cloud Lift Services for other projects. \nHow do I decide whether to use Oracle Cloud Lift Services? \nOracle Cloud Lift Services can meet customer needs when cloud engineers can help migrate a few applications to OCI; \nmigrate applications without updating software versions; configure OCI tenancies, compartments, quotas and \nidentities; perform basic reviews of network configurations and security, FastConnect setup, auditing, and assessing \nregulatory compliance; and train in-house resources on OCI. \nHow does Oracle offer support post go-live solutions as a part of Oracle Lift Services? \nFor cloud-based continuous optimization services, customers can use Oracle\u2019s partners, Oracle Consulting, and Oracle \nAdvanced Customer Services (ACS), which provides joint and fully managed 24/7 lifecycle services for database, \napplications and security. \n \n \n \n1 \nFrequently Asked Questions / Oracle Cloud Lift Services / Version 1.2 \n \n \nCopyright \u00a9 2021, Oracle and/or its affiliates / Public \n \nAre Oracle Cloud Lift Services being offered to public sector? \nAligning to the rules and regulations that govern our public sector customers, Oracle will make Cloud Lift Services \navailable to all North American public sector customers in the near-term. Oracle is committed to making these services \navailable to public sector customers globally. \nWhat are some extended benefits of using Oracle Cloud Infrastructure? \nWhen customers migrate their on-premises Oracle applications to Oracle Cloud Infrastructure, they realize the benefits",
},
],
}
def get_openai_embeddings(
texts: Union[str, List[str]], config: Dict
) -> List[List[float]]:
"""Get embeddings from the gateway via OpenAI-compatible API."""
import openai
client = openai.Client(api_key=config["api_key"], base_url=config["base_url"])
if isinstance(texts, str):
texts = [texts]
embeddings = []
for text in texts:
response = client.embeddings.create(
model=config["model_name"],
input=text,
)
embeddings.append(response.data[0].embedding)
return embeddings
def get_hf_st_embeddings(texts: Union[str, List[str]], model_path: str) -> np.ndarray:
"""Get embeddings using sentence-transformers library.
This handles the correct pooling strategy for each model automatically.
For e5-mistral, it uses last-token pooling (not mean pooling).
Uses CPU to compute reference embeddings to avoid GPU memory conflicts
with the worker being tested. This is acceptable since reference embeddings
only need to be accurate, not fast.
"""
from sentence_transformers import SentenceTransformer
if isinstance(texts, str):
texts = [texts]
# Force CPU to avoid GPU memory conflicts in CI where GPUs may be
# used by other workers. Reference embeddings just need accuracy, not speed.
model = SentenceTransformer(model_path, trust_remote_code=True, device="cpu")
embeddings = model.encode(texts, normalize_embeddings=True)
return embeddings
def get_input_texts(test_json: Dict) -> List[str]:
"""Extract document bodies from test JSON."""
return [doc["body"] for doc in test_json["sample_reference"]]
def compare_embeddings(
embeddings1: List[List[float]], embeddings2: List[List[float]]
) -> List[float]:
"""Compare two sets of embeddings using cosine similarity."""
logging.info("Comparing embeddings")
similarities = [
F.cosine_similarity(torch.tensor(e1), torch.tensor(e2), dim=0).item()
for e1, e2 in zip(embeddings1, embeddings2)
]
return similarities
class TestEmbeddingCorrectness(CustomTestCase):
"""Test embedding correctness by comparing gateway output against HuggingFace reference.
Strategy: Pre-compute HuggingFace reference embeddings on CPU, then launch the
worker on GPU and compare. Using CPU for reference avoids GPU memory conflicts
in CI where multiple workers may share limited GPU resources.
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_EMBEDDING_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
# Pre-compute all reference embeddings on CPU before launching the worker
# This avoids GPU memory conflicts in CI environments with limited GPUs
logging.info(
f"Pre-computing HuggingFace reference embeddings (CPU) for {cls.model}"
)
# Flatten all test texts for semantic similarity
all_semantic_texts = []
for text_set in SEMANTIC_TEST_SETS:
all_semantic_texts.extend(text_set)
# Get relevance test texts
query = f"Instruct: Given a search query, retrieve relevant passages that answer the query\nQuery: {RELEVANCE_TEST_DATA['sample_query']}"
docs = get_input_texts(RELEVANCE_TEST_DATA)
# Compute all reference embeddings at once
cls.hf_semantic_embeddings = get_hf_st_embeddings(all_semantic_texts, cls.model)
cls.hf_query_embedding = get_hf_st_embeddings(query, cls.model)
cls.hf_docs_embeddings = get_hf_st_embeddings(docs, cls.model)
logging.info("Reference embeddings computed on CPU")
# Now launch workers with --is-embedding flag
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
num_workers=1,
tp_size=1,
policy="round_robin",
api_key=cls.api_key,
worker_args=["--is-embedding"],
)
cls.config = {
"server_engine": "sgl-model-gateway",
"base_url": cls.base_url + "/v1",
"model_name": cls.model,
"model_path": cls.model,
"api_key": cls.api_key,
}
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def test_semantic_similarity(self, tolerance: float = 1e-2):
"""Check if gateway and HF embeddings give similar results."""
# Track position in pre-computed embeddings
embed_idx = 0
for i, input_texts in enumerate(SEMANTIC_TEST_SETS):
logging.info(f"Processing semantic similarity test set {i + 1}")
embedding_gateway = get_openai_embeddings(input_texts, self.config)
# Get pre-computed HF embeddings for this set
num_texts = len(input_texts)
embedding_hf = self.hf_semantic_embeddings[
embed_idx : embed_idx + num_texts
].tolist()
embed_idx += num_texts
logging.info(f'Comparing {self.config["server_engine"]} and HF embeddings')
similarities = compare_embeddings(embedding_gateway, embedding_hf)
logging.info(f"Similarities between embeddings: {similarities}")
# Verify similarities
for sim in similarities:
self.assertLess(
abs(sim - 1.0), tolerance, f"Similarity {sim} is not close to 1"
)
logging.info(f"Semantic similarity test set {i + 1} passed\n")
def test_relevance_scores(self, tolerance: float = 0.05):
"""Compare relevance scores between gateway and HF implementations."""
logging.info(
f'Comparing relevance scores between {self.config["server_engine"]} and HF'
)
# Format query with instruction (for e5-mistral)
query = f"Instruct: Given a search query, retrieve relevant passages that answer the query\nQuery: {RELEVANCE_TEST_DATA['sample_query']}"
docs = get_input_texts(RELEVANCE_TEST_DATA)
# Get gateway scores
query_embeddings_gateway = get_openai_embeddings(query, self.config)
docs_embeddings_gateway = get_openai_embeddings(docs, self.config)
scores_gateway = (
np.array(query_embeddings_gateway) @ np.array(docs_embeddings_gateway).T
) * 100
# Use pre-computed HF scores
scores_hf = (self.hf_query_embedding @ self.hf_docs_embeddings.T) * 100
logging.info(
f'{self.config["server_engine"]} relevance scores: {scores_gateway}'
)
logging.info(f"HF relevance scores: {scores_hf}")
self.assertTrue(
np.allclose(scores_gateway, scores_hf, atol=tolerance),
f'Scores differ beyond tolerance: \n{self.config["server_engine"]}: {scores_gateway}\nHF: {scores_hf}',
)
logging.info("Relevance scores comparison completed successfully")
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,109 @@
"""
gRPC Router E2E Test - Embedding Server
Test the embedding functionality of the gRPC router.
"""
import sys
import unittest
from pathlib import Path
import openai
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_EMBEDDING_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
kill_process_tree,
)
class TestEmbeddingServer(CustomTestCase):
"""
Test Embedding API through gRPC router.
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_EMBEDDING_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
# Launch workers with --is-embedding flag
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
num_workers=1,
tp_size=1,
policy="round_robin",
api_key=cls.api_key,
worker_args=["--is-embedding"],
)
cls.base_url += "/v1"
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def test_embedding(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
input_text = "Hello world"
response = client.embeddings.create(
model=self.model,
input=input_text,
)
assert response.object == "list"
assert len(response.data) == 1
embedding = response.data[0]
assert embedding.object == "embedding"
assert embedding.index == 0
assert len(embedding.embedding) > 0
assert isinstance(embedding.embedding[0], float)
# Verify usage statistics
assert response.usage.prompt_tokens > 0
assert response.usage.total_tokens == response.usage.prompt_tokens
def test_embedding_batch(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
input_texts = ["Hello world", "SGLang is fast"]
response = client.embeddings.create(
model=self.model,
input=input_texts,
)
assert len(response.data) == 1
assert response.data[0].index == 0
assert len(response.data[0].embedding) > 0
def test_embedding_dimensions(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response1 = client.embeddings.create(
model=self.model,
input="A short text",
)
dim1 = len(response1.data[0].embedding)
response2 = client.embeddings.create(
model=self.model,
input="A much longer text to ensure dimensions match",
)
dim2 = len(response2.data[0].embedding)
assert dim1 == dim2
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,343 @@
"""
gRPC Router E2E Test - OpenAI Server API Compatibility
This test file is REUSED from test/srt/openai_server/basic/test_openai_server.py
with minimal changes:
- Swap popen_launch_server() → popen_launch_workers_and_router()
- Update teardown to cleanup router + workers
- All test logic and assertions remain identical
Run with:
python3 -m pytest e2e_grpc/basic/test_openai_server.py -v
python3 -m unittest e2e_grpc.basic.test_openai_server.TestOpenAIServer.test_completion
"""
import json
import sys
import unittest
from pathlib import Path
import openai
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_GPT_OSS_MODEL_PATH,
DEFAULT_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
get_tokenizer,
kill_process_tree,
)
class TestOpenAIServer(CustomTestCase):
"""
Test OpenAI API through gRPC router.
REUSED from test/srt/openai_server/basic/test_openai_server.py
ONLY CHANGE: Server launch mechanism
- Launches SGLang workers with --enable-grpc
- Launches gRPC router pointing to those workers
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
num_workers=1,
tp_size=2,
policy="round_robin",
api_key=cls.api_key,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
# ALL TEST METHODS BELOW ARE UNCHANGED FROM ORIGINAL
# They validate that the router maintains OpenAI API compatibility
def run_chat_completion(self, logprobs, parallel_sample_num):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant"},
{
"role": "user",
"content": "What is the capital of France? Answer in a few words.",
},
],
temperature=0,
logprobs=logprobs is not None and logprobs > 0,
top_logprobs=logprobs,
n=parallel_sample_num,
)
if logprobs:
assert isinstance(
response.choices[0].logprobs.content[0].top_logprobs[0].token, str
)
ret_num_top_logprobs = len(
response.choices[0].logprobs.content[0].top_logprobs
)
assert (
ret_num_top_logprobs == logprobs
), f"{ret_num_top_logprobs} vs {logprobs}"
assert len(response.choices) == parallel_sample_num
assert response.choices[0].message.role == "assistant"
assert isinstance(response.choices[0].message.content, str)
assert response.id
assert response.created
assert response.usage.prompt_tokens > 0
assert response.usage.completion_tokens > 0
assert response.usage.total_tokens > 0
def run_chat_completion_stream(self, logprobs, parallel_sample_num=1):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
generator = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant"},
{"role": "user", "content": "What is the capital of France?"},
],
temperature=0,
logprobs=logprobs is not None and logprobs > 0,
top_logprobs=logprobs,
stream=True,
stream_options={"include_usage": True},
n=parallel_sample_num,
)
is_firsts = {}
is_finished = {}
finish_reason_counts = {}
for response in generator:
usage = response.usage
if usage is not None:
assert usage.prompt_tokens > 0, f"usage.prompt_tokens was zero"
assert usage.completion_tokens > 0, f"usage.completion_tokens was zero"
assert usage.total_tokens > 0, f"usage.total_tokens was zero"
continue
index = response.choices[0].index
finish_reason = response.choices[0].finish_reason
if finish_reason is not None:
is_finished[index] = True
finish_reason_counts[index] = finish_reason_counts.get(index, 0) + 1
data = response.choices[0].delta
if is_firsts.get(index, True):
assert (
data.role == "assistant"
), f"data.role was not 'assistant' for first chunk"
is_firsts[index] = False
continue
if logprobs and not is_finished.get(index, False):
assert response.choices[0].logprobs, f"logprobs was not returned"
assert isinstance(
response.choices[0].logprobs.content[0].top_logprobs[0].token, str
), f"top_logprobs token was not a string"
assert isinstance(
response.choices[0].logprobs.content[0].top_logprobs, list
), f"top_logprobs was not a list"
ret_num_top_logprobs = len(
response.choices[0].logprobs.content[0].top_logprobs
)
assert (
ret_num_top_logprobs == logprobs
), f"{ret_num_top_logprobs} vs {logprobs}"
assert (
isinstance(data.content, str)
or isinstance(data.reasoning_content, str)
or (isinstance(data.tool_calls, list) and len(data.tool_calls) > 0)
or response.choices[0].finish_reason
)
assert response.id
assert response.created
for index in [i for i in range(parallel_sample_num)]:
assert not is_firsts.get(
index, True
), f"index {index} is not found in the response"
for index in range(parallel_sample_num):
assert (
index in finish_reason_counts
), f"No finish_reason found for index {index}"
assert (
finish_reason_counts[index] == 1
), f"Expected 1 finish_reason chunk for index {index}, got {finish_reason_counts[index]}"
def test_chat_completion(self):
for logprobs in [None, 5]:
for parallel_sample_num in [1, 2]:
self.run_chat_completion(logprobs, parallel_sample_num)
def test_chat_completion_stream(self):
for logprobs in [None, 5]:
for parallel_sample_num in [1, 2]:
self.run_chat_completion_stream(logprobs, parallel_sample_num)
def test_regex(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
regex = (
r"""\{\n"""
+ r""" "name": "[\w]+",\n"""
+ r""" "population": [\d]+\n"""
+ r"""\}"""
)
response = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant"},
{"role": "user", "content": "Introduce the capital of France."},
],
temperature=0,
max_tokens=128,
extra_body={"regex": regex},
)
text = response.choices[0].message.content
try:
js_obj = json.loads(text)
except (TypeError, json.decoder.JSONDecodeError):
raise
assert isinstance(js_obj["name"], str)
assert isinstance(js_obj["population"], int)
def test_penalty(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant"},
{"role": "user", "content": "Introduce the capital of France."},
],
temperature=0,
max_tokens=32,
frequency_penalty=1.0,
)
text = response.choices[0].message.content
assert isinstance(text, str)
def test_response_prefill(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant"},
{
"role": "user",
"content": """
Extract the name, size, price, and color from this product description as a JSON object:
<description>
The SmartHome Mini is a compact smart home assistant available in black or white for only $49.99. At just 5 inches wide, it lets you control lights, thermostats, and other connected devices via voice or app—no matter where you place it in your home. This affordable little hub brings convenient hands-free control to your smart devices.
</description>
""",
},
{
"role": "assistant",
"content": "{\n",
},
],
temperature=0,
extra_body={"continue_final_message": True},
)
assert (
response.choices[0]
.message.content.strip()
.startswith('"name": "SmartHome Mini",')
)
def test_model_list(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
models = list(client.models.list().data)
assert len(models) == 1
# assert isinstance(getattr(models[0], "max_model_len", None), int)
@unittest.skip("Skipping retrieve model test as it is not supported by the router")
def test_retrieve_model(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
retrieved_model = client.models.retrieve(self.model)
self.assertEqual(retrieved_model.id, self.model)
self.assertEqual(retrieved_model.root, self.model)
with self.assertRaises(openai.NotFoundError):
client.models.retrieve("non-existent-model")
class TestOpenAIServerGptOss(TestOpenAIServer):
"""
Test OpenAI API through gRPC router with openai/gpt-oss-20b model.
Extends TestOpenAIServer and only changes the model.
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_GPT_OSS_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
num_workers=1,
tp_size=2,
policy="round_robin",
api_key=cls.api_key,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
def test_chat_completion(self):
for parallel_sample_num in [1, 2]:
self.run_chat_completion(None, parallel_sample_num)
def test_chat_completion_stream(self):
for parallel_sample_num in [1, 2]:
self.run_chat_completion_stream(None, parallel_sample_num)
@unittest.skip("Skipping for OSS models")
def test_regex(self):
super().test_regex()
@unittest.skip("Skipping for OSS models")
def test_response_prefill(self):
super().test_response_prefill()
@unittest.skip("Skipping for OSS models")
def test_penalty(self):
super().test_penalty()
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,29 @@
"""
Pytest configuration for gRPC router e2e tests.
This module provides shared fixtures that can be used across all gRPC router tests.
"""
import sys
from pathlib import Path
import pytest # noqa: F401
# Ensure router bindings/python is importable
_ROUTER_ROOT = Path(__file__).resolve().parents[2]
_ROUTER_SRC = _ROUTER_ROOT / "bindings" / "python"
if str(_ROUTER_SRC) not in sys.path:
sys.path.insert(0, str(_ROUTER_SRC))
# Ensure e2e_grpc test utilities are importable
_E2E_GRPC_DIR = Path(__file__).parent
if str(_E2E_GRPC_DIR) not in sys.path:
sys.path.insert(0, str(_E2E_GRPC_DIR))
# Pytest markers for test organization
def pytest_configure(config):
config.addinivalue_line("markers", "e2e: end-to-end tests with real workers")
config.addinivalue_line("markers", "grpc: gRPC-specific tests")
config.addinivalue_line("markers", "slow: slow-running tests")
config.addinivalue_line("markers", "pd: prefill-decode disaggregation tests")
@@ -0,0 +1,194 @@
"""
Usage:
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_chat_completion_with_reasoning
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_chat_completion_without_reasoning
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_stream_chat_completion_with_reasoning
python3 -m unittest openai_server.features.test_enable_thinking.TestEnableThinking.test_stream_chat_completion_without_reasoning
"""
import json
import sys
import unittest
from pathlib import Path
import requests
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_ENABLE_THINKING_MODEL_PATH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
kill_process_tree,
)
class TestEnableThinking(CustomTestCase):
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
cls.model = DEFAULT_ENABLE_THINKING_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-1234"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=120,
api_key=cls.api_key,
router_args=[
"--reasoning-parser",
"qwen3",
],
num_workers=1,
tp_size=4,
)
cls.additional_chat_kwargs = {}
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def test_chat_completion_with_reasoning(self):
# Test non-streaming with "enable_thinking": True, reasoning_content should not be empty
client = requests.post(
f"{self.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {self.api_key}"},
json={
"model": self.model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"chat_template_kwargs": {"enable_thinking": True},
**self.additional_chat_kwargs,
},
)
self.assertEqual(client.status_code, 200, f"Failed with: {client.text}")
data = client.json()
self.assertIn("choices", data)
self.assertTrue(len(data["choices"]) > 0)
self.assertIn("message", data["choices"][0])
self.assertIn("reasoning_content", data["choices"][0]["message"])
self.assertIsNotNone(data["choices"][0]["message"]["reasoning_content"])
def test_chat_completion_without_reasoning(self):
# Test non-streaming with "enable_thinking": False, reasoning_content should be empty
client = requests.post(
f"{self.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {self.api_key}"},
json={
"model": self.model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"chat_template_kwargs": {"enable_thinking": False},
**self.additional_chat_kwargs,
},
)
self.assertEqual(client.status_code, 200, f"Failed with: {client.text}")
data = client.json()
self.assertIn("choices", data)
self.assertTrue(len(data["choices"]) > 0)
self.assertIn("message", data["choices"][0])
if "reasoning_content" in data["choices"][0]["message"]:
self.assertIsNone(data["choices"][0]["message"]["reasoning_content"])
def test_stream_chat_completion_with_reasoning(self):
# Test streaming with "enable_thinking": True, reasoning_content should not be empty
response = requests.post(
f"{self.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {self.api_key}"},
json={
"model": self.model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"stream": True,
"chat_template_kwargs": {"enable_thinking": True},
**self.additional_chat_kwargs,
},
stream=True,
)
self.assertEqual(response.status_code, 200, f"Failed with: {response.text}")
has_reasoning = False
has_content = False
for line in response.iter_lines():
if line:
line = line.decode("utf-8")
if line.startswith("data:") and not line.startswith("data: [DONE]"):
data = json.loads(line[6:])
if "choices" in data and len(data["choices"]) > 0:
delta = data["choices"][0].get("delta", {})
if "reasoning_content" in delta and delta["reasoning_content"]:
has_reasoning = True
if "content" in delta and delta["content"]:
has_content = True
self.assertTrue(
has_reasoning,
"The reasoning content is not included in the stream response",
)
self.assertTrue(
has_content, "The stream response does not contain normal content"
)
def test_stream_chat_completion_without_reasoning(self):
# Test streaming with "enable_thinking": False, reasoning_content should be empty
response = requests.post(
f"{self.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {self.api_key}"},
json={
"model": self.model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"stream": True,
"chat_template_kwargs": {"enable_thinking": False},
**self.additional_chat_kwargs,
},
stream=True,
)
self.assertEqual(response.status_code, 200, f"Failed with: {response.text}")
has_reasoning = False
has_content = False
for line in response.iter_lines():
if line:
line = line.decode("utf-8")
if line.startswith("data:") and not line.startswith("data: [DONE]"):
data = json.loads(line[6:])
if "choices" in data and len(data["choices"]) > 0:
delta = data["choices"][0].get("delta", {})
if "reasoning_content" in delta and delta["reasoning_content"]:
has_reasoning = True
if "content" in delta and delta["content"]:
has_content = True
self.assertFalse(
has_reasoning,
"The reasoning content should not be included in the stream response",
)
self.assertTrue(
has_content, "The stream response does not contain normal content"
)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,194 @@
"""
Usage:
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_streaming_separate_reasoning_false
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_streaming_separate_reasoning_true
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_streaming_separate_reasoning_true_stream_reasoning_false
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_nonstreaming_separate_reasoning_false
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentAPI.test_nonstreaming_separate_reasoning_true
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentStartup.test_nonstreaming
python3 -m unittest openai_server.features.test_reasoning_content.TestReasoningContentStartup.test_streaming
"""
import sys
import unittest
from pathlib import Path
import openai
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_REASONING_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
kill_process_tree,
)
class TestReasoningContentAPI(CustomTestCase):
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
cls.model = DEFAULT_REASONING_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-1234"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
router_args=[
"--reasoning-parser",
"deepseek_r1",
],
num_workers=1,
tp_size=2,
)
cls.base_url += "/v1"
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def test_streaming_separate_reasoning_false(self):
# Test streaming with separate_reasoning=False, reasoning_content should be empty
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
payload = {
"model": self.model,
"messages": [
{
"role": "user",
"content": "What is 1+3?",
}
],
"max_tokens": 100,
"stream": True,
"extra_body": {"separate_reasoning": False},
}
response = client.chat.completions.create(**payload)
reasoning_content = ""
content = ""
for chunk in response:
if chunk.choices[0].delta.content:
content += chunk.choices[0].delta.content
elif chunk.choices[0].delta.reasoning_content:
reasoning_content += chunk.choices[0].delta.reasoning_content
assert len(reasoning_content) == 0
assert len(content) > 0
def test_streaming_separate_reasoning_true(self):
# Test streaming with separate_reasoning=True, reasoning_content should not be empty
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
payload = {
"model": self.model,
"messages": [
{
"role": "user",
"content": "What is 1+3?",
}
],
"max_tokens": 100,
"stream": True,
"extra_body": {"separate_reasoning": True},
}
response = client.chat.completions.create(**payload)
reasoning_content = ""
content = ""
for chunk in response:
if chunk.choices[0].delta.content:
content += chunk.choices[0].delta.content
elif chunk.choices[0].delta.reasoning_content:
reasoning_content += chunk.choices[0].delta.reasoning_content
assert len(reasoning_content) > 0
assert len(content) > 0
def test_streaming_separate_reasoning_true_stream_reasoning_false(self):
# Test streaming with separate_reasoning=True, reasoning_content should not be empty
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
payload = {
"model": self.model,
"messages": [
{
"role": "user",
"content": "What is 1+3?",
}
],
"max_tokens": 100,
"stream": True,
"extra_body": {"separate_reasoning": True, "stream_reasoning": False},
}
response = client.chat.completions.create(**payload)
reasoning_content = ""
content = ""
first_chunk = False
for chunk in response:
if chunk.choices[0].delta.reasoning_content:
reasoning_content = chunk.choices[0].delta.reasoning_content
first_chunk = True
if chunk.choices[0].delta.content:
content += chunk.choices[0].delta.content
if not first_chunk:
reasoning_content = chunk.choices[0].delta.reasoning_content
first_chunk = True
if not first_chunk:
assert (
not chunk.choices[0].delta.reasoning_content
or len(chunk.choices[0].delta.reasoning_content) == 0
)
assert len(reasoning_content) > 0
assert len(content) > 0
def test_nonstreaming_separate_reasoning_false(self):
# Test non-streaming with separate_reasoning=False, reasoning_content should be empty
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
payload = {
"model": self.model,
"messages": [
{
"role": "user",
"content": "What is 1+3?",
}
],
"max_tokens": 100,
"extra_body": {"separate_reasoning": False},
}
response = client.chat.completions.create(**payload)
assert (
not response.choices[0].message.reasoning_content
or len(response.choices[0].message.reasoning_content) == 0
)
assert len(response.choices[0].message.content) > 0
def test_nonstreaming_separate_reasoning_true(self):
# Test non-streaming with separate_reasoning=True, reasoning_content should not be empty
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
payload = {
"model": self.model,
"messages": [
{
"role": "user",
"content": "What is 1+3?",
}
],
"max_tokens": 100,
"extra_body": {"separate_reasoning": True},
}
response = client.chat.completions.create(**payload)
assert len(response.choices[0].message.reasoning_content) > 0
assert len(response.choices[0].message.content) > 0
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,342 @@
"""
Fixtures for launching gRPC router + workers for e2e testing.
This module provides fixtures for launching SGLang workers and gRPC router separately:
1. Launch N SGLang workers with gRPC enabled
2. Launch router pointing to those workers
This approach gives more control and matches production deployment patterns.
"""
import logging
import socket
import subprocess
import time
from typing import Optional
import requests
logger = logging.getLogger(__name__)
def find_free_port() -> int:
"""Find an available port on localhost."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def wait_for_workers_ready(
router_url: str,
expected_workers: int,
timeout: int = 300,
api_key: Optional[str] = None,
) -> None:
"""
Wait for router to have all workers connected.
Polls the /workers endpoint until the 'total' field matches expected_workers.
Example response from /workers endpoint:
{"workers":[],"total":0,"stats":{"prefill_count":0,"decode_count":0,"regular_count":0}}
Args:
router_url: Base URL of router (e.g., "http://127.0.0.1:30000")
expected_workers: Number of workers expected to be connected
timeout: Max seconds to wait
api_key: Optional API key for authentication
"""
start_time = time.time()
last_error = None
attempt = 0
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
with requests.Session() as session:
while time.time() - start_time < timeout:
attempt += 1
elapsed = int(time.time() - start_time)
# Log progress every 10 seconds
if elapsed > 0 and elapsed % 10 == 0 and attempt % 10 == 0:
logger.info(
f" Still waiting for workers... ({elapsed}/{timeout}s elapsed)"
)
try:
response = session.get(
f"{router_url}/workers", headers=headers, timeout=5
)
if response.status_code == 200:
data = response.json()
total_workers = data.get("total", 0)
if total_workers == expected_workers:
logger.info(
f" All {expected_workers} workers connected after {elapsed}s"
)
return
else:
last_error = f"Workers: {total_workers}/{expected_workers}"
else:
last_error = f"HTTP {response.status_code}"
except requests.ConnectionError:
last_error = "Connection refused (router not ready yet)"
except requests.Timeout:
last_error = "Timeout"
except requests.RequestException as e:
last_error = str(e)
except (ValueError, KeyError) as e:
last_error = f"Invalid response: {e}"
time.sleep(1)
raise TimeoutError(
f"Router at {router_url} did not get {expected_workers} workers within {timeout}s.\n"
f"Last status: {last_error}\n"
f"Hint: Run with SHOW_ROUTER_LOGS=1 to see startup logs"
)
def popen_launch_workers_and_router(
model: str,
base_url: str,
timeout: int = 300,
num_workers: int = 2,
policy: str = "round_robin",
api_key: Optional[str] = None,
worker_args: Optional[list] = None,
router_args: Optional[list] = None,
tp_size: int = 1,
env: Optional[dict] = None,
stdout=None,
stderr=None,
) -> dict:
"""
Launch SGLang workers and gRPC router separately.
This approach:
1. Starts N SGLang workers with --grpc-mode flag
2. Waits for workers to initialize (process startup)
3. Starts a gRPC router pointing to those workers
4. Waits for router health check to pass (router validates worker connectivity)
This matches production deployment patterns better than the integrated approach.
Args:
model: Model path (e.g., /home/ubuntu/models/llama-3.1-8b-instruct)
base_url: Base URL for router (e.g., "http://127.0.0.1:8080")
timeout: Timeout for server startup (default: 300s)
num_workers: Number of workers to launch
policy: Routing policy (round_robin, random, power_of_two, cache_aware)
api_key: Optional API key for router
worker_args: Additional arguments for workers (e.g., ["--context-len", "8192"])
router_args: Additional arguments for router (e.g., ["--max-total-token", "1536"])
tp_size: Tensor parallelism size for workers (default: 1)
env: Optional environment variables for workers (e.g., {"SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION": "256"})
stdout: Optional file handle for worker stdout (default: subprocess.PIPE)
stderr: Optional file handle for worker stderr (default: subprocess.PIPE)
Returns:
dict with:
- workers: list of worker process objects
- worker_urls: list of gRPC worker URLs
- router: router process object
- base_url: router URL (HTTP endpoint)
Example:
>>> cluster = popen_launch_workers_and_router(model, base_url, num_workers=2)
>>> # Use cluster['base_url'] for HTTP requests
>>> # Cleanup:
>>> for worker in cluster['workers']:
>>> kill_process_tree(worker.pid)
>>> kill_process_tree(cluster['router'].pid)
"""
import os
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
# Note: timeout parameter is used for router health check below
# Parse router port from base_url
if ":" in base_url.split("//")[-1]:
router_port = int(base_url.split(":")[-1])
else:
router_port = find_free_port()
logger.info(f"\n{'='*70}")
logger.info(f"Launching gRPC cluster (separate workers + router)")
logger.info(f"{'='*70}")
logger.info(f" Model: {model}")
logger.info(f" Router port: {router_port}")
logger.info(f" Workers: {num_workers}")
logger.info(f" TP size: {tp_size}")
logger.info(f" Policy: {policy}")
# Step 1: Launch workers with gRPC enabled
workers = []
worker_urls = []
for i in range(num_workers):
worker_port = find_free_port()
worker_url = f"grpc://127.0.0.1:{worker_port}"
worker_urls.append(worker_url)
logger.info(f"\n[Worker {i+1}/{num_workers}]")
logger.info(f" Port: {worker_port}")
logger.info(f" URL: {worker_url}")
# Build worker command
worker_cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
"127.0.0.1",
"--port",
str(worker_port),
"--grpc-mode", # Enable gRPC for this worker
"--mem-fraction-static",
"0.8",
]
# Add TP size
if tp_size > 1:
worker_cmd.extend(["--tp-size", str(tp_size)])
# Add worker-specific args
if worker_args:
worker_cmd.extend(worker_args)
# Launch worker with optional environment variables
if show_output:
worker_proc = subprocess.Popen(
worker_cmd,
env=env,
stdout=stdout,
stderr=stderr,
)
else:
worker_proc = subprocess.Popen(
worker_cmd,
stdout=stdout if stdout is not None else subprocess.PIPE,
stderr=stderr if stderr is not None else subprocess.PIPE,
env=env,
)
workers.append(worker_proc)
logger.info(f" PID: {worker_proc.pid}")
# Give workers a moment to start binding to ports
# The router will check worker health when it starts
logger.info(f"\nWaiting for {num_workers} workers to initialize (20s)...")
time.sleep(20)
# Quick check: make sure worker processes are still alive
for i, worker in enumerate(workers):
if worker.poll() is not None:
logger.error(
f" ✗ Worker {i+1} died during startup (exit code: {worker.poll()})"
)
# Cleanup: kill all workers
for w in workers:
try:
w.kill()
except:
pass
raise RuntimeError(f"Worker {i+1} failed to start")
logger.info(
f"✓ All {num_workers} workers started (router will verify connectivity)"
)
# Step 2: Launch router pointing to workers
logger.info(f"\n[Router]")
logger.info(f" Port: {router_port}")
logger.info(f" Worker URLs: {', '.join(worker_urls)}")
# Build router command
router_cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(router_port),
"--prometheus-port",
"9321",
"--policy",
policy,
"--model-path",
model,
"--log-level",
"warn",
]
# Add worker URLs
router_cmd.append("--worker-urls")
router_cmd.extend(worker_urls)
# Add API key
if api_key:
router_cmd.extend(["--api-key", api_key])
# Add router-specific args
if router_args:
router_cmd.extend(router_args)
if show_output:
logger.info(f" Command: {' '.join(router_cmd)}")
# Launch router
if show_output:
router_proc = subprocess.Popen(router_cmd)
else:
router_proc = subprocess.Popen(
router_cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
logger.info(f" PID: {router_proc.pid}")
# Wait for router to be ready
router_url = f"http://127.0.0.1:{router_port}"
logger.info(f"\nWaiting for router to start at {router_url}...")
try:
wait_for_workers_ready(
router_url, expected_workers=num_workers, timeout=180, api_key=api_key
)
logger.info(f"✓ Router ready at {router_url}")
except TimeoutError:
logger.error(f"✗ Router failed to start")
# Cleanup: kill router and all workers
try:
router_proc.kill()
except:
pass
for worker in workers:
try:
worker.kill()
except:
pass
raise
logger.info(f"\n{'='*70}")
logger.info(f"✓ gRPC cluster ready!")
logger.info(f" Router: {router_url}")
logger.info(f" Workers: {len(workers)}")
logger.info(f"{'='*70}\n")
return {
"workers": workers,
"worker_urls": worker_urls,
"router": router_proc,
"base_url": router_url,
}
@@ -0,0 +1,950 @@
"""
gRPC Router E2E Test - Test Openai Function Calling
This test file is REUSED from test/srt/openai_server/function_call/test_openai_function_calling.py
with minimal changes:
num_workers=2,
- Swap popen_launch_server() → popen_launch_workers_and_router()
- Update teardown to cleanup router + workers
- All test logic and assertions remain identical
Run with:
pytest e2e_test/e2e_grpc/e2e_grpc/function_call/test_openai_function_calling.py -v
"""
import json
import sys
import unittest
from pathlib import Path
import openai
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_MODEL_PATH,
DEFAULT_SMALL_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
get_tokenizer,
kill_process_tree,
)
class TestOpenAIServerFunctionCalling(CustomTestCase):
# NOTE: this system_message is for Llama3.2 system prompt. Without this,
# sometimes Llama3.2 gives a different tool call format such as:
# '<|python_tag|>{"type": "function", "function": "add", "parameters": {"a": "3", "b": "5"}}'
SYSTEM_MESSAGE = (
"You are a helpful assistant with tool calling capabilities. "
"Only reply with a tool call if the function exists in the library provided by the user. "
"If it doesn't exist, just reply directly in natural language. "
"When you receive a tool call response, use the output to format an answer to the original user question. "
"You have access to the following functions. "
"To call a function, please respond with JSON for a function call. "
'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. '
"Do not use variables.\n\n"
)
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
# Using small model for function calling tests
cls.model = DEFAULT_SMALL_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
# Start the local OpenAI Server. If necessary, you can add other parameters such as --enable-tools.
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
router_args=[
# If your server needs extra parameters to test function calling, please add them here.
"--tool-call-parser",
"llama",
],
num_workers=1,
tp_size=2,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def test_function_calling_format(self):
"""
Test: Whether the function call format returned by the AI is correct.
When returning a tool call, message.content should be None, and tool_calls should be a list.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "add",
"description": "Compute the sum of two numbers",
"parameters": {
"type": "object",
"properties": {
"a": {
"type": "integer",
"description": "A number",
},
"b": {
"type": "integer",
"description": "A number",
},
},
"required": ["a", "b"],
},
},
}
]
messages = [
{"role": "system", "content": self.SYSTEM_MESSAGE},
{"role": "user", "content": "Compute (3+5)"},
]
response = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=False,
tools=tools,
)
tool_calls = response.choices[0].message.tool_calls
assert (
isinstance(tool_calls, list) and len(tool_calls) > 0
), "tool_calls should be a non-empty list"
function_name = tool_calls[0].function.name
assert function_name == "add", "Function name should be 'add'"
# This unit test is too difficult for default model. Mark it as optional unit tests so it won't trigger unless specified.
def _test_function_calling_multiturn(self):
"""
Test: Whether the function call format returned by the AI is correct.
When returning a tool call, message.content should be None, and tool_calls should be a list.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "add",
"description": "Compute the sum of two numbers",
"parameters": {
"type": "object",
"properties": {
"a": {
"type": "integer",
"description": "A number",
},
"b": {
"type": "integer",
"description": "A number",
},
},
"required": ["a", "b"],
},
},
}
]
messages = [{"role": "user", "content": "Compute (3+5)"}]
response = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=False,
tools=tools,
)
tool_call = response.choices[0].message.tool_calls[0]
function_name = tool_call.function.name
assert function_name == "add", "Function name should be 'add'"
function_arguments = tool_call.function.arguments
function_arguments = json.loads(tool_call.function.arguments)
assert function_arguments in [
{"a": 3, "b": 5},
{"a": "3", "b": "5"},
], f"Unexpected function arguments: {function_arguments}"
messages.append(response.choices[0].message)
messages.append(
{
"role": "tool",
"tool_call_id": tool_call.id,
"content": "8",
"name": function_name,
}
)
final_response = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=False,
tools=tools,
)
assert (
"8" in final_response.choices[0].message.content
), "tool_call response should have the sum 8 in the content"
def test_function_calling_streaming_simple(self):
"""
Test: Whether the function name can be correctly recognized in streaming mode.
- Expect a function call to be found, and the function name to be correct.
- Verify that streaming mode returns at least multiple chunks.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The city to find the weather for",
},
"unit": {
"type": "string",
"description": "Weather unit (celsius or fahrenheit)",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["city", "unit"],
},
},
}
]
messages = [
{"role": "system", "content": self.SYSTEM_MESSAGE},
{
"role": "user",
"content": "What is the temperature in Paris in celsius??",
},
]
response_stream = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=True,
tools=tools,
)
chunks = list(response_stream)
self.assertTrue(len(chunks) > 0, "Streaming should return at least one chunk")
found_function_name = False
for chunk in chunks:
choice = chunk.choices[0]
# Check whether the current chunk contains tool_calls
if choice.delta.tool_calls:
tool_call = choice.delta.tool_calls[0]
if tool_call.function.name:
self.assertEqual(
tool_call.function.name,
"get_current_weather",
"Function name should be 'get_current_weather'",
)
found_function_name = True
break
self.assertTrue(
found_function_name,
"Target function name 'get_current_weather' was not found in the streaming chunks",
)
finish_reason = chunks[-1].choices[0].finish_reason
self.assertEqual(
finish_reason,
"tool_calls",
"Final response of function calling should have finish_reason 'tool_calls'",
)
def test_function_calling_streaming_args_parsing(self):
"""
Test: Whether the function call arguments returned in streaming mode can be correctly concatenated into valid JSON.
- The user request requires multiple parameters.
- AI may return the arguments in chunks that need to be concatenated.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "add",
"description": "Compute the sum of two integers",
"parameters": {
"type": "object",
"properties": {
"a": {
"type": "integer",
"description": "First integer",
},
"b": {
"type": "integer",
"description": "Second integer",
},
},
"required": ["a", "b"],
},
"strict": True, # Llama-3.2-1B is flaky in tool call. It won't always respond with parameters unless we set strict.
},
}
]
messages = [
{"role": "system", "content": self.SYSTEM_MESSAGE},
{"role": "user", "content": "Please sum 5 and 7, just call the function."},
]
response_stream = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.9,
top_p=0.9,
stream=True,
tools=tools,
)
argument_fragments = []
chunks = list(response_stream)
function_name = None
for chunk in chunks:
choice = chunk.choices[0]
if choice.delta.tool_calls:
tool_call = choice.delta.tool_calls[0]
# Record the function name on first occurrence
function_name = tool_call.function.name or function_name
# In case of multiple chunks, JSON fragments may need to be concatenated
if tool_call.function.arguments is not None:
argument_fragments.append(tool_call.function.arguments)
self.assertEqual(function_name, "add", "Function name should be 'add'")
joined_args = "".join(argument_fragments)
self.assertTrue(
len(joined_args) > 0,
"No parameter fragments were returned in the function call",
)
finish_reason = chunks[-1].choices[0].finish_reason
self.assertEqual(
finish_reason,
"tool_calls",
"Final response of function calling should have finish_reason 'tool_calls'",
)
# Check whether the concatenated JSON is valid
try:
args_obj = json.loads(joined_args)
except json.JSONDecodeError:
self.fail(
"The concatenated tool call arguments are not valid JSON, parsing failed"
)
self.assertIn("a", args_obj, "Missing parameter 'a'")
self.assertIn("b", args_obj, "Missing parameter 'b'")
self.assertEqual(str(args_obj["a"]), "5", "Parameter a should be 5")
self.assertEqual(str(args_obj["b"]), "7", "Parameter b should be 7")
@unittest.skip(
"Skipping function call strict test as it is not supported by the router"
)
def test_function_call_strict(self):
"""
Test: Whether the strict mode of function calling works as expected.
- When strict mode is enabled, the AI should not return a function call if the function name is not recognized.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "sub",
"description": "Compute the difference of two integers",
"parameters": {
"type": "object",
"properties": {
"int_a": {
"type": "integer",
"description": "First integer",
},
"int_b": {
"type": "integer",
"description": "Second integer",
},
},
"required": ["int_a", "int_b"],
},
"strict": True,
},
}
]
messages = [
{"role": "user", "content": "Please compute 5 - 7, using your tool."}
]
response = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=False,
tools=tools,
)
tool_calls = response.choices[0].message.tool_calls
function_name = tool_calls[0].function.name
arguments = tool_calls[0].function.arguments
args_obj = json.loads(arguments)
self.assertEqual(function_name, "sub", "Function name should be 'sub'")
self.assertEqual(str(args_obj["int_a"]), "5", "Parameter int_a should be 5")
self.assertEqual(str(args_obj["int_b"]), "7", "Parameter int_b should be 7")
def test_function_call_required(self):
"""
Test: Whether tool_choice: "required" works as expected
- When tool_choice == "required", the model should return one or more tool_calls.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "sub",
"description": "Compute the difference of two integers",
"parameters": {
"type": "object",
"properties": {
"int_a": {
"type": "integer",
"description": "First integer",
},
"int_b": {
"type": "integer",
"description": "Second integer",
},
},
"required": ["int_a", "int_b"],
},
"strict": True,
},
},
{
"type": "function",
"function": {
"name": "get_weather",
"description": "use this to get latest weather information for a city given its name",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "name of the city to get weather for",
}
},
"required": ["city"],
},
},
},
]
messages = [{"role": "user", "content": "What is the capital of France?"}]
response = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=False,
tools=tools,
tool_choice="required",
)
tool_calls = response.choices[0].message.tool_calls
self.assertIsNotNone(tool_calls, "No tool_calls in the response")
function_name = tool_calls[0].function.name
arguments = tool_calls[0].function.arguments
args_obj = json.loads(arguments)
self.assertEqual(
function_name,
"get_weather",
f"Function name should be 'get_weather', got: {function_name}",
)
self.assertIn(
"city", args_obj, f"Function arguments should have 'city', got: {args_obj}"
)
# Make the test more robust by checking type and accepting valid responses
city_value = args_obj["city"]
self.assertIsInstance(
city_value,
str,
f"Parameter city should be a string, got: {type(city_value)}",
)
self.assertTrue(
"Paris" in city_value or "France" in city_value,
f"Parameter city should contain either 'Paris' or 'France', got: {city_value}",
)
def test_function_call_specific(self):
"""
Test: Whether tool_choice: ToolChoice works as expected
- When tool_choice is a specific ToolChoice, the model should return one or more tool_calls.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "sub",
"description": "Compute the difference of two integers",
"parameters": {
"type": "object",
"properties": {
"int_a": {
"type": "integer",
"description": "First integer",
},
"int_b": {
"type": "integer",
"description": "Second integer",
},
},
"required": ["int_a", "int_b"],
},
"strict": True,
},
},
{
"type": "function",
"function": {
"name": "get_weather",
"description": "use this to get latest weather information for a city given its name",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "name of the city to get weather for",
}
},
"required": ["city"],
},
},
},
]
messages = [{"role": "user", "content": "What is the capital of France?"}]
response = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=False,
tools=tools,
tool_choice={"type": "function", "function": {"name": "get_weather"}},
)
tool_calls = response.choices[0].message.tool_calls
self.assertIsNotNone(tool_calls, "No tool_calls in the response")
function_name = tool_calls[0].function.name
arguments = tool_calls[0].function.arguments
args_obj = json.loads(arguments)
self.assertEqual(
function_name, "get_weather", "Function name should be 'get_weather'"
)
self.assertIn("city", args_obj, "Function arguments should have 'city'")
def test_streaming_multiple_choices_finish_reason(self):
"""
Test: Verify that each choice gets its own finish_reason chunk in streaming mode with n > 1.
This tests the fix for the bug where only the last index got a finish_reason chunk.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
},
}
]
messages = [
{"role": "user", "content": "What is the weather like in Los Angeles?"}
]
# Request with n=2 to get multiple choices
response_stream = client.chat.completions.create(
model=self.model,
messages=messages,
max_tokens=2048,
temperature=0.8,
stream=True,
tools=tools,
tool_choice="required", # Force tool calls
n=2, # Multiple choices
)
chunks = list(response_stream)
# Track finish_reason chunks for each index
finish_reason_chunks = {}
for chunk in chunks:
if chunk.choices:
for choice in chunk.choices:
if choice.finish_reason is not None:
index = choice.index
if index not in finish_reason_chunks:
finish_reason_chunks[index] = []
finish_reason_chunks[index].append(choice.finish_reason)
# Verify we got finish_reason chunks for both indices
self.assertEqual(
len(finish_reason_chunks),
2,
f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}",
)
# Verify both index 0 and 1 have finish_reason
self.assertIn(
0, finish_reason_chunks, "Missing finish_reason chunk for index 0"
)
self.assertIn(
1, finish_reason_chunks, "Missing finish_reason chunk for index 1"
)
# Verify the finish_reason is "tool_calls" since we forced tool calls
for index, reasons in finish_reason_chunks.items():
self.assertEqual(
reasons[-1], # Last finish_reason for this index
"tool_calls",
f"Expected finish_reason 'tool_calls' for index {index}, got {reasons[-1]}",
)
def test_function_calling_streaming_no_tool_call(self):
"""
Test: Whether the finish_reason is stop in streaming mode when no tool call is given.
- Expect no function call to be found.
- Verify that finish_reason is stop
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The city to find the weather for",
},
"unit": {
"type": "string",
"description": "Weather unit (celsius or fahrenheit)",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["city", "unit"],
},
},
}
]
messages = [{"role": "user", "content": "Who are you?"}]
response_stream = client.chat.completions.create(
model=self.model,
max_tokens=2048,
messages=messages,
temperature=0.8,
top_p=0.8,
stream=True,
tools=tools,
tool_choice="none",
)
chunks = list(response_stream)
self.assertTrue(len(chunks) > 0, "Streaming should return at least one chunk")
found_tool_call = False
for chunk in chunks:
choice = chunk.choices[0]
# Check whether the current chunk contains tool_calls
found_tool_call = choice.delta.tool_calls is not None
self.assertFalse(
found_tool_call,
"Shouldn't have any tool_call in the streaming chunks",
)
finish_reason = chunks[-1].choices[0].finish_reason
self.assertEqual(
finish_reason,
"stop",
"Final response of no function calling should have finish_reason 'stop'",
)
def test_streaming_multiple_choices_without_tools(self):
"""
Test: Verify that each choice gets its own finish_reason chunk without tool calls.
This tests the fix for regular content streaming with multiple choices.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
messages = [{"role": "user", "content": "Say hello in one word."}]
# Request with n=2 to get multiple choices, no tools
response_stream = client.chat.completions.create(
model=self.model,
messages=messages,
temperature=0.8,
stream=True,
max_tokens=10, # Keep it short
n=2, # Multiple choices
)
chunks = list(response_stream)
# Track finish_reason chunks for each index
finish_reason_chunks = {}
for chunk in chunks:
if chunk.choices:
for choice in chunk.choices:
if choice.finish_reason is not None:
index = choice.index
if index not in finish_reason_chunks:
finish_reason_chunks[index] = []
finish_reason_chunks[index].append(choice.finish_reason)
# Verify we got finish_reason chunks for both indices
self.assertEqual(
len(finish_reason_chunks),
2,
f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}",
)
# Verify both index 0 and 1 have finish_reason
self.assertIn(
0, finish_reason_chunks, "Missing finish_reason chunk for index 0"
)
self.assertIn(
1, finish_reason_chunks, "Missing finish_reason chunk for index 1"
)
# Verify the finish_reason is "stop" (regular completion)
for index, reasons in finish_reason_chunks.items():
self.assertIn(
reasons[-1],
["stop", "length"], # Could be either depending on how model responds
f"Expected finish_reason 'stop' or 'length' for index {index}, got {reasons[-1]}",
)
class TestOpenAIPythonicFunctionCalling(CustomTestCase):
PYTHONIC_TOOLS = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a given location.",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The name of the city or location.",
}
},
"required": ["location"],
},
},
},
{
"type": "function",
"function": {
"name": "get_tourist_attractions",
"description": "Get a list of top tourist attractions for a given city.",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The name of the city to find attractions for.",
}
},
"required": ["city"],
},
},
},
]
PYTHONIC_MESSAGES = [
{
"role": "system",
"content": (
"You are a travel assistant. "
"When asked to call functions, ALWAYS respond ONLY with a python list of function calls, "
"using this format: [func_name1(param1=value1, param2=value2), func_name2(param=value)]. "
"Do NOT use JSON, do NOT use variables, do NOT use any other format. "
"Here is an example:\n"
'[get_weather(location="Paris"), get_tourist_attractions(city="Paris")]'
),
},
{
"role": "user",
"content": (
"I'm planning a trip to Tokyo next week. What's the weather like and what are some top tourist attractions? "
"Propose parallel tool calls at once, using the python list of function calls format as shown above."
),
},
]
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
cls.model = DEFAULT_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
router_args=[
"--tool-call-parser",
"pythonic",
],
num_workers=1,
tp_size=2,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def test_pythonic_tool_call_prompt(self):
"""
Test: Explicit prompt for pythonic tool call format without chat template.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response = client.chat.completions.create(
model=self.model,
messages=self.PYTHONIC_MESSAGES,
tools=self.PYTHONIC_TOOLS,
temperature=0.1,
stream=False,
)
tool_calls = response.choices[0].message.tool_calls
self.assertIsInstance(tool_calls, list, "No tool_calls found")
self.assertGreaterEqual(len(tool_calls), 1)
names = [tc.function.name for tc in tool_calls]
self.assertTrue(
"get_weather" in names or "get_tourist_attractions" in names,
f"Function name '{names}' should container either 'get_weather' or 'get_tourist_attractions'",
)
def test_pythonic_tool_call_streaming(self):
"""
Test: Streaming pythonic tool call format; assert tool_call index is present.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response_stream = client.chat.completions.create(
model=self.model,
messages=self.PYTHONIC_MESSAGES,
tools=self.PYTHONIC_TOOLS,
temperature=0.1,
stream=True,
)
found_tool_calls = False
found_index = False
found_names = set()
for chunk in response_stream:
choice = chunk.choices[0]
if getattr(choice.delta, "tool_calls", None):
found_tool_calls = True
tool_call = choice.delta.tool_calls[0]
if hasattr(tool_call, "index") or (
isinstance(tool_call, dict) and "index" in tool_call
):
found_index = True
found_names.add(str(tool_call.function.name))
self.assertTrue(found_tool_calls, "No tool_calls found in streaming response")
self.assertTrue(found_index, "No index field found in any streamed tool_call")
self.assertTrue(
"get_weather" in found_names or "get_tourist_attractions" in found_names,
f"Function name '{found_names}' should container either 'get_weather' or 'get_tourist_attractions'",
)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,795 @@
"""
Test script for tool_choice functionality in SGLang
Tests: required, auto, and specific function choices in both streaming and non-streaming modes
# To run the tests, use the following command:
#
# python3 -m unittest openai_server.function_call.test_tool_choice
"""
import json
import sys
import unittest
from pathlib import Path
import openai
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_SMALL_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
get_tokenizer,
kill_process_tree,
)
class TestToolChoiceLlama32(CustomTestCase):
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
# Mark flaky tests for this model
cls.flaky_tests = {
"test_multi_tool_scenario_auto",
"test_multi_tool_scenario_required",
}
# Use a model that supports function calling
cls.model = DEFAULT_SMALL_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
# Start the local OpenAI Server with tool calling support
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
router_args=[
"--tool-call-parser",
"llama", # Default parser for the test model
],
num_workers=1,
tp_size=2,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def setUp(self):
self.client = openai.Client(base_url=self.base_url, api_key=self.api_key)
self.model_name = self.client.models.list().data[0].id
def _is_flaky_test(self):
"""Check if the current test is marked as flaky for this class"""
return (
hasattr(self.__class__, "flaky_tests")
and self._testMethodName in self.__class__.flaky_tests
)
def get_test_tools(self):
"""Get the test tools for function calling"""
return [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "use this to get latest weather information for a city given its name",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "name of the city to get weather for",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["city"],
},
},
},
{
"type": "function",
"function": {
"name": "get_pokemon_info",
"description": "get detailed information about a pokemon given its name",
"parameters": {
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "name of the pokemon to get info for",
}
},
"required": ["name"],
},
},
},
{
"type": "function",
"function": {
"name": "make_next_step_decision",
"description": "You will be given a trace of thinking process in the following format.\n\nQuestion: the input question you must answer\nTOOL: think about what to do, and choose a tool to use ONLY IF there are defined tools. \n You should never call the same tool with the same input twice in a row.\n If the previous conversation history already contains the information that can be retrieved from the tool, you should not call the tool again.\nOBSERVATION: the result of the tool call, NEVER include this in your response, this information will be provided\n... (this TOOL/OBSERVATION can repeat N times)\nANSWER: If you know the answer to the original question, require for more information,\n or you don't know the answer and there are no defined tools or all available tools are not helpful, respond with the answer without mentioning anything else.\n If the previous conversation history already contains the answer, respond with the answer right away.\n\n If no tools are configured, naturally mention this limitation while still being helpful. Briefly note that adding tools in the agent configuration would expand capabilities.\n\nYour task is to respond with the next step to take, based on the traces, \nor answer the question if you have enough information.",
"parameters": {
"type": "object",
"properties": {
"decision": {
"type": "string",
"description": 'The next step to take, it must be either "TOOL" or "ANSWER". If the previous conversation history already contains the information that can be retrieved from the tool, you should not call the tool again. If there are no defined tools, you should not return "TOOL" in your response.',
},
"content": {
"type": "string",
"description": 'The content of the next step. If the decision is "TOOL", this should be a short and concise reasoning of why you chose the tool, MUST include the tool name. If the decision is "ANSWER", this should be the answer to the question. If no tools are available, integrate this limitation conversationally without sounding scripted.',
},
},
"required": ["decision", "content"],
},
},
},
]
def get_test_messages(self):
"""Get test messages that should trigger tool usage"""
return [
{
"role": "user",
"content": "Answer the following questions as best you can:\n\nYou will be given a trace of thinking process in the following format.\n\nQuestion: the input question you must answer\nTOOL: think about what to do, and choose a tool to use ONLY IF there are defined tools\nOBSERVATION: the result of the tool call or the observation of the current task, NEVER include this in your response, this information will be provided\n... (this TOOL/OBSERVATION can repeat N times)\nANSWER: If you know the answer to the original question, require for more information, \nif the previous conversation history already contains the answer, \nor you don't know the answer and there are no defined tools or all available tools are not helpful, respond with the answer without mentioning anything else.\nYou may use light Markdown formatting to improve clarity (e.g. lists, **bold**, *italics*), but keep it minimal and unobtrusive.\n\nYour task is to respond with the next step to take, based on the traces, \nor answer the question if you have enough information.\n\nQuestion: what is the weather in top 5 populated cities in the US in celsius?\n\nTraces:\n\n\nThese are some additional instructions that you should follow:",
}
]
def get_travel_tools(self):
"""Get tools for travel assistant scenario that should trigger multiple tool calls"""
return [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a given location.",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The name of the city or location.",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
},
},
{
"type": "function",
"function": {
"name": "get_tourist_attractions",
"description": "Get a list of top tourist attractions for a given city.",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The name of the city to find attractions for.",
}
},
"required": ["city"],
},
},
},
]
def get_travel_messages(self):
"""Get travel assistant messages that should trigger multiple tool calls"""
return [
{
"content": "You are a travel assistant providing real-time weather updates and top tourist attractions.",
"role": "system",
},
{
"content": "I'm planning a trip to Tokyo next week. What's the weather like? What are the most amazing sights?",
"role": "user",
},
]
def test_tool_choice_auto_non_streaming(self):
"""Test tool_choice='auto' in non-streaming mode"""
tools = self.get_test_tools()
messages = self.get_test_messages()
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
tools=tools,
tool_choice="auto",
stream=False,
)
self.assertIsNotNone(response.choices[0].message)
# With auto, tool calls are optional
def test_tool_choice_auto_streaming(self):
"""Test tool_choice='auto' in streaming mode"""
tools = self.get_test_tools()
messages = self.get_test_messages()
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
tools=tools,
tool_choice="auto",
stream=True,
)
# Collect streaming response
content_chunks = []
tool_call_chunks = []
for chunk in response:
if chunk.choices[0].delta.content:
content_chunks.append(chunk.choices[0].delta.content)
elif chunk.choices[0].delta.tool_calls:
tool_call_chunks.extend(chunk.choices[0].delta.tool_calls)
# Should complete without errors
self.assertIsInstance(content_chunks, list)
self.assertIsInstance(tool_call_chunks, list)
def test_tool_choice_required_non_streaming(self):
"""Test tool_choice='required' in non-streaming mode"""
tools = self.get_test_tools()
messages = self.get_test_messages()
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
temperature=0.2,
tools=tools,
tool_choice="required",
stream=False,
)
# With required, we should get tool calls
tool_calls = response.choices[0].message.tool_calls
self.assertIsNotNone(tool_calls)
self.assertGreater(len(tool_calls), 0)
def test_tool_choice_required_streaming(self):
"""Test tool_choice='required' in streaming mode"""
tools = self.get_test_tools()
messages = self.get_test_messages()
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
tools=tools,
tool_choice="required",
stream=True,
)
# Collect streaming response
tool_call_chunks = []
for chunk in response:
if chunk.choices[0].delta.tool_calls:
tool_call_chunks.extend(chunk.choices[0].delta.tool_calls)
# With required, we should get tool call chunks
self.assertGreater(len(tool_call_chunks), 0)
def test_tool_choice_specific_function_non_streaming(self):
"""Test tool_choice with specific function in non-streaming mode"""
tools = self.get_test_tools()
messages = self.get_test_messages()
tool_choice = {"type": "function", "function": {"name": "get_weather"}}
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
tools=tools,
tool_choice=tool_choice,
stream=False,
)
# Should call the specific function
tool_calls = response.choices[0].message.tool_calls
self.assertIsNotNone(tool_calls)
# Our messages ask the top 5 populated cities in the US, so the model could get 5 tool calls
self.assertGreaterEqual(len(tool_calls), 1)
for tool_call in tool_calls:
self.assertEqual(tool_call.function.name, "get_weather")
def test_tool_choice_specific_function_streaming(self):
"""Test tool_choice with specific function in streaming mode"""
tools = self.get_test_tools()
messages = self.get_test_messages()
tool_choice = {"type": "function", "function": {"name": "get_weather"}}
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
tools=tools,
tool_choice=tool_choice,
stream=True,
)
# Collect streaming response
tool_call_chunks = []
for chunk in response:
if chunk.choices[0].delta.tool_calls:
tool_call_chunks.extend(chunk.choices[0].delta.tool_calls)
# Should get tool call chunks for the specific function
self.assertGreater(len(tool_call_chunks), 0)
# Find function name in chunks
found_name = None
for chunk in tool_call_chunks:
if chunk.function and chunk.function.name:
found_name = chunk.function.name
break
self.assertEqual(found_name, "get_weather")
def test_required_streaming_arguments_chunks_json(self):
"""In streaming required mode, complete tool call arguments should be valid JSON when all chunks are combined"""
tools = self.get_test_tools()
messages = self.get_test_messages()
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=1024,
temperature=0.1,
tools=tools,
tool_choice="required",
stream=True,
)
# Collect all tool call chunks and reconstruct complete tool calls
tool_calls_by_index = {}
for chunk in response:
if chunk.choices[0].delta.tool_calls:
for tool_call_delta in chunk.choices[0].delta.tool_calls:
tool_index = tool_call_delta.index
# Initialize tool call if not seen before
if tool_index not in tool_calls_by_index:
tool_calls_by_index[tool_index] = {
"id": tool_call_delta.id,
"type": "function",
"function": {"name": "", "arguments": ""},
}
# Update function name if present (first chunk)
if tool_call_delta.function and tool_call_delta.function.name:
tool_calls_by_index[tool_index]["function"][
"name"
] = tool_call_delta.function.name
# Accumulate arguments (all chunks)
if tool_call_delta.function and tool_call_delta.function.arguments:
tool_calls_by_index[tool_index]["function"][
"arguments"
] += tool_call_delta.function.arguments
self.assertGreater(len(tool_calls_by_index), 0)
# Validate that complete tool calls have valid JSON arguments
for tool_call in tool_calls_by_index.values():
self.assertIsNotNone(tool_call["function"]["name"])
self.assertIsNotNone(tool_call["function"]["arguments"])
# The complete arguments should be valid JSON
try:
args = json.loads(tool_call["function"]["arguments"])
self.assertIsInstance(args, dict)
except json.JSONDecodeError:
self.fail(
f"Invalid JSON in complete tool call arguments: {tool_call['function']['arguments']}"
)
def test_complex_parameters_required_non_streaming(self):
"""Validate complex nested parameter schemas in non-streaming required mode"""
complex_tools = [
{
"type": "function",
"function": {
"name": "analyze_data",
"description": "Analyze complex data structures",
"parameters": {
"type": "object",
"properties": {
"data": {
"type": "object",
"properties": {
"metrics": {
"type": "array",
"items": {"type": "string"},
},
"config": {
"type": "object",
"properties": {
"threshold": {"type": "number"},
"enabled": {"type": "boolean"},
},
},
},
"required": ["metrics"],
},
"options": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": {"type": "string"},
"value": {"type": "string"},
},
},
},
},
"required": ["data"],
},
},
}
]
messages = [
{
"role": "user",
"content": "Analyze some data with metrics and configuration",
}
]
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=1024,
temperature=0.1,
tools=complex_tools,
tool_choice="required",
stream=False,
)
tool_calls = response.choices[0].message.tool_calls
self.assertIsNotNone(tool_calls)
self.assertGreater(len(tool_calls), 0)
for tool_call in tool_calls:
self.assertEqual(tool_call.function.name, "analyze_data")
try:
args = json.loads(tool_call.function.arguments)
self.assertIsInstance(args, dict)
self.assertIn("data", args)
self.assertIsInstance(args["data"], dict)
except json.JSONDecodeError:
self.fail(
f"Invalid JSON in complex tool call arguments: {tool_call.function.arguments}"
)
def test_multi_tool_scenario_auto(self):
"""Test multi-tool scenario with tool_choice='auto'"""
tools = self.get_travel_tools()
messages = self.get_travel_messages()
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
temperature=0.2,
tools=tools,
tool_choice="auto",
stream=False,
)
# Should complete without errors
self.assertIsNotNone(response.choices[0].message)
tool_calls = response.choices[0].message.tool_calls
expected_functions = {"get_weather", "get_tourist_attractions"}
if self._is_flaky_test():
# For flaky tests, just verify all called functions are available tools
if tool_calls:
available_names = [tool["function"]["name"] for tool in tools]
for call in tool_calls:
self.assertIn(call.function.name, available_names)
else:
# For non-flaky tests, enforce strict requirements
self.assertIsNotNone(tool_calls, "Expected tool calls but got none")
self.assertEqual(
len(tool_calls), 2, f"Expected 2 tool calls, got {len(tool_calls)}"
)
called_functions = {call.function.name for call in tool_calls}
self.assertEqual(
called_functions,
expected_functions,
f"Expected functions {expected_functions}, got {called_functions}",
)
def test_multi_tool_scenario_required(self):
"""Test multi-tool scenario with tool_choice='required'"""
tools = self.get_travel_tools()
messages = self.get_travel_messages()
response = self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
temperature=0.2,
tools=tools,
tool_choice="required",
stream=False,
)
# With required, we should get at least one tool call
tool_calls = response.choices[0].message.tool_calls
self.assertIsNotNone(tool_calls)
self.assertGreater(len(tool_calls), 0)
# Verify all called functions are available tools
available_names = [tool["function"]["name"] for tool in tools]
expected_functions = {"get_weather", "get_tourist_attractions"}
for tool_call in tool_calls:
self.assertIsNotNone(tool_call.function.name)
self.assertIsNotNone(tool_call.function.arguments)
if self._is_flaky_test():
# For flaky tests, just ensure basic functionality works
self.assertGreater(
len(tool_calls),
0,
f"Expected at least 1 tool call, got {len(tool_calls)}",
)
for call in tool_calls:
self.assertIn(call.function.name, available_names)
else:
# For non-flaky tests, enforce strict requirements
self.assertEqual(
len(tool_calls), 2, f"Expected 2 tool calls, got {len(tool_calls)}"
)
called_functions = {call.function.name for call in tool_calls}
self.assertEqual(
called_functions,
expected_functions,
f"Expected functions {expected_functions}, got {called_functions}",
)
def test_error_handling_invalid_tool_choice(self):
"""Test error handling for invalid tool_choice"""
tools = self.get_test_tools()
messages = self.get_test_messages()
# Test with invalid function name
tool_choice = {"type": "function", "function": {"name": "nonexistent_function"}}
# Expect a 400 BadRequestError to be raised for invalid tool_choice
with self.assertRaises(openai.BadRequestError) as context:
self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=2048,
tools=tools,
tool_choice=tool_choice,
stream=False,
)
# Verify the error message contains the expected text
self.assertIn(
"function 'nonexistent_function' not found in",
str(context.exception),
)
def test_invalid_tool_missing_name(self):
"""Test what happens when user doesn't provide a tool name in request"""
# Test with malformed JSON in tool parameters - missing required "name" field
invalid_tools = [
{
"type": "function",
"function": {
# Missing required "name" field
"description": "Test function with invalid schema",
"parameters": {
"type": "object",
"properties": {
"test_field": {
"type": "string",
"description": "Test field",
}
},
"required": ["test_field"],
},
},
}
]
messages = [
{
"role": "user",
"content": "Test the function",
}
]
# Should raise BadRequestError due to missing required 'name' field
with self.assertRaises(openai.BadRequestError) as context:
self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=100,
temperature=0.1,
tools=invalid_tools,
tool_choice="required",
stream=False,
)
# Verify the error message indicates missing name field
error_msg = str(context.exception).lower()
self.assertIn("name", error_msg)
def test_conflicting_defs_required_tool_choice(self):
"""Test that conflicting $defs with required tool_choice returns 400 error"""
conflicting_tools = [
{
"type": "function",
"function": {
"name": "tool1",
"description": "Tool 1 with conflicting $defs",
"parameters": {
"type": "object",
"properties": {
"data": {"$ref": "#/$defs/DataType"},
},
"required": ["data"],
"$defs": {
"DataType": {
"type": "object",
"properties": {"value": {"type": "string"}},
"required": ["value"],
},
},
},
},
},
{
"type": "function",
"function": {
"name": "tool2",
"description": "Tool 2 with conflicting $defs",
"parameters": {
"type": "object",
"properties": {
"data": {"$ref": "#/$defs/DataType"},
},
"required": ["data"],
"$defs": {
"DataType": { # Different definition for DataType
"type": "object",
"properties": {"value": {"type": "number"}},
"required": ["value"],
},
},
},
},
},
]
messages = [
{
"role": "user",
"content": "Test the conflicting tools",
}
]
# Should raise BadRequestError due to conflicting $defs
with self.assertRaises(openai.BadRequestError) as context:
self.client.chat.completions.create(
model=self.model_name,
messages=messages,
max_tokens=100,
temperature=0.1,
tools=conflicting_tools,
tool_choice="required",
stream=False,
)
# Verify the error message indicates conflicting tool definitions
error_msg = str(context.exception).lower()
self.assertIn("invalid tool configuration", error_msg)
self.assertIn("not supported", error_msg)
class TestToolChoiceQwen25(TestToolChoiceLlama32):
"""Test tool_choice functionality with Qwen2.5 model"""
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
cls.flaky_tests = {}
cls.model = DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
router_args=[
"--tool-call-parser",
"qwen",
],
num_workers=1,
tp_size=2,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
class TestToolChoiceMistral(TestToolChoiceLlama32):
"""Test tool_choice functionality with Mistral model"""
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
# Mark flaky tests for this model
cls.flaky_tests = {
"test_multi_tool_scenario_auto",
"test_multi_tool_scenario_required",
}
cls.model = DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
router_args=[
"--tool-call-parser",
"mistral",
],
num_workers=1,
tp_size=2,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
@unittest.skip("Fails due to whitespace issue with Mistral - skipping")
def test_complex_parameters_required_non_streaming(self):
"""Validate complex nested parameter schemas in non-streaming required mode"""
super().test_complex_parameters_required_non_streaming()
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,13 @@
[pytest]
# Show print statements and logs
log_cli = true
log_cli_level = INFO
log_cli_format = %(asctime)s [%(levelname)8s] %(message)s
log_cli_date_format = %Y-%m-%d %H:%M:%S
# Show stdout/stderr
addopts = -v -s --tb=short
# Capture settings
# -s means don't capture stdout (show print statements)
# --tb=short means short traceback format
+263
View File
@@ -0,0 +1,263 @@
"""
Standalone utilities for e2e_grpc tests.
This module provides all necessary utilities without depending on sglang Python package.
Extracted and adapted from:
- sglang.srt.utils.kill_process_tree
- sglang.srt.utils.hf_transformers_utils.get_tokenizer
- sglang.test.test_utils (constants and CustomTestCase)
"""
import logging
import os
import signal
import threading
import unittest
from pathlib import Path
from typing import Optional, Union
import psutil
logger = logging.getLogger(__name__)
try:
from transformers import (
AutoTokenizer,
PreTrainedTokenizer,
PreTrainedTokenizerBase,
PreTrainedTokenizerFast,
)
except ImportError:
raise ImportError(
"transformers is required for tokenizer utilities. "
"Install with: pip install transformers"
)
# ============================================================================
# Constants
# ============================================================================
# Server and timeout constants
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH = 600
DEFAULT_PORT_FOR_SRT_TEST_RUNNER = 20000
DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000}"
# File name constants for test output
STDOUT_FILENAME = "/tmp/sglang_test_stdout.txt"
STDERR_FILENAME = "/tmp/sglang_test_stderr.txt"
# Model base path - can be overridden via environment variable
# By default, use HuggingFace model identifiers (no local path prefix)
# Set ROUTER_LOCAL_MODEL_PATH to use local models (e.g., "/home/ubuntu/models")
ROUTER_LOCAL_MODEL_PATH = os.environ.get("ROUTER_LOCAL_MODEL_PATH", "")
# Helper function to build model paths
def _get_model_path(model_identifier: str) -> str:
"""
Build model path from base path and model identifier.
If ROUTER_LOCAL_MODEL_PATH is set, prepend it to the identifier.
Otherwise, return the identifier as-is (for HuggingFace download).
"""
if ROUTER_LOCAL_MODEL_PATH:
return os.path.join(ROUTER_LOCAL_MODEL_PATH, model_identifier)
return model_identifier
# Model paths used in e2e_grpc tests
# These can be either HuggingFace identifiers or local paths (depending on ROUTER_LOCAL_MODEL_PATH)
# Main test model - Llama 3.1 8B Instruct
DEFAULT_MODEL_PATH = _get_model_path("meta-llama/Llama-3.1-8B-Instruct")
# Small models for function calling tests
DEFAULT_SMALL_MODEL_PATH = _get_model_path("meta-llama/Llama-3.2-1B-Instruct")
# Reasoning models
DEFAULT_REASONING_MODEL_PATH = _get_model_path(
"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
)
# Thinking-enabled models
DEFAULT_ENABLE_THINKING_MODEL_PATH = _get_model_path("Qwen/Qwen3-30B-A3B")
# Function calling models
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH = _get_model_path("Qwen/Qwen2.5-7B-Instruct")
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH = _get_model_path(
"mistralai/Mistral-7B-Instruct-v0.3"
)
# GPT-OSS models
DEFAULT_GPT_OSS_MODEL_PATH = _get_model_path("openai/gpt-oss-20b")
# Embedding models
DEFAULT_EMBEDDING_MODEL_PATH = _get_model_path("intfloat/e5-mistral-7b-instruct")
# ============================================================================
# Process Management
# ============================================================================
def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
"""
Kill the process and all its child processes.
Args:
parent_pid: PID of the parent process
include_parent: Whether to kill the parent process itself
skip_pid: Optional PID to skip during cleanup
"""
# Remove sigchld handler to avoid spammy logs
if threading.current_thread() is threading.main_thread():
signal.signal(signal.SIGCHLD, signal.SIG_DFL)
if parent_pid is None:
parent_pid = os.getpid()
include_parent = False
try:
itself = psutil.Process(parent_pid)
except psutil.NoSuchProcess:
return
children = itself.children(recursive=True)
for child in children:
if child.pid == skip_pid:
continue
try:
child.kill()
except psutil.NoSuchProcess:
pass
if include_parent:
try:
itself.kill()
except psutil.NoSuchProcess:
pass
# ============================================================================
# Tokenizer Utilities
# ============================================================================
def check_gguf_file(model_path: str) -> bool:
"""Check if the model path points to a GGUF file."""
if not isinstance(model_path, str):
return False
return model_path.endswith(".gguf")
def is_remote_url(path: str) -> bool:
"""Check if the path is a remote URL."""
if not isinstance(path, str):
return False
return path.startswith("http://") or path.startswith("https://")
def get_tokenizer(
tokenizer_name: str,
*args,
tokenizer_mode: str = "auto",
trust_remote_code: bool = False,
tokenizer_revision: Optional[str] = None,
**kwargs,
) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
"""
Gets a tokenizer for the given model name via Huggingface.
Args:
tokenizer_name: Name or path of the tokenizer
tokenizer_mode: Mode for tokenizer loading ("auto", "slow")
trust_remote_code: Whether to trust remote code
tokenizer_revision: Specific revision to use
**kwargs: Additional arguments passed to AutoTokenizer.from_pretrained
Returns:
Loaded tokenizer instance
"""
if tokenizer_mode == "slow":
if kwargs.get("use_fast", False):
raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
kwargs["use_fast"] = False
# Handle special model name mapping
if tokenizer_name == "mistralai/Devstral-Small-2505":
tokenizer_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
is_gguf = check_gguf_file(tokenizer_name)
if is_gguf:
kwargs["gguf_file"] = tokenizer_name
tokenizer_name = Path(tokenizer_name).parent
# Note: Removed remote URL handling and local directory download
# as they depend on sglang-specific utilities
try:
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_name,
*args,
trust_remote_code=trust_remote_code,
tokenizer_revision=tokenizer_revision,
**kwargs,
)
except TypeError as e:
# Handle specific errors
err_msg = (
"Failed to load the tokenizer. If you are running a model with "
"a custom tokenizer, please set the --trust-remote-code flag."
)
raise RuntimeError(err_msg) from e
if not isinstance(tokenizer, PreTrainedTokenizerFast):
logger.warning(
f"Using a slow tokenizer. This might cause a performance "
f"degradation. Consider using a fast tokenizer instead."
)
return tokenizer
def get_tokenizer_from_processor(processor):
"""Extract tokenizer from a processor object."""
if isinstance(processor, PreTrainedTokenizerBase):
return processor
return processor.tokenizer
# ============================================================================
# Test Utilities
# ============================================================================
class CustomTestCase(unittest.TestCase):
"""
Custom test case base class with retry support.
This provides automatic test retry functionality based on environment variables.
"""
def _callTestMethod(self, method):
"""Override to add retry logic."""
max_retry = int(os.environ.get("SGLANG_TEST_MAX_RETRY", "0"))
if max_retry == 0:
# No retry, just run once
return super(CustomTestCase, self)._callTestMethod(method)
# Retry logic
for attempt in range(max_retry + 1):
try:
return super(CustomTestCase, self)._callTestMethod(method)
except Exception as e:
if attempt < max_retry:
logger.info(
f"Test failed on attempt {attempt + 1}/{max_retry + 1}, retrying..."
)
continue
else:
# Last attempt, re-raise the exception
raise
@@ -0,0 +1,115 @@
"""
python3 -m unittest openai_server.validation.test_large_max_new_tokens.TestLargeMaxNewTokens.test_chat_completion
"""
import os
import sys
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import openai
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
STDERR_FILENAME,
STDOUT_FILENAME,
CustomTestCase,
get_tokenizer,
kill_process_tree,
)
class TestLargeMaxNewTokens(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.stdout = open(STDOUT_FILENAME, "w")
cls.stderr = open(STDERR_FILENAME, "w")
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
worker_args=(
"--max-total-token",
"1536",
"--context-len",
"8192",
"--decode-log-interval",
"2",
),
num_workers=1,
tp_size=2,
env={"SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION": "256", **os.environ},
stdout=cls.stdout,
stderr=cls.stderr,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
cls.stdout.close()
cls.stderr.close()
os.remove(STDOUT_FILENAME)
os.remove(STDERR_FILENAME)
def run_chat_completion(self):
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant"},
{
"role": "user",
"content": "Please repeat the world 'hello' for 10000 times.",
},
],
temperature=0,
)
return response
def test_chat_completion(self):
num_requests = 4
all_requests_running = False
futures = []
with ThreadPoolExecutor(num_requests) as executor:
# Send multiple requests
for i in range(num_requests):
futures.append(executor.submit(self.run_chat_completion))
# Ensure that they are running concurrently
pt = 0
while pt >= 0:
time.sleep(5)
# Flush stderr to ensure logs are written
self.stderr.flush()
lines = open(STDERR_FILENAME).readlines()
for line in lines[pt:]:
if f"#running-req: {num_requests}" in line:
all_requests_running = True
pt = -1
break
pt += 1
assert all_requests_running
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,109 @@
"""
gRPC Router E2E Test - Test Openai Server Ignore Eos
This test file is REUSED from test/srt/openai_server/validation/test_openai_server_ignore_eos.py
with minimal changes:
num_workers=2,
- Swap popen_launch_server() → popen_launch_workers_and_router()
- Update teardown to cleanup router + workers
- All test logic and assertions remain identical
Run with:
pytest e2e_test/e2e_grpc/e2e_grpc/validation/test_openai_server_ignore_eos.py -v
"""
import sys
from pathlib import Path
import openai
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR.parent))
from fixtures import popen_launch_workers_and_router
from util import (
DEFAULT_MODEL_PATH,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
get_tokenizer,
kill_process_tree,
)
class TestOpenAIServerIgnoreEOS(CustomTestCase):
@classmethod
def setUpClass(cls):
# CHANGE: Launch gRPC router with integrated workers (single command)
cls.model = DEFAULT_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
cls.api_key = "sk-123456"
cls.cluster = popen_launch_workers_and_router(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
num_workers=1,
tp_size=2,
)
cls.base_url += "/v1"
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
# Cleanup router and workers
kill_process_tree(cls.cluster["router"].pid)
for worker in cls.cluster.get("workers", []):
kill_process_tree(worker.pid)
def test_ignore_eos(self):
"""
Test that ignore_eos=True allows generation to continue beyond EOS token
and reach the max_tokens limit.
"""
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
max_tokens = 200
response_default = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Count from 1 to 20."},
],
temperature=0,
max_tokens=max_tokens,
extra_body={"ignore_eos": False},
)
response_ignore_eos = client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Count from 1 to 20."},
],
temperature=0,
max_tokens=max_tokens,
extra_body={"ignore_eos": True},
)
default_tokens = len(
self.tokenizer.encode(response_default.choices[0].message.content)
)
ignore_eos_tokens = len(
self.tokenizer.encode(response_ignore_eos.choices[0].message.content)
)
# Check if ignore_eos resulted in more tokens or exactly max_tokens
# The ignore_eos response should either:
# 1. Have more tokens than the default response (if default stopped at EOS before max_tokens)
# 2. Have exactly max_tokens (if it reached the max_tokens limit)
self.assertTrue(
ignore_eos_tokens > default_tokens or ignore_eos_tokens >= max_tokens,
f"ignore_eos did not generate more tokens: {ignore_eos_tokens} vs {default_tokens}",
)
self.assertEqual(
response_ignore_eos.choices[0].finish_reason,
"length",
f"Expected finish_reason='length' for ignore_eos=True, got {response_ignore_eos.choices[0].finish_reason}",
)
@@ -0,0 +1,807 @@
import json
import logging
import os
import shutil
import signal
import socket
import subprocess
import time
from pathlib import Path
from types import SimpleNamespace
from typing import Callable, Optional
from urllib.parse import urlparse
import pytest
import requests
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
)
logger = logging.getLogger(__name__)
def _find_available_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _parse_url(base_url: str) -> tuple[str, str]:
"""Parse a base URL and return (host, port) as strings.
This is more robust than simple string splitting and supports different schemes
and URL shapes like trailing paths.
"""
parsed = urlparse(base_url)
return parsed.hostname or "127.0.0.1", (
str(parsed.port) if parsed.port is not None else ""
)
def _wait_router_health(base_url: str, timeout: float) -> None:
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{base_url}/health", timeout=5)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(2)
raise TimeoutError("Router failed to become healthy in time")
def _popen_launch_router(
model: str,
base_url: str,
dp_size: int,
timeout: float,
policy: str = "cache_aware",
) -> subprocess.Popen:
host, port = _parse_url(base_url)
prom_port = _find_available_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
"--dp",
str(dp_size),
"--router-policy",
policy,
"--allow-auto-truncate",
"--router-prometheus-port",
str(prom_port),
"--router-prometheus-host",
"127.0.0.1",
"--router-log-level",
"warn",
]
proc = subprocess.Popen(cmd)
_wait_router_health(base_url, timeout)
return proc
def _popen_launch_worker(
model: str,
base_url: str,
*,
dp_size: int | None = None,
api_key: str | None = None,
base_gpu_id: int | None = 0,
) -> subprocess.Popen:
host, port = _parse_url(base_url)
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
"--base-gpu-id",
str(base_gpu_id or 0),
"--log-level",
"warning",
]
if dp_size is not None:
cmd += ["--dp-size", str(dp_size)]
if api_key is not None:
cmd += ["--api-key", api_key]
return subprocess.Popen(cmd)
def _popen_launch_router_only(
base_url: str,
policy: str = "round_robin",
timeout: float = 120.0,
*,
dp_aware: bool = False,
enable_igw: bool = False,
api_key: str | None = None,
) -> subprocess.Popen:
host, port = _parse_url(base_url)
prom_port = _find_available_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
host,
"--port",
port,
"--policy",
policy,
]
if dp_aware:
cmd += ["--dp-aware"]
if enable_igw:
cmd += ["--enable-igw"]
if api_key is not None:
cmd += ["--api-key", api_key]
cmd += [
"--prometheus-port",
str(prom_port),
"--prometheus-host",
"127.0.0.1",
"--log-level",
"warn",
]
proc = subprocess.Popen(cmd)
_wait_router_health(base_url, timeout)
return proc
def _terminate(proc: subprocess.Popen, timeout: float = 120) -> None:
if proc is None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
def _which(cmd: str) -> Optional[str]:
try:
return shutil.which(cmd)
except Exception as e:
logger.warning("shutil.which(%r) failed: %s", cmd, e)
return None
def _graceful_stop_popen(p: subprocess.Popen) -> None:
if p is None:
return
try:
if p.poll() is None:
p.terminate()
for _ in range(5):
if p.poll() is not None:
break
time.sleep(1)
if p.poll() is None:
p.kill()
except Exception as e:
logger.warning("Exception during graceful stop of popen: %s", e)
def _pid_alive(pid: int) -> bool:
try:
os.kill(pid, 0)
return True
except Exception:
return False
def _graceful_stop_pid(pid: int) -> None:
try:
if _pid_alive(pid):
try:
os.kill(pid, signal.SIGTERM)
except Exception:
pass
for _ in range(5):
if not _pid_alive(pid):
break
time.sleep(1)
if _pid_alive(pid):
try:
os.kill(pid, signal.SIGKILL)
except Exception:
pass
except Exception:
pass
def _graceful_stop_any(obj) -> None:
try:
if isinstance(obj, subprocess.Popen):
_graceful_stop_popen(obj)
return
if isinstance(obj, int):
_graceful_stop_pid(obj)
return
proc_obj = getattr(obj, "proc", None)
if isinstance(proc_obj, subprocess.Popen):
_graceful_stop_popen(proc_obj)
except Exception:
pass
def _gpu_monitor_should_run(thresholds: Optional[dict]) -> bool:
"""Decide whether to enable the GPU monitor.
Runs if thresholds request GPU checks or if GPU_UTIL_LOG is truthy.
"""
want = False
try:
mean_th = None if thresholds is None else thresholds.get("gpu_util_mean_min")
p50_th = None if thresholds is None else thresholds.get("gpu_util_p50_min")
want = bool(mean_th is not None or p50_th is not None)
except Exception:
want = False
if not want:
env_flag = os.environ.get("GPU_UTIL_LOG", "").lower() in ("1", "true", "yes")
want = want or env_flag
return want
def _gpu_monitor_path(experiment_folder: str) -> str:
"""Return the JSON path for storing GPU monitor results."""
base = Path.cwd() / experiment_folder
return str(base / "gpu_utilization.json")
def _launch_gpu_monitor(bench_pid: int, experiment_folder: str, interval: float):
"""Start the GPU monitor process. Returns (proc, path) or (None, None)."""
try:
from multiprocessing import Process
out_path = _gpu_monitor_path(experiment_folder)
proc = Process(
target=_gpu_monitor_proc_entry,
args=(bench_pid, out_path, interval),
daemon=True,
)
proc.start()
return proc, out_path
except Exception as e:
logger.warning("Failed to launch GPU monitor: %s", e)
return None, None
def _read_gpu_monitor_result(path: Optional[str]) -> Optional[dict]:
try:
if path and os.path.exists(path):
with open(path, "r") as f:
return json.load(f)
except Exception as e:
logger.warning("Failed to read GPU monitor result from %r: %s", path, e)
return None
def _log_and_assert_gpu_thresholds(
result: Optional[dict], thresholds: Optional[dict]
) -> None:
if not result or not isinstance(result, dict) or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples.")
return
overall = result.get("overall", {}) if isinstance(result, dict) else {}
count = int(result.get("count", 0))
mean_th = None if thresholds is None else thresholds.get("gpu_util_mean_min")
p50_th = None if thresholds is None else thresholds.get("gpu_util_p50_min")
mean_v = float(overall.get("mean", 0.0))
p50_v = overall.get("p50")
logger.info(
"GPU utilization overall: mean=%.2f%% p50=%s (samples=%d)",
mean_v,
(f"{float(p50_v):.2f}%" if p50_v is not None else "n/a"),
count,
)
if mean_th is not None:
assert mean_v >= float(
mean_th
), f"GPU utilization mean below threshold: {mean_v:.2f}% < {mean_th}%"
if p50_th is not None and p50_v is not None:
p50_f = float(p50_v)
assert p50_f >= float(
p50_th
), f"GPU utilization p50 below threshold: {p50_f:.2f}% < {p50_th}%"
def _gpu_monitor_proc_entry(bench_pid: int, out_file: str, interval: float) -> None:
"""Low-impact GPU utilization monitor using NVML in a separate process.
Writes JSON to out_file that includes overall and per-GPU raw samples and summary stats.
"""
try:
try:
os.nice(10)
except Exception:
pass
total = 0.0
n = 0
try:
import pynvml # type: ignore
pynvml.nvmlInit()
except Exception:
with open(out_file, "w") as f:
os.makedirs(os.path.dirname(out_file), exist_ok=True)
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {},
},
f,
)
return
try:
import pynvml # type: ignore
count = pynvml.nvmlDeviceGetCount()
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)]
except Exception:
with open(out_file, "w") as f:
os.makedirs(os.path.dirname(out_file), exist_ok=True)
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {},
},
f,
)
return
# Prepare per-GPU and overall raw collectors
per_gpu_samples: dict[str, list[float]] = {}
overall_samples: list[float] = []
while True:
if not os.path.exists(f"/proc/{bench_pid}"):
break
try:
vals = []
import pynvml # type: ignore
for idx, h in enumerate(handles):
try:
util = pynvml.nvmlDeviceGetUtilizationRates(h).gpu
vals.append(float(util))
key = str(idx)
per_gpu_samples.setdefault(key, []).append(float(util))
except Exception:
continue
if vals:
avg = sum(vals) / len(vals)
overall_samples.append(avg)
total += avg
n += 1
except Exception:
pass
time.sleep(interval)
finally:
try:
os.makedirs(os.path.dirname(out_file), exist_ok=True)
with open(out_file, "w") as f:
def pct_from(samples: list[float], p: float) -> float:
if not samples:
return 0.0
srt = sorted(samples)
i = max(
0, min(len(srt) - 1, int(round((p / 100.0) * (len(srt) - 1))))
)
return float(srt[i])
overall_mean = (total / n) if n > 0 else 0.0
per_gpu_summary: dict[str, dict] = {}
for key, arr in per_gpu_samples.items():
per_gpu_summary[key] = {
"mean": float(sum(arr) / len(arr)) if arr else 0.0,
"p5": pct_from(arr, 5),
"p10": pct_from(arr, 10),
"p25": pct_from(arr, 25),
"p50": pct_from(arr, 50),
"p75": pct_from(arr, 75),
"p90": pct_from(arr, 90),
"p95": pct_from(arr, 95),
"min": float(min(arr)) if arr else 0.0,
"max": float(max(arr)) if arr else 0.0,
"count": len(arr),
}
out_payload = {
"bench_pid": bench_pid,
"interval_sec": interval,
"count": n,
"overall": {
"mean": float(overall_mean),
"p5": pct_from(overall_samples, 5),
"p10": pct_from(overall_samples, 10),
"p25": pct_from(overall_samples, 25),
"p50": pct_from(overall_samples, 50),
"p75": pct_from(overall_samples, 75),
"p90": pct_from(overall_samples, 90),
"p95": pct_from(overall_samples, 95),
"min": float(min(overall_samples)) if overall_samples else 0.0,
"max": float(max(overall_samples)) if overall_samples else 0.0,
},
"per_gpu": per_gpu_summary,
"raw": {
"overall": overall_samples,
"per_gpu": per_gpu_samples,
},
}
json.dump(out_payload, f)
except Exception:
pass
try:
import pynvml # type: ignore
pynvml.nvmlShutdown()
except Exception:
pass
@pytest.fixture(scope="session")
def genai_bench_runner() -> Callable[..., None]:
"""Provide a callable to run genai-bench and validate metrics.
Usage in tests:
def test(..., genai_bench_runner):
genai_bench_runner(router_url=..., model_path=..., experiment_folder=...)
"""
def _run(
*,
router_url: str,
model_path: str,
experiment_folder: str,
timeout_sec: int | None = None,
thresholds: dict | None = None,
extra_env: dict | None = None,
num_concurrency: int = 32,
traffic_scenario: str = "D(4000,100)",
max_requests_per_run: int | None = None,
clean_experiment: bool = True,
kill_procs: list | None = None,
drain_delay_sec: int = 6,
) -> None:
cli = _which("genai-bench")
if not cli:
pytest.fail(
"genai-bench CLI not found; please install it to run benchmarks"
)
# Clean previous experiment folder under current working directory
if clean_experiment:
exp_dir = Path.cwd() / experiment_folder
if exp_dir.exists():
shutil.rmtree(exp_dir, ignore_errors=True)
# Default requests per run if not provided
mrr = (
max_requests_per_run
if max_requests_per_run is not None
else num_concurrency * 5
)
cmd = [
cli,
"benchmark",
"--api-backend",
"openai",
"--api-base",
router_url,
"--api-key",
"dummy-token",
"--api-model-name",
model_path,
"--model-tokenizer",
model_path,
"--task",
"text-to-text",
"--num-concurrency",
str(num_concurrency),
"--traffic-scenario",
traffic_scenario,
"--max-requests-per-run",
str(mrr),
"--max-time-per-run",
"3",
"--experiment-folder-name",
experiment_folder,
"--experiment-base-dir",
str(Path.cwd()),
]
env = os.environ.copy()
if extra_env:
env.update(extra_env)
to = timeout_sec or int(os.environ.get("GENAI_BENCH_TEST_TIMEOUT", "120"))
proc = subprocess.Popen(
cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
)
# Optional GPU utilization monitor in a low-priority child process (pynvml only)
# Enabled only when gpu_util_mean_min is provided in thresholds.
monitor_path = None
monitor_proc = None
gpu_util_result: dict | None = None
want_gpu_monitor = _gpu_monitor_should_run(thresholds)
if want_gpu_monitor:
interval = float(os.environ.get("GPU_UTIL_SAMPLE_INTERVAL", "2.0"))
monitor_proc, monitor_path = _launch_gpu_monitor(
bench_pid=proc.pid,
experiment_folder=experiment_folder,
interval=interval,
)
stdout = stderr = ""
rc = None
try:
try:
stdout, stderr = proc.communicate(timeout=to)
except subprocess.TimeoutExpired:
# Simple: kill the CLI process if it doesn't exit in time
try:
proc.kill()
except Exception:
pass
stdout, stderr = proc.communicate()
rc = proc.returncode
# Prefer exact path under cwd; fallback to rglob search
base = Path.cwd()
direct = base / experiment_folder
candidates = [direct] if direct.is_dir() else []
if not candidates:
for p in base.rglob(experiment_folder):
if p.is_dir() and p.name == experiment_folder:
candidates = [p]
break
if not candidates:
raise AssertionError(
"Benchmark failed: experiment folder not found: "
f"{experiment_folder}\nExit code: {rc}\nSTDOUT (tail):\n{stdout[-1000:]}\nSTDERR (tail):\n{stderr[-1000:]}"
)
actual_folder = candidates[0]
json_files = []
for _ in range(10):
json_files = [
p
for p in actual_folder.rglob("*.json")
if "experiment_metadata" not in p.name
]
if json_files:
break
time.sleep(1)
if not json_files:
raise AssertionError(
"Benchmark failed: no JSON results found\n"
f"Exit code: {rc}\nSTDOUT (tail):\n{stdout[-1000:]}\nSTDERR (tail):\n{stderr[-1000:]}"
)
th = thresholds # None means "log only", no validation
for jf in json_files:
with jf.open("r") as f:
data = json.load(f)
stats = data.get("aggregated_metrics", {}).get("stats", {})
ttft_mean = float(stats.get("ttft", {}).get("mean", float("inf")))
e2e_latency_mean = float(
stats.get("e2e_latency", {}).get("mean", float("inf"))
)
input_tp_mean = float(
stats.get("input_throughput", {}).get("mean", 0.0)
)
output_tp_mean = float(
stats.get("output_throughput", {}).get("mean", 0.0)
)
logger.info(
"genai-bench[%s] %s ttft_mean=%.3fs e2e_latency_mean=%.3fs input_tp_mean=%.1f tok/s output_tp_mean=%.1f tok/s",
experiment_folder,
jf.name,
ttft_mean,
e2e_latency_mean,
input_tp_mean,
output_tp_mean,
)
if th is not None:
assert (
ttft_mean <= th["ttft_mean_max"]
), f"TTFT validation failed: {ttft_mean} > {th['ttft_mean_max']} (file={jf.name})"
assert (
e2e_latency_mean <= th["e2e_latency_mean_max"]
), f"E2E latency validation failed: {e2e_latency_mean} > {th['e2e_latency_mean_max']} (file={jf.name})"
assert (
input_tp_mean >= th["input_throughput_mean_min"]
), f"Input throughput validation failed: {input_tp_mean} < {th['input_throughput_mean_min']} (file={jf.name})"
assert (
output_tp_mean >= th["output_throughput_mean_min"]
), f"Output throughput validation failed: {output_tp_mean} < {th['output_throughput_mean_min']} (file={jf.name})"
# Validate optional GPU utilization threshold if provided
if want_gpu_monitor:
try:
if monitor_proc is not None:
monitor_proc.join(timeout=5)
except Exception:
pass
gpu_util_result = _read_gpu_monitor_result(monitor_path)
_log_and_assert_gpu_thresholds(gpu_util_result, thresholds)
finally:
# Always attempt to stop workers to avoid resource leakage
if kill_procs:
# Give router/workers a small grace period to finish any last drains
if drain_delay_sec > 0:
try:
time.sleep(drain_delay_sec)
except Exception:
pass
for p in kill_procs:
_graceful_stop_any(p)
try:
time.sleep(2)
except Exception:
pass
# Ensure GPU monitor process is cleaned up
if monitor_proc is not None and monitor_proc.is_alive():
try:
monitor_proc.terminate()
except Exception:
pass
return _run
def pytest_configure(config):
config.addinivalue_line("markers", "e2e: mark as end-to-end test")
@pytest.fixture(scope="session")
def e2e_model() -> str:
# Always use the default test model
return os.getenv("E2E_PRIMARY_MODEL", DEFAULT_MODEL_NAME_FOR_TEST)
@pytest.fixture
def e2e_router(e2e_model: str):
# Keep this available but tests below use router-only to avoid GPU contention
base_url = DEFAULT_URL_FOR_TEST
proc = _popen_launch_router(
e2e_model, base_url, dp_size=2, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture
def e2e_router_only_rr():
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_router_only(base_url, policy="round_robin")
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_embedding_model() -> str:
"""Embedding model to use for E2E tests.
Defaults to an E5 Mistral model, can be overridden via E2E_EMBEDDING_MODEL env var.
"""
import os
return os.getenv("E2E_EMBEDDING_MODEL", "intfloat/e5-mistral-7b-instruct")
@pytest.fixture
def e2e_primary_embedding_worker(e2e_embedding_model: str):
"""Launch a single embedding worker using the specified model."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_worker(e2e_embedding_model, base_url)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_primary_worker(e2e_model: str):
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
proc = _popen_launch_worker(e2e_model, base_url)
# Router health gate will handle worker readiness
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture
def e2e_router_only_rr_dp_aware_api():
"""Router-only with dp-aware enabled and an API key."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
api_key = "secret"
proc = _popen_launch_router_only(
base_url, policy="round_robin", timeout=180.0, dp_aware=True, api_key=api_key
)
try:
yield SimpleNamespace(proc=proc, url=base_url, api_key=api_key)
finally:
_terminate(proc)
@pytest.fixture
def e2e_worker_dp2_api(e2e_model: str, e2e_router_only_rr_dp_aware_api):
"""Worker with dp-size=2 and the same API key as the dp-aware router."""
port = _find_available_port()
base_url = f"http://127.0.0.1:{port}"
api_key = e2e_router_only_rr_dp_aware_api.api_key
proc = _popen_launch_worker(e2e_model, base_url, dp_size=2, api_key=api_key)
try:
yield SimpleNamespace(proc=proc, url=base_url)
finally:
_terminate(proc)
@pytest.fixture(scope="session")
def e2e_two_workers_dp2(e2e_model: str):
"""Launch two workers, each with dp_size=2, mapped to GPUs [0,1] and [2,3]."""
workers = []
try:
# Worker A on GPUs 0-1
port_a = _find_available_port()
url_a = f"http://127.0.0.1:{port_a}"
proc_a = _popen_launch_worker(e2e_model, url_a, dp_size=2, base_gpu_id=0)
workers.append(SimpleNamespace(proc=proc_a, url=url_a))
# Worker B on GPUs 2-3
port_b = _find_available_port()
url_b = f"http://127.0.0.1:{port_b}"
proc_b = _popen_launch_worker(e2e_model, url_b, dp_size=2, base_gpu_id=2)
workers.append(SimpleNamespace(proc=proc_b, url=url_b))
yield workers
finally:
for w in workers:
_terminate(w.proc)
@@ -0,0 +1,62 @@
import time
import pytest
import requests
def _wait_for_workers(
base_url: str, expected_count: int, timeout: float = 60.0, headers: dict = None
) -> None:
"""Poll /workers endpoint until expected number of workers are registered."""
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{base_url}/workers", headers=headers, timeout=5)
if r.status_code == 200:
workers = r.json().get("workers", [])
if len(workers) >= expected_count:
return
except requests.RequestException:
pass
time.sleep(0.5)
raise TimeoutError(
f"Expected {expected_count} workers at {base_url}, timed out after {timeout}s"
)
@pytest.mark.e2e
def test_embeddings_basic(
e2e_router_only_rr, e2e_primary_embedding_worker, e2e_embedding_model
):
base = e2e_router_only_rr.url
worker_url = e2e_primary_embedding_worker.url
# Attach embedding worker to router-only instance
r = requests.post(f"{base}/workers", json={"url": worker_url}, timeout=180)
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
# Wait for worker to be registered
_wait_for_workers(base, expected_count=1, timeout=60.0)
# Simple embedding request with two inputs
payload = {
"model": e2e_embedding_model,
"input": [
"the quick brown fox",
"jumps over the lazy dog",
],
}
r = requests.post(f"{base}/v1/embeddings", json=payload, timeout=120)
assert r.status_code == 200, f"unexpected status: {r.status_code} {r.text}"
data = r.json()
assert "data" in data and isinstance(data["data"], list)
assert len(data["data"]) == 2
# Validate shape of embedding objects
for item in data["data"]:
assert "embedding" in item and isinstance(item["embedding"], list)
# Ensure non-empty vectors
assert len(item["embedding"]) > 0
@@ -0,0 +1,264 @@
import logging
import socket
import subprocess
import time
from types import SimpleNamespace
from typing import Optional
import pytest
import requests
from sglang.test.run_eval import run_eval
logger = logging.getLogger(__name__)
def _find_available_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def _wait_health(url: str, timeout: float = 180.0) -> None:
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{url}/health", timeout=5)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(1)
raise TimeoutError(f"Service at {url} failed to become healthy in time")
def _detect_ib_device() -> Optional[str]:
"""Return first active IB device name (e.g., mlx5_0) or None if unavailable."""
# Fast check that ibv_devinfo exists
try:
subprocess.run(
["ibv_devinfo", "-l"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=1,
)
except (FileNotFoundError, subprocess.TimeoutExpired):
return None
for i in range(12):
dev = f"mlx5_{i}"
try:
res = subprocess.run(
["ibv_devinfo", dev],
capture_output=True,
text=True,
timeout=2,
)
if res.returncode == 0 and ("state:" in res.stdout):
for line in res.stdout.splitlines():
if "state:" in line and "PORT_ACTIVE" in line:
return dev
except Exception:
pass
return None
def _popen_launch_prefill_worker(
model: str,
bootstrap_port: int,
ib_device: Optional[str] = None,
base_gpu_id: int = 0,
) -> SimpleNamespace:
port = _find_available_port()
url = f"http://127.0.0.1:{port}"
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--disaggregation-mode",
"prefill",
"--host",
"127.0.0.1",
"--port",
str(port),
"--disaggregation-bootstrap-port",
str(bootstrap_port),
"--base-gpu-id",
str(base_gpu_id),
]
if ib_device:
cmd += ["--disaggregation-ib-device", ib_device]
proc = subprocess.Popen(cmd)
_wait_health(url, timeout=300.0)
return SimpleNamespace(proc=proc, url=url, bootstrap_port=bootstrap_port)
def _popen_launch_decode_worker(
model: str, ib_device: Optional[str] = None, base_gpu_id: int = 0
) -> SimpleNamespace:
port = _find_available_port()
url = f"http://127.0.0.1:{port}"
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--disaggregation-mode",
"decode",
"--host",
"127.0.0.1",
"--port",
str(port),
"--base-gpu-id",
str(base_gpu_id),
]
if ib_device:
cmd += ["--disaggregation-ib-device", ib_device]
proc = subprocess.Popen(cmd)
_wait_health(url, timeout=300.0)
return SimpleNamespace(proc=proc, url=url)
def _terminate(proc: subprocess.Popen, timeout: float = 120) -> None:
if proc is None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
@pytest.fixture(scope="module")
def pd_cluster(e2e_model: str):
"""Start 2 prefill + 2 decode workers and one PD router, once per module."""
# Environment capability checks: require sgl_kernel and GPU backend
try:
import sgl_kernel # noqa: F401
except Exception as e: # pragma: no cover - environment dependent
pytest.fail(f"PD e2e requires sgl_kernel but it is not available: {e}")
try:
import torch # noqa: F401
except Exception as e: # pragma: no cover - environment dependent
pytest.fail(
f"PD e2e requires torch but it is not available or misconfigured: {e}"
)
if not torch.cuda.is_available(): # pragma: no cover - environment dependent
pytest.fail("PD e2e requires CUDA backend, but CUDA is not available")
workers: list[SimpleNamespace] = []
router_proc = None
try:
ib_device = _detect_ib_device()
# Launch 4 workers across 4 GPUs: prefill on 0,1 and decode on 2,3
pf1 = _popen_launch_prefill_worker(
e2e_model,
bootstrap_port=_find_available_port(),
ib_device=ib_device,
base_gpu_id=0,
)
pf2 = _popen_launch_prefill_worker(
e2e_model,
bootstrap_port=_find_available_port(),
ib_device=ib_device,
base_gpu_id=1,
)
dc1 = _popen_launch_decode_worker(e2e_model, ib_device=ib_device, base_gpu_id=2)
dc2 = _popen_launch_decode_worker(e2e_model, ib_device=ib_device, base_gpu_id=3)
prefills = [pf1, pf2]
decodes = [dc1, dc2]
workers.extend(prefills + decodes)
# PD router with two prefill and two decode endpoints
rport = _find_available_port()
router_url = f"http://127.0.0.1:{rport}"
pport = _find_available_port()
prefill = [(pf.url, pf.bootstrap_port) for pf in prefills]
decode = [dc.url for dc in decodes]
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(rport),
"--policy",
"round_robin",
"--pd-disaggregation",
"--log-level",
"warn",
]
for url, bport in prefill:
cmd += ["--prefill", url, str(bport)]
for url in decode:
cmd += ["--decode", url]
cmd += [
"--prometheus-port",
str(pport),
"--prometheus-host",
"127.0.0.1",
]
router_proc = subprocess.Popen(cmd)
_wait_health(router_url, timeout=180.0)
yield SimpleNamespace(
router_url=router_url, workers=workers, router_proc=router_proc
)
finally:
if router_proc is not None:
_terminate(router_proc)
for w in workers:
_terminate(w.proc)
@pytest.mark.e2e
def test_pd_mmlu(e2e_model: str, pd_cluster):
"""
Launch 4 workers, start a PD router (2 prefill + 2 decode), then run MMLU.
"""
args = SimpleNamespace(
base_url=pd_cluster.router_url,
model=e2e_model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert metrics["score"] >= 0.65
@pytest.mark.e2e
def test_pd_genai_bench(e2e_model: str, pd_cluster, genai_bench_runner):
"""
Launch 4 workers, start a PD router (2 prefill + 2 decode), then run a
short genai-bench benchmark and validate aggregate metrics.
"""
# Run genai-bench against the shared router
policy_label = "benchmark_round_robin_pd"
genai_bench_runner(
router_url=pd_cluster.router_url,
model_path=e2e_model,
experiment_folder=policy_label,
thresholds={
"ttft_mean_max": 13,
"e2e_latency_mean_max": 16,
"input_throughput_mean_min": 350,
"output_throughput_mean_min": 18,
"gpu_util_p50_min": 99,
},
kill_procs=pd_cluster.workers,
)
@@ -0,0 +1,238 @@
import threading
import time
from types import SimpleNamespace
import pytest
import requests
from sglang.test.run_eval import run_eval
def _wait_for_workers(
base_url: str, expected_count: int, timeout: float = 60.0, headers: dict = None
) -> None:
"""Poll /workers endpoint until expected number of workers are registered."""
start = time.perf_counter()
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
r = session.get(f"{base_url}/workers", headers=headers, timeout=5)
if r.status_code == 200:
workers = r.json().get("workers", [])
if len(workers) >= expected_count:
return
except requests.RequestException:
pass
time.sleep(0.5)
raise TimeoutError(
f"Expected {expected_count} workers at {base_url}, timed out after {timeout}s"
)
def _get_worker_id_by_url(base_url: str, worker_url: str, headers: dict = None) -> str:
r = requests.get(f"{base_url}/workers", headers=headers, timeout=10)
r.raise_for_status()
workers = r.json().get("workers", [])
worker_id = next((w.get("id") for w in workers if w.get("url") == worker_url), None)
if not worker_id:
raise RuntimeError(
f"Could not find worker_id for url={worker_url}. workers={workers}"
)
return worker_id
@pytest.mark.e2e
def test_mmlu(e2e_router_only_rr, e2e_two_workers_dp2, e2e_model):
# Attach two dp=2 workers (total 4 GPUs) to a fresh router-only instance
base = e2e_router_only_rr.url
for w in e2e_two_workers_dp2:
r = requests.post(f"{base}/workers", json={"url": w.url}, timeout=180)
assert (
r.status_code == 202
), f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
# Wait for workers to be registered
_wait_for_workers(base, expected_count=2, timeout=60.0)
args = SimpleNamespace(
base_url=base,
model=e2e_model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert metrics["score"] >= 0.65
@pytest.mark.e2e
def test_genai_bench(
e2e_router_only_rr, e2e_two_workers_dp2, e2e_model, genai_bench_runner
):
"""Attach a worker to the regular router and run a short genai-bench."""
base = e2e_router_only_rr.url
for w in e2e_two_workers_dp2:
r = requests.post(f"{base}/workers", json={"url": w.url}, timeout=180)
assert (
r.status_code == 202
), f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
# Wait for workers to be registered
_wait_for_workers(base, expected_count=2, timeout=60.0)
genai_bench_runner(
router_url=base,
model_path=e2e_model,
experiment_folder="benchmark_round_robin_regular",
thresholds={
"ttft_mean_max": 6,
"e2e_latency_mean_max": 14,
"input_throughput_mean_min": 800, # temp relax from 1000 to 800 for now
"output_throughput_mean_min": 12,
# Enforce GPU utilization p50 >= 99% during the run.
"gpu_util_p50_min": 99,
},
kill_procs=e2e_two_workers_dp2,
)
@pytest.mark.e2e
def test_add_and_remove_worker_live(e2e_router_only_rr, e2e_primary_worker, e2e_model):
base = e2e_router_only_rr.url
worker_url = e2e_primary_worker.url
r = requests.post(f"{base}/workers", json={"url": worker_url}, timeout=180)
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
# Wait for worker to be registered
_wait_for_workers(base, expected_count=1, timeout=60.0)
with requests.Session() as s:
for i in range(8):
r = s.post(
f"{base}/v1/completions",
json={
"model": e2e_model,
"prompt": f"x{i}",
"max_tokens": 1,
"stream": False,
},
timeout=120,
)
r.raise_for_status()
# Remove the worker
worker_id = _get_worker_id_by_url(base, worker_url)
r = requests.delete(f"{base}/workers/{worker_id}", timeout=60)
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
@pytest.mark.e2e
def test_lazy_fault_tolerance_live(e2e_router_only_rr, e2e_primary_worker, e2e_model):
base = e2e_router_only_rr.url
worker = e2e_primary_worker
r = requests.post(f"{base}/workers", json={"url": worker.url}, timeout=180)
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
# Wait for worker to be registered
_wait_for_workers(base, expected_count=1, timeout=60.0)
def killer():
time.sleep(10)
try:
worker.proc.terminate()
except Exception:
pass
t = threading.Thread(target=killer, daemon=True)
t.start()
args = SimpleNamespace(
base_url=base,
model=e2e_model,
eval_name="mmlu",
num_examples=32,
num_threads=16,
temperature=0.0,
)
metrics = run_eval(args)
assert 0.0 <= metrics["score"] <= 1.0
@pytest.mark.e2e
def test_dp_aware_worker_expansion_and_api_key(
e2e_model,
e2e_router_only_rr_dp_aware_api,
e2e_worker_dp2_api,
):
"""
Launch a router-only instance in dp_aware mode and a single worker with dp_size=2
and API key protection. Verify expansion, auth enforcement, and basic eval.
"""
import os
router_url = e2e_router_only_rr_dp_aware_api.url
worker_url = e2e_worker_dp2_api.url
api_key = e2e_router_only_rr_dp_aware_api.api_key
# Attach worker; router should expand to dp_size logical workers
r = requests.post(
f"{router_url}/workers",
json={"url": worker_url, "api_key": api_key},
headers={"Authorization": f"Bearer {api_key}"},
timeout=180,
)
assert r.status_code == 202, f"Expected 202 ACCEPTED, got {r.status_code}: {r.text}"
# Wait for workers to be registered and expanded
_wait_for_workers(
router_url,
expected_count=2,
timeout=60.0,
headers={"Authorization": f"Bearer {api_key}"},
)
# Verify the expanded workers have correct URLs
r = requests.get(
f"{router_url}/workers",
headers={"Authorization": f"Bearer {api_key}"},
timeout=30,
)
r.raise_for_status()
workers = r.json().get("workers", [])
urls = [w["url"] for w in workers]
assert len(urls) == 2
assert set(urls) == {f"{worker_url}@0", f"{worker_url}@1"}
# Verify API key enforcement
# 1) Without Authorization -> Should get 401 Unauthorized
r = requests.post(
f"{router_url}/v1/completions",
json={"model": e2e_model, "prompt": "hi", "max_tokens": 1},
timeout=60,
)
assert r.status_code == 401
# 2) With correct Authorization -> 200
r = requests.post(
f"{router_url}/v1/completions",
json={"model": e2e_model, "prompt": "hi", "max_tokens": 1},
headers={"Authorization": f"Bearer {api_key}"},
timeout=60,
)
assert r.status_code == 200
# Finally, run MMLU eval through the router with auth
os.environ["OPENAI_API_KEY"] = api_key
args = SimpleNamespace(
base_url=router_url,
model=e2e_model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert metrics["score"] >= 0.65
@@ -0,0 +1,153 @@
import datetime
import os
import socket
import subprocess
import sys
import time
from contextlib import closing
import requests
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID
from urllib3.exceptions import InsecureRequestWarning
# Suppress insecure request warnings due to self-signed cert
requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
def find_free_port() -> int:
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
s.bind(("", 0))
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
return s.getsockname()[1]
def generate_self_signed_cert(cert_path: str, key_path: str) -> None:
"""Generate a self-signed certificate and private key for localhost."""
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
subject = issuer = x509.Name(
[
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "California"),
x509.NameAttribute(NameOID.LOCALITY_NAME, "San Francisco"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
x509.NameAttribute(NameOID.COMMON_NAME, "localhost"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(issuer)
.public_key(key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=10))
.add_extension(
x509.SubjectAlternativeName([x509.DNSName("localhost")]), critical=False
)
.sign(key, hashes.SHA256())
)
with open(key_path, "wb") as f:
f.write(
key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.TraditionalOpenSSL,
encryption_algorithm=serialization.NoEncryption(),
)
)
with open(cert_path, "wb") as f:
f.write(cert.public_bytes(serialization.Encoding.PEM))
def test_tls_server() -> None:
"""End-to-end test for TLS-enabled router startup and basic endpoints."""
cert_path = "cert.pem"
key_path = "key.pem"
generate_self_signed_cert(cert_path, key_path)
port = find_free_port()
cmd = [
sys.executable,
"-m",
"sglang_router.launch_router",
"--worker-urls",
"http://127.0.0.1:9999", # Dummy worker
"--host",
"127.0.0.1",
"--port",
str(port),
"--tls-cert-path",
cert_path,
"--tls-key-path",
key_path,
"--log-level",
"info",
]
proc = subprocess.Popen(
cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
)
try:
# Wait for server to start and respond to health check
start_time = time.time()
while time.time() - start_time < 15:
try:
response = requests.get(
f"https://localhost:{port}/health", verify=False, timeout=2
)
if response.status_code == 200:
break
except requests.RequestException:
pass
if proc.poll() is not None:
stdout, stderr = proc.communicate()
raise RuntimeError(
f"Router process died early.\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}"
)
time.sleep(0.5)
else:
raise TimeoutError("Server did not become healthy within 15 seconds")
# Verify basic endpoints work over TLS
models_resp = requests.get(
f"https://localhost:{port}/v1/models", verify=False, timeout=2
)
assert models_resp.status_code in (
200,
503,
) # 503 expected with no healthy workers
# Minimal generate request (should be rejected or queued)
gen_payload = {"model": "dummy", "prompt": "test", "max_new_tokens": 1}
gen_resp = requests.post(
f"https://localhost:{port}/generate",
json=gen_payload,
verify=False,
timeout=2,
)
assert gen_resp.status_code in (
200,
400,
503,
) # Various valid responses with dummy worker
finally:
proc.terminate()
try:
proc.wait(timeout=5)
except subprocess.TimeoutExpired:
proc.kill()
for path in (cert_path, key_path):
if os.path.exists(path):
os.remove(path)
@@ -0,0 +1,132 @@
"""
pytest configuration for e2e_response_api tests.
This configures pytest to not collect base test classes that are meant to be inherited.
"""
import os
import openai
import pytest # noqa: F401
from router_fixtures import (
popen_launch_openai_xai_router,
popen_launch_workers_and_router,
)
from util import kill_process_tree
# ------------------------------
# Backend Configuration Map
# ------------------------------
BACKENDS = {
"openai": {
"model": "gpt-5-nano",
"base_url_port": "http://127.0.0.1:30010",
"launcher": popen_launch_openai_xai_router,
"launcher_kwargs": {
"backend": "openai",
"history_backend": "memory",
},
"api_key_env": "OPENAI_API_KEY",
"needs_workers": False,
},
"xai": {
"model": "grok-4-fast",
"base_url_port": "http://127.0.0.1:30023",
"launcher": popen_launch_openai_xai_router,
"launcher_kwargs": {
"backend": "xai",
"history_backend": "memory",
},
"api_key_env": "XAI_API_KEY",
"needs_workers": False,
},
"grpc": {
"model": "/home/ubuntu/models/Qwen/Qwen2.5-14B-Instruct",
"base_url_port": "http://127.0.0.1:30030",
"launcher": popen_launch_workers_and_router,
"launcher_kwargs": {
"timeout": 90,
"num_workers": 1,
"tp_size": 2,
"policy": "round_robin",
"worker_args": ["--context-length=1000"],
"router_args": [
"--history-backend",
"memory",
"--tool-call-parser",
"qwen",
],
},
"api_key_env": None, # grpc does not use API keys
"needs_workers": True,
},
"grpc_harmony": {
"model": "/home/ubuntu/models/openai/gpt-oss-20b",
"base_url_port": "http://127.0.0.1:30030",
"launcher": popen_launch_workers_and_router,
"launcher_kwargs": {
"timeout": 90,
"num_workers": 1,
"tp_size": 2,
"policy": "round_robin",
"worker_args": ["--reasoning-parser=gpt-oss"],
"router_args": ["--history-backend", "memory"],
},
"api_key_env": None,
"needs_workers": True,
},
"oracle_store": {
"model": "gpt-5-nano",
"base_url_port": "http://127.0.0.1:30040",
"launcher": popen_launch_openai_xai_router,
"launcher_kwargs": {
"backend": "openai",
"history_backend": "oracle",
},
"api_key_env": "OPENAI_API_KEY",
"needs_workers": False,
},
}
@pytest.fixture(scope="class")
def setup_backend(request):
backend = request.param
if backend not in BACKENDS:
raise RuntimeError(f"Unknown backend {backend}")
cfg = BACKENDS[backend]
# Launch cluster
cluster = (
cfg["launcher"](
cfg["model"],
cfg["base_url_port"],
**cfg["launcher_kwargs"],
)
if cfg["launcher"] is popen_launch_workers_and_router
else cfg["launcher"](
backend=cfg["launcher_kwargs"]["backend"],
base_url=cfg["base_url_port"],
history_backend=cfg["launcher_kwargs"]["history_backend"],
)
)
# Build client
api_key = os.environ.get(cfg["api_key_env"]) if cfg["api_key_env"] else None
client = openai.Client(
api_key=api_key,
base_url=cluster["base_url"] + "/v1",
)
# Yield data to test
try:
yield backend, cfg["model"], client
finally:
# Always kill router
kill_process_tree(cluster["router"].pid)
# If workers exist, kill them as well
if cfg["needs_workers"]:
for w in cluster.get("workers", []):
kill_process_tree(w.pid)
@@ -0,0 +1,241 @@
"""
Base test class for Response API e2e tests.
This module provides base test classes that can be reused across different backends
(OpenAI, XAI, gRPC) with common test logic.
"""
import sys
import time
from pathlib import Path
import openai
import pytest
from openai import OpenAI
from openai.types import responses
# Add current directory for local imports
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR))
@pytest.mark.parametrize("setup_backend", ["openai", "oracle_store"], indirect=True)
class TestResponseCRUD:
"""Base class for Response API CRUD tests."""
def test_create_and_get_response(self, setup_backend):
"""Test creating response and retrieving it."""
_, model, client = setup_backend
# Create response
create_resp = client.responses.create(model=model, input="Hello, world!")
assert create_resp.id is not None
assert create_resp.error is None
assert create_resp.status == "completed"
assert len(create_resp.output_text) > 0
response_id = create_resp.id
# Get response
get_resp = client.responses.retrieve(response_id=response_id)
assert get_resp.error is None
assert get_resp.id == response_id
assert get_resp.status == "completed"
input_resp = client.responses.input_items.list(response_id=get_resp.id)
assert input_resp.data is not None
assert len(input_resp.data) > 0
@pytest.mark.skip(reason="TODO: Add delete response feature")
def test_delete_response(self, setup_backend):
"""Test deleting response."""
_, model, client = setup_backend
# Create response
create_resp = client.responses.create(model=model, input="Test deletion")
assert create_resp.id is not None
assert create_resp.error is None
assert create_resp.status == "completed"
assert len(create_resp.output_text) > 0
response_id = create_resp.id
# Delete response
client.responses.delete(response_id=response_id)
# Verify it's deleted (should return 404)
with pytest.raises(openai.NotFoundError):
client.responses.retrieve(response_id=response_id)
@pytest.mark.skip(reason="TODO: Add background response feature")
def test_background_response(self, setup_backend):
"""Test background response execution."""
_, model, client = setup_backend
# Create background response
create_resp = client.responses.create(
model=model,
input="Write a short story",
background=True,
max_output_tokens=100,
)
assert create_resp.id is not None
assert create_resp.error is None
assert create_resp.status in ["in_progress", "queued"]
response_id = create_resp.id
# Wait for completion
final_data = wait_for_background_task(client, response_id, timeout=60)
assert final_data.status == "completed"
@pytest.mark.parametrize("setup_backend", ["openai", "oracle_store"], indirect=True)
class TestConversationCRUD:
"""Base class for Conversation API CRUD tests."""
def test_create_and_get_conversation(self, setup_backend):
"""Test creating and retrieving conversation."""
_, model, client = setup_backend
# Create conversation
create_resp = client.conversations.create(metadata={"user": "test_user"})
assert create_resp.id is not None
assert create_resp.created_at is not None
create_data = create_resp.metadata
assert create_data["user"] == "test_user"
conversation_id = create_resp.id
# Get conversation
get_resp = client.conversations.retrieve(conversation_id=conversation_id)
assert get_resp.id is not None
assert get_resp.created_at is not None
get_data = get_resp.metadata
assert get_resp.id == conversation_id
assert get_data["user"] == "test_user"
def test_update_conversation(self, setup_backend):
"""Test updating conversation metadata."""
_, model, client = setup_backend
# Create conversation
create_resp = client.conversations.create(metadata={"key1": "value1"})
assert create_resp.id is not None
assert create_resp.created_at is not None
create_data = create_resp.metadata
assert create_data["key1"] == "value1"
assert "key2" not in create_data
conversation_id = create_resp.id
# Update conversation
update_resp = client.conversations.update(
conversation_id=conversation_id,
metadata={"key1": "value1", "key2": "value2"},
)
assert update_resp.id == conversation_id
update_data = update_resp.metadata
assert update_data["key1"] == "value1"
assert update_data["key2"] == "value2"
# Verify update
get_resp = client.conversations.retrieve(conversation_id=conversation_id)
get_data = get_resp.metadata
assert update_data["key1"] == "value1"
assert update_data["key2"] == "value2"
def test_delete_conversation(self, setup_backend):
"""Test deleting conversation."""
_, model, client = setup_backend
# Create conversation
create_resp = client.conversations.create()
assert create_resp.id is not None
assert create_resp.created_at is not None
conversation_id = create_resp.id
# Delete conversation
delete_resp = client.conversations.delete(conversation_id=conversation_id)
assert delete_resp.id is not None
assert delete_resp.deleted
# Verify deletion
with pytest.raises(openai.NotFoundError):
client.conversations.retrieve(conversation_id=conversation_id)
def test_list_conversation_items(self, setup_backend):
"""Test listing conversation items."""
_, model, client = setup_backend
# Create conversation
conv_resp = client.conversations.create()
assert conv_resp.id is not None
conversation_id = conv_resp.id
# Create response with conversation
resp1 = client.responses.create(
model=model,
input="First message",
conversation=conversation_id,
max_output_tokens=50,
)
assert resp1.error is None
resp2 = client.responses.create(
model=model,
input="Second message",
conversation=conversation_id,
max_output_tokens=50,
)
assert resp2.error is None
# List items
list_resp = client.conversations.items.list(conversation_id=conversation_id)
assert list_resp is not None
assert list_resp.data is not None
list_data = list_resp.data
# Should have at least 4 items (2 inputs + 2 outputs)
assert len(list_data) >= 4
def wait_for_background_task(
client: OpenAI, response_id: str, timeout: int = 30, poll_interval: float = 0.5
) -> responses.Response:
"""
Wait for background task to complete.
Args:
client: openai client
response_id: Response ID to poll
timeout: Max seconds to wait
poll_interval: Seconds between polls
Returns:
Final response data
Raises:
TimeoutError: If task doesn't complete in time
AssertionError: If task fails
"""
start_time = time.time()
while time.time() - start_time < timeout:
resp = client.responses.retrieve(response_id=response_id)
assert resp.error is None
assert resp.id == response_id
status = resp.status
if status == "completed":
return resp
elif status == "failed":
raise AssertionError(f"Background task failed: {resp.error}")
elif status == "cancelled":
raise AssertionError("Background task was cancelled")
time.sleep(poll_interval)
raise TimeoutError(
f"Background task {response_id} did not complete within {timeout}s"
)
@@ -0,0 +1,161 @@
"""
State management tests for Response API.
Tests both previous_response_id and conversation-based state management.
These tests should work across all backends (OpenAI, XAI, gRPC).
"""
import openai
import pytest
@pytest.mark.parametrize(
"setup_backend", ["openai", "xai", "grpc", "grpc_harmony"], indirect=True
)
class TestStateManagement:
"""Tests for state management using previous_response_id and conversation."""
def test_basic_response_creation(self, setup_backend):
"""Test basic response creation without state."""
_, model, client = setup_backend
resp = client.responses.create(model=model, input="What is 2+2?")
assert resp.id is not None
assert resp.error is None
assert resp.status == "completed"
assert len(resp.output_text) > 0
assert resp.usage is not None
def test_streaming_response(self, setup_backend):
"""Test streaming response."""
_, model, client = setup_backend
resp = client.responses.create(
model=model, input="Count to 5", stream=True, max_output_tokens=50
)
# Check for response.created event
events = [event for event in resp]
created_events = [event for event in events if event.type == "response.created"]
assert len(created_events) > 0
# Check for final completed event or in_progress events
assert any(
event.type in ["response.completed", "response.in_progress"]
for event in events
)
def test_previous_response_id_chaining(self, setup_backend):
"""Test chaining responses using previous_response_id."""
_, model, client = setup_backend
# First response
resp1 = client.responses.create(
model=model, input="My name is Alice and my friend is Bob. Remember it."
)
assert resp1.error is None
assert resp1.status == "completed"
response1_id = resp1.id
# Second response referencing first
resp2 = client.responses.create(
model=model, input="What is my name", previous_response_id=response1_id
)
assert resp2.error is None
assert resp2.status == "completed"
# The model should remember the name from previous response
assert "Alice" in resp2.output_text
# Third response referencing second
resp3 = client.responses.create(
model=model,
input="What is my friend name?",
previous_response_id=resp2.id,
)
assert resp3.error is None
assert resp3.status == "completed"
assert "Bob" in resp3.output_text
@pytest.mark.skip(reason="TODO: Add the invalid previous_response_id check")
def test_previous_response_id_invalid(self, setup_backend):
"""Test using invalid previous_response_id."""
_, model, client = setup_backend
with pytest.raises(openai.BadRequestError):
client.responses.create(
model=model,
input="Test",
previous_response_id="resp_invalid123",
max_output_tokens=50,
)
def test_conversation_with_multiple_turns(self, setup_backend):
"""Test state management using conversation ID."""
backend, model, client = setup_backend
if backend in ["grpc", "grpc_harmony"]:
pytest.skip("TODO: 501 Not Implemented")
# Create conversation
conv_resp = client.conversations.create(metadata={"topic": "math"})
assert conv_resp.id is not None
assert conv_resp.created_at is not None
conversation_id = conv_resp.id
# First response in conversation
resp1 = client.responses.create(
model=model, input="I have 5 apples.", conversation=conversation_id
)
assert resp1.error is None
assert resp1.status == "completed"
# Second response in same conversation
resp2 = client.responses.create(
model=model,
input="How many apples do I have?",
conversation=conversation_id,
)
assert resp2.error is None
assert resp2.status == "completed"
output_text = resp2.output_text
# Should remember "5 apples"
assert "5" in output_text or "five" in output_text.lower()
# Third response in same conversation
resp3 = client.responses.create(
model=model,
input="If I get 3 more, how many total?",
conversation=conversation_id,
)
assert resp3.error is None
assert resp3.status == "completed"
output_text = resp3.output_text
# Should calculate 5 + 3 = 8
assert "8" in output_text or "eight" in output_text.lower()
list_resp = client.conversations.items.list(conversation_id)
assert list_resp.data is not None
items = list_resp.data
# Should have at least 6 items (3 inputs + 3 outputs)
assert len(items) >= 6
def test_mutually_exclusive_parameters(self, setup_backend):
"""Test that previous_response_id and conversation are mutually exclusive."""
_, model, client = setup_backend
# TODO: Remove this once the conversation API is implemented for GRPC backend
conversation_id = "conv_123"
resp1 = client.responses.create(model=model, input="Test")
response1_id = resp1.id
# Try to use both parameters
with pytest.raises(openai.BadRequestError):
client.responses.create(
model=model,
input="This should fail",
previous_response_id=response1_id,
conversation=conversation_id,
)
@@ -0,0 +1,152 @@
"""
Streaming events tests for Response API.
Tests for streaming event validation including:
- Zero-based output_index for reasoning content
- OutputItemDone event emission and output array construction
"""
import pytest
@pytest.mark.parametrize("setup_backend", ["grpc", "grpc_harmony"], indirect=True)
class TestStreamingEvents:
"""Tests for streaming event validation."""
def test_output_item_event_emitted(self, setup_backend):
"""
Test that output_index is zero-based in streaming responses.
Verifies that the first output item has output_index: 0.
"""
_, model, client = setup_backend
resp = client.responses.create(
model=model,
input="Count from 1 to 3",
stream=True,
max_output_tokens=50,
)
events = [event for event in resp]
assert len(events) > 0
# Find output_item.added events
output_item_added_events = [
event for event in events if event.type == "response.output_item.added"
]
assert len(output_item_added_events) > 0, "Should have output_item.added events"
# Verify first output item has output_index: 0
first_item_event = output_item_added_events[0]
assert first_item_event.item is not None
assert first_item_event.output_index is not None
assert (
first_item_event.output_index == 0
), "First output item must have output_index: 0 (zero-based indexing)"
# Verify subsequent items increment correctly
for i, event in enumerate(output_item_added_events):
assert (
event.output_index == i
), f"Output item {i} should have output_index: {i}"
# Verify output_item.done event exists
output_item_done_events = [
event for event in events if event.type == "response.output_item.done"
]
assert len(output_item_done_events) > 0
# Verify output_item.done event structure
for event in output_item_done_events:
assert event.item is not None
assert event.output_index is not None
assert event.item.type is not None
# Find response.completed event
completed_events = [
event for event in events if event.type == "response.completed"
]
assert len(completed_events) == 1, "Should have exactly one completed event"
# Verify output array exists and contains items
completed_event = completed_events[0]
assert completed_event.response.output is not None
output_array = completed_event.response.output
assert isinstance(output_array, list)
assert len(output_array) > 0, "Output array should contain at least one item"
# Verify each item in output array has proper structure
for i, item in enumerate(output_array):
assert item.type is not None
# Verify output_item.added events match items in final output array
output_item_added_events = [
event for event in events if event.type == "response.output_item.added"
]
assert len(output_item_added_events) == len(
output_array
), "Number of output_item.added events should match output array length"
def test_reasoning_content(self, setup_backend):
"""
Test that reasoning content has correct zero-based output_index.
Specifically tests that reasoning item has output_index: 0
and message item has output_index: 1.
"""
backend, model, client = setup_backend
if backend in ["grpc"]:
pytest.skip("skip test_reasoning_content for grpc")
resp = client.responses.create(
model=model,
input="What is the capital of France? Think step by step.",
stream=True,
max_output_tokens=200,
)
events = [event for event in resp]
assert len(events) > 0
# Find output_item.added events
output_item_added_events = [
event for event in events if event.type == "response.output_item.added"
]
assert len(output_item_added_events) > 0
reasoning_items = [
item for item in output_item_added_events if item.item.type == "reasoning"
]
message_items = [
item for item in output_item_added_events if item.item.type == "message"
]
# If reasoning is present, verify it has output_index: 0
if reasoning_items:
reasoning_item = reasoning_items[0]
assert (
reasoning_item.output_index == 0
), "Reasoning item should have output_index: 0"
# If message is present after reasoning, verify it has output_index: 1
if reasoning_items and message_items:
message_item = message_items[0]
assert (
message_item.output_index == 1
), "Message item after reasoning should have output_index: 1"
# Find response.completed event
completed_events = [
event for event in events if event.type == "response.completed"
]
assert len(completed_events) == 1
# Get output array from completed event
output_array = completed_events[0].response.output
assert len(output_array) > 0
# Check if reasoning items are in output array
reasoning_items_in_output = [
item for item in output_array if item.type == "reasoning"
]
assert len(reasoning_items_in_output) > 0
@@ -0,0 +1,172 @@
"""
Structured output tests for Response API.
Tests for text.format field with json_object and json_schema formats.
"""
import json
import sys
from pathlib import Path
import pytest
# Add current directory for local imports
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR))
@pytest.mark.parametrize("setup_backend", ["openai", "grpc_harmony"], indirect=True)
class TestStructuredOutput:
def test_structured_output_json_schema(self, setup_backend):
"""Test structured output with json_schema format."""
_, model, client = setup_backend
# Create response with structured output
params = {
"model": model,
"input": [
{
"role": "system",
"content": "You are a helpful math tutor. Guide the user through the solution step by step.",
},
{"role": "user", "content": "how can I solve 8x + 7 = -23"},
],
"text": {
"format": {
"type": "json_schema",
"name": "math_reasoning",
"schema": {
"type": "object",
"properties": {
"steps": {
"type": "array",
"items": {
"type": "object",
"properties": {
"explanation": {"type": "string"},
"output": {"type": "string"},
},
"required": ["explanation", "output"],
"additionalProperties": False,
},
},
"final_answer": {"type": "string"},
},
"required": ["steps", "final_answer"],
"additionalProperties": False,
},
"strict": True,
}
},
}
create_resp = client.responses.create(**params)
assert create_resp.error is None
assert create_resp.id is not None
assert create_resp.output is not None
assert create_resp.text is not None
# Verify text format was echoed back correctly
assert create_resp.text.format is not None
assert create_resp.text.format.type == "json_schema"
assert create_resp.text.format.name == "math_reasoning"
assert create_resp.text.format.schema_ is not None
assert create_resp.text.format.strict
# Find the message output (output[0] may be reasoning, output[1] is message)
output_text = next(
(
content.text
for item in create_resp.output
if item.type == "message"
for content in item.content
if content.type == "output_text"
),
None,
)
assert output_text is not None, "No output_text found in response"
assert output_text.strip(), "output_text is empty"
# Parse JSON output
output_json = json.loads(output_text)
# Verify schema structure
assert "steps" in output_json
assert "final_answer" in output_json
assert isinstance(output_json["steps"], list)
assert len(output_json["steps"]) > 0
# Verify each step has required fields
for step in output_json["steps"]:
assert "explanation" in step
assert "output" in step
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestSimpleSchemaStructuredOutput:
def test_structured_output_json_schema(self, setup_backend):
"""Override with simpler schema for Llama model (complex schemas not well supported)."""
_, model, client = setup_backend
params = {
"model": model,
"input": [
{
"role": "system",
"content": "You are a math solver. Return ONLY a JSON object that matches the schema—no extra text.",
},
{
"role": "user",
"content": "What is 1 + 1?",
},
],
"text": {
"format": {
"type": "json_schema",
"name": "math_answer",
"schema": {
"type": "object",
"properties": {"answer": {"type": "string"}},
"required": ["answer"],
},
}
},
}
create_resp = client.responses.create(**params)
assert create_resp.error is None
assert create_resp.id is not None
assert create_resp.output is not None
assert create_resp.text is not None
# Verify text format was echoed back correctly
assert create_resp.text.format is not None
assert create_resp.text.format.type == "json_schema"
assert create_resp.text.format.name == "math_answer"
assert create_resp.text.format.schema_ is not None
# Find the message output
output_text = next(
(
content.text
for item in create_resp.output
if item.type == "message"
for content in item.content
if content.type == "output_text"
),
None,
)
assert output_text is not None, "No output_text found in response"
assert output_text.strip(), "output_text is empty"
# Parse JSON output
output_json = json.loads(output_text)
# Verify simple schema structure (just answer field)
assert "answer" in output_json
assert isinstance(output_json["answer"], str)
assert output_json["answer"], "Answer is empty"
@@ -0,0 +1,765 @@
"""
Test class for tool calling tests.
This module provides test cases for function calling functionality, tool choices
and mcp calling functionality across different backends.
"""
import json
import sys
import time
from pathlib import Path
import pytest
# Add current directory for local imports
_TEST_DIR = Path(__file__).parent
sys.path.insert(0, str(_TEST_DIR))
@pytest.mark.parametrize(
"setup_backend", ["openai", "grpc", "grpc_harmony"], indirect=True
)
class TestToolCalling:
# Shared function tool definitions
SYSTEM_DIAGNOSTICS_FUNCTION = {
"type": "function",
"name": "get_system_diagnostics",
"description": "Retrieve real-time diagnostics for a spacecraft system.",
"parameters": {
"type": "object",
"properties": {
"system_name": {
"type": "string",
"description": "Name of the spacecraft system to query. "
"Example: 'Astra-7 Core Reactor'.",
}
},
"required": ["system_name"],
},
}
GET_WEATHER_FUNCTION = {
"type": "function",
"name": "get_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city name, e.g., San Francisco",
}
},
"required": ["location"],
},
}
CALCULATE_FUNCTION = {
"type": "function",
"name": "calculate",
"description": "Perform a mathematical calculation",
"parameters": {
"type": "object",
"properties": {
"expression": {
"type": "string",
"description": "The mathematical expression to evaluate",
}
},
"required": ["expression"],
},
}
SEARCH_WEB_FUNCTION = {
"type": "function",
"name": "search_web",
"description": "Search the web for information",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
}
LOCAL_SEARCH_FUNCTION = {
"type": "function",
"name": "local_search",
"description": "Search local database",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
}
# Shared constants for MCP tests
BRAVE_MCP_TOOL = {
"type": "mcp",
"server_label": "brave",
"server_description": "A Tool to do web search",
"server_url": "http://localhost:8001/sse",
"require_approval": "never",
}
DEEPWIKI_MCP_TOOL = {
"type": "mcp",
"server_label": "deepwiki",
"server_url": "https://mcp.deepwiki.com/mcp",
"require_approval": "never",
}
MCP_TEST_PROMPT = (
"show me some news about sglang router, use the tool to just search "
"one result and return one sentence response"
)
# Test cases for basic function calling functionality
def test_basic_function_call(self, setup_backend):
"""
Test basic function calling workflow.
This test follows the pattern from function_call_test.py:
1. Define a function tool (get_horoscope)
2. Send user message asking for horoscope
3. Model should return function_call
4. Execute function locally and provide output
5. Model should generate final response using the function output
"""
backend, model, client = setup_backend
if backend in ["grpc"]:
pytest.skip("skip for grpc")
# 1. Define a list of callable tools for the model
tools = [
{
"type": "function",
"name": "get_horoscope",
"description": "Get today's horoscope for an astrological sign.",
"parameters": {
"type": "object",
"properties": {
"sign": {
"type": "string",
"description": "An astrological sign like Taurus or Aquarius",
},
},
"required": ["sign"],
},
},
]
system_prompt = (
"You are a helpful assistant that can call functions. "
"When a user asks for horoscope information, call the function. "
"IMPORTANT: Don't reply directly to the user, only call the function. "
)
# Create a running input list we will add to over time
input_list = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": "What is my horoscope? I am an Aquarius."},
]
# 2. Prompt the model with tools defined
resp = client.responses.create(model=model, input=input_list, tools=tools)
# Should successfully make the request
assert resp.error is None
# Basic response structure
assert resp.id is not None
assert resp.status == "completed"
assert resp.output is not None
# Verify output array is not empty
output = resp.output
assert isinstance(output, list)
assert len(output) > 0
# Check for function_call in output
function_calls = [item for item in output if item.type == "function_call"]
assert (
len(function_calls) > 0
), "Response should contain at least one function_call"
# Verify function_call structure
function_call = function_calls[0]
assert function_call.call_id is not None
assert function_call.name is not None
assert function_call.name == "get_horoscope"
assert function_call.arguments is not None
# Parse arguments
args = json.loads(function_call.arguments)
assert "sign" in args
assert args["sign"].lower() == "aquarius"
# 3. Save function call outputs for subsequent requests
input_list.append(function_call)
# 4. Execute the function logic for get_horoscope
horoscope = f"{args['sign']}: Next Tuesday you will befriend a baby otter."
# 5. Provide function call results to the model
input_list.append(
{
"type": "function_call_output",
"call_id": function_call.call_id,
"output": json.dumps({"horoscope": horoscope}),
}
)
# 6. Make second request with function output
resp2 = client.responses.create(
model=model,
input=input_list,
instructions="Respond only with a horoscope generated by a tool.",
tools=tools,
)
assert resp2.error is None
assert resp2.status == "completed"
# The model should be able to give a response using the function output
output2 = resp2.output
assert len(output2) > 0
# Find message output
messages = [item for item in output2 if item.type == "message"]
assert len(messages) > 0, "Response should contain at least one message"
# Verify message contains the horoscope
message = messages[0]
assert message.content is not None
content_parts = message.content
assert len(content_parts) > 0
# Get text from content
text_parts = [part.text for part in content_parts if part.type == "output_text"]
full_text = " ".join(text_parts).lower()
# Should mention the horoscope or baby otter
assert (
"baby otter" in full_text or "aquarius" in full_text
), "Response should reference the horoscope content"
# Test cases for tool_choice parameter support, these tests require --reasoning-parser
def test_tool_choice_auto(self, setup_backend):
"""
Test tool_choice="auto" allows model to decide whether to use tools.
The model should be able to choose to call a tool or not.
"""
backend, model, client = setup_backend
if backend in ["openai"]:
pytest.skip("skip for openai")
tools = [self.GET_WEATHER_FUNCTION]
# Query that should trigger tool use
resp = client.responses.create(
model=model,
input="What is the weather in Seattle?",
tools=tools,
tool_choice="auto",
stream=False,
)
assert resp.id is not None
assert resp.error is None
output = resp.output
assert len(output) > 0
# With auto, model should choose to call get_weather for this query
function_calls = [item for item in output if item.type == "function_call"]
assert (
len(function_calls) > 0
), "Model should choose to call function with tool_choice='auto'"
def test_tool_choice_required(self, setup_backend):
"""
Test tool_choice="required" forces the model to call at least one tool.
The model must make at least one function call.
"""
backend, model, client = setup_backend
if backend in ["openai"]:
pytest.skip("skip for openai")
tools = [self.CALCULATE_FUNCTION]
resp = client.responses.create(
model=model,
input="What is 15 * 23?",
tools=tools,
tool_choice="required",
stream=False,
)
assert resp.id is not None
assert resp.error is None
output = resp.output
# Must have at least one function call
function_calls = [item for item in output if item.type == "function_call"]
assert (
len(function_calls) > 0
), "tool_choice='required' must force at least one function call"
def test_tool_choice_specific_function(self, setup_backend):
"""
Test tool_choice with specific function name forces that function to be called.
The model must call the specified function.
"""
backend, model, client = setup_backend
if backend in ["openai"]:
pytest.skip("skip for openai")
tools = [self.SEARCH_WEB_FUNCTION, self.GET_WEATHER_FUNCTION]
# Force specific function call
resp = client.responses.create(
model=model,
input="What's happening in the news today?",
tools=tools,
tool_choice={"type": "function", "function": {"name": "search_web"}},
stream=False,
)
assert resp.id is not None
assert resp.error is None
output = resp.output
# Must have function call
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0, "Must call the specified function"
# Must be the specified function
called_function = function_calls[0]
assert (
called_function.name == "search_web"
), "Must call the function specified in tool_choice"
def test_tool_choice_streaming(self, setup_backend):
"""
Test tool_choice parameter works correctly with streaming.
Verifies that tool_choice constraints are applied in streaming mode.
"""
backend, model, client = setup_backend
if backend in ["openai", "grpc"]:
pytest.skip("skip for openai")
tools = [self.CALCULATE_FUNCTION]
resp = client.responses.create(
model=model,
input="Calculate 42 * 17",
tools=tools,
tool_choice="required",
stream=True,
)
events = [event for event in resp]
assert len(events) > 0
event_types = [e.type for e in events]
# Should have function call events
assert (
"response.function_call_arguments.delta" in event_types
), "Should have function_call_arguments.delta events"
# Verify completed event has function call
completed_events = [e for e in events if e.type == "response.completed"]
assert len(completed_events) == 1
output = completed_events[0].response.output
function_calls = [item for item in output if item.type == "function_call"]
assert (
len(function_calls) > 0
), "Streaming with tool_choice='required' must produce function call"
def test_tool_choice_with_mcp_tools(self, setup_backend):
"""
Test tool_choice parameter works with MCP tools.
Verifies that tool_choice can control MCP tool usage.
"""
backend, model, client = setup_backend
if backend in ["openai"]:
pytest.skip("skip for openai")
tools = [self.DEEPWIKI_MCP_TOOL]
# With tool_choice="auto", should allow MCP tool calls
resp = client.responses.create(
model=model,
input="What transport protocols does the 2025-03-26 version of the MCP spec (modelcontextprotocol/modelcontextprotocol) support?",
tools=tools,
tool_choice="auto",
stream=False,
)
assert resp.id is not None
assert resp.error is None
output = resp.output
# Should have mcp_call with auto
mcp_calls = [item for item in output if item.type == "mcp_call"]
assert len(mcp_calls) > 0, "tool_choice='auto' should allow MCP tool calls"
def test_tool_choice_mixed_function_and_mcp(self, setup_backend):
"""
Test tool_choice with mixed function and MCP tools.
Verifies tool_choice can select specific tools when both function and MCP tools are available.
"""
backend, model, client = setup_backend
if backend in ["openai"]:
pytest.skip("skip for openai")
tools = [self.DEEPWIKI_MCP_TOOL, self.LOCAL_SEARCH_FUNCTION]
# Force specific function call
resp = client.responses.create(
model=model,
input="Search for information about Python",
tools=tools,
tool_choice={"type": "function", "function": {"name": "local_search"}},
stream=False,
)
assert resp.id is not None
assert resp.error is None
output = resp.output
# Must call local_search, not MCP
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0
assert function_calls[0].name == "local_search"
# Should not have mcp_call
mcp_calls = [item for item in output if item.type == "mcp_call"]
assert len(mcp_calls) == 0, "Should only call specified function, not MCP tools"
# Tests for MCP tool calling in both streaming and non-streaming modes.
def test_mcp_basic_tool_call(self, setup_backend):
"""
Test basic MCP tool call (non-streaming).
Validation strictness is controlled by parameter `backend` from setup_backend fixture.
Set to "strict" if backend is http.
"""
backend, model, client = setup_backend
# To avoid being rate-limited by brave search server
time.sleep(2)
resp = client.responses.create(
model=model,
input=self.MCP_TEST_PROMPT,
tools=[self.BRAVE_MCP_TOOL],
stream=False,
reasoning={"effort": "low"},
)
# Should successfully make the request
assert resp.error is None
# Basic response structure
assert resp.id is not None
assert resp.status == "completed"
assert resp.model is not None
assert resp.output is not None
# Verify output array is not empty
assert len(resp.output_text) > 0
# Check for MCP-specific output types
output_types = [item.type for item in resp.output]
# Should have mcp_list_tools - tools are listed before calling
assert (
"mcp_list_tools" in output_types
), "Response should contain mcp_list_tools"
# Should have at least one mcp_call
mcp_calls = [item for item in resp.output if item.type == "mcp_call"]
assert len(mcp_calls) > 0, "Response should contain at least one mcp_call"
# Verify mcp_call structure
for mcp_call in mcp_calls:
assert mcp_call.id is not None
assert mcp_call.error is None
assert mcp_call.status == "completed"
assert mcp_call.server_label == "brave"
assert mcp_call.name is not None
assert mcp_call.arguments is not None
assert mcp_call.output is not None
# Strict mode: additional validation for HTTP backends
if backend == "openai":
# Should have final message output
messages = [item for item in resp.output if item.type == "message"]
assert len(messages) > 0, "Response should contain at least one message"
# Verify message structure
for msg in messages:
assert msg.content is not None
assert isinstance(msg.content, list)
# Check content has text
for content_item in msg.content:
if content_item.type == "output_text":
assert content_item.text is not None
assert isinstance(content_item.text, str)
assert len(content_item.text) > 0
def test_mcp_basic_tool_call_streaming(self, setup_backend):
"""Test basic MCP tool call (streaming).
Validation strictness is controlled by the class attribute `mcp_validation_mode`.
Set to "strict" in subclasses for additional HTTP-specific validation.
"""
backend, model, client = setup_backend
# To avoid being rate-limited by brave search server
time.sleep(2)
resp = client.responses.create(
model=model,
input=self.MCP_TEST_PROMPT,
tools=[self.BRAVE_MCP_TOOL],
stream=True,
reasoning={"effort": "low"},
)
# Should successfully make the request
events = [event for event in resp]
assert len(events) > 0
event_types = [event.type for event in events]
# Check for lifecycle events
assert "response.created" in event_types, "Should have response.created event"
assert (
"response.completed" in event_types
), "Should have response.completed event"
# Check for MCP list tools events
assert (
"response.output_item.added" in event_types
), "Should have output_item.added events"
assert (
"response.mcp_list_tools.in_progress" in event_types
), "Should have mcp_list_tools.in_progress event"
assert (
"response.mcp_list_tools.completed" in event_types
), "Should have mcp_list_tools.completed event"
# Check for MCP call events
assert (
"response.mcp_call.in_progress" in event_types
), "Should have mcp_call.in_progress event"
assert (
"response.mcp_call_arguments.delta" in event_types
), "Should have mcp_call_arguments.delta event"
assert (
"response.mcp_call_arguments.done" in event_types
), "Should have mcp_call_arguments.done event"
assert (
"response.mcp_call.completed" in event_types
), "Should have mcp_call.completed event"
# Verify final completed event has full response
completed_events = [e for e in events if e.type == "response.completed"]
assert len(completed_events) == 1
final_response = completed_events[0].response
assert final_response.id is not None
assert final_response.status == "completed"
assert final_response.output is not None
# Verify final output contains expected items
final_output = final_response.output
final_output_types = [item.type for item in final_output]
assert "mcp_list_tools" in final_output_types
assert "mcp_call" in final_output_types
# Verify mcp_call items in final output
mcp_calls = [item for item in final_output if item.type == "mcp_call"]
assert len(mcp_calls) > 0
for mcp_call in mcp_calls:
assert mcp_call.error is None
assert mcp_call.status == "completed"
assert mcp_call.server_label == "brave"
assert mcp_call.name is not None
assert mcp_call.arguments is not None
assert mcp_call.output is not None
# Strict mode: additional validation for HTTP backends
if backend == "openai":
# Check for text output events
assert (
"response.content_part.added" in event_types
), "Should have content_part.added event"
assert (
"response.output_text.delta" in event_types
), "Should have output_text.delta events"
assert (
"response.output_text.done" in event_types
), "Should have output_text.done event"
assert (
"response.content_part.done" in event_types
), "Should have content_part.done event"
assert "message" in final_output_types
# Verify text deltas combine to final message
text_deltas = [
e.delta for e in events if e.type == "response.output_text.delta"
]
assert len(text_deltas) > 0, "Should have text deltas"
# Get final text from output_text.done event
text_done_events = [
e for e in events if e.type == "response.output_text.done"
]
assert len(text_done_events) > 0
final_text = text_done_events[0].text
assert len(final_text) > 0, "Final text should not be empty"
def test_mixed_mcp_and_function_tools(self, setup_backend):
"""Test mixed MCP and function tools (non-streaming)."""
backend, model, client = setup_backend
if backend in ["openai"]:
pytest.skip(
"Requires external MCP server (deepwiki) - may not be accessible in CI"
)
resp = client.responses.create(
model=model,
input="Give me diagnostics for the Astra-7 Core Reactor.",
tools=[self.BRAVE_MCP_TOOL, self.SYSTEM_DIAGNOSTICS_FUNCTION],
stream=False,
tool_choice="auto",
)
# Should successfully make the request
assert resp.error is None
# Basic response structure
assert resp.id is not None
assert resp.status is not None
assert resp.output is not None
# Verify output array is not empty
output = resp.output
assert isinstance(output, list)
assert len(output) > 0
# Check for function_call (not mcp_call for get_system_diagnostics)
function_calls = [item for item in output if item.type == "function_call"]
assert (
len(function_calls) > 0
), "Response should contain at least one function_call"
# Verify function_call structure for get_system_diagnostics
system_diagnostics_call = function_calls[0]
assert system_diagnostics_call.name == "get_system_diagnostics"
assert system_diagnostics_call.call_id is not None
assert system_diagnostics_call.arguments is not None
assert system_diagnostics_call.status is not None
# Parse and verify arguments
args = json.loads(system_diagnostics_call.arguments)
assert "system_name" in args
assert "astra-7" in args["system_name"].lower()
def test_mixed_mcp_and_function_tools_streaming(self, setup_backend):
"""Test mixed MCP and function tools (streaming)."""
backend, model, client = setup_backend
if backend in ["openai"]:
pytest.skip(
"Requires external MCP server (deepwiki) - may not be accessible in CI"
)
resp = client.responses.create(
model=model,
input="Give me diagnostics for the Astra-7 Core Reactor.",
tools=[self.BRAVE_MCP_TOOL, self.SYSTEM_DIAGNOSTICS_FUNCTION],
stream=True,
tool_choice="auto", # Encourage tool usage
)
# Should successfully make the request
events = [event for event in resp]
assert len(events) > 0
event_types = [e.type for e in events]
# Check for lifecycle events
assert "response.created" in event_types, "Should have response.created event"
# Should have mcp_list_tools events
assert (
"response.mcp_list_tools.completed" in event_types
), "Should have mcp_list_tools.completed event"
# Should have function_call_arguments events (not mcp_call_arguments)
assert (
"response.function_call_arguments.delta" in event_types
), "Should have function_call_arguments.delta event for function tools"
assert (
"response.function_call_arguments.done" in event_types
), "Should have function_call_arguments.done event for function tools"
# Should NOT have mcp_call_arguments events for function tools
# (get_system_diagnostics should use function_call_arguments, not mcp_call_arguments)
mcp_call_arg_events = [
e
for e in events
if e.type == "response.mcp_call_arguments.delta"
and "get_system_diagnostics" in str(e.delta)
]
assert (
len(mcp_call_arg_events) == 0
), "Should NOT emit mcp_call_arguments.delta for function tools (get_system_diagnostics)"
# Verify function_call_arguments.delta event structure
func_arg_deltas = [
e for e in events if e.type == "response.function_call_arguments.delta"
]
assert (
len(func_arg_deltas) > 0
), "Should have function_call_arguments.delta events"
# Check that delta event contains system_name arguments
full_delta_event = ""
for event in func_arg_deltas:
full_delta_event += event.delta
assert (
"system_name" in full_delta_event.lower()
and "astra-7" in full_delta_event.lower()
), "function_call_arguments.delta should contain system_name and astra-7"
@@ -0,0 +1,565 @@
"""
Fixtures for launching OpenAI/XAI router for response API e2e testing.
This module provides fixtures for launching SGLang router with OpenAI or XAI backends:
1. Launch router with --backend openai pointing to OpenAI or XAI API
2. Configure history backend (memory or oracle)
This supports testing the Response API against real cloud providers.
"""
import logging
import os
import socket
import subprocess
import time
from typing import Optional
import requests
logger = logging.getLogger(__name__)
def wait_for_workers_ready(
router_url: str,
expected_workers: int,
timeout: int = 300,
api_key: Optional[str] = None,
) -> None:
"""
Wait for router to have all workers connected.
Polls the /workers endpoint until the 'total' field matches expected_workers.
Example response from /workers endpoint:
{"workers":[],"total":0,"stats":{"prefill_count":0,"decode_count":0,"regular_count":0}}
Args:
router_url: Base URL of router (e.g., "http://127.0.0.1:30000")
expected_workers: Number of workers expected to be connected
timeout: Max seconds to wait
api_key: Optional API key for authentication
"""
start_time = time.time()
last_error = None
attempt = 0
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
with requests.Session() as session:
while time.time() - start_time < timeout:
attempt += 1
elapsed = int(time.time() - start_time)
# Log progress every 10 seconds
if elapsed > 0 and elapsed % 10 == 0 and attempt % 10 == 0:
logger.info(
f" Still waiting for workers... ({elapsed}/{timeout}s elapsed)"
)
try:
response = session.get(
f"{router_url}/workers", headers=headers, timeout=5
)
if response.status_code == 200:
data = response.json()
total_workers = data.get("total", 0)
if total_workers == expected_workers:
logger.info(
f" All {expected_workers} workers connected after {elapsed}s"
)
return
else:
last_error = f"Workers: {total_workers}/{expected_workers}"
else:
last_error = f"HTTP {response.status_code}"
except requests.ConnectionError:
last_error = "Connection refused (router not ready yet)"
except requests.Timeout:
last_error = "Timeout"
except requests.RequestException as e:
last_error = str(e)
except (ValueError, KeyError) as e:
last_error = f"Invalid response: {e}"
time.sleep(1)
raise TimeoutError(
f"Router at {router_url} did not get {expected_workers} workers within {timeout}s.\n"
f"Last status: {last_error}\n"
f"Hint: Run with SHOW_ROUTER_LOGS=1 to see startup logs"
)
def find_free_port() -> int:
"""Find an available port on localhost."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def wait_for_router_ready(
router_url: str,
timeout: int = 60,
api_key: Optional[str] = None,
) -> None:
"""
Wait for router to be ready.
Polls the /health endpoint until it returns 200.
Args:
router_url: Base URL of router (e.g., "http://127.0.0.1:30000")
timeout: Max seconds to wait
api_key: Optional API key for authentication
"""
start_time = time.time()
last_error = None
attempt = 0
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
with requests.Session() as session:
while time.time() - start_time < timeout:
attempt += 1
elapsed = int(time.time() - start_time)
# Log progress every 10 seconds
if elapsed > 0 and elapsed % 10 == 0 and attempt % 10 == 0:
logger.info(
f" Still waiting for router... ({elapsed}/{timeout}s elapsed)"
)
try:
response = session.get(
f"{router_url}/health", headers=headers, timeout=5
)
if response.status_code == 200:
logger.info(f" Router ready after {elapsed}s")
return
else:
last_error = f"HTTP {response.status_code}"
except requests.ConnectionError:
last_error = "Connection refused (router not ready yet)"
except requests.Timeout:
last_error = "Timeout"
except requests.RequestException as e:
last_error = str(e)
time.sleep(1)
raise TimeoutError(
f"Router at {router_url} did not become ready within {timeout}s.\n"
f"Last status: {last_error}\n"
f"Hint: Run with SHOW_ROUTER_LOGS=1 to see startup logs"
)
def popen_launch_openai_xai_router(
backend: str, # "openai" or "xai"
base_url: str,
timeout: int = 60,
history_backend: str = "memory",
api_key: Optional[str] = None,
router_args: Optional[list] = None,
stdout=None,
stderr=None,
prometheus_port: Optional[int] = None,
) -> dict:
"""
Launch SGLang router with OpenAI or XAI backend.
This approach:
1. Starts router with --backend openai
2. Points to OpenAI or XAI API via --worker-urls
3. Configures history backend (memory or oracle)
4. Waits for router health check to pass
Args:
backend: "openai" or "xai"
base_url: Base URL for router (e.g., "http://127.0.0.1:30000")
timeout: Timeout for router startup (default: 60s)
history_backend: "memory" or "oracle" (default: memory)
api_key: Optional API key for router authentication
router_args: Additional arguments for router
stdout: Optional file handle for router stdout
stderr: Optional file handle for router stderr
Returns:
dict with:
- router: router process object
- base_url: router URL (HTTP endpoint)
Example:
>>> cluster = popen_launch_openai_xai_router(
... "openai", "http://127.0.0.1:30000"
... )
>>> # Use cluster['base_url'] for HTTP requests
>>> # Cleanup:
>>> kill_process_tree(cluster['router'].pid)
"""
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
# Parse router port from base_url
if ":" in base_url.split("//")[-1]:
router_port = int(base_url.split(":")[-1])
else:
router_port = find_free_port()
logger.info(f"\n{'='*70}")
logger.info(f"Launching {backend.upper()} router")
logger.info(f"{'='*70}")
logger.info(f" Backend: {backend}")
logger.info(f" Router port: {router_port}")
logger.info(f" History backend: {history_backend}")
# Determine worker URL based on backend
if backend == "openai":
worker_url = "https://api.openai.com"
# Get API key from environment
backend_api_key = os.environ.get("OPENAI_API_KEY")
if not backend_api_key:
raise ValueError(
"OPENAI_API_KEY environment variable must be set for OpenAI backend"
)
elif backend == "xai":
worker_url = "https://api.x.ai"
# Get API key from environment
backend_api_key = os.environ.get("XAI_API_KEY")
if not backend_api_key:
raise ValueError(
"XAI_API_KEY environment variable must be set for XAI backend"
)
else:
raise ValueError(f"Unsupported backend: {backend}")
logger.info(f" Worker URL: {worker_url}")
# Build router command
router_cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(router_port),
"--backend",
"openai",
"--worker-urls",
worker_url,
"--history-backend",
history_backend,
"--log-level",
"warn",
]
# Note: Not adding --api-key to router command for local testing
# The router will not require authentication
# Add Prometheus port to avoid conflicts (use unique port or disable)
if prometheus_port is None:
# Auto-assign a unique prometheus port based on router port
prometheus_port = router_port + 1000
router_cmd.extend(["--prometheus-port", str(prometheus_port)])
# Add router-specific args
if router_args:
router_cmd.extend(router_args)
if show_output:
logger.info(f" Command: {' '.join(router_cmd)}")
# Set up environment with backend API key
env = os.environ.copy()
if backend == "openai":
env["OPENAI_API_KEY"] = backend_api_key
else:
env["XAI_API_KEY"] = backend_api_key
# Launch router
if show_output:
router_proc = subprocess.Popen(
router_cmd,
env=env,
stdout=stdout,
stderr=stderr,
)
else:
router_proc = subprocess.Popen(
router_cmd,
stdout=stdout if stdout is not None else subprocess.PIPE,
stderr=stderr if stderr is not None else subprocess.PIPE,
env=env,
)
print(f" PID: {router_proc.pid}")
# Wait for router to be ready
router_url = f"http://127.0.0.1:{router_port}"
print(f"\nWaiting for router to start at {router_url}...")
try:
wait_for_router_ready(router_url, timeout=timeout, api_key=None)
logger.info(f"✓ Router ready at {router_url}")
except TimeoutError:
logger.error(f"✗ Router failed to start")
# Cleanup: kill router
try:
router_proc.kill()
except:
pass
raise
logger.info(f"\n{'='*70}")
logger.info(f"{backend.upper()} router ready!")
logger.info(f" Router: {router_url}")
logger.info(f"{'='*70}\n")
return {
"router": router_proc,
"base_url": router_url,
}
def popen_launch_workers_and_router(
model: str,
base_url: str,
timeout: int = 300,
num_workers: int = 2,
policy: str = "round_robin",
api_key: Optional[str] = None,
worker_args: Optional[list] = None,
router_args: Optional[list] = None,
tp_size: int = 1,
env: Optional[dict] = None,
stdout=None,
stderr=None,
) -> dict:
"""
Launch SGLang workers and gRPC router separately.
This approach:
1. Starts N SGLang workers with --grpc-mode flag
2. Waits for workers to initialize (process startup)
3. Starts a gRPC router pointing to those workers
4. Waits for router health check to pass (router validates worker connectivity)
This matches production deployment patterns better than the integrated approach.
Args:
model: Model path (e.g., /home/ubuntu/models/llama-3.1-8b-instruct)
base_url: Base URL for router (e.g., "http://127.0.0.1:8080")
timeout: Timeout for server startup (default: 300s)
num_workers: Number of workers to launch
policy: Routing policy (round_robin, random, power_of_two, cache_aware)
api_key: Optional API key for router
worker_args: Additional arguments for workers (e.g., ["--context-len", "8192"])
router_args: Additional arguments for router (e.g., ["--max-total-token", "1536"])
tp_size: Tensor parallelism size for workers (default: 1)
env: Optional environment variables for workers (e.g., {"SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION": "256"})
stdout: Optional file handle for worker stdout (default: subprocess.PIPE)
stderr: Optional file handle for worker stderr (default: subprocess.PIPE)
Returns:
dict with:
- workers: list of worker process objects
- worker_urls: list of gRPC worker URLs
- router: router process object
- base_url: router URL (HTTP endpoint)
Example:
>>> cluster = popen_launch_workers_and_router(model, base_url, num_workers=2)
>>> # Use cluster['base_url'] for HTTP requests
>>> # Cleanup:
>>> for worker in cluster['workers']:
>>> kill_process_tree(worker.pid)
>>> kill_process_tree(cluster['router'].pid)
"""
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
# Parse router port from base_url
if ":" in base_url.split("//")[-1]:
router_port = int(base_url.split(":")[-1])
else:
router_port = find_free_port()
logger.info(f"\n{'='*70}")
logger.info(f"Launching gRPC cluster (separate workers + router)")
logger.info(f"{'='*70}")
logger.info(f" Model: {model}")
logger.info(f" Router port: {router_port}")
logger.info(f" Workers: {num_workers}")
logger.info(f" TP size: {tp_size}")
logger.info(f" Policy: {policy}")
# Step 1: Launch workers with gRPC enabled
workers = []
worker_urls = []
for i in range(num_workers):
worker_port = find_free_port()
worker_url = f"grpc://127.0.0.1:{worker_port}"
worker_urls.append(worker_url)
logger.info(f"\n[Worker {i+1}/{num_workers}]")
logger.info(f" Port: {worker_port}")
logger.info(f" URL: {worker_url}")
# Build worker command
worker_cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
"127.0.0.1",
"--port",
str(worker_port),
"--grpc-mode", # Enable gRPC for this worker
"--mem-fraction-static",
"0.8",
]
# Add TP size
if tp_size > 1:
worker_cmd.extend(["--tp-size", str(tp_size)])
# Add worker-specific args
if worker_args:
worker_cmd.extend(worker_args)
# Launch worker with optional environment variables
if show_output:
worker_proc = subprocess.Popen(
worker_cmd,
env=env,
stdout=stdout,
stderr=stderr,
)
else:
worker_proc = subprocess.Popen(
worker_cmd,
stdout=stdout if stdout is not None else subprocess.PIPE,
stderr=stderr if stderr is not None else subprocess.PIPE,
env=env,
)
workers.append(worker_proc)
logger.info(f" PID: {worker_proc.pid}")
# Give workers a moment to start binding to ports
# The router will check worker health when it starts
logger.info(f"\nWaiting for {num_workers} workers to initialize (20s)...")
time.sleep(20)
# Quick check: make sure worker processes are still alive
for i, worker in enumerate(workers):
if worker.poll() is not None:
logger.error(
f" ✗ Worker {i+1} died during startup (exit code: {worker.poll()})"
)
# Cleanup: kill all workers
for w in workers:
try:
w.kill()
except:
pass
raise RuntimeError(f"Worker {i+1} failed to start")
logger.info(
f"✓ All {num_workers} workers started (router will verify connectivity)"
)
# Step 2: Launch router pointing to workers
logger.info(f"\n[Router]")
logger.info(f" Port: {router_port}")
logger.info(f" Worker URLs: {', '.join(worker_urls)}")
# Build router command
router_cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(router_port),
"--prometheus-port",
"9321",
"--policy",
policy,
"--model-path",
model,
"--log-level",
"warn",
]
# Add worker URLs
router_cmd.append("--worker-urls")
router_cmd.extend(worker_urls)
# Add API key
if api_key:
router_cmd.extend(["--api-key", api_key])
# Add router-specific args
if router_args:
router_cmd.extend(router_args)
if show_output:
logger.info(f" Command: {' '.join(router_cmd)}")
# Launch router
if show_output:
router_proc = subprocess.Popen(router_cmd)
else:
router_proc = subprocess.Popen(
router_cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
logger.info(f" PID: {router_proc.pid}")
# Wait for router to be ready
router_url = f"http://127.0.0.1:{router_port}"
logger.info(f"\nWaiting for router to start at {router_url}...")
try:
wait_for_workers_ready(
router_url, expected_workers=num_workers, timeout=180, api_key=api_key
)
logger.info(f"✓ Router ready at {router_url}")
except TimeoutError:
logger.error(f"✗ Router failed to start")
# Cleanup: kill router and all workers
try:
router_proc.kill()
except:
pass
for worker in workers:
try:
worker.kill()
except:
pass
raise
logger.info(f"\n{'='*70}")
logger.info(f"✓ gRPC cluster ready!")
logger.info(f" Router: {router_url}")
logger.info(f" Workers: {len(workers)}")
logger.info(f"{'='*70}\n")
return {
"workers": workers,
"worker_urls": worker_urls,
"router": router_proc,
"base_url": router_url,
}
@@ -0,0 +1,81 @@
"""
Utility functions for Response API e2e tests.
"""
import logging
import os
import signal
import threading
import unittest
import psutil
logger = logging.getLogger(__name__)
def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
"""
Kill the process and all its child processes.
Args:
parent_pid: PID of the parent process
include_parent: Whether to kill the parent process itself
skip_pid: Optional PID to skip during cleanup
"""
# Remove sigchld handler to avoid spammy logs
if threading.current_thread() is threading.main_thread():
signal.signal(signal.SIGCHLD, signal.SIG_DFL)
if parent_pid is None:
parent_pid = os.getpid()
include_parent = False
try:
itself = psutil.Process(parent_pid)
except psutil.NoSuchProcess:
return
children = itself.children(recursive=True)
for child in children:
if child.pid == skip_pid:
continue
try:
child.kill()
except psutil.NoSuchProcess:
pass
if include_parent:
try:
itself.kill()
except psutil.NoSuchProcess:
pass
class CustomTestCase(unittest.TestCase):
"""
Custom test case base class with retry support.
This provides automatic test retry functionality based on environment variables.
"""
def _callTestMethod(self, method):
"""Override to add retry logic."""
max_retry = int(os.environ.get("SGLANG_TEST_MAX_RETRY", "0"))
if max_retry == 0:
# No retry, just run once
return super(CustomTestCase, self)._callTestMethod(method)
# Retry logic
for attempt in range(max_retry + 1):
try:
return super(CustomTestCase, self)._callTestMethod(method)
except Exception as e:
if attempt < max_retry:
logger.info(
f"Test failed on attempt {attempt + 1}/{max_retry + 1}, retrying..."
)
continue
else:
# Last attempt, re-raise the exception
raise
@@ -0,0 +1 @@
"""Shared fixtures for router integration tests."""
@@ -0,0 +1,236 @@
"""
Generate self-signed certificates for mTLS integration testing.
Creates a Certificate Authority (CA), server certificates, and client certificates.
"""
import datetime
import ipaddress
from pathlib import Path
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID
def generate_private_key():
"""Generate an RSA private key."""
return rsa.generate_private_key(
public_exponent=65537,
key_size=2048,
)
def generate_ca_certificate():
"""Generate a self-signed CA certificate."""
private_key = generate_private_key()
subject = issuer = x509.Name(
[
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
x509.NameAttribute(NameOID.COMMON_NAME, "Test CA"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(issuer)
.public_key(private_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
.add_extension(
x509.BasicConstraints(ca=True, path_length=None),
critical=True,
)
.add_extension(
x509.KeyUsage(
digital_signature=True,
key_cert_sign=True,
crl_sign=True,
key_encipherment=False,
content_commitment=False,
data_encipherment=False,
key_agreement=False,
encipher_only=False,
decipher_only=False,
),
critical=True,
)
.sign(private_key, hashes.SHA256())
)
return private_key, cert
def generate_server_certificate(ca_key, ca_cert):
"""Generate a server certificate signed by the CA."""
private_key = generate_private_key()
subject = x509.Name(
[
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
x509.NameAttribute(NameOID.COMMON_NAME, "localhost"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(ca_cert.subject)
.public_key(private_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
.add_extension(
x509.SubjectAlternativeName(
[
x509.DNSName("localhost"),
x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
]
),
critical=False,
)
.add_extension(
x509.KeyUsage(
digital_signature=True,
key_encipherment=True,
key_cert_sign=False,
crl_sign=False,
content_commitment=False,
data_encipherment=False,
key_agreement=False,
encipher_only=False,
decipher_only=False,
),
critical=True,
)
.add_extension(
x509.ExtendedKeyUsage(
[
x509.oid.ExtendedKeyUsageOID.SERVER_AUTH,
]
),
critical=False,
)
.sign(ca_key, hashes.SHA256())
)
return private_key, cert
def generate_client_certificate(ca_key, ca_cert):
"""Generate a client certificate signed by the CA."""
private_key = generate_private_key()
subject = x509.Name(
[
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
x509.NameAttribute(NameOID.COMMON_NAME, "test-client"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(ca_cert.subject)
.public_key(private_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
.add_extension(
x509.KeyUsage(
digital_signature=True,
key_encipherment=True,
key_cert_sign=False,
crl_sign=False,
content_commitment=False,
data_encipherment=False,
key_agreement=False,
encipher_only=False,
decipher_only=False,
),
critical=True,
)
.add_extension(
x509.ExtendedKeyUsage(
[
x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH,
]
),
critical=False,
)
.sign(ca_key, hashes.SHA256())
)
return private_key, cert
def save_key(key, path: Path):
"""Save private key to PEM file."""
with open(path, "wb") as f:
f.write(
key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.TraditionalOpenSSL,
encryption_algorithm=serialization.NoEncryption(),
)
)
def save_cert(cert, path: Path):
"""Save certificate to PEM file."""
with open(path, "wb") as f:
f.write(cert.public_bytes(serialization.Encoding.PEM))
def generate_all_certificates(output_dir: Path):
"""Generate all certificates and keys for mTLS testing."""
output_dir.mkdir(parents=True, exist_ok=True)
print("==> Generating CA certificate...")
ca_key, ca_cert = generate_ca_certificate()
save_key(ca_key, output_dir / "ca-key.pem")
save_cert(ca_cert, output_dir / "ca-cert.pem")
print("==> Generating server certificate...")
server_key, server_cert = generate_server_certificate(ca_key, ca_cert)
save_key(server_key, output_dir / "server-key.pem")
save_cert(server_cert, output_dir / "server-cert.pem")
print("==> Generating client certificate...")
client_key, client_cert = generate_client_certificate(ca_key, ca_cert)
save_key(client_key, output_dir / "client-key.pem")
save_cert(client_cert, output_dir / "client-cert.pem")
print(f"==> Certificates generated successfully in {output_dir}")
print()
print("Files created:")
print(" - ca-cert.pem : CA certificate (for verifying server/client certs)")
print(" - ca-key.pem : CA private key")
print(" - server-cert.pem : Server certificate")
print(" - server-key.pem : Server private key")
print(" - client-cert.pem : Client certificate")
print(" - client-key.pem : Client private key")
print()
print("Test server can use: server-cert.pem + server-key.pem")
print("Test router can use: client-cert.pem + client-key.pem + ca-cert.pem")
if __name__ == "__main__":
script_dir = Path(__file__).parent
certs_dir = script_dir / "test_certs"
generate_all_certificates(certs_dir)
@@ -0,0 +1,288 @@
"""
Lightweight mock worker HTTP server for router integration tests.
Implements minimal endpoints used by the router:
- GET /health, /health_generate
- POST /generate, /v1/completions, /v1/chat/completions
- POST /flush_cache
- GET /get_server_info, /get_model_info, /v1/models
Behavior knobs are controlled via CLI flags to simulate failures, latency, and load.
"""
import argparse
import asyncio
import json
import os
import random
import signal
import sys
import time
from contextlib import asynccontextmanager
from typing import Optional
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
# Global state (per-process)
_inflight = 0
_failures_seen = 0
def _parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--host", default="127.0.0.1")
p.add_argument("--port", type=int, required=True)
p.add_argument("--worker-id", default=None)
p.add_argument("--latency-ms", type=int, default=0)
p.add_argument("--timeout", action="store_true")
p.add_argument("--status-code", type=int, default=200)
p.add_argument("--fail-first-n", type=int, default=0)
p.add_argument("--random-fail-rate", type=float, default=0.0)
p.add_argument("--require-api-key", action="store_true")
p.add_argument("--api-key", default=None)
p.add_argument("--max-payload-bytes", type=int, default=10 * 1024 * 1024)
p.add_argument("--stream", action="store_true")
p.add_argument("--dp-size", type=int, default=1)
p.add_argument("--crash-on-request", action="store_true")
p.add_argument("--health-fail-after-ms", type=int, default=0)
# TLS/mTLS configuration
p.add_argument(
"--ssl-certfile", type=str, default=None, help="Path to SSL certificate file"
)
p.add_argument("--ssl-keyfile", type=str, default=None, help="Path to SSL key file")
p.add_argument(
"--ssl-ca-certs",
type=str,
default=None,
help="Path to CA certificates for client verification",
)
return p.parse_args()
def _extract_worker_id(args: argparse.Namespace) -> str:
if args.worker_id:
return str(args.worker_id)
# default to port (unique enough for tests)
return f"worker-{args.port}"
def create_app(args: argparse.Namespace) -> FastAPI:
app = FastAPI()
worker_id = _extract_worker_id(args)
start_ts = time.time()
crashed = {"done": False}
async def maybe_delay():
if args.latency_ms > 0:
await asyncio.sleep(args.latency_ms / 1000.0)
def should_fail() -> Optional[int]:
global _failures_seen
# Fail first N requests (500)
if args.fail_first_n > 0 and _failures_seen < args.fail_first_n:
_failures_seen += 1
return 500
# Random failure probability (500)
if args.random_fail_rate > 0.0 and random.random() < args.random_fail_rate:
return 500
# Forced status code override (non-200) for all responses
if args.status_code != 200:
return int(args.status_code)
return None
def check_api_key(request: Request):
if not args.require_api_key:
return
auth = request.headers.get("Authorization")
if not auth or not auth.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Unauthorized")
key = auth.split(" ", 1)[1]
if args.api_key and key != args.api_key:
raise HTTPException(status_code=401, detail="Unauthorized")
@asynccontextmanager
async def track_inflight():
global _inflight
_inflight += 1
try:
yield
finally:
_inflight -= 1
@app.get("/health")
async def health():
if (
args.health_fail_after_ms
and (time.time() - start_ts) * 1000.0 >= args.health_fail_after_ms
):
return PlainTextResponse("bad", status_code=500)
return PlainTextResponse("ok", status_code=200)
@app.get("/health_generate")
async def health_generate():
return PlainTextResponse("ok", status_code=200)
@app.post("/flush_cache")
async def flush_cache():
return PlainTextResponse("ok", status_code=200)
@app.get("/get_model_info")
async def get_model_info():
return JSONResponse({"model": "mock", "vocab_size": 32000})
@app.get("/v1/models")
async def list_models():
return JSONResponse({"data": [{"id": "mock", "object": "model"}]})
@app.get("/get_server_info")
async def get_server_info(request: Request):
# Enforce API key on server info when required (used by dp_aware probing)
check_api_key(request)
return JSONResponse(
{
"worker_id": worker_id,
"load_in_flight": _inflight,
"cache": {"size": 0, "hit_rate": 0.0},
"dp_size": int(args.dp_size),
}
)
@app.get("/get_load")
async def get_load(request: Request):
check_api_key(request)
# Return format matching real workers: array of load info per DP rank
return JSONResponse(
[
{
"dp_rank": 0,
"num_reqs": _inflight,
"num_waiting_reqs": 0,
"num_tokens": _inflight,
}
]
)
def make_json_response(obj: dict, status_code: int = 200) -> JSONResponse:
resp = JSONResponse(obj, status_code=status_code)
resp.headers["X-Worker-Id"] = worker_id
return resp
async def handle_text_request(request: Request):
# Authorization
check_api_key(request)
# Payload limit
body = await request.body()
if len(body) > args.max_payload_bytes:
return make_json_response({"error": "payload too large"}, status_code=413)
# Simulate crash on first request
if args.crash_on_request and not crashed["done"]:
crashed["done"] = True
os._exit(1)
# Optional timeout (simulate hang)
if args.timeout:
await asyncio.sleep(3600)
# Optional latency
await maybe_delay()
# Optional failures
fail_code = should_fail()
if fail_code is not None and fail_code != 200:
return make_json_response(
{"error": f"mock failure {fail_code}"}, status_code=fail_code
)
# Build response echoing minimal shape
try:
data = await request.json()
except (json.JSONDecodeError, ValueError):
data = {}
received_headers = {k.lower(): v for k, v in request.headers.items()}
now = time.time()
ret = {
"id": f"cmpl-{int(now*1000)}",
"object": "text_completion",
"created": int(now),
"model": "mock",
"choices": [
{
"text": "ok",
"index": 0,
"finish_reason": "stop",
}
],
"worker_id": worker_id,
"echo": data,
"received_headers": received_headers,
}
return make_json_response(ret, status_code=200)
async def handle_stream_request(request: Request):
check_api_key(request)
async def gen():
# minimal 2-chunk stream then [DONE]
for i in range(2):
await asyncio.sleep(0.01)
chunk = {
"choices": [{"delta": {"content": "x"}}],
"worker_id": worker_id,
}
yield f"data: {json.dumps(chunk)}\n\n"
yield "data: [DONE]\n\n"
headers = {"X-Worker-Id": worker_id}
return StreamingResponse(gen(), media_type="text/event-stream", headers=headers)
@app.post("/generate")
async def generate(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
@app.post("/v1/completions")
async def completions(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
return app
def main() -> None:
args = _parse_args()
app = create_app(args)
# Handle SIGTERM gracefully for fast test teardown
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
# Configure SSL if certificates are provided
ssl_config = {}
if args.ssl_certfile and args.ssl_keyfile:
ssl_config["ssl_certfile"] = args.ssl_certfile
ssl_config["ssl_keyfile"] = args.ssl_keyfile
# If CA certs provided, require client certificates (mTLS)
if args.ssl_ca_certs:
ssl_config["ssl_ca_certs"] = args.ssl_ca_certs
ssl_config["ssl_cert_reqs"] = 2 # ssl.CERT_REQUIRED
uvicorn.run(app, host=args.host, port=args.port, log_level="warning", **ssl_config)
if __name__ == "__main__":
main()
@@ -0,0 +1,8 @@
import socket
def find_free_port() -> int:
"""Return an available TCP port on localhost."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
@@ -0,0 +1,248 @@
import subprocess
import time
from dataclasses import dataclass
from typing import Dict, List, Optional
import requests
from .ports import find_free_port
@dataclass
class ProcHandle:
process: subprocess.Popen
url: str
class RouterManager:
"""Helper to spawn a router process and interact with admin endpoints."""
def __init__(self):
self._children: List[subprocess.Popen] = []
def start_router(
self,
worker_urls: Optional[List[str]] = None,
policy: str = "round_robin",
port: Optional[int] = None,
extra: Optional[Dict] = None,
# PD options
pd_disaggregation: bool = False,
prefill_urls: Optional[List[tuple]] = None,
decode_urls: Optional[List[str]] = None,
prefill_policy: Optional[str] = None,
decode_policy: Optional[str] = None,
) -> ProcHandle:
worker_urls = worker_urls or []
port = port or find_free_port()
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(port),
"--policy",
policy,
]
# Avoid Prometheus port collisions by assigning a free port per router
prom_port = find_free_port()
cmd.extend(
["--prometheus-port", str(prom_port), "--prometheus-host", "127.0.0.1"]
)
if worker_urls:
cmd.extend(["--worker-urls", *worker_urls])
# PD routing configuration
if pd_disaggregation:
cmd.append("--pd-disaggregation")
if prefill_urls:
for url, bport in prefill_urls:
if bport is None:
cmd.extend(["--prefill", url, "none"])
else:
cmd.extend(["--prefill", url, str(bport)])
if decode_urls:
for url in decode_urls:
cmd.extend(["--decode", url])
if prefill_policy:
cmd.extend(["--prefill-policy", prefill_policy])
if decode_policy:
cmd.extend(["--decode-policy", decode_policy])
# Map supported extras to CLI flags (subset for integration)
if extra:
flag_map = {
"max_payload_size": "--max-payload-size",
"dp_aware": "--dp-aware",
"api_key": "--api-key",
# Health/monitoring
"worker_startup_check_interval": "--worker-startup-check-interval",
# Cache-aware tuning
"cache_threshold": "--cache-threshold",
"balance_abs_threshold": "--balance-abs-threshold",
"balance_rel_threshold": "--balance-rel-threshold",
# Retry
"retry_max_retries": "--retry-max-retries",
"retry_initial_backoff_ms": "--retry-initial-backoff-ms",
"retry_max_backoff_ms": "--retry-max-backoff-ms",
"retry_backoff_multiplier": "--retry-backoff-multiplier",
"retry_jitter_factor": "--retry-jitter-factor",
"disable_retries": "--disable-retries",
# Circuit breaker
"cb_failure_threshold": "--cb-failure-threshold",
"cb_success_threshold": "--cb-success-threshold",
"cb_timeout_duration_secs": "--cb-timeout-duration-secs",
"cb_window_duration_secs": "--cb-window-duration-secs",
"disable_circuit_breaker": "--disable-circuit-breaker",
# Rate limiting
"max_concurrent_requests": "--max-concurrent-requests",
"queue_size": "--queue-size",
"queue_timeout_secs": "--queue-timeout-secs",
"rate_limit_tokens_per_second": "--rate-limit-tokens-per-second",
# mTLS configuration
"client_cert_path": "--client-cert-path",
"client_key_path": "--client-key-path",
"ca_cert_paths": "--ca-cert-paths",
}
for k, v in extra.items():
if v is None:
continue
flag = flag_map.get(k)
if not flag:
continue
if isinstance(v, bool):
if v:
cmd.append(flag)
elif isinstance(v, list):
# Handle list arguments (e.g., ca_cert_paths)
if v: # Only add if list is not empty
cmd.append(flag)
cmd.extend([str(item) for item in v])
else:
cmd.extend([flag, str(v)])
proc = subprocess.Popen(cmd)
self._children.append(proc)
url = f"http://127.0.0.1:{port}"
self._wait_health(url)
return ProcHandle(process=proc, url=url)
def _wait_health(self, base_url: str, timeout: float = 30.0):
start = time.time()
with requests.Session() as s:
while time.time() - start < timeout:
try:
r = s.get(f"{base_url}/health", timeout=2)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(0.2)
raise TimeoutError(f"Router at {base_url} did not become healthy")
def add_worker(self, base_url: str, worker_url: str, timeout: float = 30.0) -> None:
r = requests.post(f"{base_url}/workers", json={"url": worker_url})
assert (
r.status_code == 202
), f"add_worker failed: {r.status_code} {r.text}" # ACCEPTED status
payload = r.json()
worker_id = payload.get("worker_id")
assert worker_id, f"add_worker did not return worker_id: {payload}"
# Poll until worker is actually added and healthy
start = time.time()
with requests.Session() as s:
while time.time() - start < timeout:
try:
r = s.get(f"{base_url}/workers/{worker_id}", timeout=2)
if r.status_code == 200:
data = r.json()
# Check if registration job failed
job_status = data.get("job_status")
if job_status and job_status.get("state") == "failed":
raise RuntimeError(
f"Worker registration failed: {job_status.get('message', 'Unknown error')}"
)
# Check if worker is healthy and registered (not just in job queue)
if data.get("is_healthy", False):
return
# Worker not ready yet, continue polling
except requests.RequestException:
pass
time.sleep(0.1)
raise TimeoutError(
f"Worker {worker_url} was not added and healthy after {timeout}s"
)
def remove_worker(
self, base_url: str, worker_url: str, timeout: float = 30.0
) -> None:
# Resolve worker_id from the current registry snapshot
r_list = requests.get(f"{base_url}/workers")
assert (
r_list.status_code == 200
), f"list_workers failed: {r_list.status_code} {r_list.text}"
workers = r_list.json().get("workers", [])
worker_id = next(
(w.get("id") for w in workers if w.get("url") == worker_url), None
)
assert (
worker_id
), f"could not find worker_id for url={worker_url}. workers={workers}"
r = requests.delete(f"{base_url}/workers/{worker_id}")
assert (
r.status_code == 202
), f"remove_worker failed: {r.status_code} {r.text}" # ACCEPTED status
# Poll until worker is actually removed (GET returns 404) or timeout
start = time.time()
last_status = None
with requests.Session() as s:
while time.time() - start < timeout:
try:
r = s.get(f"{base_url}/workers/{worker_id}", timeout=2)
if r.status_code == 404:
# Worker successfully removed
return
elif r.status_code == 200:
# Check if removal job failed
data = r.json()
job_status = data.get("job_status")
if job_status:
last_status = job_status
if job_status.get("state") == "failed":
raise RuntimeError(
f"Worker removal failed: {job_status.get('message', 'Unknown error')}"
)
# Worker still being processed, continue polling
except requests.RequestException:
pass
time.sleep(0.1)
# Provide detailed timeout error with last known status
error_msg = f"Worker {worker_url} was not removed after {timeout}s"
if last_status:
error_msg += f". Last job status: {last_status}"
raise TimeoutError(error_msg)
def list_workers(self, base_url: str) -> list[str]:
r = requests.get(f"{base_url}/workers")
assert r.status_code == 200, f"list_workers failed: {r.status_code} {r.text}"
data = r.json()
# Extract URLs from WorkerInfo objects
workers = data.get("workers", [])
return [w["url"] for w in workers]
def stop_all(self):
for p in self._children:
if p.poll() is None:
p.terminate()
try:
p.wait(timeout=5)
except subprocess.TimeoutExpired:
p.kill()
self._children.clear()