[model-gateway] rename py_test to e2e_test (#16454)
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Test package root for router Python tests."""
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user