chore: vendor sglang v0.5.10 snapshot

This commit is contained in:
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions

View File

@@ -0,0 +1 @@
"""Test package root for router Python tests."""

View File

@@ -0,0 +1,222 @@
"""Benchmark-specific fixtures."""
from __future__ import annotations
import logging
import os
import shutil
import subprocess
import time
from pathlib import Path
import pytest
from infra import GPUMonitor, should_monitor_gpu, terminate_process
from .results import BenchmarkResult
logger = logging.getLogger(__name__)
def _build_command(
cli: str,
router_url: str,
model_path: str,
experiment_folder: str,
num_concurrency: int,
traffic_scenario: str,
max_requests: int,
) -> list[str]:
"""Build genai-bench command."""
return [
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(max_requests),
"--max-time-per-run",
"3",
"--experiment-folder-name",
experiment_folder,
"--experiment-base-dir",
str(Path.cwd()),
]
def _find_results(experiment_folder: str, timeout: int = 10) -> list[Path]:
"""Find benchmark result JSON files."""
base = Path.cwd()
folder = base / experiment_folder
if not folder.is_dir():
# Search for folder
for p in base.rglob(experiment_folder):
if p.is_dir() and p.name == experiment_folder:
folder = p
break
if not folder.is_dir():
raise AssertionError(f"Experiment folder not found: {experiment_folder}")
# Wait for JSON results
for _ in range(timeout):
files = [
p
for p in folder.rglob("*.json")
if "experiment_metadata" not in p.name and "gpu_utilization" not in p.name
]
if files:
return files
time.sleep(1)
raise AssertionError(f"No JSON results found in {folder}")
def _cleanup_procs(procs: list, drain_delay: int) -> None:
"""Terminate processes gracefully."""
if not procs:
return
if drain_delay > 0:
time.sleep(drain_delay)
for p in procs:
try:
proc = getattr(p, "proc", p) if hasattr(p, "proc") else p
if isinstance(proc, subprocess.Popen):
terminate_process(proc)
except Exception:
pass
time.sleep(2)
@pytest.fixture(scope="session")
def genai_bench_runner():
"""Run genai-bench and validate metrics.
Usage:
def test_perf(setup_backend, genai_bench_runner):
backend, model_path, client, gateway = setup_backend
genai_bench_runner(
router_url=gateway.base_url,
model_path=model_path,
experiment_folder="benchmark_results",
thresholds={"ttft_mean_max": 5, "gpu_util_p50_min": 99},
)
"""
def _run(
*,
router_url: str,
model_path: str,
experiment_folder: str,
thresholds: dict | None = None,
timeout_sec: int | None = None,
num_concurrency: int = 32,
traffic_scenario: str = "D(4000,100)",
max_requests_per_run: int | None = None,
kill_procs: list | None = None,
drain_delay_sec: int = 6,
) -> None:
cli = shutil.which("genai-bench")
if not cli:
pytest.fail("genai-bench CLI not found")
# Clean previous results
exp_dir = Path.cwd() / experiment_folder
if exp_dir.exists():
shutil.rmtree(exp_dir, ignore_errors=True)
# Build and run command
max_requests = max_requests_per_run or num_concurrency * 5
cmd = _build_command(
cli,
router_url,
model_path,
experiment_folder,
num_concurrency,
traffic_scenario,
max_requests,
)
timeout = timeout_sec or int(os.environ.get("GENAI_BENCH_TEST_TIMEOUT", "120"))
try:
proc = subprocess.Popen(
cmd,
env=os.environ.copy(),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
)
except FileNotFoundError:
pytest.fail(f"genai-bench executable not found at {cli}")
except PermissionError:
pytest.fail(f"Permission denied executing {cli}")
except OSError as e:
pytest.fail(f"Failed to start genai-bench: {e}")
# Start GPU monitor if needed
gpu_monitor: GPUMonitor | None = None
if should_monitor_gpu(thresholds):
interval = float(os.environ.get("GPU_UTIL_SAMPLE_INTERVAL", "2.0"))
gpu_monitor = GPUMonitor(output_dir=exp_dir, interval=interval)
gpu_monitor.start(target_pid=proc.pid)
try:
stdout, stderr = proc.communicate(timeout=timeout)
except subprocess.TimeoutExpired:
proc.kill()
stdout, stderr = proc.communicate()
logger.error("genai-bench timed out after %ds", timeout)
# Log output if process failed or for debugging
if proc.returncode != 0:
logger.error(
"genai-bench exited with code %d\nstdout:\n%s\nstderr:\n%s",
proc.returncode,
stdout or "(empty)",
stderr or "(empty)",
)
try:
# Parse and validate results
for path in _find_results(experiment_folder):
result = BenchmarkResult.from_json(path)
result.log(experiment_folder, logger)
if thresholds:
result.validate(thresholds)
# Validate GPU utilization
if gpu_monitor:
gpu_monitor.stop()
gpu_monitor.log_summary()
gpu_monitor.assert_thresholds(thresholds)
except AssertionError:
# Log genai-bench output when results not found
logger.error(
"genai-bench output (returncode=%d):\nstdout:\n%s\nstderr:\n%s",
proc.returncode,
stdout or "(empty)",
stderr or "(empty)",
)
raise
finally:
_cleanup_procs(kill_procs, drain_delay_sec)
if gpu_monitor:
gpu_monitor.stop(timeout=2)
return _run

View File

@@ -0,0 +1,98 @@
"""Benchmark result dataclasses for parsing genai-bench and GPU monitor output."""
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
@dataclass
class BenchmarkResult:
"""Parsed benchmark metrics from genai-bench output."""
ttft_mean: float
e2e_latency_mean: float
input_throughput_mean: float
output_throughput_mean: float
file_name: str
@classmethod
def from_json(cls, path: Path) -> "BenchmarkResult":
"""Parse benchmark results from JSON file."""
with path.open() as f:
data = json.load(f)
stats = data.get("aggregated_metrics", {}).get("stats", {})
return cls(
ttft_mean=float(stats.get("ttft", {}).get("mean", float("inf"))),
e2e_latency_mean=float(
stats.get("e2e_latency", {}).get("mean", float("inf"))
),
input_throughput_mean=float(
stats.get("input_throughput", {}).get("mean", 0.0)
),
output_throughput_mean=float(
stats.get("output_throughput", {}).get("mean", 0.0)
),
file_name=path.name,
)
def log(self, experiment: str, logger) -> None:
"""Log benchmark results."""
logger.info(
"genai-bench[%s] %s ttft=%.3fs e2e=%.3fs input=%.1f tok/s output=%.1f tok/s",
experiment,
self.file_name,
self.ttft_mean,
self.e2e_latency_mean,
self.input_throughput_mean,
self.output_throughput_mean,
)
def validate(self, thresholds: dict) -> None:
"""Validate metrics against thresholds."""
checks = [
("ttft_mean_max", self.ttft_mean, "<=", "TTFT"),
("e2e_latency_mean_max", self.e2e_latency_mean, "<=", "E2E latency"),
(
"input_throughput_mean_min",
self.input_throughput_mean,
">=",
"Input throughput",
),
(
"output_throughput_mean_min",
self.output_throughput_mean,
">=",
"Output throughput",
),
]
for key, value, op, name in checks:
if key not in thresholds:
continue
threshold = thresholds[key]
if op == "<=" and value > threshold:
raise AssertionError(f"{name}: {value:.2f} > {threshold}")
if op == ">=" and value < threshold:
raise AssertionError(f"{name}: {value:.2f} < {threshold}")
@dataclass
class GPUUtilization:
"""Parsed GPU utilization metrics from gpu_monitor output."""
overall_mean: float
per_gpu: dict[str, dict[str, float]]
@classmethod
def from_json(cls, path: Path) -> "GPUUtilization | None":
"""Parse GPU utilization from JSON file."""
try:
with path.open() as f:
data = json.load(f)
return cls(
overall_mean=float(data.get("overall", {}).get("mean", 0)),
per_gpu=data.get("per_gpu", {}),
)
except Exception:
return None

View File

@@ -0,0 +1,119 @@
"""Generate benchmark summary for GitHub Actions."""
from __future__ import annotations
import os
import sys
from pathlib import Path
from results import BenchmarkResult, GPUUtilization
def discover_benchmarks(base_dir: Path) -> list[tuple[Path, str]]:
"""Auto-discover benchmark folders and their result JSON files.
Returns list of (json_path, label) tuples sorted by folder name.
"""
results = []
for folder in base_dir.rglob("benchmark_*"):
if not folder.is_dir():
continue
# Find result JSON (exclude metadata and gpu files)
for json_file in folder.glob("*.json"):
if (
"experiment_metadata" not in json_file.name
and "gpu_utilization" not in json_file.name
):
# Generate label from folder name: benchmark_cache_aware_pd_grpc -> cache_aware pd grpc
label = folder.name.replace("benchmark_", "").replace("_", " ")
results.append((json_file, label))
break # One JSON per folder
return sorted(results, key=lambda x: x[0].parent.name)
def find_gpu_utilization(result_path: Path) -> Path | None:
"""Find GPU utilization JSON in same folder as result."""
gpu_json = result_path.parent / "gpu_utilization.json"
return gpu_json if gpu_json.exists() else None
def generate_summary(base_dir: Path) -> str:
"""Generate markdown summary."""
benchmarks = discover_benchmarks(base_dir)
if not benchmarks:
return (
"## Gateway E2E Genai-Bench Results Summary\n\nNo benchmark results found."
)
lines = [
"## Gateway E2E Genai-Bench Results Summary",
"",
"| Scenario | Status | TTFT (s) | E2E Latency (s) | Input Throughput (tok/s) | Output Throughput (tok/s) |",
"|----------|--------|----------|-----------------|--------------------------|---------------------------|",
]
gpu_sections = []
for result_path, label in benchmarks:
try:
result = BenchmarkResult.from_json(result_path)
except Exception as e:
print(f"Warning: Failed to parse {result_path}: {e}", file=sys.stderr)
lines.append(f"| {label} | ❌ Failed | - | - | - | - |")
continue
lines.append(
f"| {label} | ✅ Success | "
f"{result.ttft_mean:.2f} | "
f"{result.e2e_latency_mean:.2f} | "
f"{result.input_throughput_mean:.0f} | "
f"{result.output_throughput_mean:.0f} |"
)
# GPU utilization
gpu_path = find_gpu_utilization(result_path)
if gpu_path:
gpu = GPUUtilization.from_json(gpu_path)
if gpu and gpu.per_gpu:
gpu_lines = [
f"### GPU Utilization — {label}",
"",
f"Overall mean: {gpu.overall_mean:.2f}%",
"",
"| GPU | Mean (%) | p5 | p10 | p25 | p50 | p75 | p90 | p95 |",
"|-----|----------|----|-----|-----|-----|-----|-----|-----|",
]
for gpu_id, stats in sorted(
gpu.per_gpu.items(), key=lambda x: int(x[0])
):
gpu_lines.append(
f"| {gpu_id} | {stats.get('mean', 0):.2f} | "
f"{stats.get('p5', 0):.2f} | {stats.get('p10', 0):.2f} | "
f"{stats.get('p25', 0):.2f} | {stats.get('p50', 0):.2f} | "
f"{stats.get('p75', 0):.2f} | {stats.get('p90', 0):.2f} | "
f"{stats.get('p95', 0):.2f} |"
)
gpu_sections.append("\n".join(gpu_lines))
return "\n".join(lines) + "\n" + "\n\n".join(gpu_sections)
def main() -> None:
"""Main entry point."""
base_dir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.cwd()
summary = generate_summary(base_dir)
# Write to GITHUB_STEP_SUMMARY if available
summary_file = os.environ.get("GITHUB_STEP_SUMMARY")
if summary_file:
with open(summary_file, "a") as f:
f.write(summary)
f.write("\n")
print(f"Summary written to {summary_file}")
else:
print(summary)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,26 @@
"""PD (prefill/decode disaggregation) router performance benchmark test."""
import pytest
@pytest.mark.e2e
@pytest.mark.workers(prefill=2, decode=2)
@pytest.mark.parametrize("setup_backend", ["pd"], indirect=True)
class TestPDPerf:
"""Performance benchmark for PD disaggregation router."""
def test_pd_perf(self, setup_backend, genai_bench_runner):
"""Run genai-bench against PD router and validate metrics."""
backend, model_path, client, gateway = setup_backend
genai_bench_runner(
router_url=gateway.base_url,
model_path=model_path,
experiment_folder="benchmark_round_robin_pd",
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,
},
)

View File

@@ -0,0 +1,27 @@
"""Regular router performance benchmark test."""
import pytest
@pytest.mark.e2e
@pytest.mark.workers(count=4)
@pytest.mark.gateway(policy="cache_aware")
@pytest.mark.parametrize("setup_backend", ["http", "grpc"], indirect=True)
class TestRegularPerf:
"""Performance benchmark for regular (non-PD) router."""
def test_regular_perf(self, setup_backend, genai_bench_runner):
"""Run genai-bench against regular router and validate metrics."""
backend, model_path, client, gateway = setup_backend
genai_bench_runner(
router_url=gateway.base_url,
model_path=model_path,
experiment_folder=f"benchmark_cache_aware_regular_{backend}",
thresholds={
"ttft_mean_max": 6,
"e2e_latency_mean_max": 14,
"input_throughput_mean_min": 800,
"output_throughput_mean_min": 12,
"gpu_util_p50_min": 99,
},
)

View File

@@ -0,0 +1,168 @@
"""Enable Thinking E2E Tests.
Tests for chat completions with enable_thinking feature (Qwen3 reasoning).
Source: Migrated from e2e_grpc/features/test_enable_thinking.py
"""
from __future__ import annotations
import json
import logging
import pytest
import requests
logger = logging.getLogger(__name__)
# API key is not validated by the gateway, but required for OpenAI-compatible headers
API_KEY = "not-used"
# =============================================================================
# Enable Thinking Tests (Qwen 30B)
# =============================================================================
@pytest.mark.model("qwen-30b")
@pytest.mark.gateway(
extra_args=["--reasoning-parser", "qwen3", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestEnableThinking:
"""Tests for enable_thinking feature with Qwen3 reasoning parser."""
def test_chat_completion_with_reasoning(self, setup_backend):
"""Test non-streaming with enable_thinking=True, reasoning_content should not be empty."""
_, model, client, gateway = setup_backend
response = requests.post(
f"{gateway.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {API_KEY}"},
json={
"model": model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"chat_template_kwargs": {"enable_thinking": True},
},
)
assert response.status_code == 200, f"Failed with: {response.text}"
data = response.json()
assert "choices" in data
assert len(data["choices"]) > 0
assert "message" in data["choices"][0]
assert "reasoning_content" in data["choices"][0]["message"]
assert data["choices"][0]["message"]["reasoning_content"] is not None
def test_chat_completion_without_reasoning(self, setup_backend):
"""Test non-streaming with enable_thinking=False, reasoning_content should be empty."""
_, model, client, gateway = setup_backend
response = requests.post(
f"{gateway.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {API_KEY}"},
json={
"model": model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"chat_template_kwargs": {"enable_thinking": False},
},
)
assert response.status_code == 200, f"Failed with: {response.text}"
data = response.json()
assert "choices" in data
assert len(data["choices"]) > 0
assert "message" in data["choices"][0]
if "reasoning_content" in data["choices"][0]["message"]:
assert data["choices"][0]["message"]["reasoning_content"] is None
def test_stream_chat_completion_with_reasoning(self, setup_backend):
"""Test streaming with enable_thinking=True, reasoning_content should not be empty."""
_, model, client, gateway = setup_backend
response = requests.post(
f"{gateway.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {API_KEY}"},
json={
"model": model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"stream": True,
"chat_template_kwargs": {"enable_thinking": True},
},
stream=True,
)
assert 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
assert (
has_reasoning
), "The reasoning content is not included in the stream response"
assert has_content, "The stream response does not contain normal content"
def test_stream_chat_completion_without_reasoning(self, setup_backend):
"""Test streaming with enable_thinking=False, reasoning_content should be empty."""
_, model, client, gateway = setup_backend
response = requests.post(
f"{gateway.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {API_KEY}"},
json={
"model": model,
"messages": [{"role": "user", "content": "Hello"}],
"temperature": 0,
"separate_reasoning": True,
"stream": True,
"chat_template_kwargs": {"enable_thinking": False},
},
stream=True,
)
assert 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
assert (
not has_reasoning
), "The reasoning content should not be included in the stream response"
assert has_content, "The stream response does not contain normal content"

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,316 @@
"""Chat Completions API E2E Tests - OpenAI Server Compatibility.
Tests for OpenAI-compatible chat completions API through the gateway.
Source: Migrated from e2e_grpc/basic/test_openai_server.py
"""
from __future__ import annotations
import json
import logging
import pytest
logger = logging.getLogger(__name__)
# =============================================================================
# Chat Completion Tests (Llama 8B)
# =============================================================================
@pytest.mark.model("llama-8b")
@pytest.mark.gateway(extra_args=["--history-backend", "memory"])
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestChatCompletion:
"""Tests for OpenAI-compatible chat completions API."""
@pytest.mark.parametrize("logprobs", [None, 5])
@pytest.mark.parametrize("parallel_sample_num", [1, 2])
def test_chat_completion(self, setup_backend, logprobs, parallel_sample_num):
"""Test non-streaming chat completion with logprobs and parallel sampling."""
_, model, client, gateway = setup_backend
self._run_chat_completion(client, model, logprobs, parallel_sample_num)
@pytest.mark.parametrize("logprobs", [None, 5])
@pytest.mark.parametrize("parallel_sample_num", [1, 2])
def test_chat_completion_stream(self, setup_backend, logprobs, parallel_sample_num):
"""Test streaming chat completion with logprobs and parallel sampling."""
_, model, client, gateway = setup_backend
self._run_chat_completion_stream(client, model, logprobs, parallel_sample_num)
def test_regex(self, setup_backend):
"""Test structured output with regex constraint."""
_, model, client, gateway = setup_backend
regex = (
r"""\{\n"""
+ r""" "name": "[\w]+",\n"""
+ r""" "population": [\d]+\n"""
+ r"""\}"""
)
response = client.chat.completions.create(
model=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, setup_backend):
"""Test frequency penalty parameter."""
_, model, client, gateway = setup_backend
response = client.chat.completions.create(
model=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, setup_backend):
"""Test assistant message prefill with continue_final_message."""
_, model, client, gateway = setup_backend
response = client.chat.completions.create(
model=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, setup_backend):
"""Test listing available models."""
_, model, client, gateway = setup_backend
models = list(client.models.list().data)
assert len(models) == 1
@pytest.mark.skip(
reason="Skipping retrieve model test as it is not supported by the router"
)
def test_retrieve_model(self, setup_backend):
"""Test retrieving a specific model."""
import openai
_, model, client, gateway = setup_backend
retrieved_model = client.models.retrieve(model)
assert retrieved_model.id == model
assert retrieved_model.root == model
with pytest.raises(openai.NotFoundError):
client.models.retrieve("non-existent-model")
# -------------------------------------------------------------------------
# Helper methods
# -------------------------------------------------------------------------
def _run_chat_completion(self, client, model, logprobs, parallel_sample_num):
"""Run a non-streaming chat completion and verify response."""
response = client.chat.completions.create(
model=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, client, model, logprobs, parallel_sample_num=1
):
"""Run a streaming chat completion and verify response chunks."""
generator = client.chat.completions.create(
model=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, "usage.prompt_tokens was zero"
assert usage.completion_tokens > 0, "usage.completion_tokens was zero"
assert usage.total_tokens > 0, "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"
), "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, "logprobs was not returned"
assert isinstance(
response.choices[0].logprobs.content[0].top_logprobs[0].token, str
), "top_logprobs token was not a string"
assert isinstance(
response.choices[0].logprobs.content[0].top_logprobs, list
), "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 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}, "
f"got {finish_reason_counts[index]}"
)
# =============================================================================
# Chat Completion Tests (GPT-OSS)
#
# NOTE: Some tests are skipped because they don't work with OSS models:
# - test_regex: OSS models don't support regex constraints
# - test_penalty: OSS models don't support frequency_penalty
# - test_response_prefill: OSS models don't support continue_final_message
# =============================================================================
@pytest.mark.model("gpt-oss")
@pytest.mark.gateway(
extra_args=["--reasoning-parser=gpt-oss", "--history-backend", "memory"]
)
class TestChatCompletionGptOss(TestChatCompletion):
"""Tests for chat completions API with GPT-OSS model.
Inherits from TestChatCompletion and overrides tests that don't work
with OSS models.
"""
@pytest.mark.parametrize("logprobs", [None]) # No logprobs for OSS
@pytest.mark.parametrize("parallel_sample_num", [1, 2])
def test_chat_completion(self, setup_backend, logprobs, parallel_sample_num):
"""Test non-streaming chat completion with parallel sampling (no logprobs)."""
super().test_chat_completion(setup_backend, logprobs, parallel_sample_num)
@pytest.mark.parametrize("logprobs", [None]) # No logprobs for OSS
@pytest.mark.parametrize("parallel_sample_num", [1, 2])
def test_chat_completion_stream(self, setup_backend, logprobs, parallel_sample_num):
"""Test streaming chat completion with parallel sampling (no logprobs)."""
super().test_chat_completion_stream(
setup_backend, logprobs, parallel_sample_num
)
@pytest.mark.skip(reason="OSS models don't support regex constraints")
def test_regex(self, setup_backend):
pass
@pytest.mark.skip(reason="OSS models don't support frequency_penalty")
def test_penalty(self, setup_backend):
pass
@pytest.mark.skip(reason="OSS models don't support continue_final_message")
def test_response_prefill(self, setup_backend):
pass

View File

@@ -0,0 +1,165 @@
"""Reasoning Content E2E Tests.
Tests for chat completions with reasoning content (DeepSeek R1 reasoning parser).
Source: Migrated from e2e_grpc/features/test_reasoning_content.py
"""
from __future__ import annotations
import logging
import pytest
logger = logging.getLogger(__name__)
# =============================================================================
# Reasoning Content API Tests (DeepSeek 7B)
# =============================================================================
@pytest.mark.model("deepseek-7b")
@pytest.mark.gateway(
extra_args=["--reasoning-parser", "deepseek_r1", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestReasoningContentAPI:
"""Tests for reasoning content API with DeepSeek R1 reasoning parser."""
def test_streaming_separate_reasoning_false(self, setup_backend):
"""Test streaming with separate_reasoning=False, reasoning_content should be empty."""
_, model, client, gateway = setup_backend
response = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": "What is 1+3?",
}
],
max_tokens=100,
stream=True,
extra_body={"separate_reasoning": False},
)
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, setup_backend):
"""Test streaming with separate_reasoning=True, reasoning_content should not be empty."""
_, model, client, gateway = setup_backend
response = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": "What is 1+3?",
}
],
max_tokens=100,
stream=True,
extra_body={"separate_reasoning": True},
)
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, setup_backend
):
"""Test streaming with separate_reasoning=True and stream_reasoning=False."""
_, model, client, gateway = setup_backend
response = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": "What is 1+3?",
}
],
max_tokens=100,
stream=True,
extra_body={"separate_reasoning": True, "stream_reasoning": False},
)
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, setup_backend):
"""Test non-streaming with separate_reasoning=False, reasoning_content should be empty."""
_, model, client, gateway = setup_backend
response = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": "What is 1+3?",
}
],
max_tokens=100,
extra_body={"separate_reasoning": False},
)
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, setup_backend):
"""Test non-streaming with separate_reasoning=True, reasoning_content should not be empty."""
_, model, client, gateway = setup_backend
response = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": "What is 1+3?",
}
],
max_tokens=100,
extra_body={"separate_reasoning": True},
)
assert len(response.choices[0].message.reasoning_content) > 0
assert len(response.choices[0].message.content) > 0

View File

@@ -0,0 +1,167 @@
"""Validation E2E Tests.
Tests for validation features like ignore_eos and large token handling.
Source: Migrated from e2e_grpc/validation/test_openai_server_ignore_eos.py
and e2e_grpc/validation/test_large_max_new_tokens.py
"""
from __future__ import annotations
import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor
import pytest
logger = logging.getLogger(__name__)
# Lazy load tokenizer to avoid import errors if transformers not installed
_tokenizer_cache: dict = {}
_tokenizer_lock = threading.Lock()
def get_tokenizer(model_path: str):
"""Get tokenizer for a model, with caching."""
if model_path not in _tokenizer_cache:
with _tokenizer_lock:
# Re-check after acquiring the lock to handle race conditions
if model_path not in _tokenizer_cache:
from transformers import AutoTokenizer
_tokenizer_cache[model_path] = AutoTokenizer.from_pretrained(model_path)
return _tokenizer_cache[model_path]
# =============================================================================
# Ignore EOS Tests (Llama 8B)
# =============================================================================
@pytest.mark.model("llama-8b")
@pytest.mark.gateway(extra_args=["--history-backend", "memory"])
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestIgnoreEOS:
"""Tests for ignore_eos feature."""
def test_ignore_eos(self, setup_backend):
"""Test that ignore_eos=True allows generation to continue beyond EOS token.
When ignore_eos=True, the model should generate until max_tokens is reached,
even if it encounters an EOS token.
"""
_, model, client, _ = setup_backend
tokenizer = get_tokenizer(model)
max_tokens = 200
# Request without ignore_eos (default behavior - stops at EOS)
response_default = client.chat.completions.create(
model=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},
)
# Request with ignore_eos=True (continues past EOS until max_tokens)
response_ignore_eos = client.chat.completions.create(
model=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(
tokenizer.encode(response_default.choices[0].message.content)
)
ignore_eos_tokens = len(
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)
assert (
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}"
assert response_ignore_eos.choices[0].finish_reason == "length", (
f"Expected finish_reason='length' for ignore_eos=True, "
f"got {response_ignore_eos.choices[0].finish_reason}"
)
# =============================================================================
# Large Max New Tokens Tests (Llama 8B)
#
# NOTE: This test verifies concurrent request handling with large token limits.
# The original test monitored server logs to verify concurrency, which is not
# possible with the pool-based infrastructure. This simplified version verifies
# that concurrent requests complete successfully.
# =============================================================================
@pytest.mark.model("llama-8b")
@pytest.mark.gateway(extra_args=["--history-backend", "memory"])
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestLargeMaxNewTokens:
"""Tests for handling large max_new_tokens with concurrent requests."""
def test_concurrent_chat_completions(self, setup_backend):
"""Test that multiple concurrent requests with large token generation complete.
This test sends multiple requests that ask for long outputs concurrently
to verify the server can handle concurrent long-running requests.
"""
_, model, client, _ = setup_backend
num_requests = 4
def run_chat_completion():
response = client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant"},
{
"role": "user",
"content": "Please repeat the word 'hello' for 100 times.",
},
],
temperature=0,
max_tokens=256, # Reasonable limit for concurrent test
)
return response
# Send concurrent requests
start_time = time.time()
futures = []
with ThreadPoolExecutor(max_workers=num_requests) as executor:
for _ in range(num_requests):
futures.append(executor.submit(run_chat_completion))
# Wait for all to complete and collect results
responses = [f.result() for f in futures]
elapsed = time.time() - start_time
logger.info("Completed %d concurrent requests in %.2fs", num_requests, elapsed)
# Verify all requests completed successfully
assert len(responses) == num_requests
for i, response in enumerate(responses):
assert response.choices[
0
].message.content, f"Request {i} returned empty content"
assert response.choices[0].finish_reason in ("stop", "length"), (
f"Request {i} had unexpected finish_reason: "
f"{response.choices[0].finish_reason}"
)

View File

@@ -0,0 +1,225 @@
"""Pytest configuration for E2E tests.
Parallel Execution
------------------
Tests can run in parallel using pytest-parallel with shared worker processes.
Use --workers 1 --tests-per-worker N for N concurrent test threads:
pytest --workers 1 --tests-per-worker 4 e2e_test/router/
This leverages the thread-safe ModelPool and GPUAllocator classes to enable
true shared-worker parallelism where all threads share the same session-scoped
model_pool fixture. Tests marked with @pytest.mark.thread_unsafe will be
automatically skipped in parallel mode.
Markers
-------
This module defines several pytest markers for configuring E2E tests:
@pytest.mark.model(name)
Specify which model to use for the test.
Args:
name: Model ID from MODEL_SPECS (e.g., "llama-8b", "qwen-7b")
GPU Resource Management:
When GPUs are limited (e.g., 4 GPUs, 6 models), the model pool uses
MRU (Most Recently Used) eviction:
1. Models are pre-launched until GPUs are full
2. When a test needs a model that isn't running, MRU model is evicted
(models just used are likely done, models not yet used are waiting)
3. The needed model is then launched on-demand
Examples:
@pytest.mark.model("llama-8b")
@pytest.mark.model("qwen-72b")
@pytest.mark.workers(count=1, prefill=None, decode=None)
Configure worker topology for the test.
Args:
count: Number of regular workers (default: 1)
prefill: Number of prefill workers for PD disaggregation
decode: Number of decode workers for PD disaggregation
Examples:
@pytest.mark.workers(count=3) # 3 regular workers
@pytest.mark.workers(prefill=2, decode=2) # PD mode
@pytest.mark.gateway(policy="round_robin", timeout=None, extra_args=None)
Configure the gateway/router.
Args:
policy: Routing policy ("round_robin", "random", etc.)
timeout: Startup timeout in seconds
extra_args: Additional CLI arguments for the router
Examples:
@pytest.mark.gateway(policy="random")
@pytest.mark.gateway(extra_args=["--cache-routing"])
@pytest.mark.e2e
Mark test as an end-to-end test requiring GPU workers.
@pytest.mark.slow
Mark test as slow-running.
@pytest.mark.thread_unsafe(reason=None)
Mark test as incompatible with parallel thread execution.
Tests with this marker are automatically skipped when running
with --tests-per-worker > 1.
Args:
reason: Optional explanation of why the test is thread-unsafe.
Examples:
@pytest.mark.thread_unsafe
@pytest.mark.thread_unsafe(reason="Modifies global state")
Fixtures
--------
model_pool: Session-scoped fixture managing SGLang worker processes.
setup_backend: Class-scoped fixture that launches gateway + provides client.
Usage Examples
--------------
Basic test with default model:
@pytest.mark.e2e
@pytest.mark.parametrize("setup_backend", ["http"], indirect=True)
class TestBasic:
def test_chat(self, setup_backend):
backend, model, client, gateway = setup_backend
response = client.chat.completions.create(...)
Test with specific model and multiple backends:
@pytest.mark.e2e
@pytest.mark.model("qwen-7b")
@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True)
class TestQwen:
def test_generate(self, setup_backend):
...
PD disaggregation mode:
@pytest.mark.e2e
@pytest.mark.workers(prefill=1, decode=1)
@pytest.mark.parametrize("setup_backend", ["pd"], indirect=True)
class TestPD:
def test_pd_inference(self, setup_backend):
...
"""
from __future__ import annotations
import logging
import sys
from importlib.util import find_spec
from pathlib import Path
# ---------------------------------------------------------------------------
# Path setup (must happen before other imports)
# ---------------------------------------------------------------------------
_ROOT = Path(__file__).resolve().parents[1] # sgl-model-gateway/
_E2E_TEST = Path(__file__).resolve().parent # e2e_test/
_SRC = _ROOT / "bindings" / "python"
# Add e2e_test to path so "from infra import ..." works
if str(_E2E_TEST) not in sys.path:
sys.path.insert(0, str(_E2E_TEST))
# Add bindings/python to path if the wheel is not installed (for local development)
_wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None
if not _wheel_installed and str(_SRC) not in sys.path:
sys.path.insert(0, str(_SRC))
# ---------------------------------------------------------------------------
# Logging setup (clean output without pytest's "---- live log ----" dividers)
# ---------------------------------------------------------------------------
def _setup_logging() -> None:
"""Configure clean logging to stdout with timestamps and thread info.
In parallel mode (--tests-per-worker > 1), logs from different threads
would be interleaved. Including thread name helps identify which test
produced each log line.
"""
# Include thread name for parallel execution readability
# MainThread for sequential, Thread-N for parallel workers
fmt = "%(asctime)s.%(msecs)03d [%(threadName)s] [%(name)s] %(message)s"
datefmt = "%H:%M:%S"
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(logging.Formatter(fmt, datefmt))
for logger_name in ("e2e_test", "infra", "fixtures"):
log = logging.getLogger(logger_name)
log.setLevel(logging.INFO)
log.addHandler(handler)
log.propagate = False
for logger_name in ("openai", "httpx", "httpcore", "numexpr"):
logging.getLogger(logger_name).setLevel(logging.WARNING)
_setup_logging()
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Test visibility hooks
# ---------------------------------------------------------------------------
def pytest_runtest_logstart(nodeid: str, location: tuple) -> None:
"""Print clear test header at start of each test."""
import threading
from infra import LOG_SEPARATOR_WIDTH
test_name = nodeid.split("::")[-1] if "::" in nodeid else nodeid
thread_name = threading.current_thread().name
print(f"\n{'=' * LOG_SEPARATOR_WIDTH}")
print(f"[{thread_name}] TEST: {test_name}")
print(f"{'=' * LOG_SEPARATOR_WIDTH}")
# ---------------------------------------------------------------------------
# Import pytest hooks and fixtures from fixtures/ package
# ---------------------------------------------------------------------------
# Import fixtures - pytest discovers these by name
# Import hooks - pytest discovers these by name
from fixtures import (
backend_router,
model_base_url,
model_client,
model_pool,
pytest_collection_finish,
pytest_collection_modifyitems,
pytest_configure,
pytest_runtest_setup,
setup_backend,
)
# Re-export for pytest discovery
__all__ = [
# Hooks
"pytest_runtest_logstart",
"pytest_collection_modifyitems",
"pytest_collection_finish",
"pytest_configure",
"pytest_runtest_setup",
# Fixtures
"model_pool",
"model_client",
"model_base_url",
"setup_backend",
"backend_router",
]

View File

@@ -0,0 +1,143 @@
"""Basic embedding API tests.
Tests the embedding functionality through the router with both gRPC and HTTP backends.
Source: Migrated from e2e_grpc/basic/test_embedding_server.py
Usage:
pytest e2e_test/embeddings/test_basic.py -v
pytest e2e_test/embeddings/test_basic.py -v -k "grpc"
"""
from __future__ import annotations
import logging
import pytest
logger = logging.getLogger(__name__)
@pytest.mark.e2e
@pytest.mark.model("embedding")
@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True)
class TestEmbeddingBasic:
"""Basic embedding API tests using local workers (gRPC and HTTP)."""
def test_embedding_single(self, setup_backend):
"""Test single text embedding.
Verifies that:
- Response object structure is correct
- Embedding is a non-empty list of floats
- Usage statistics are present
"""
backend, model, client, gateway = setup_backend
input_text = "Hello world"
response = client.embeddings.create(
model=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
logger.info(
"Single embedding: %d dimensions, %d tokens",
len(embedding.embedding),
response.usage.prompt_tokens,
)
def test_embedding_batch(self, setup_backend):
"""Test batch embedding with multiple texts.
Note: The original test expected len(response.data) == 1 for batch,
which seems incorrect. This might be model-specific behavior.
"""
backend, model, client, gateway = setup_backend
input_texts = ["Hello world", "SGLang is fast"]
response = client.embeddings.create(
model=model,
input=input_texts,
)
# Note: Original test had len(response.data) == 1, which seems like
# a bug or model-specific behavior. Standard behavior should return
# one embedding per input text.
assert len(response.data) >= 1
assert response.data[0].index == 0
assert len(response.data[0].embedding) > 0
logger.info("Batch embedding: %d results", len(response.data))
def test_embedding_dimensions_consistent(self, setup_backend):
"""Test that embedding dimensions are consistent across different inputs.
Verifies that different length inputs produce embeddings with
the same dimensionality.
"""
backend, model, client, gateway = setup_backend
response1 = client.embeddings.create(
model=model,
input="A short text",
)
dim1 = len(response1.data[0].embedding)
response2 = client.embeddings.create(
model=model,
input="A much longer text to ensure dimensions match regardless of input length",
)
dim2 = len(response2.data[0].embedding)
assert dim1 == dim2, f"Dimensions differ: {dim1} vs {dim2}"
logger.info("Embedding dimensions: %d (consistent)", dim1)
def test_embedding_empty_string(self, setup_backend):
"""Test embedding with empty string input.
Some models may handle empty strings differently.
This test verifies the API doesn't crash on empty input.
"""
backend, model, client, gateway = setup_backend
try:
response = client.embeddings.create(
model=model,
input="",
)
# If it succeeds, verify structure
assert len(response.data) >= 1
logger.info("Empty string embedding succeeded")
except Exception as e:
# Some models may reject empty strings - that's acceptable
logger.info("Empty string embedding rejected: %s", e)
def test_embedding_unicode(self, setup_backend):
"""Test embedding with unicode characters.
Verifies that the API handles non-ASCII characters correctly.
"""
backend, model, client, gateway = setup_backend
input_text = "Hello 世界! 🚀 Привет мир"
response = client.embeddings.create(
model=model,
input=input_text,
)
assert len(response.data) == 1
assert len(response.data[0].embedding) > 0
logger.info("Unicode embedding: %d dimensions", len(response.data[0].embedding))

View File

@@ -0,0 +1,262 @@
"""Embedding correctness tests.
Tests that embeddings from the router match HuggingFace reference embeddings.
Validates numerical correctness including tokenization and inference.
Source: Migrated from e2e_grpc/basic/test_embedding_correctness.py
Usage:
pytest e2e_test/embeddings/test_correctness.py -v
pytest e2e_test/embeddings/test_correctness.py -v -k "grpc"
Requirements:
- sentence-transformers (for reference embeddings)
- torch
- numpy
"""
from __future__ import annotations
import logging
import threading
from typing import Any
import numpy as np
import pytest
import torch
import torch.nn.functional as F
logger = logging.getLogger(__name__)
# Thread-safe storage for HF reference embeddings
_hf_embeddings_cache: dict[str, Any] | None = None
_hf_embeddings_lock = threading.Lock()
# 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[str, Any] = {
"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.",
},
{
"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.",
},
{
"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's 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.",
},
{
"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.",
},
],
}
def get_openai_embeddings(
texts: str | list[str],
client,
model: str,
) -> list[list[float]]:
"""Get embeddings from the gateway via OpenAI-compatible API."""
if isinstance(texts, str):
texts = [texts]
embeddings = []
for text in texts:
response = client.embeddings.create(
model=model,
input=text,
)
embeddings.append(response.data[0].embedding)
return embeddings
def get_hf_st_embeddings(texts: 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.
"""
from sentence_transformers import SentenceTransformer
if isinstance(texts, str):
texts = [texts]
# Force CPU to avoid GPU memory conflicts in CI
model = SentenceTransformer(model_path, trust_remote_code=True, device="cpu")
embeddings = model.encode(texts, normalize_embeddings=True)
return embeddings
def compare_embeddings(
embeddings1: list[list[float]], embeddings2: list[list[float]]
) -> list[float]:
"""Compare two sets of embeddings using cosine similarity."""
similarities = [
F.cosine_similarity(torch.tensor(e1), torch.tensor(e2), dim=0).item()
for e1, e2 in zip(embeddings1, embeddings2)
]
return similarities
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"]]
@pytest.fixture(scope="session")
def hf_reference_embeddings(request):
"""Pre-compute HuggingFace reference embeddings on CPU.
This is done once per session with thread-safe initialization to support
pytest-parallel execution. Uses CPU to avoid GPU memory conflicts.
"""
global _hf_embeddings_cache
# Thread-safe initialization - only one thread computes embeddings
with _hf_embeddings_lock:
if _hf_embeddings_cache is not None:
return _hf_embeddings_cache
from infra.model_specs import MODEL_SPECS
# Get model path from MODEL_SPECS for the embedding model
model_path = MODEL_SPECS.get("embedding", {}).get("model")
if model_path is None:
pytest.skip("Embedding model not found in MODEL_SPECS")
logger.info(
"Pre-computing HuggingFace reference embeddings (CPU) for %s", model_path
)
# 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
hf_semantic = get_hf_st_embeddings(all_semantic_texts, model_path)
hf_query = get_hf_st_embeddings(query, model_path)
hf_docs = get_hf_st_embeddings(docs, model_path)
logger.info("Reference embeddings computed on CPU")
_hf_embeddings_cache = {
"semantic": hf_semantic,
"query": hf_query,
"docs": hf_docs,
}
return _hf_embeddings_cache
@pytest.mark.e2e
@pytest.mark.model("embedding")
@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True)
class TestEmbeddingCorrectness:
"""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.
"""
def test_semantic_similarity(self, setup_backend, hf_reference_embeddings):
"""Check if gateway and HF embeddings give similar results.
For each text in the semantic test sets, the gateway embedding should
have >0.98 cosine similarity with the HuggingFace reference embedding.
"""
backend, model_path, client, gateway = setup_backend
tolerance = 1e-2
# Track position in pre-computed embeddings
embed_idx = 0
for i, input_texts in enumerate(SEMANTIC_TEST_SETS):
logger.info("Processing semantic similarity test set %d", i + 1)
embedding_gateway = get_openai_embeddings(input_texts, client, model_path)
# Get pre-computed HF embeddings for this set
num_texts = len(input_texts)
embedding_hf = hf_reference_embeddings["semantic"][
embed_idx : embed_idx + num_texts
].tolist()
embed_idx += num_texts
similarities = compare_embeddings(embedding_gateway, embedding_hf)
logger.info("Similarities: %s", similarities)
# Verify all similarities are close to 1.0
for j, sim in enumerate(similarities):
assert (
abs(sim - 1.0) < tolerance
), f"Set {i+1}, text {j+1}: similarity {sim:.4f} not close to 1.0"
logger.info("Semantic similarity test set %d passed", i + 1)
def test_relevance_scores(self, setup_backend, hf_reference_embeddings):
"""Compare relevance scores between gateway and HF implementations.
The relevance scores (query @ docs) should match between the gateway
and HuggingFace implementations within tolerance.
"""
backend, model_path, client, gateway = setup_backend
tolerance = 0.05
# 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, client, model_path)
docs_embeddings_gateway = get_openai_embeddings(docs, client, model_path)
scores_gateway = (
np.array(query_embeddings_gateway) @ np.array(docs_embeddings_gateway).T
) * 100
# Use pre-computed HF scores
scores_hf = (
hf_reference_embeddings["query"] @ hf_reference_embeddings["docs"].T
) * 100
logger.info("Gateway relevance scores: %s", scores_gateway)
logger.info("HF relevance scores: %s", scores_hf)
assert np.allclose(
scores_gateway, scores_hf, atol=tolerance
), f"Scores differ beyond tolerance:\nGateway: {scores_gateway}\nHF: {scores_hf}"
logger.info("Relevance scores comparison passed")

View File

@@ -0,0 +1,51 @@
"""Fixtures for E2E tests.
This package contains modular pytest fixtures split by responsibility:
- hooks.py: Pytest collection hooks and marker registration
- pool.py: Model pool fixtures (session-scoped worker management)
- setup_backend.py: Backend setup fixtures (class/function-scoped)
- markers.py: Helper utilities for marker extraction
Legacy modules (to be removed during e2e_response_api migration):
- ports.py: Use infra.get_open_port() instead
- router_manager.py: Use infra.Gateway instead
"""
# Pytest hooks (imported by conftest.py via pytest_plugins)
from .hooks import (
get_pool_requirements,
is_parallel_execution,
pytest_collection_finish,
pytest_collection_modifyitems,
pytest_configure,
pytest_runtest_setup,
validate_gpu_requirements,
)
# Marker helpers
from .markers import get_marker_kwargs, get_marker_value
# Fixtures (imported by conftest.py)
from .pool import model_base_url, model_client, model_pool
from .setup_backend import backend_router, setup_backend
__all__ = [
# Hooks
"pytest_collection_modifyitems",
"pytest_collection_finish",
"pytest_configure",
"pytest_runtest_setup",
"get_pool_requirements",
"validate_gpu_requirements",
"is_parallel_execution",
# Pool fixtures
"model_pool",
"model_client",
"model_base_url",
# Backend fixtures
"setup_backend",
"backend_router",
# Marker helpers
"get_marker_value",
"get_marker_kwargs",
]

View File

@@ -0,0 +1,430 @@
"""Pytest hooks for E2E test collection and validation.
This module handles:
- Test collection: Scanning markers to determine required workers
- GPU validation: Ensuring sufficient GPUs for test requirements
- Marker registration: Defining custom pytest markers
"""
from __future__ import annotations
import logging
import os
from typing import TYPE_CHECKING
import pytest
if TYPE_CHECKING:
from infra import ConnectionMode, WorkerIdentity, WorkerType
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Test collection state
# ---------------------------------------------------------------------------
# Track max worker counts: (model_id, mode, worker_type) -> max_count
_worker_counts: dict[tuple["ConnectionMode", "WorkerType"], int] = {}
# Track first-seen order to preserve test collection order
_first_seen_order: list[tuple[str, "ConnectionMode", "WorkerType"]] = []
# Track max GPU requirement for any single test (for validation)
_max_test_gpu_requirement: int = 0
_max_test_name: str = ""
_needs_default_model: bool = False
def reset_collection_state() -> None:
"""Reset collection state (useful for testing)."""
global _worker_counts, _first_seen_order
global _max_test_gpu_requirement, _max_test_name, _needs_default_model
_worker_counts = {}
_first_seen_order = []
_max_test_gpu_requirement = 0
_max_test_name = ""
_needs_default_model = False
def get_worker_counts() -> dict:
"""Get the worker counts dictionary."""
return _worker_counts
def get_first_seen_order() -> list:
"""Get the first-seen order list."""
return _first_seen_order
def get_max_gpu_requirement() -> tuple[int, str]:
"""Get the max GPU requirement and test name."""
return _max_test_gpu_requirement, _max_test_name
def needs_default_model() -> bool:
"""Check if any test needs the default model."""
return _needs_default_model
# ---------------------------------------------------------------------------
# Test collection hook
# ---------------------------------------------------------------------------
def pytest_collection_modifyitems(
session: pytest.Session,
config: pytest.Config,
items: list[pytest.Item],
) -> None:
"""Scan collected tests to determine required workers.
This runs after test collection but before tests execute.
It extracts worker requirements from markers in test collection order,
tracking the max count needed for each (model, mode, worker_type) combination.
"""
global _worker_counts, _first_seen_order, _needs_default_model
global _max_test_gpu_requirement, _max_test_name
from infra import (
DEFAULT_MODEL,
LOG_SEPARATOR_WIDTH,
MODEL_SPECS,
PARAM_MODEL,
PARAM_SETUP_BACKEND,
ConnectionMode,
WorkerType,
)
def track_worker(
model_id: str, mode: ConnectionMode, worker_type: WorkerType, count: int
) -> None:
"""Track a worker requirement, updating max count if needed."""
key = (model_id, mode, worker_type)
if key not in _worker_counts:
_first_seen_order.append(key)
_worker_counts[key] = count
else:
_worker_counts[key] = max(_worker_counts[key], count)
def calculate_test_gpus(
model_id: str, prefill: int, decode: int, regular: int
) -> int:
"""Calculate GPU requirement for a single test."""
if model_id not in MODEL_SPECS:
return 0
tp = MODEL_SPECS[model_id].get("tp", 1)
return tp * (prefill + decode + regular)
for item in items:
# Extract model from marker or use default
# First check the class directly (handles inheritance correctly)
model_id = None
if hasattr(item, "cls") and item.cls is not None:
for marker in (
item.cls.pytestmark if hasattr(item.cls, "pytestmark") else []
):
if marker.name == PARAM_MODEL and marker.args:
model_id = marker.args[0]
break
# Fall back to get_closest_marker for method-level markers
if model_id is None:
model_marker = item.get_closest_marker(PARAM_MODEL)
model_id = (
model_marker.args[0] if model_marker and model_marker.args else None
)
# Check parametrize for model
if model_id is None:
for marker in item.iter_markers("parametrize"):
if marker.args and len(marker.args) >= 2:
param_name = marker.args[0]
if param_name == PARAM_MODEL or PARAM_MODEL in param_name:
param_values = marker.args[1]
if isinstance(param_values, (list, tuple)) and param_values:
model_id = param_values[0]
break
# Extract backends from parametrize
backends: list[str] = []
for marker in item.iter_markers("parametrize"):
if marker.args and len(marker.args) >= 2:
param_name = marker.args[0]
param_values = marker.args[1]
if param_name == PARAM_SETUP_BACKEND:
if isinstance(param_values, (list, tuple)):
backends.extend(param_values)
# Check for workers marker
workers_marker = item.get_closest_marker("workers")
prefill_count = 0
decode_count = 0
regular_count = 1
if workers_marker:
prefill_count = workers_marker.kwargs.get("prefill") or 0
decode_count = workers_marker.kwargs.get("decode") or 0
regular_count = workers_marker.kwargs.get("count") or 1
# Track if this test needs default model
is_e2e = item.get_closest_marker("e2e") is not None
if model_id is None and is_e2e:
_needs_default_model = True
model_id = DEFAULT_MODEL
# Track worker requirements
test_gpus = 0
if model_id and backends:
for backend in backends:
if backend == "pd":
mode = ConnectionMode.HTTP
p_count = prefill_count if prefill_count > 0 else 1
d_count = decode_count if decode_count > 0 else 1
track_worker(model_id, mode, WorkerType.PREFILL, p_count)
track_worker(model_id, mode, WorkerType.DECODE, d_count)
test_gpus = max(
test_gpus, calculate_test_gpus(model_id, p_count, d_count, 0)
)
else:
try:
mode = ConnectionMode(backend)
except ValueError:
continue
if prefill_count > 0 or decode_count > 0:
track_worker(model_id, mode, WorkerType.PREFILL, prefill_count)
track_worker(model_id, mode, WorkerType.DECODE, decode_count)
test_gpus = max(
test_gpus,
calculate_test_gpus(
model_id, prefill_count, decode_count, 0
),
)
else:
track_worker(model_id, mode, WorkerType.REGULAR, regular_count)
test_gpus = max(
test_gpus,
calculate_test_gpus(model_id, 0, 0, regular_count),
)
elif model_id and is_e2e:
track_worker(model_id, ConnectionMode.HTTP, WorkerType.REGULAR, 1)
test_gpus = calculate_test_gpus(model_id, 0, 0, 1)
if test_gpus > _max_test_gpu_requirement:
_max_test_gpu_requirement = test_gpus
_max_test_name = item.nodeid
# Log results
if _worker_counts:
summary = []
for key in _first_seen_order:
model_id, mode, worker_type = key
count = _worker_counts[key]
if worker_type == WorkerType.REGULAR:
summary.append(f"{model_id}:{mode.value}x{count}")
else:
summary.append(f"{model_id}:{mode.value}:{worker_type.value}x{count}")
logger.info("Scanned worker requirements (in test order): %s", summary)
logger.info(
"Max GPU requirement for single test: %d (%s)",
_max_test_gpu_requirement,
_max_test_name,
)
else:
logger.info("Scanned worker requirements: (none)")
# ---------------------------------------------------------------------------
# Pool requirements
# ---------------------------------------------------------------------------
def get_pool_requirements() -> list["WorkerIdentity"]:
"""Build pool requirements from scanned test markers.
Returns:
List of WorkerIdentity objects to pre-launch.
"""
from infra import DEFAULT_MODEL, ConnectionMode, WorkerIdentity, WorkerType
# Track which models have PD workers as their first requirement
models_with_pd_first: set[str] = set()
first_worker_type_per_model: dict[str, WorkerType] = {}
for model_id, mode, worker_type in _first_seen_order:
if model_id not in first_worker_type_per_model:
first_worker_type_per_model[model_id] = worker_type
if worker_type in (WorkerType.PREFILL, WorkerType.DECODE):
models_with_pd_first.add(model_id)
logger.info(
"Model %s has PD test first - skipping regular worker pre-launch",
model_id,
)
# Generate individual WorkerIdentity objects in first-seen order
requirements: list[WorkerIdentity] = []
for model_id, mode, worker_type in _first_seen_order:
if model_id in models_with_pd_first and worker_type == WorkerType.REGULAR:
continue
count = _worker_counts.get((model_id, mode, worker_type), 1)
for i in range(count):
requirements.append(WorkerIdentity(model_id, mode, worker_type, i))
if not requirements:
requirements.append(WorkerIdentity(DEFAULT_MODEL, ConnectionMode.HTTP))
return requirements
# ---------------------------------------------------------------------------
# GPU validation
# ---------------------------------------------------------------------------
def _count_gpus_without_cuda() -> int:
"""Count available GPUs without initializing CUDA.
Uses nvidia-smi to avoid CUDA initialization, which is critical for
pytest-parallel compatibility. CUDA cannot be re-initialized after a fork.
"""
import subprocess
try:
result = subprocess.run(
["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
capture_output=True,
text=True,
timeout=10,
)
if result.returncode == 0:
return len([line for line in result.stdout.strip().split("\n") if line])
except (subprocess.SubprocessError, FileNotFoundError, OSError):
pass
return 0
def validate_gpu_requirements() -> tuple[int, int]:
"""Check if there are enough GPUs for any single test.
Uses nvidia-smi instead of torch.cuda to avoid CUDA initialization,
which would break pytest-parallel (CUDA cannot be re-initialized after fork).
Returns:
Tuple of (max_required_gpus, available_gpus).
"""
available_gpus = _count_gpus_without_cuda()
return _max_test_gpu_requirement, available_gpus
def pytest_collection_finish(session: pytest.Session) -> None:
"""Validate GPU requirements after test collection."""
from infra import ENV_SKIP_MODEL_POOL, LOG_SEPARATOR_WIDTH
if not _worker_counts:
return
if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
return
max_required, available_gpus = validate_gpu_requirements()
if max_required > available_gpus:
sep = "=" * LOG_SEPARATOR_WIDTH
raise pytest.UsageError(
f"\n{sep}\n"
f"GPU REQUIREMENTS EXCEEDED\n"
f"{sep}\n"
f"Test '{_max_test_name}' requires {max_required} GPUs\n"
f"Available: {available_gpus} GPUs\n"
f"\nOptions:\n"
f" 1. Run tests that fit: pytest -k 'not {_max_test_name.split('::')[0]}'\n"
f" 2. Reduce workers: @pytest.mark.workers(prefill=1, decode=1)\n"
f" 3. Skip GPU tests: SKIP_MODEL_POOL=1 pytest\n"
f"{sep}"
)
logger.info(
"GPU validation passed: max %d required (by %s), %d available",
max_required,
_max_test_name,
available_gpus,
)
# ---------------------------------------------------------------------------
# Marker registration
# ---------------------------------------------------------------------------
def pytest_configure(config: pytest.Config) -> None:
"""Register custom markers."""
config.addinivalue_line(
"markers",
"model(name): mark test to use a specific model from MODEL_SPECS",
)
config.addinivalue_line(
"markers",
"backend(name): mark test to use a specific backend (grpc, http, openai, etc.)",
)
config.addinivalue_line(
"markers",
"workers(count=1, prefill=None, decode=None): "
"worker configuration - use count for regular workers, "
"or prefill/decode for PD disaggregation mode",
)
config.addinivalue_line(
"markers",
"gateway(policy='round_robin', timeout=None, extra_args=None): "
"gateway/router configuration",
)
config.addinivalue_line(
"markers",
"e2e: mark test as an end-to-end test requiring GPU workers",
)
config.addinivalue_line(
"markers",
"slow: mark test as slow-running",
)
config.addinivalue_line(
"markers",
"thread_unsafe: mark test as incompatible with parallel thread execution",
)
config.addinivalue_line(
"markers",
"storage(backend): mark test to use a specific history storage backend "
"(memory, oracle). Default is memory.",
)
# ---------------------------------------------------------------------------
# Parallel execution support
# ---------------------------------------------------------------------------
def is_parallel_execution(config: pytest.Config) -> bool:
"""Check if tests are running in parallel mode (pytest-parallel).
Returns True if --tests-per-worker > 1, indicating concurrent thread execution.
"""
# pytest-parallel adds the 'tests_per_worker' option
tests_per_worker = getattr(config.option, "tests_per_worker", None)
if tests_per_worker is None:
return False
if tests_per_worker == "auto":
return True
try:
return int(tests_per_worker) > 1
except (ValueError, TypeError):
return False
def pytest_runtest_setup(item: pytest.Item) -> None:
"""Skip thread_unsafe tests when running in parallel mode."""
if is_parallel_execution(item.config):
marker = item.get_closest_marker("thread_unsafe")
if marker:
reason = marker.kwargs.get("reason", "Test is not thread-safe")
pytest.skip(f"Skipping in parallel mode: {reason}")

View File

@@ -0,0 +1,57 @@
"""Marker helper utilities for E2E tests.
This module provides helper functions for extracting values from pytest markers.
"""
from __future__ import annotations
from typing import Any
import pytest
def get_marker_value(
request: pytest.FixtureRequest,
marker_name: str,
arg_index: int = 0,
default: Any = None,
) -> Any:
"""Get a value from a pytest marker.
Args:
request: The pytest fixture request.
marker_name: Name of the marker to look for.
arg_index: Index of positional argument to extract.
default: Default value if marker not found.
Returns:
The marker argument value or default.
"""
marker = request.node.get_closest_marker(marker_name)
if marker is None:
return default
if marker.args and len(marker.args) > arg_index:
return marker.args[arg_index]
return default
def get_marker_kwargs(
request: pytest.FixtureRequest,
marker_name: str,
defaults: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Get keyword arguments from a pytest marker.
Args:
request: The pytest fixture request.
marker_name: Name of the marker to look for.
defaults: Default values if marker not found or missing kwargs.
Returns:
Dict of keyword arguments merged with defaults.
"""
result = dict(defaults) if defaults else {}
marker = request.node.get_closest_marker(marker_name)
if marker is not None:
result.update(marker.kwargs)
return result

View File

@@ -0,0 +1,242 @@
"""Model pool fixtures for E2E tests.
This module provides session-scoped fixtures for managing SGLang worker processes.
Workers are expensive to start (~30-60s each), so they're kept running across tests.
"""
from __future__ import annotations
import atexit
import logging
import os
import threading
from typing import TYPE_CHECKING
import pytest
if TYPE_CHECKING:
from infra import ModelPool
from .hooks import get_pool_requirements
logger = logging.getLogger(__name__)
# Global model pool instance with thread-safe initialization
_model_pool: "ModelPool | None" = None
_model_pool_lock = threading.Lock()
_shutdown_registered = False
def _shutdown_model_pool() -> None:
"""Shutdown the global model pool at process exit.
This is registered with atexit to ensure cleanup happens after all tests
complete, which is important for pytest-parallel where multiple threads
share the session-scoped fixture.
"""
global _model_pool
if _model_pool is not None:
logger.info("Shutting down model pool at process exit")
_model_pool.shutdown()
_model_pool = None
@pytest.fixture(scope="session")
def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
"""Session-scoped fixture that manages SGLang worker processes.
Workers (sglang.launch_server) are expensive to start (~30-60s each due to
model loading). This fixture starts them ONCE per session and keeps them
running across all tests. The setup_backend fixture then launches cheap
routers (~1-2s) pointing to these workers.
Startup behavior:
- Scans test markers to determine required workers (model, mode, type, count)
- Launches workers in test collection order
- Waits for all workers to become healthy before returning
Test requirements are auto-detected from:
- @pytest.mark.parametrize("setup_backend", ["grpc", "http", "pd"])
- @pytest.mark.model("model-name")
- @pytest.mark.workers(count=N) for regular workers
- @pytest.mark.workers(prefill=N, decode=N) for PD workers
Environment variable overrides:
- E2E_MODELS: Comma-separated model IDs (e.g., "llama-8b,qwen-7b")
- E2E_BACKENDS: Comma-separated backends (e.g., "grpc,http")
- SKIP_MODEL_POOL: Set to "1" to skip worker startup
"""
global _model_pool
from infra import (
DEFAULT_MODEL,
ENV_BACKENDS,
ENV_MODELS,
ENV_SKIP_MODEL_POOL,
ENV_STARTUP_TIMEOUT,
LOCAL_MODES,
MODEL_SPECS,
ConnectionMode,
GPUAllocator,
ModelPool,
WorkerIdentity,
WorkerType,
)
# Thread-safe initialization: use lock to ensure only one thread creates the pool
# This is critical for pytest-parallel which runs tests as concurrent threads
with _model_pool_lock:
if _model_pool is not None:
return _model_pool
# Check if we should skip model startup
if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
logger.info("%s is set, skipping model pool startup", ENV_SKIP_MODEL_POOL)
_model_pool = ModelPool(GPUAllocator(gpus=[]))
return _model_pool
# Determine requirements from scanned tests or env vars
models_env = os.environ.get(ENV_MODELS, "")
backends_env = os.environ.get(ENV_BACKENDS, "")
if models_env or backends_env:
# Use env var overrides
models = (
{m.strip() for m in models_env.split(",") if m.strip()}
if models_env
else {DEFAULT_MODEL}
)
# Parse backend strings to ConnectionMode enums
backend_modes: set[ConnectionMode] = set()
if backends_env:
for b in backends_env.split(","):
b = b.strip()
if b:
try:
mode = ConnectionMode(b)
if mode in LOCAL_MODES:
backend_modes.add(mode)
except ValueError:
logger.warning("Unknown backend '%s', skipping", b)
# Default to HTTP if no valid backends
if not backend_modes:
backend_modes = {ConnectionMode.HTTP}
# Create WorkerIdentity objects (regular workers only from env vars)
requirements = [
WorkerIdentity(m, b, WorkerType.REGULAR, 0)
for m in models
for b in backend_modes
]
logger.info(
"Using env var requirements: %s", [str(r) for r in requirements]
)
else:
# Use scanned requirements from test markers
requirements = get_pool_requirements()
logger.info(
"Using scanned requirements: %s", [str(r) for r in requirements]
)
# Filter to valid models
requirements = [r for r in requirements if r.model_id in MODEL_SPECS]
if not requirements:
logger.warning("No valid requirements, model pool will be empty")
_model_pool = ModelPool(GPUAllocator(gpus=[]))
return _model_pool
# Create and start the pool
allocator = GPUAllocator()
_model_pool = ModelPool(allocator)
startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300"))
_model_pool.startup(
requirements=requirements,
startup_timeout=startup_timeout,
)
# Log final GPU allocation summary
logger.info(_model_pool.allocator.summary())
# Register cleanup with atexit instead of request.addfinalizer
# This is critical for pytest-parallel where multiple threads share
# the session-scoped fixture - addfinalizer can fire too early
global _shutdown_registered
if not _shutdown_registered:
atexit.register(_shutdown_model_pool)
_shutdown_registered = True
return _model_pool
@pytest.fixture
def model_client(request: pytest.FixtureRequest, model_pool: "ModelPool"):
"""Get OpenAI client for the model specified by @pytest.mark.model().
Usage:
@pytest.mark.model("llama-8b")
def test_chat(model_client):
response = model_client.chat.completions.create(...)
"""
import openai
from infra import PARAM_MODEL
marker = request.node.get_closest_marker(PARAM_MODEL)
if marker is None:
pytest.fail(
f"Test must be marked with @pytest.mark.{PARAM_MODEL}('model-id') "
"to use model_client fixture"
)
model_id = marker.args[0]
try:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id)
except KeyError:
pytest.skip(f"Model {model_id} not available in model pool")
client = openai.OpenAI(
base_url=f"{instance.base_url}/v1",
api_key="not-used",
)
yield client
# Release reference to allow eviction
instance.release()
@pytest.fixture
def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> str:
"""Get the base URL for the model specified by @pytest.mark.model().
Usage:
@pytest.mark.model("llama-8b")
def test_direct_http(model_base_url):
response = httpx.get(f"{model_base_url}/health")
"""
from infra import PARAM_MODEL
marker = request.node.get_closest_marker(PARAM_MODEL)
if marker is None:
pytest.fail(
f"Test must be marked with @pytest.mark.{PARAM_MODEL}('model-id') "
"to use model_base_url fixture"
)
model_id = marker.args[0]
try:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id)
except KeyError:
pytest.skip(f"Model {model_id} not available in model pool")
yield instance.base_url
# Release reference to allow eviction
instance.release()

View File

@@ -0,0 +1,14 @@
"""Legacy port utilities.
DEPRECATED: This module will be removed during e2e_response_api migration.
Use infra.get_open_port() instead.
"""
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]

View File

@@ -0,0 +1,449 @@
"""Backend setup fixtures for E2E tests.
This module provides fixtures for launching gateways/routers for different backends.
"""
from __future__ import annotations
import logging
import os
from typing import TYPE_CHECKING
import pytest
if TYPE_CHECKING:
from infra import ModelPool
from .markers import get_marker_kwargs, get_marker_value
logger = logging.getLogger(__name__)
@pytest.fixture(scope="class")
def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
"""Class-scoped fixture that launches a router for each test class.
Routers are cheap to start (~1-2s) compared to workers (~30-60s), so we
launch a fresh router per test class for isolation while reusing the
expensive workers from model_pool.
Backend types:
- "http", "grpc": Gets existing worker from model_pool, launches router
- "pd": Launches prefill/decode workers via model_pool, launches PD router
- "openai", "xai", etc.: Launches cloud router (no local workers)
Configuration via markers:
- @pytest.mark.model("model-id"): Override default model
- @pytest.mark.workers(count=1): Number of regular workers behind router
- @pytest.mark.workers(prefill=1, decode=1): PD worker configuration
- @pytest.mark.gateway(policy="round_robin", timeout=60): Gateway configuration
Returns:
Tuple of (backend_name, model_path, openai_client, gateway)
Usage:
@pytest.mark.parametrize("setup_backend", ["http"], indirect=True)
class TestBasic:
def test_chat(self, setup_backend):
backend, model, client, gateway = setup_backend
"""
import openai
from infra import (
DEFAULT_MODEL,
DEFAULT_ROUTER_TIMEOUT,
ENV_MODEL,
ENV_SKIP_BACKEND_SETUP,
LOCAL_MODES,
ConnectionMode,
Gateway,
WorkerIdentity,
WorkerType,
)
backend_name = request.param
# Skip if requested
if os.environ.get(ENV_SKIP_BACKEND_SETUP, "").lower() in ("1", "true", "yes"):
pytest.skip(f"{ENV_SKIP_BACKEND_SETUP} is set")
# Get model from marker or env var or default
model_id = get_marker_value(request, "model")
if model_id is None:
model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL)
# Get worker configuration from marker
workers_config = get_marker_kwargs(
request, "workers", defaults={"count": 1, "prefill": None, "decode": None}
)
# Get gateway configuration from marker
gateway_config = get_marker_kwargs(
request,
"gateway",
defaults={
"policy": "round_robin",
"timeout": DEFAULT_ROUTER_TIMEOUT,
"extra_args": None,
},
)
# PD disaggregation backend
if backend_name == "pd":
yield from _setup_pd_backend(
request, model_pool, model_id, workers_config, gateway_config
)
return
# Check if this is a local backend (grpc, http)
try:
connection_mode = ConnectionMode(backend_name)
is_local = connection_mode in LOCAL_MODES
except ValueError:
is_local = False
connection_mode = None
# Local backends: use worker from pool + launch gateway
if is_local:
yield from _setup_local_backend(
request,
model_pool,
backend_name,
model_id,
connection_mode,
workers_config,
gateway_config,
)
return
# Get storage backend from marker (default: memory)
storage_backend = get_marker_value(request, "storage", default="memory")
# Cloud backends: launch cloud router
yield from _setup_cloud_backend(backend_name, storage_backend, gateway_config)
def _setup_pd_backend(
request: pytest.FixtureRequest,
model_pool: "ModelPool",
model_id: str,
workers_config: dict,
gateway_config: dict,
):
"""Setup PD disaggregation backend."""
import openai
from infra import ConnectionMode, Gateway, WorkerIdentity, WorkerType
logger.info("Setting up PD backend for model %s", model_id)
# Get PD configuration from workers marker
num_prefill = workers_config.get("prefill") or 1
num_decode = workers_config.get("decode") or 1
logger.info("PD config: %d prefill, %d decode workers", num_prefill, num_decode)
# Try to use pre-launched PD workers, or launch additional ones if needed
# get_workers_by_type auto-acquires all returned workers
existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL)
existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE)
# Calculate how many more we need
missing_prefill = max(0, num_prefill - len(existing_prefills))
missing_decode = max(0, num_decode - len(existing_decodes))
if missing_prefill == 0 and missing_decode == 0:
prefills = existing_prefills[:num_prefill]
decodes = existing_decodes[:num_decode]
# Release excess workers we won't use
for w in existing_prefills[num_prefill:]:
w.release()
for w in existing_decodes[num_decode:]:
w.release()
logger.info(
"Using pre-launched PD workers: %d prefill, %d decode",
len(prefills),
len(decodes),
)
else:
# Build WorkerIdentity list for missing workers
workers_to_launch: list[WorkerIdentity] = []
for i in range(missing_prefill):
workers_to_launch.append(
WorkerIdentity(
model_id,
ConnectionMode.HTTP,
WorkerType.PREFILL,
len(existing_prefills) + i,
)
)
for i in range(missing_decode):
workers_to_launch.append(
WorkerIdentity(
model_id,
ConnectionMode.HTTP,
WorkerType.DECODE,
len(existing_decodes) + i,
)
)
logger.info(
"Have %d/%d prefill, %d/%d decode. Launching %d more workers",
len(existing_prefills),
num_prefill,
len(existing_decodes),
num_decode,
len(workers_to_launch),
)
new_instances = model_pool.launch_workers(
workers_to_launch, startup_timeout=300
)
if not new_instances:
# Release any existing workers we acquired
for w in existing_prefills + existing_decodes:
w.release()
pytest.fail(
f"Failed to launch PD workers: needed {len(workers_to_launch)} workers "
f"but could not allocate GPUs (all in use or timeout)"
)
# Acquire newly launched instances (launch_workers doesn't auto-acquire)
for inst in new_instances:
inst.acquire()
new_prefills = [w for w in new_instances if w.worker_type == WorkerType.PREFILL]
new_decodes = [w for w in new_instances if w.worker_type == WorkerType.DECODE]
prefills = existing_prefills + new_prefills
decodes = existing_decodes + new_decodes
# All workers in prefills and decodes are now acquired
if not prefills or not decodes:
# This shouldn't happen but guard against it
for w in prefills + decodes:
w.release()
pytest.fail(
f"PD setup incomplete: have {len(prefills)} prefill, {len(decodes)} decode "
f"(need {num_prefill} prefill, {num_decode} decode)"
)
model_path = prefills[0].model_path
# Launch PD gateway
gateway = Gateway()
gateway.start(
prefill_workers=prefills,
decode_workers=decodes,
policy=gateway_config["policy"],
timeout=gateway_config["timeout"],
extra_args=gateway_config["extra_args"],
)
client = openai.OpenAI(
base_url=f"{gateway.base_url}/v1",
api_key="not-used",
)
logger.info(
"Setup PD backend: model=%s, %d prefill + %d decode workers, "
"gateway=%s, policy=%s",
model_id,
len(prefills),
len(decodes),
gateway.base_url,
gateway_config["policy"],
)
try:
yield "pd", model_path, client, gateway
finally:
logger.info("Tearing down PD gateway")
gateway.shutdown()
# Release references to allow eviction
for worker in prefills + decodes:
worker.release()
def _setup_local_backend(
request: pytest.FixtureRequest,
model_pool: "ModelPool",
backend_name: str,
model_id: str,
connection_mode,
workers_config: dict,
gateway_config: dict,
):
"""Setup local backend (grpc, http)."""
import openai
from infra import Gateway, WorkerIdentity, WorkerType
num_workers = workers_config.get("count") or 1
instances: list = [] # Track instances for reference counting
try:
if num_workers > 1:
# get_workers_by_type auto-acquires all returned workers
all_existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR)
existing_for_mode = [w for w in all_existing if w.mode == connection_mode]
# Release workers we won't use (wrong mode)
for w in all_existing:
if w not in existing_for_mode:
w.release()
if len(existing_for_mode) >= num_workers:
instances = existing_for_mode[:num_workers]
# Release excess workers we won't use
for w in existing_for_mode[num_workers:]:
w.release()
else:
missing = num_workers - len(existing_for_mode)
workers_to_launch = [
WorkerIdentity(
model_id,
connection_mode,
WorkerType.REGULAR,
len(existing_for_mode) + i,
)
for i in range(missing)
]
new_instances = model_pool.launch_workers(
workers_to_launch, startup_timeout=300
)
# Acquire newly launched instances
for inst in new_instances:
inst.acquire()
instances = existing_for_mode + new_instances
if not instances:
pytest.fail(f"Failed to get {num_workers} workers for {model_id}")
worker_urls = [inst.worker_url for inst in instances]
model_path = instances[0].model_path
else:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id, connection_mode)
instances = [instance]
worker_urls = [instance.worker_url]
model_path = instance.model_path
except RuntimeError as e:
pytest.fail(str(e))
# Launch gateway
gateway = Gateway()
gateway.start(
worker_urls=worker_urls,
model_path=model_path,
policy=gateway_config["policy"],
timeout=gateway_config["timeout"],
extra_args=gateway_config["extra_args"],
)
client = openai.OpenAI(
base_url=f"{gateway.base_url}/v1",
api_key="not-used",
)
logger.info(
"Setup %s backend: model=%s, workers=%d, gateway=%s, policy=%s",
backend_name,
model_id,
num_workers,
gateway.base_url,
gateway_config["policy"],
)
try:
yield backend_name, model_path, client, gateway
finally:
logger.info("Tearing down gateway for %s backend", backend_name)
gateway.shutdown()
# Release references to allow eviction
for inst in instances:
inst.release()
def _setup_cloud_backend(
backend_name: str,
storage_backend: str = "memory",
gateway_config: dict | None = None,
):
"""Setup cloud backend (openai, xai, etc.).
Args:
backend_name: Cloud backend name (openai, xai).
storage_backend: History storage backend (memory, oracle).
gateway_config: Gateway configuration from marker.
"""
import openai
from infra import THIRD_PARTY_MODELS, launch_cloud_gateway
if backend_name not in THIRD_PARTY_MODELS:
pytest.fail(f"Unknown cloud runtime: {backend_name}")
cfg = THIRD_PARTY_MODELS[backend_name]
api_key_env = cfg.get("api_key_env")
if api_key_env and not os.environ.get(api_key_env):
pytest.skip(f"{api_key_env} not set, skipping {backend_name} tests")
extra_args = gateway_config.get("extra_args") if gateway_config else None
logger.info(
"Launching cloud backend: %s with storage=%s", backend_name, storage_backend
)
gateway = launch_cloud_gateway(
backend_name,
history_backend=storage_backend,
extra_args=extra_args,
)
api_key = os.environ.get(api_key_env) if api_key_env else "not-used"
client = openai.OpenAI(
base_url=f"{gateway.base_url}/v1",
api_key=api_key,
)
try:
yield backend_name, cfg["model"], client, gateway
finally:
logger.info("Tearing down cloud backend: %s", backend_name)
gateway.shutdown()
@pytest.fixture
def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
"""Function-scoped fixture for launching a fresh router per test.
This launches a new Gateway for each test, pointing to workers from the pool.
Use for tests that need isolated router state.
Usage:
@pytest.mark.parametrize("backend_router", ["grpc", "http"], indirect=True)
def test_router_state(backend_router):
gateway = backend_router
"""
from infra import DEFAULT_MODEL, ENV_MODEL, ConnectionMode, Gateway
backend_name = request.param
model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL)
connection_mode = ConnectionMode(backend_name)
try:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id, connection_mode)
except KeyError:
pytest.skip(f"Model {model_id}:{backend_name} not available in pool")
except RuntimeError as e:
pytest.fail(str(e))
gateway = Gateway()
gateway.start(
worker_urls=[instance.worker_url],
model_path=instance.model_path,
)
try:
yield gateway
finally:
gateway.shutdown()
# Release reference to allow eviction
instance.release()

View File

@@ -0,0 +1,144 @@
"""Infrastructure for parallel GPU test execution."""
from .constants import ( # Enums; Convenience sets; Fixture parameters; Defaults; Environment variables
CLOUD_RUNTIMES,
DEFAULT_HOST,
DEFAULT_MODEL,
DEFAULT_ROUTER_TIMEOUT,
DEFAULT_STARTUP_TIMEOUT,
ENV_BACKENDS,
ENV_MODEL,
ENV_MODELS,
ENV_SHOW_ROUTER_LOGS,
ENV_SHOW_WORKER_LOGS,
ENV_SKIP_BACKEND_SETUP,
ENV_SKIP_MODEL_POOL,
ENV_STARTUP_TIMEOUT,
HEALTH_CHECK_INTERVAL,
LOCAL_MODES,
LOCAL_RUNTIMES,
LOG_SEPARATOR_WIDTH,
MAX_RETRY_ATTEMPTS,
PARAM_BACKEND_ROUTER,
PARAM_MODEL,
PARAM_SETUP_BACKEND,
ConnectionMode,
Runtime,
WorkerType,
)
from .gateway import Gateway, WorkerInfo, launch_cloud_gateway
from .gpu_allocator import (
GPUAllocator,
GPUInfo,
GPUSlot,
get_gpu_memory_usage,
get_open_port,
get_physical_device_indices,
nvml_context,
wait_for_gpu_memory_to_clear,
)
from .gpu_monitor import GPUMonitor
from .gpu_monitor import should_monitor as should_monitor_gpu
from .model_pool import ModelInstance, ModelPool, WorkerIdentity
from .model_specs import ( # Default model paths; Model groups
CHAT_MODELS,
DEFAULT_EMBEDDING_MODEL_PATH,
DEFAULT_ENABLE_THINKING_MODEL_PATH,
DEFAULT_GPT_OSS_MODEL_PATH,
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_MODEL_PATH,
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_REASONING_MODEL_PATH,
DEFAULT_SMALL_MODEL_PATH,
EMBEDDING_MODELS,
FUNCTION_CALLING_MODELS,
MODEL_SPECS,
REASONING_MODELS,
THIRD_PARTY_MODELS,
)
from .process_utils import (
detect_ib_device,
kill_process_tree,
terminate_process,
wait_for_health,
wait_for_workers_ready,
)
from .run_eval import run_eval
__all__ = [
# Enums and Identity
"ConnectionMode",
"WorkerType",
"Runtime",
"WorkerIdentity",
# Convenience sets
"LOCAL_MODES",
"LOCAL_RUNTIMES",
"CLOUD_RUNTIMES",
# Fixture params
"PARAM_SETUP_BACKEND",
"PARAM_BACKEND_ROUTER",
"PARAM_MODEL",
# Defaults
"DEFAULT_MODEL",
"DEFAULT_HOST",
"DEFAULT_STARTUP_TIMEOUT",
"DEFAULT_ROUTER_TIMEOUT",
"HEALTH_CHECK_INTERVAL",
"MAX_RETRY_ATTEMPTS",
"LOG_SEPARATOR_WIDTH",
# Env vars
"ENV_MODELS",
"ENV_BACKENDS",
"ENV_MODEL",
"ENV_STARTUP_TIMEOUT",
"ENV_SKIP_MODEL_POOL",
"ENV_SKIP_BACKEND_SETUP",
"ENV_SHOW_ROUTER_LOGS",
"ENV_SHOW_WORKER_LOGS",
# GPU allocation
"GPUAllocator",
"GPUInfo",
"GPUSlot",
# GPU utilities
"nvml_context",
"get_open_port",
"get_physical_device_indices",
"get_gpu_memory_usage",
"wait_for_gpu_memory_to_clear",
# Process utilities
"kill_process_tree",
"terminate_process",
"wait_for_health",
"wait_for_workers_ready",
"detect_ib_device",
# GPU monitoring
"GPUMonitor",
"should_monitor_gpu",
# Model management
"ModelInstance",
"ModelPool",
"MODEL_SPECS",
# Gateway
"Gateway",
"WorkerInfo",
"launch_cloud_gateway",
# Default model paths
"DEFAULT_MODEL_PATH",
"DEFAULT_SMALL_MODEL_PATH",
"DEFAULT_REASONING_MODEL_PATH",
"DEFAULT_ENABLE_THINKING_MODEL_PATH",
"DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH",
"DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH",
"DEFAULT_GPT_OSS_MODEL_PATH",
"DEFAULT_EMBEDDING_MODEL_PATH",
# Model groups
"CHAT_MODELS",
"EMBEDDING_MODELS",
"REASONING_MODELS",
"FUNCTION_CALLING_MODELS",
# Third-party models
"THIRD_PARTY_MODELS",
# Evaluation
"run_eval",
]

View File

@@ -0,0 +1,74 @@
"""Constants and enums for E2E test infrastructure."""
from enum import Enum
class ConnectionMode(str, Enum):
"""Worker connection protocol."""
HTTP = "http"
GRPC = "grpc"
class WorkerType(str, Enum):
"""Worker specialization type."""
REGULAR = "regular"
PREFILL = "prefill"
DECODE = "decode"
class Runtime(str, Enum):
"""Inference runtime/backend."""
SGLANG = "sglang"
VLLM = "vllm"
OPENAI = "openai"
XAI = "xai"
GEMINI = "gemini"
# Convenience sets
LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC})
LOCAL_RUNTIMES = frozenset({Runtime.SGLANG, Runtime.VLLM})
CLOUD_RUNTIMES = frozenset({Runtime.OPENAI, Runtime.XAI, Runtime.GEMINI})
# Fixture parameter names (used in @pytest.mark.parametrize)
PARAM_SETUP_BACKEND = "setup_backend"
PARAM_BACKEND_ROUTER = "backend_router"
PARAM_MODEL = "model"
# Default model
DEFAULT_MODEL = "llama-8b"
# Environment variable names
ENV_MODELS = "E2E_MODELS"
ENV_BACKENDS = "E2E_BACKENDS"
ENV_MODEL = "E2E_MODEL"
ENV_STARTUP_TIMEOUT = "E2E_STARTUP_TIMEOUT"
ENV_SKIP_MODEL_POOL = "SKIP_MODEL_POOL"
ENV_SKIP_BACKEND_SETUP = "SKIP_BACKEND_SETUP"
ENV_SHOW_ROUTER_LOGS = "SHOW_ROUTER_LOGS"
ENV_SHOW_WORKER_LOGS = "SHOW_WORKER_LOGS"
# Network
DEFAULT_HOST = "127.0.0.1"
# Timeouts (seconds)
DEFAULT_STARTUP_TIMEOUT = 300
DEFAULT_ROUTER_TIMEOUT = 60
HEALTH_CHECK_INTERVAL = 2 # Check every 2s (was 5s)
# Model loading configuration
INITIAL_GRACE_PERIOD = 30 # Wait before first health check (model loading time)
LAUNCH_STAGGER_DELAY = (
10 # Delay between launching multiple workers (avoid I/O contention)
)
# Retry configuration
MAX_RETRY_ATTEMPTS = (
6 # Max retries with exponential backoff (total ~63s: 1+2+4+8+16+32)
)
# Display formatting
LOG_SEPARATOR_WIDTH = 60 # Width for log separator lines (e.g., "="*60)

View File

@@ -0,0 +1,595 @@
"""Gateway class for managing sgl-model-gateway router instances."""
from __future__ import annotations
import logging
import os
import subprocess
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
import httpx
from .constants import DEFAULT_HOST, DEFAULT_ROUTER_TIMEOUT, ENV_SHOW_ROUTER_LOGS
from .gpu_allocator import get_open_port
from .process_utils import kill_process_tree, wait_for_health, wait_for_workers_ready
if TYPE_CHECKING:
from .model_pool import ModelInstance
logger = logging.getLogger(__name__)
@dataclass
class WorkerInfo:
"""Information about a worker connected to the gateway."""
id: str
url: str
model: str | None = None
status: str = "unknown"
pending_requests: int = 0
metadata: dict[str, Any] = field(default_factory=dict)
class Gateway:
"""Manages a sgl-model-gateway router instance.
Provides lifecycle management and API access for:
- Starting/stopping the router
- Worker management (list, add, remove)
- Health and metrics endpoints
Four startup modes:
1. Regular mode: Start with worker URLs
2. PD mode: Start with prefill/decode workers
3. IGW mode: Start empty, add workers via API
4. Cloud mode: Start with cloud backend (OpenAI, xAI)
Example (regular mode):
gateway = Gateway()
gateway.start(
worker_urls=["http://127.0.0.1:30000"],
model_path="/path/to/model",
)
Example (PD disaggregation mode):
gateway = Gateway()
gateway.start(
prefill_workers=prefill_instances,
decode_workers=decode_instances,
)
Example (IGW mode):
gateway = Gateway()
gateway.start(igw_mode=True)
gateway.add_worker("http://127.0.0.1:30000")
gateway.add_worker("http://127.0.0.1:30001")
# Use gateway
workers = gateway.list_workers()
health = gateway.health()
# Cleanup
gateway.shutdown()
Example (cloud mode):
gateway = Gateway()
gateway.start(cloud_backend="openai") # or "xai"
# Requires OPENAI_API_KEY or XAI_API_KEY env var
"""
def __init__(
self,
host: str = DEFAULT_HOST,
port: int | None = None,
prometheus_port: int | None = None,
):
"""Initialize gateway configuration.
Args:
host: Host to bind the router to.
port: Port for the router. If None, auto-assigns.
prometheus_port: Port for prometheus metrics. If None, auto-assigns.
"""
self.host = host
self.port = port or get_open_port()
self.prometheus_port = prometheus_port or get_open_port()
self.base_url = f"http://{self.host}:{self.port}"
self.metrics_url = f"http://{self.host}:{self.prometheus_port}"
self.process: subprocess.Popen | None = None
self.model_path: str | None = None
self.policy: str = "round_robin"
self.pd_mode: bool = False
self.igw_mode: bool = False
self.cloud_mode: bool = False
self.cloud_backend: str | None = None
self._started: bool = False
self._env: dict[str, str] | None = None # Custom env for subprocess
@property
def is_running(self) -> bool:
"""Check if the gateway process is running."""
return self.process is not None and self.process.poll() is None
def start(
self,
*,
# Regular mode arguments
worker_urls: list[str] | None = None,
model_path: str | None = None,
# PD mode arguments
prefill_workers: list["ModelInstance"] | None = None,
decode_workers: list["ModelInstance"] | None = None,
# IGW mode arguments
igw_mode: bool = False,
# Cloud mode arguments
cloud_backend: str | None = None,
history_backend: str = "memory",
# Common arguments
policy: str = "round_robin",
timeout: float = DEFAULT_ROUTER_TIMEOUT,
show_output: bool | None = None,
extra_args: list[str] | None = None,
) -> None:
"""Start the gateway.
Can be started in four modes:
1. Regular mode: Provide worker_urls and model_path
2. PD mode: Provide prefill_workers and decode_workers
3. IGW mode: Set igw_mode=True, add workers later via add_worker()
4. Cloud mode: Provide cloud_backend ("openai" or "xai")
Args:
worker_urls: List of worker URLs for regular mode.
model_path: Model path for regular mode.
prefill_workers: List of prefill ModelInstance objects for PD mode.
decode_workers: List of decode ModelInstance objects for PD mode.
igw_mode: Start in IGW mode (no workers, add via API).
cloud_backend: Cloud backend type ("openai" or "xai").
history_backend: History backend for cloud mode ("memory" or "oracle").
policy: Routing policy (round_robin, random, etc.)
timeout: Startup timeout in seconds.
show_output: Show subprocess output (env var override).
extra_args: Additional router arguments.
Raises:
RuntimeError: If gateway is already started.
ValueError: If arguments are invalid for the mode.
"""
if self._started:
raise RuntimeError("Gateway already started")
# Determine mode based on arguments
is_pd_mode = prefill_workers is not None or decode_workers is not None
is_regular_mode = worker_urls is not None
is_igw_mode = igw_mode
is_cloud_mode = cloud_backend is not None
# Validate mode exclusivity
modes_specified = sum([is_pd_mode, is_regular_mode, is_igw_mode, is_cloud_mode])
if modes_specified > 1:
raise ValueError(
"Cannot specify multiple modes. Choose one of: "
"worker_urls (regular), prefill/decode_workers (PD), "
"igw_mode, or cloud_backend"
)
if modes_specified == 0:
raise ValueError(
"Must specify one mode: worker_urls (regular), "
"prefill/decode_workers (PD), igw_mode=True, or cloud_backend"
)
if show_output is None:
show_output = os.environ.get(ENV_SHOW_ROUTER_LOGS, "0") == "1"
self.policy = policy
if is_igw_mode:
# IGW mode: start empty, add workers via API
self.pd_mode = False
self.igw_mode = True
self._launch(
mode_args=["--enable-igw"],
timeout=timeout,
show_output=show_output,
extra_args=extra_args,
log_msg="IGW gateway (no workers)",
)
elif is_pd_mode:
# PD mode: prefill/decode disaggregation
self.pd_mode = True
self.igw_mode = False
prefills = prefill_workers or []
decodes = decode_workers or []
mode_args = ["--pd-disaggregation"]
for pf in prefills:
mode_args += ["--prefill", pf.base_url, str(pf.bootstrap_port)]
for dc in decodes:
mode_args += ["--decode", dc.base_url]
self._launch(
mode_args=mode_args,
timeout=timeout,
show_output=show_output,
extra_args=extra_args,
log_msg=f"PD gateway ({len(prefills)} prefill, {len(decodes)} decode)",
)
elif is_cloud_mode:
# Cloud mode: OpenAI/xAI backend
self.pd_mode = False
self.igw_mode = False
self.cloud_mode = True
self.cloud_backend = cloud_backend
# Get worker URL and API key based on backend
if cloud_backend == "openai":
worker_url = "https://api.openai.com"
api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
raise ValueError("OPENAI_API_KEY environment variable required")
self._env = os.environ.copy()
self._env["OPENAI_API_KEY"] = api_key
elif cloud_backend == "xai":
worker_url = "https://api.x.ai"
api_key = os.environ.get("XAI_API_KEY")
if not api_key:
raise ValueError("XAI_API_KEY environment variable required")
self._env = os.environ.copy()
self._env["XAI_API_KEY"] = api_key
else:
raise ValueError(f"Unsupported cloud backend: {cloud_backend}")
mode_args = [
"--backend",
"openai", # Both OpenAI and xAI use openai backend type
"--worker-urls",
worker_url,
"--history-backend",
history_backend,
]
self._launch(
mode_args=mode_args,
timeout=timeout,
show_output=show_output,
extra_args=extra_args,
log_msg=f"{cloud_backend} cloud gateway",
)
else:
# Regular mode: worker URLs
if model_path is None:
raise ValueError("model_path is required for regular mode")
self.model_path = model_path
self.pd_mode = False
self.igw_mode = False
self._launch(
mode_args=["--model-path", model_path, "--worker-urls", *worker_urls],
timeout=timeout,
show_output=show_output,
extra_args=extra_args,
num_workers=len(worker_urls),
log_msg=f"gateway with {len(worker_urls)} worker(s)",
)
def _launch(
self,
mode_args: list[str],
timeout: float,
show_output: bool,
extra_args: list[str] | None,
num_workers: int | None = None,
log_msg: str = "",
) -> None:
"""Launch the gateway process.
Args:
mode_args: Mode-specific CLI arguments.
timeout: Startup timeout in seconds.
show_output: Show subprocess output.
extra_args: Additional router arguments.
num_workers: If set, wait for this many workers to be ready.
If None, just wait for health check.
log_msg: Log message describing the startup.
"""
cmd = self._build_base_cmd()
cmd.extend(mode_args)
if extra_args:
cmd.extend(extra_args)
logger.info("Starting %s on port %d", log_msg or "gateway", self.port)
logger.debug("Gateway command: %s", " ".join(cmd))
self.process = subprocess.Popen(
cmd,
env=self._env, # Use custom env if set (e.g., for cloud mode API keys)
stdout=None if show_output else subprocess.PIPE,
stderr=None if show_output else subprocess.PIPE,
start_new_session=True,
)
try:
if num_workers is not None:
wait_for_workers_ready(self.base_url, num_workers, timeout=timeout)
else:
wait_for_health(self.base_url, timeout=timeout)
except TimeoutError:
self.shutdown()
raise
self._started = True
logger.info("Gateway ready at %s", self.base_url)
def shutdown(self) -> None:
"""Shutdown the gateway process."""
if self.process is not None:
logger.info("Shutting down gateway (PID %d)", self.process.pid)
kill_process_tree(self.process.pid)
self.process = None
self._started = False
def _build_base_cmd(self) -> list[str]:
"""Build the base command for launching the router."""
return [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
self.host,
"--port",
str(self.port),
"--prometheus-port",
str(self.prometheus_port),
"--prometheus-host",
self.host,
"--policy",
self.policy,
"--log-level",
"warn",
]
# -------------------------------------------------------------------------
# Health & Metrics APIs
# -------------------------------------------------------------------------
def health(self, timeout: float = 5.0) -> bool:
"""Check gateway health.
Returns:
True if healthy, False otherwise.
"""
try:
resp = httpx.get(f"{self.base_url}/health", timeout=timeout)
return resp.status_code == 200
except (httpx.RequestError, httpx.TimeoutException):
return False
def get_metrics(self, timeout: float = 5.0) -> str | None:
"""Get Prometheus metrics.
Returns:
Metrics text or None if unavailable.
"""
try:
resp = httpx.get(f"{self.metrics_url}/metrics", timeout=timeout)
if resp.status_code == 200:
return resp.text
return None
except (httpx.RequestError, httpx.TimeoutException):
return None
# -------------------------------------------------------------------------
# Worker Management APIs
# -------------------------------------------------------------------------
def _worker_from_api_response(self, w: dict) -> WorkerInfo:
"""Convert API response dict to WorkerInfo.
Args:
w: Worker dict from API response.
Returns:
WorkerInfo object.
"""
status = "healthy" if w.get("is_healthy", False) else "unhealthy"
return WorkerInfo(
id=w.get("id", ""),
url=w.get("url", ""),
model=w.get("model_id"),
status=status,
pending_requests=w.get("load", 0),
metadata={
"worker_type": w.get("worker_type"),
"connection_mode": w.get("connection_mode"),
"priority": w.get("priority"),
"cost": w.get("cost"),
},
)
def list_workers(self, timeout: float = 5.0) -> list[WorkerInfo]:
"""List all workers connected to the gateway.
Returns:
List of WorkerInfo objects.
"""
try:
resp = httpx.get(f"{self.base_url}/workers", timeout=timeout)
if resp.status_code == 200:
data = resp.json()
return [
self._worker_from_api_response(w) for w in data.get("workers", [])
]
return []
except (httpx.RequestError, httpx.TimeoutException):
return []
def get_worker(self, worker_id: str, timeout: float = 5.0) -> WorkerInfo | None:
"""Get information about a specific worker.
Args:
worker_id: The worker ID.
Returns:
WorkerInfo or None if not found.
"""
try:
resp = httpx.get(f"{self.base_url}/workers/{worker_id}", timeout=timeout)
if resp.status_code == 200:
return self._worker_from_api_response(resp.json())
return None
except (httpx.RequestError, httpx.TimeoutException):
return None
def add_worker(
self,
worker_url: str,
timeout: float = 10.0,
wait_ready: bool = True,
ready_timeout: float = 60.0,
) -> tuple[bool, str | None]:
"""Add a worker to the gateway.
Args:
worker_url: URL of the worker to add.
timeout: HTTP request timeout.
wait_ready: If True, wait for worker to become ready.
ready_timeout: Timeout for waiting for worker to be ready.
Returns:
Tuple of (success, worker_id or error message).
"""
try:
resp = httpx.post(
f"{self.base_url}/workers",
json={"url": worker_url},
timeout=timeout,
)
# API returns 200 OK or 202 Accepted for async processing
if resp.status_code in (200, 202):
data = resp.json()
worker_id = data.get("worker_id")
if wait_ready and worker_id:
# Wait for worker to appear in list
import time
start = time.time()
while time.time() - start < ready_timeout:
workers = self.list_workers()
for w in workers:
if w.id == worker_id:
return True, worker_id
time.sleep(1.0)
return (
False,
f"Worker {worker_id} not ready within {ready_timeout}s",
)
return True, worker_id
return False, resp.text
except (httpx.RequestError, httpx.TimeoutException) as e:
return False, str(e)
def remove_worker(self, worker_url: str, timeout: float = 10.0) -> tuple[bool, str]:
"""Remove a worker from the gateway by URL.
Args:
worker_url: URL of the worker to remove.
Returns:
Tuple of (success, message).
"""
# Find worker_id by URL
workers = self.list_workers(timeout=timeout)
worker_id = None
for w in workers:
if w.url == worker_url:
worker_id = w.id
break
if not worker_id:
return False, f"Worker with URL {worker_url} not found"
try:
resp = httpx.delete(
f"{self.base_url}/workers/{worker_id}",
timeout=timeout,
)
if resp.status_code == 200:
return True, "Worker removed"
return False, resp.text
except (httpx.RequestError, httpx.TimeoutException) as e:
return False, str(e)
# -------------------------------------------------------------------------
# Model APIs
# -------------------------------------------------------------------------
def list_models(self, timeout: float = 5.0) -> list[dict]:
"""List available models (OpenAI-compatible).
Returns:
List of model info dicts.
"""
try:
resp = httpx.get(f"{self.base_url}/v1/models", timeout=timeout)
if resp.status_code == 200:
data = resp.json()
return data.get("data", [])
return []
except (httpx.RequestError, httpx.TimeoutException):
return []
# -------------------------------------------------------------------------
# Context manager support
# -------------------------------------------------------------------------
def __enter__(self) -> "Gateway":
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
self.shutdown()
def launch_cloud_gateway(
runtime: str, # "openai" or "xai"
*,
history_backend: str = "memory",
extra_args: list[str] | None = None,
timeout: float = 60,
show_output: bool | None = None,
) -> Gateway:
"""Launch gateway with cloud API runtime.
Args:
runtime: Cloud runtime ("openai" or "xai")
history_backend: History storage backend ("memory" or "oracle")
extra_args: Additional router arguments
timeout: Startup timeout in seconds
show_output: Show subprocess output
Returns:
Gateway instance with running router
"""
from .model_specs import THIRD_PARTY_MODELS
if runtime not in THIRD_PARTY_MODELS:
raise ValueError(
f"Unknown cloud runtime: {runtime}. "
f"Available: {list(THIRD_PARTY_MODELS.keys())}"
)
gateway = Gateway()
gateway.start(
cloud_backend=runtime,
history_backend=history_backend,
timeout=timeout,
show_output=show_output,
extra_args=extra_args,
)
return gateway

View File

@@ -0,0 +1,437 @@
"""GPU detection and slot allocation for parallel test execution."""
from __future__ import annotations
import logging
import os
import socket
import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass
logger = logging.getLogger(__name__)
# Try to import nvidia-ml-py for GPU detection
try:
import pynvml
NVML_AVAILABLE = True
except ImportError:
NVML_AVAILABLE = False
logger.debug("nvidia-ml-py not available, GPU detection will be limited")
@contextmanager
def nvml_context():
"""Context manager for NVML initialization/shutdown.
Usage:
with nvml_context():
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
...
"""
if not NVML_AVAILABLE:
yield
return
try:
pynvml.nvmlInit()
yield
finally:
pynvml.nvmlShutdown()
@dataclass
class GPUInfo:
"""Information about a single GPU."""
id: int
name: str
memory_mb: int
@property
def memory_gb(self) -> float:
return self.memory_mb / 1024
@dataclass
class GPUSlot:
"""A slot representing one or more GPUs allocated for a model."""
gpu_ids: list[int]
total_memory_mb: int
assigned_model: str | None = None
port: int | None = None
@property
def total_memory_gb(self) -> float:
return self.total_memory_mb / 1024
def cuda_visible_devices(self) -> str:
"""Return CUDA_VISIBLE_DEVICES string for this slot."""
return ",".join(str(g) for g in self.gpu_ids)
def get_open_port() -> int:
"""Get an available port by binding to port 0 and reading the assigned port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
s.listen(1)
port = s.getsockname()[1]
return port
def get_physical_device_indices(devices: list[int]) -> list[int]:
"""Map logical device indices to physical indices based on CUDA_VISIBLE_DEVICES."""
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
if visible_devices is None:
return devices
visible_indices = [int(x) for x in visible_devices.split(",")]
index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
return [index_mapping[i] for i in devices if i in index_mapping]
def get_gpu_memory_usage(device_id: int) -> tuple[float, float]:
"""Get GPU memory usage in GB (used, total).
Args:
device_id: Physical GPU device ID
Returns:
Tuple of (used_gb, total_gb)
"""
if not NVML_AVAILABLE:
return (0.0, 0.0)
with nvml_context():
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
return (mem_info.used / (1024**3), mem_info.total / (1024**3))
def wait_for_gpu_memory_to_clear(
*,
devices: list[int],
threshold_bytes: int | None = None,
threshold_ratio: float | None = None,
timeout_s: float = 120,
) -> None:
"""Wait for GPU memory to be freed below a threshold.
Args:
devices: List of logical GPU device IDs to check
threshold_bytes: Memory threshold in bytes (used <= threshold)
threshold_ratio: Memory threshold as ratio (used/total <= ratio)
timeout_s: Timeout in seconds
Raises:
ValueError: If memory doesn't clear within timeout
"""
if not NVML_AVAILABLE:
logger.warning("nvidia-ml-py not available, skipping memory wait")
return
if threshold_bytes is None and threshold_ratio is None:
raise ValueError("Must specify threshold_bytes or threshold_ratio")
physical_devices = get_physical_device_indices(devices)
start_time = time.time()
with nvml_context():
while True:
output: dict[int, str] = {}
output_raw: dict[int, tuple[float, float]] = {}
for device in physical_devices:
handle = pynvml.nvmlDeviceGetHandleByIndex(device)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
gb_used = mem_info.used / (1024**3)
gb_total = mem_info.total / (1024**3)
output_raw[device] = (gb_used, gb_total)
output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
logger.debug(
"GPU memory used/total (GiB): %s",
" ".join(f"{k}={v}" for k, v in output.items()),
)
if threshold_bytes is not None:
def is_free(used: float, total: float) -> bool:
return used <= threshold_bytes / (1024**3)
threshold_desc = f"{threshold_bytes / (1024**3):.1f} GiB"
else:
def is_free(used: float, total: float) -> bool:
return used / total <= threshold_ratio # type: ignore[operator]
threshold_desc = f"{threshold_ratio:.2%}" # type: ignore[str-format]
dur_s = time.time() - start_time
if all(is_free(used, total) for used, total in output_raw.values()):
logger.info(
"GPU memory cleared on devices %s (threshold=%s) in %.1fs",
devices,
threshold_desc,
dur_s,
)
return
if dur_s >= timeout_s:
raise ValueError(
f"GPU memory on devices {devices} not freed after {dur_s:.1f}s "
f"(threshold={threshold_desc})"
)
time.sleep(5)
class GPUAllocator:
"""Detects GPUs and assigns them to model slots using bin-packing."""
def __init__(self, gpus: list[GPUInfo] | None = None):
"""Initialize the allocator.
Args:
gpus: Optional list of GPUs. If None, auto-detects via nvidia-ml-py.
"""
self.gpus = gpus if gpus is not None else self._detect_gpus()
self.slots: list[GPUSlot] = []
self._used_gpus: set[int] = set() # Track GPUs used across all allocations
self._lock = threading.RLock() # Protects slots and _used_gpus
def _detect_gpus(self) -> list[GPUInfo]:
"""Auto-detect available GPUs via nvidia-ml-py (NVML)."""
if not NVML_AVAILABLE:
logger.warning("nvidia-ml-py not available - no GPUs detected")
return []
# Check for CUDA_VISIBLE_DEVICES restriction
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
allowed_ids: set[int] | None = None
if visible_devices:
allowed_ids = set(int(x) for x in visible_devices.split(",") if x.strip())
try:
with nvml_context():
device_count = pynvml.nvmlDeviceGetCount()
gpus = []
for idx in range(device_count):
# Skip GPUs not in CUDA_VISIBLE_DEVICES if set
if allowed_ids is not None and idx not in allowed_ids:
continue
handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
name = pynvml.nvmlDeviceGetName(handle)
# Handle bytes vs string return type (varies by pynvml version)
if isinstance(name, bytes):
name = name.decode("utf-8", errors="replace")
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
# Convert bytes to MB
memory_mb = mem_info.total // (1024 * 1024)
gpus.append(GPUInfo(idx, name, memory_mb))
logger.info("Detected %d GPUs: %s", len(gpus), [g.name for g in gpus])
return gpus
except pynvml.NVMLError as e:
logger.warning("NVML error during GPU detection: %s", e)
return []
except Exception as e:
logger.warning("Failed to detect GPUs: %s", e)
return []
def allocate_slots(
self, model_specs: dict[str, dict], preserve_order: bool = False
) -> list[GPUSlot]:
"""Allocate GPU slots based on model memory requirements.
Uses a first-fit decreasing bin-packing algorithm by default:
1. Sort models by memory requirement (largest first)
2. For each model, find the first GPU(s) that can fit it
3. For multi-GPU models, find consecutive GPUs
When preserve_order=True, processes models in dict insertion order
(test collection order) instead of sorting by memory. This ensures
models needed by earlier tests are allocated first.
Note: This method tracks used GPUs across multiple calls, so subsequent
allocations will use different GPUs than previous ones.
Thread-safe: Protected by internal lock.
Args:
model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys
preserve_order: If True, allocate in dict order (test order) instead
of sorting by memory size. Default False.
Returns:
List of GPUSlots with assigned models (only the newly allocated slots)
"""
with self._lock:
return self._allocate_slots_unlocked(model_specs, preserve_order)
def _allocate_slots_unlocked(
self, model_specs: dict[str, dict], preserve_order: bool = False
) -> list[GPUSlot]:
"""Internal allocation logic. Caller must hold _lock."""
if not self.gpus:
logger.warning("No GPUs available for allocation")
return []
if preserve_order:
# Process in dict insertion order (test collection order)
ordered_models = list(model_specs.items())
else:
# Sort models by memory requirement (largest first for better packing)
ordered_models = sorted(
model_specs.items(),
key=lambda x: x[1].get("memory_gb", 0),
reverse=True,
)
# Track new slots allocated in this call
new_slots: list[GPUSlot] = []
for model_id, spec in ordered_models:
memory_gb = spec.get("memory_gb", 16)
tp_size = spec.get("tp", 1)
# Find available GPUs (not used by any previous allocation)
available = [g for g in self.gpus if g.id not in self._used_gpus]
if tp_size == 1:
# Single GPU - find one with enough memory
for gpu in available:
if gpu.memory_gb >= memory_gb:
slot = GPUSlot(
gpu_ids=[gpu.id],
total_memory_mb=gpu.memory_mb,
assigned_model=model_id,
port=get_open_port(),
)
new_slots.append(slot)
self._used_gpus.add(gpu.id)
logger.info(
"Allocated GPU %d (%s, %.1fGB) for %s",
gpu.id,
gpu.name,
gpu.memory_gb,
model_id,
)
break
else:
logger.warning(
"No GPU with %.1fGB available for %s (used: %s)",
memory_gb,
model_id,
self._used_gpus,
)
else:
# Multi-GPU - find consecutive GPUs with enough total memory
# Sort available by ID for consecutive allocation
available_sorted = sorted(available, key=lambda g: g.id)
for i in range(len(available_sorted) - tp_size + 1):
candidate_gpus = available_sorted[i : i + tp_size]
total_mem = sum(g.memory_mb for g in candidate_gpus)
if total_mem >= memory_gb * 1024:
gpu_ids = [g.id for g in candidate_gpus]
slot = GPUSlot(
gpu_ids=gpu_ids,
total_memory_mb=total_mem,
assigned_model=model_id,
port=get_open_port(),
)
new_slots.append(slot)
self._used_gpus.update(gpu_ids)
logger.info(
"Allocated GPUs %s (%.1fGB total) for %s (tp=%d)",
gpu_ids,
total_mem / 1024,
model_id,
tp_size,
)
break
else:
logger.warning(
"No %d consecutive GPUs with %.1fGB available for %s (used: %s)",
tp_size,
memory_gb,
model_id,
self._used_gpus,
)
# Add new slots to existing slots list
self.slots.extend(new_slots)
return new_slots
def get_slot_for_model(self, model_id: str) -> GPUSlot | None:
"""Get the slot assigned to a specific model.
Thread-safe: Protected by internal lock.
"""
with self._lock:
for slot in self.slots:
if slot.assigned_model == model_id:
return slot
return None
def release_gpus(self, gpu_ids: list[int]) -> None:
"""Release GPUs back to the available pool.
Thread-safe: Protected by internal lock.
Args:
gpu_ids: List of GPU IDs to release.
"""
with self._lock:
for gpu_id in gpu_ids:
self._used_gpus.discard(gpu_id)
# Remove slots that used these GPUs
self.slots = [
s for s in self.slots if not any(g in gpu_ids for g in s.gpu_ids)
]
logger.info("Released GPUs %s, now used: %s", gpu_ids, self._used_gpus)
def release_slot(self, slot: GPUSlot) -> None:
"""Release a GPU slot back to the available pool.
Args:
slot: The GPUSlot to release.
"""
self.release_gpus(slot.gpu_ids)
def available_gpus(self) -> list[int]:
"""Get list of available (unused) GPU IDs.
Thread-safe: Protected by internal lock.
Returns:
List of GPU IDs that are not currently allocated.
"""
with self._lock:
return [g.id for g in self.gpus if g.id not in self._used_gpus]
def summary(self) -> str:
"""Return a summary of GPU allocations.
Thread-safe: Protected by internal lock.
"""
with self._lock:
lines = ["GPU Allocation Summary:"]
lines.append(f" Total GPUs: {len(self.gpus)}")
lines.append(f" Used GPUs: {sorted(self._used_gpus)}")
lines.append(f" Allocated Slots: {len(self.slots)}")
for slot in self.slots:
lines.append(
f" - {slot.assigned_model}: GPUs {slot.gpu_ids} "
f"({slot.total_memory_gb:.1f}GB) port={slot.port}"
)
return "\n".join(lines)

View File

@@ -0,0 +1,329 @@
"""GPU utilization monitoring for benchmarks.
This module provides a low-impact GPU monitor that runs in a separate process
and collects utilization samples using NVML.
"""
from __future__ import annotations
import json
import logging
import os
import time
from multiprocessing import Process
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
logger = logging.getLogger(__name__)
def _percentile(samples: list[float], p: float) -> float:
"""Calculate percentile from sorted samples."""
if not samples:
return 0.0
sorted_samples = sorted(samples)
idx = max(
0,
min(
len(sorted_samples) - 1, int(round((p / 100.0) * (len(sorted_samples) - 1)))
),
)
return float(sorted_samples[idx])
def _compute_stats(samples: list[float]) -> dict[str, float]:
"""Compute statistics for a list of samples."""
if not samples:
return {
"mean": 0.0,
"min": 0.0,
"max": 0.0,
"p5": 0.0,
"p10": 0.0,
"p25": 0.0,
"p50": 0.0,
"p75": 0.0,
"p90": 0.0,
"p95": 0.0,
"count": 0,
}
return {
"mean": sum(samples) / len(samples),
"min": min(samples),
"max": max(samples),
"p5": _percentile(samples, 5),
"p10": _percentile(samples, 10),
"p25": _percentile(samples, 25),
"p50": _percentile(samples, 50),
"p75": _percentile(samples, 75),
"p90": _percentile(samples, 90),
"p95": _percentile(samples, 95),
"count": len(samples),
}
def _monitor_loop(pid: int, output_path: str, interval: float) -> None:
"""Main monitoring loop - runs in separate process.
Monitors GPU utilization until the target process exits, then writes
results to output_path as JSON.
"""
# Lower process priority to minimize impact on benchmark
try:
os.nice(10)
except Exception:
pass
# Initialize NVML
try:
import pynvml
pynvml.nvmlInit()
except Exception as e:
logger.warning("Failed to initialize NVML: %s", e)
_write_empty_result(output_path)
return
# Get GPU handles
try:
device_count = pynvml.nvmlDeviceGetCount()
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(device_count)]
except Exception as e:
logger.warning("Failed to get GPU handles: %s", e)
_write_empty_result(output_path)
_shutdown_nvml()
return
# Collect samples
per_gpu_samples: dict[str, list[float]] = {str(i): [] for i in range(device_count)}
overall_samples: list[float] = []
try:
while _process_alive(pid):
try:
gpu_utils = []
for idx, handle in enumerate(handles):
try:
util = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
gpu_utils.append(float(util))
per_gpu_samples[str(idx)].append(float(util))
except Exception:
continue
if gpu_utils:
avg = sum(gpu_utils) / len(gpu_utils)
overall_samples.append(avg)
except Exception:
pass
time.sleep(interval)
finally:
# Write results
_write_result(output_path, pid, interval, overall_samples, per_gpu_samples)
_shutdown_nvml()
def _process_alive(pid: int) -> bool:
"""Check if process is still running."""
try:
os.kill(pid, 0)
return True
except (OSError, ProcessLookupError):
return False
def _write_empty_result(path: str) -> None:
"""Write empty result file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {"overall": [], "per_gpu": {}},
},
f,
)
except Exception:
pass
def _write_result(
path: str,
pid: int,
interval: float,
overall_samples: list[float],
per_gpu_samples: dict[str, list[float]],
) -> None:
"""Write monitoring results to JSON file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"bench_pid": pid,
"interval_sec": interval,
"count": len(overall_samples),
"overall": _compute_stats(overall_samples),
"per_gpu": {
k: _compute_stats(v) for k, v in per_gpu_samples.items()
},
"raw": {
"overall": overall_samples,
"per_gpu": per_gpu_samples,
},
},
f,
)
except Exception as e:
logger.warning("Failed to write GPU monitor results: %s", e)
def _shutdown_nvml() -> None:
"""Shutdown NVML."""
try:
import pynvml
pynvml.nvmlShutdown()
except Exception:
pass
class GPUMonitor:
"""GPU utilization monitor for benchmarks.
Usage:
monitor = GPUMonitor(output_dir="benchmark_results")
monitor.start(target_pid=12345)
# ... run benchmark ...
result = monitor.stop()
monitor.assert_thresholds({"gpu_util_p50_min": 99})
"""
def __init__(
self,
output_dir: str | Path = ".",
interval: float = 2.0,
):
self.output_dir = Path(output_dir)
self.interval = interval
self._process: Process | None = None
self._output_path: str | None = None
self._result: dict[str, Any] | None = None
@property
def output_path(self) -> str | None:
"""Path to the GPU utilization JSON file."""
return self._output_path
def start(self, target_pid: int) -> None:
"""Start monitoring GPU utilization for the target process."""
self._output_path = str(self.output_dir / "gpu_utilization.json")
self._result = None
self._process = Process(
target=_monitor_loop,
args=(target_pid, self._output_path, self.interval),
daemon=True,
)
self._process.start()
logger.debug("Started GPU monitor for PID %d", target_pid)
def stop(self, timeout: float = 5.0) -> dict[str, Any] | None:
"""Stop monitoring and return results."""
if self._process is None:
return None
try:
self._process.join(timeout=timeout)
except Exception:
pass
if self._process.is_alive():
try:
self._process.terminate()
except Exception:
pass
self._process = None
self._result = self._read_result()
return self._result
def _read_result(self) -> dict[str, Any] | None:
"""Read results from output file."""
if not self._output_path or not os.path.exists(self._output_path):
return None
try:
with open(self._output_path) as f:
return json.load(f)
except Exception as e:
logger.warning("Failed to read GPU monitor result: %s", e)
return None
def log_summary(self) -> None:
"""Log a summary of GPU utilization."""
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
logger.info(
"GPU utilization: mean=%.2f%% p50=%.2f%% (samples=%d)",
overall.get("mean", 0.0),
overall.get("p50", 0.0),
result.get("count", 0),
)
def assert_thresholds(self, thresholds: dict[str, float] | None) -> None:
"""Assert GPU utilization meets thresholds.
Supported thresholds:
- gpu_util_mean_min: Minimum mean GPU utilization %
- gpu_util_p50_min: Minimum p50 GPU utilization %
"""
if not thresholds:
return
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
mean_threshold = thresholds.get("gpu_util_mean_min")
if mean_threshold is not None:
mean_value = overall.get("mean", 0.0)
assert (
mean_value >= mean_threshold
), f"GPU utilization mean below threshold: {mean_value:.2f}% < {mean_threshold}%"
p50_threshold = thresholds.get("gpu_util_p50_min")
if p50_threshold is not None:
p50_value = overall.get("p50")
if p50_value is not None:
assert (
p50_value >= p50_threshold
), f"GPU utilization p50 below threshold: {p50_value:.2f}% < {p50_threshold}%"
def should_monitor(thresholds: dict[str, Any] | None) -> bool:
"""Check if GPU monitoring should be enabled.
Returns True if:
- thresholds contains gpu_util_mean_min or gpu_util_p50_min, OR
- GPU_UTIL_LOG environment variable is truthy
"""
if thresholds:
if thresholds.get("gpu_util_mean_min") is not None:
return True
if thresholds.get("gpu_util_p50_min") is not None:
return True
return os.environ.get("GPU_UTIL_LOG", "").lower() in ("1", "true", "yes")

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,151 @@
"""Model specifications for E2E tests.
Each model spec defines:
- model: HuggingFace model path or local path
- memory_gb: Estimated GPU memory required
- tp: Tensor parallelism size (number of GPUs needed)
- features: List of features this model supports (for test filtering)
"""
from __future__ import annotations
import os
# Environment variable for local model paths (CI uses local copies for speed)
ROUTER_LOCAL_MODEL_PATH = os.environ.get("ROUTER_LOCAL_MODEL_PATH", "")
def _resolve_model_path(hf_path: str) -> str:
"""Resolve model path, preferring local path if available."""
if ROUTER_LOCAL_MODEL_PATH:
local_path = os.path.join(ROUTER_LOCAL_MODEL_PATH, hf_path)
if os.path.exists(local_path):
return local_path
return hf_path
MODEL_SPECS: dict[str, dict] = {
# Primary chat model - used for most tests
"llama-8b": {
"model": _resolve_model_path("meta-llama/Llama-3.1-8B-Instruct"),
"memory_gb": 16,
"tp": 1,
"features": ["chat", "streaming", "function_calling"],
},
# Small model for quick tests
"llama-1b": {
"model": _resolve_model_path("meta-llama/Llama-3.2-1B-Instruct"),
"memory_gb": 4,
"tp": 1,
"features": ["chat", "streaming", "tool_choice"],
},
# Function calling specialist
"qwen-7b": {
"model": _resolve_model_path("Qwen/Qwen2.5-7B-Instruct"),
"memory_gb": 14,
"tp": 1,
"features": ["chat", "streaming", "function_calling", "pythonic_tools"],
},
# Function calling specialist (larger, for Response API tests)
"qwen-14b": {
"model": _resolve_model_path("Qwen/Qwen2.5-14B-Instruct"),
"memory_gb": 28,
"tp": 2,
"features": ["chat", "streaming", "function_calling", "pythonic_tools"],
"worker_args": [
"--context-length=1000"
], # Faster startup, prevents memory issues
},
# Reasoning model
"deepseek-7b": {
"model": _resolve_model_path("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"),
"memory_gb": 14,
"tp": 1,
"features": ["chat", "streaming", "reasoning"],
},
# Thinking/reasoning model (larger)
"qwen-30b": {
"model": _resolve_model_path("Qwen/Qwen3-30B-A3B"),
"memory_gb": 60,
"tp": 4,
"features": ["chat", "streaming", "thinking", "reasoning"],
},
# Mistral for function calling
"mistral-7b": {
"model": _resolve_model_path("mistralai/Mistral-7B-Instruct-v0.3"),
"memory_gb": 14,
"tp": 1,
"features": ["chat", "streaming", "function_calling"],
},
# Embedding model
"embedding": {
"model": _resolve_model_path("intfloat/e5-mistral-7b-instruct"),
"memory_gb": 14,
"tp": 1,
"features": ["embedding"],
},
# GPT-OSS model (Harmony)
"gpt-oss": {
"model": _resolve_model_path("openai/gpt-oss-20b"),
"memory_gb": 40,
"tp": 2,
"features": ["chat", "streaming", "reasoning", "harmony"],
},
}
def get_models_with_feature(feature: str) -> list[str]:
"""Get list of model IDs that support a specific feature."""
return [
model_id
for model_id, spec in MODEL_SPECS.items()
if feature in spec.get("features", [])
]
def get_model_spec(model_id: str) -> dict:
"""Get spec for a specific model, raising KeyError if not found."""
if model_id not in MODEL_SPECS:
raise KeyError(
f"Unknown model: {model_id}. Available: {list(MODEL_SPECS.keys())}"
)
return MODEL_SPECS[model_id]
# Convenience groupings for test parametrization
CHAT_MODELS = get_models_with_feature("chat")
EMBEDDING_MODELS = get_models_with_feature("embedding")
REASONING_MODELS = get_models_with_feature("reasoning")
FUNCTION_CALLING_MODELS = get_models_with_feature("function_calling")
# =============================================================================
# Default model path constants (for backward compatibility with existing tests)
# =============================================================================
DEFAULT_MODEL_PATH = MODEL_SPECS["llama-8b"]["model"]
DEFAULT_SMALL_MODEL_PATH = MODEL_SPECS["llama-1b"]["model"]
DEFAULT_REASONING_MODEL_PATH = MODEL_SPECS["deepseek-7b"]["model"]
DEFAULT_ENABLE_THINKING_MODEL_PATH = MODEL_SPECS["qwen-30b"]["model"]
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH = MODEL_SPECS["qwen-7b"]["model"]
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH = MODEL_SPECS["mistral-7b"]["model"]
DEFAULT_GPT_OSS_MODEL_PATH = MODEL_SPECS["gpt-oss"]["model"]
DEFAULT_EMBEDDING_MODEL_PATH = MODEL_SPECS["embedding"]["model"]
# =============================================================================
# Third-party model configurations (cloud APIs)
# =============================================================================
THIRD_PARTY_MODELS: dict[str, dict] = {
"openai": {
"description": "OpenAI API",
"model": "gpt-5-nano",
"api_key_env": "OPENAI_API_KEY",
},
"xai": {
"description": "xAI API",
"model": "grok-4-fast",
"api_key_env": "XAI_API_KEY",
},
}

View File

@@ -0,0 +1,161 @@
"""Process management utilities for E2E tests."""
from __future__ import annotations
import logging
import os
import signal
import subprocess
import time
import requests
logger = logging.getLogger(__name__)
def kill_process_tree(pid: int, sig: int = signal.SIGTERM) -> None:
"""Kill a process and all its children.
Args:
pid: Process ID to kill
sig: Signal to send (default: SIGTERM)
"""
try:
import psutil
parent = psutil.Process(pid)
children = parent.children(recursive=True)
for child in children:
try:
child.send_signal(sig)
except psutil.NoSuchProcess:
pass
parent.send_signal(sig)
except ImportError:
# Fallback if psutil not available
os.kill(pid, sig)
except Exception as e:
logger.warning("Failed to kill process tree for PID %d: %s", pid, e)
def terminate_process(proc: subprocess.Popen, timeout: float = 30) -> None:
"""Gracefully terminate a process, kill if needed.
Args:
proc: Process to terminate
timeout: Seconds to wait before force-killing
"""
if proc is None or proc.poll() is not 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 wait_for_health(
url: str,
timeout: float = 60,
api_key: str | None = None,
check_interval: float = 1.0,
) -> None:
"""Wait for a server's /health endpoint to return 200.
Args:
url: Base URL of the server
timeout: Seconds to wait before timing out
api_key: Optional API key for auth header
check_interval: Seconds between health checks
"""
start = time.perf_counter()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
resp = session.get(f"{url}/health", headers=headers, timeout=5)
if resp.status_code == 200:
logger.info("Service healthy at %s", url)
return
except requests.RequestException:
pass
time.sleep(check_interval)
raise TimeoutError(f"Server at {url} did not become healthy within {timeout}s")
def wait_for_workers_ready(
router_url: str,
expected_workers: int,
timeout: float = 300,
api_key: str | None = None,
) -> None:
"""Wait for router to have all workers connected.
Args:
router_url: Base URL of the router
expected_workers: Number of workers to wait for
timeout: Seconds to wait before timing out
api_key: Optional API key for auth header
"""
start = time.perf_counter()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
while time.perf_counter() - start < timeout:
try:
resp = requests.get(f"{router_url}/workers", headers=headers, timeout=5)
if resp.status_code == 200:
data = resp.json()
total = data.get("total", len(data.get("workers", [])))
if total >= expected_workers:
logger.info(
"All %d workers connected after %.1fs",
expected_workers,
time.perf_counter() - start,
)
return
except requests.RequestException:
pass
time.sleep(2)
raise TimeoutError(
f"Router at {router_url} did not get {expected_workers} workers within {timeout}s"
)
def detect_ib_device() -> str | None:
"""Detect first active InfiniBand device (e.g., mlx5_0).
Returns:
Device name if found (e.g., "mlx5_0"), None otherwise.
"""
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:
logger.info("Detected IB device: %s", dev)
return dev
except Exception:
pass
return None

View File

@@ -0,0 +1,138 @@
"""MMLU evaluation runner for E2E tests.
Simplified evaluation runner that uses local eval implementations
with cleaner logging for CI/CD environments.
Usage:
from infra.run_eval import run_eval
from types import SimpleNamespace
args = SimpleNamespace(
base_url="http://127.0.0.1:30000",
model="meta-llama/Llama-3.1-8B-Instruct",
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
"""
from __future__ import annotations
import logging
import os
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from .simple_eval_common import Eval
from .simple_eval_common import ChatCompletionSampler, set_ulimit
from .simple_eval_mmlu import MMLU_DATASET_URL
logger = logging.getLogger(__name__)
@dataclass
class EvalConfig:
"""Configuration for running an evaluation."""
base_url: str
model: str | None = None
eval_name: str = "mmlu"
num_examples: int = 64
num_threads: int = 32
temperature: float = 0.0
max_tokens: int = 2048
host: str = "127.0.0.1"
port: int = 30000
def _get_eval(eval_name: str, num_examples: int, num_threads: int) -> "Eval":
"""Get the evaluation object by name."""
if eval_name == "mmlu":
from .simple_eval_mmlu import MMLUEval
return MMLUEval(MMLU_DATASET_URL, num_examples, num_threads)
else:
raise ValueError(f"Unknown eval: {eval_name}. Supported: mmlu")
def run_eval(args: Any) -> dict:
"""Run an evaluation and return metrics.
Args:
args: Configuration object with attributes:
- base_url: Base URL of the server (e.g., "http://127.0.0.1:30000")
- model: Model name/path (optional, will be auto-detected)
- eval_name: Evaluation name ("mmlu")
- num_examples: Number of examples to evaluate
- num_threads: Number of parallel threads
- temperature: Sampling temperature
Returns:
Dict with metrics including 'score' key.
"""
set_ulimit()
if "OPENAI_API_KEY" not in os.environ:
os.environ["OPENAI_API_KEY"] = "EMPTY"
# Build base URL
base_url = getattr(args, "base_url", None)
if base_url:
base_url = base_url.rstrip("/") # Remove trailing slashes
if not base_url.endswith("/v1"):
base_url = f"{base_url}/v1"
else:
host = getattr(args, "host", "127.0.0.1")
port = getattr(args, "port", 30000)
base_url = f"http://{host}:{port}/v1"
eval_name = getattr(args, "eval_name", "mmlu")
num_examples = getattr(args, "num_examples", 64)
num_threads = getattr(args, "num_threads", 32)
temperature = getattr(args, "temperature", 0.0)
max_tokens = getattr(args, "max_tokens", 2048)
model = getattr(args, "model", None)
logger.info(
"Starting %s eval: %d examples, %d threads, temp=%.2f",
eval_name,
num_examples,
num_threads,
temperature,
)
# Create sampler
sampler = ChatCompletionSampler(
model=model,
max_tokens=max_tokens,
base_url=base_url,
temperature=temperature,
)
# Get eval object
eval_obj = _get_eval(eval_name, num_examples, num_threads)
# Run evaluation
start_time = time.perf_counter()
result = eval_obj(sampler)
latency = time.perf_counter() - start_time
# Build metrics
metrics = result.metrics.copy() if result.metrics else {}
metrics["score"] = result.score
metrics["latency"] = latency
logger.info(
"%s eval complete: score=%.3f, latency=%.1fs, model=%s",
eval_name,
result.score,
latency,
sampler.model,
)
return metrics

View File

@@ -0,0 +1,492 @@
# Adapted from https://github.com/openai/simple-evals/
"""Common utilities for simple evaluations."""
from __future__ import annotations
import logging
import os
import resource
import time
from collections import defaultdict
from dataclasses import dataclass, field
from multiprocessing.pool import ThreadPool
from typing import Any
import httpx
import jinja2
import numpy as np
import openai
import requests
from openai import OpenAI
from tqdm import tqdm
from .constants import MAX_RETRY_ATTEMPTS
logger = logging.getLogger(__name__)
OPENAI_SYSTEM_MESSAGE_API = "You are a helpful assistant."
OPENAI_SYSTEM_MESSAGE_CHATGPT = (
"You are ChatGPT, a large language model trained by OpenAI, based on the GPT-4 architecture."
+ "\nKnowledge cutoff: 2023-12\nCurrent date: 2024-04-01"
)
Message = dict[str, Any] # keys role, content
MessageList = list[Message]
class SamplerBase:
"""
Base class for defining a sampling model, which can be evaluated,
or used as part of the grading process.
"""
def __call__(self, message_list: MessageList) -> str:
raise NotImplementedError()
@dataclass
class EvalResult:
"""Result of running an evaluation (usually consisting of many samples)."""
score: float | None # top-line metric
metrics: dict[str, float] | None # other metrics
htmls: list[str] # strings of valid HTML
convos: list[MessageList] # sampled conversations
@dataclass
class SingleEvalResult:
"""Result of evaluating a single sample."""
score: float | None
metrics: dict[str, float] = field(default_factory=dict)
html: str | None = None
convo: MessageList | None = None # sampled conversation
class Eval:
"""
Base class for defining an evaluation.
"""
def __call__(self, sampler: SamplerBase) -> EvalResult:
raise NotImplementedError()
class LargerHttpxClient(httpx.Client):
def __init__(self):
timeout_config = httpx.Timeout(3600)
limits = httpx.Limits(
max_keepalive_connections=3600,
max_connections=3600,
)
super().__init__(timeout=timeout_config, limits=limits)
class ChatCompletionSampler(SamplerBase):
"""Sample from OpenAI's chat completion API."""
def __init__(
self,
base_url: str | None = None,
model: str | None = None,
system_message: str | None = None,
temperature: float = 0.0,
reasoning_effort: str | None = None,
max_tokens: int = 2048,
extra_body: dict[str, Any] | None = None,
):
self.client = OpenAI(base_url=base_url, http_client=LargerHttpxClient())
if model is None:
model = self.client.models.list().data[0].id
self.model = model
self.system_message = system_message
self.temperature = temperature
self.max_tokens = max_tokens
self.reasoning_effort = reasoning_effort
self.extra_body = extra_body
self.image_format = "url"
logger.debug(
"ChatCompletionSampler: model=%s, temp=%.2f, max_tokens=%d",
self.model,
self.temperature,
self.max_tokens,
)
def _handle_image(
self,
image: str,
encoding: str = "base64",
format: str = "png",
):
new_image = {
"type": "image_url",
"image_url": {
"url": f"data:image/{format};{encoding},{image}",
},
}
return new_image
def _handle_text(self, text: str):
return {"type": "text", "text": text}
def _pack_message(self, role: str, content: Any):
return {"role": str(role), "content": content}
def __call__(self, message_list: MessageList) -> str:
if self.system_message:
message_list = [
self._pack_message("system", self.system_message)
] + message_list
trial = 0
while trial < MAX_RETRY_ATTEMPTS:
try:
response = self.client.chat.completions.create(
model=self.model,
messages=message_list,
temperature=self.temperature,
max_tokens=self.max_tokens,
reasoning_effort=self.reasoning_effort,
extra_body=self.extra_body,
)
return response.choices[0].message.content or ""
except openai.BadRequestError as e:
logger.warning("Bad request error: %s", e)
return ""
except Exception as e:
exception_backoff = 2**trial # exponential back off
# Log first few retries at debug, later ones at warning
log_fn = logger.warning if trial >= 3 else logger.debug
log_fn(
"Request failed (retry %d/%d, backoff %ds): %s",
trial + 1,
MAX_RETRY_ATTEMPTS,
exception_backoff,
e,
)
time.sleep(exception_backoff)
trial += 1
logger.warning(
"All retry attempts exhausted after %d retries, returning empty response",
MAX_RETRY_ATTEMPTS,
)
return ""
QUERY_TEMPLATE_MULTICHOICE = """
Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering.
{Question}
A) {A}
B) {B}
C) {C}
D) {D}
""".strip()
ANSWER_PATTERN_MULTICHOICE = r"(?i)Answer\s*:\s*([A-D])"
ANSWER_PATTERN = r"(?i)Answer\s*:\s*([^\n]+)"
EQUALITY_TEMPLATE = r"""
Look at the following two expressions (answers to a math problem) and judge whether they are equivalent. Only perform trivial simplifications
Examples:
Expression 1: $2x+3$
Expression 2: $3+2x$
Yes
Expression 1: 3/2
Expression 2: 1.5
Yes
Expression 1: $x^2+2x+1$
Expression 2: $y^2+2y+1$
No
Expression 1: $x^2+2x+1$
Expression 2: $(x+1)^2$
Yes
Expression 1: 3245/5
Expression 2: 649
No
(these are actually equal, don't mark them equivalent if you need to do nontrivial simplifications)
Expression 1: 2/(-3)
Expression 2: -2/3
Yes
(trivial simplifications are allowed)
Expression 1: 72 degrees
Expression 2: 72
Yes
(give benefit of the doubt to units)
Expression 1: 64
Expression 2: 64 square feet
Yes
(give benefit of the doubt to units)
---
YOUR TASK
Respond with only "Yes" or "No" (without quotes). Do not include a rationale.
Expression 1: %(expression1)s
Expression 2: %(expression2)s
""".strip()
HTML_JINJA = """
<h3>Prompt conversation</h3>
{% for message in prompt_messages %}
{{ message_to_html(message) | safe }}
{% endfor %}
<h3>Sampled message</h3>
{{ message_to_html(next_message) | safe }}
<h3>Results</h3>
<p>Correct Answer: {{ correct_answer }}</p>
<p>Extracted Answer: {{ extracted_answer }}</p>
<p>Score: {{ score }}</p>
"""
def format_multichoice_question(row):
return QUERY_TEMPLATE_MULTICHOICE.format(**row)
def check_equality(sampler: SamplerBase, expr1: str, expr2: str):
prompt = EQUALITY_TEMPLATE % {"expression1": expr1, "expression2": expr2}
response = sampler([dict(content=prompt, role="user")])
return (response or "").lower().strip() == "yes"
def _compute_stat(values: list, stat: str):
if stat == "mean":
return np.mean(values)
elif stat == "std":
return np.std(values)
elif stat == "min":
return np.min(values)
elif stat == "max":
return np.max(values)
else:
raise ValueError(f"Unknown {stat =}")
def aggregate_results(
single_eval_results: list[SingleEvalResult],
default_stats: tuple[str, ...] = ("mean", "std"),
name2stats: dict[str, tuple[str, ...]] | None = None,
) -> EvalResult:
"""
Aggregate results from multiple evaluations into a single EvalResult.
"""
name2stats = name2stats or {}
name2values = defaultdict(list)
htmls = []
convos = []
for single_eval_result in single_eval_results:
# Skip None results
if single_eval_result is None:
continue
for name, value in single_eval_result.metrics.items():
name2values[name].append(value)
if single_eval_result.score is not None:
name2values["score"].append(single_eval_result.score)
htmls.append(single_eval_result.html)
convos.append(single_eval_result.convo)
final_metrics = {}
for name, values in name2values.items():
stats = name2stats.get(name, default_stats)
for stat in stats:
key = name if stat == "mean" else f"{name}:{stat}"
final_metrics[key] = _compute_stat(values, stat)
return EvalResult(
score=final_metrics.pop("score", None),
metrics=final_metrics,
htmls=htmls,
convos=convos,
)
def map_with_progress(f: callable, xs: list[Any], num_threads: int) -> list[Any]:
"""Apply f to each element of xs, using a ThreadPool, and show progress."""
# Use quiet progress bar that doesn't pollute logs
if os.getenv("debug"):
return list(map(f, tqdm(xs, total=len(xs), leave=False)))
else:
with ThreadPool(min(num_threads, len(xs))) as pool:
return list(tqdm(pool.imap(f, xs), total=len(xs), leave=False))
jinja_env = jinja2.Environment(
loader=jinja2.BaseLoader(),
undefined=jinja2.StrictUndefined,
autoescape=jinja2.select_autoescape(["html", "xml"]),
)
_message_template = """
<div class="message {{ role }}">
<div class="role">
{{ role }}
{% if variant %}<span class="variant">({{ variant }})</span>{% endif %}
</div>
<div class="content">
<pre>{{ content }}</pre>
</div>
</div>
"""
def message_to_html(message: Message) -> str:
"""
Generate HTML snippet (inside a <div>) for a message.
"""
return jinja_env.from_string(_message_template).render(
role=message["role"],
content=message["content"],
variant=message.get("variant", None),
)
jinja_env.globals["message_to_html"] = message_to_html
_report_template = """<!DOCTYPE html>
<html>
<head>
<style>
.message {
padding: 8px 16px;
margin-bottom: 8px;
border-radius: 4px;
}
.message.user {
background-color: #B2DFDB;
color: #00695C;
}
.message.assistant {
background-color: #B39DDB;
color: #4527A0;
}
.message.system {
background-color: #EEEEEE;
color: #212121;
}
.role {
font-weight: bold;
margin-bottom: 4px;
}
.variant {
color: #795548;
}
table, th, td {
border: 1px solid black;
}
pre {
white-space: pre-wrap;
}
</style>
</head>
<body>
{% if metrics %}
<h1>Metrics</h1>
<table>
<tr>
<th>Metric</th>
<th>Value</th>
</tr>
<tr>
<td><b>Score</b></td>
<td>{{ score | float | round(3) }}</td>
</tr>
{% for name, value in metrics.items() %}
<tr>
<td>{{ name }}</td>
<td>{{ value }}</td>
</tr>
{% endfor %}
</table>
{% endif %}
<h1>Examples</h1>
{% for html in htmls %}
{{ html | safe }}
<hr>
{% endfor %}
</body>
</html>
"""
def make_report(eval_result: EvalResult) -> str:
"""
Create a standalone HTML report from an EvalResult.
"""
return jinja_env.from_string(_report_template).render(
score=eval_result.score,
metrics=eval_result.metrics,
htmls=eval_result.htmls,
)
def make_report_from_example_htmls(htmls: list[str]):
"""
Create a standalone HTML report from a list of example htmls
"""
return jinja_env.from_string(_report_template).render(
score=None, metrics={}, htmls=htmls
)
def download_dataset(path: str, url: str) -> None:
"""Download a dataset from URL to path."""
logger.info("Downloading dataset from %s", url)
try:
response = requests.get(url, stream=True, timeout=30)
response.raise_for_status()
total_size = int(response.headers.get("content-length", 0))
block_size = 8192
with open(path, "wb") as f, tqdm(
desc="Downloading",
total=total_size,
unit="iB",
unit_scale=True,
unit_divisor=1024,
leave=False,
) as progress_bar:
for data in response.iter_content(block_size):
size = f.write(data)
progress_bar.update(size)
logger.debug("Dataset saved to %s", path)
except requests.RequestException as e:
raise RuntimeError(f"Failed to download dataset: {e}") from e
def set_ulimit(target_soft_limit: int = 65535) -> None:
"""Set the file descriptor limit for parallel requests."""
resource_type = resource.RLIMIT_NOFILE
current_soft, current_hard = resource.getrlimit(resource_type)
if current_soft < target_soft_limit:
try:
resource.setrlimit(resource_type, (target_soft_limit, current_hard))
except ValueError as e:
logger.debug("Could not set RLIMIT_NOFILE: %s", e)

View File

@@ -0,0 +1,132 @@
# Adapted from https://github.com/openai/simple-evals/
"""
MMLU Evaluation - Measuring Massive Multitask Language Understanding
Dan Hendrycks et al. https://arxiv.org/abs/2009.03300
"""
from __future__ import annotations
import random
import re
from typing import TYPE_CHECKING
import pandas
from . import simple_eval_common as common
from .simple_eval_common import (
ANSWER_PATTERN_MULTICHOICE,
HTML_JINJA,
Eval,
EvalResult,
SingleEvalResult,
format_multichoice_question,
)
if TYPE_CHECKING:
from .simple_eval_common import SamplerBase
# MMLU dataset URL (hosted by OpenAI)
MMLU_DATASET_URL = "https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv"
SUBJECT_TO_CATEGORY = {
"abstract_algebra": "stem",
"anatomy": "other",
"astronomy": "stem",
"business_ethics": "other",
"clinical_knowledge": "other",
"college_biology": "stem",
"college_chemistry": "stem",
"college_computer_science": "stem",
"college_mathematics": "stem",
"college_medicine": "other",
"college_physics": "stem",
"computer_security": "stem",
"conceptual_physics": "stem",
"econometrics": "social_sciences",
"electrical_engineering": "stem",
"elementary_mathematics": "stem",
"formal_logic": "humanities",
"global_facts": "other",
"high_school_biology": "stem",
"high_school_chemistry": "stem",
"high_school_computer_science": "stem",
"high_school_european_history": "humanities",
"high_school_geography": "social_sciences",
"high_school_government_and_politics": "social_sciences",
"high_school_macroeconomics": "social_sciences",
"high_school_mathematics": "stem",
"high_school_microeconomics": "social_sciences",
"high_school_physics": "stem",
"high_school_psychology": "social_sciences",
"high_school_statistics": "stem",
"high_school_us_history": "humanities",
"high_school_world_history": "humanities",
"human_aging": "other",
"human_sexuality": "social_sciences",
"international_law": "humanities",
"jurisprudence": "humanities",
"logical_fallacies": "humanities",
"machine_learning": "stem",
"management": "other",
"marketing": "other",
"medical_genetics": "other",
"miscellaneous": "other",
"moral_disputes": "humanities",
"moral_scenarios": "humanities",
"nutrition": "other",
"philosophy": "humanities",
"prehistory": "humanities",
"professional_accounting": "other",
"professional_law": "humanities",
"professional_medicine": "other",
"professional_psychology": "social_sciences",
"public_relations": "social_sciences",
"security_studies": "social_sciences",
"sociology": "social_sciences",
"us_foreign_policy": "social_sciences",
"virology": "other",
"world_religions": "humanities",
}
class MMLUEval(Eval):
"""MMLU benchmark evaluation."""
def __init__(self, filename: str, num_examples: int | None, num_threads: int):
if "://" in filename:
df = pandas.read_csv(filename, storage_options={"timeout": 30})
else:
df = pandas.read_csv(filename)
examples = [row.to_dict() for _, row in df.iterrows()]
if num_examples:
examples = random.Random(0).sample(examples, num_examples)
self.examples = examples
self.num_threads = num_threads
def __call__(self, sampler: "SamplerBase") -> EvalResult:
def fn(row: dict) -> SingleEvalResult:
prompt_messages = [
sampler._pack_message(
content=format_multichoice_question(row), role="user"
)
]
response_text = sampler(prompt_messages)
response_text = response_text or ""
match = re.search(ANSWER_PATTERN_MULTICHOICE, response_text)
extracted_answer = match.group(1) if match else None
score = 1.0 if extracted_answer == row["Answer"] else 0.0
html = common.jinja_env.from_string(HTML_JINJA).render(
prompt_messages=prompt_messages,
next_message=dict(content=response_text, role="assistant"),
score=score,
correct_answer=row["Answer"],
extracted_answer=extracted_answer,
)
convo = prompt_messages + [dict(content=response_text, role="assistant")]
category = SUBJECT_TO_CATEGORY.get(row["Subject"], "other")
return SingleEvalResult(
html=html, score=score, metrics={category: score}, convo=convo
)
results = common.map_with_progress(fn, self.examples, self.num_threads)
return common.aggregate_results(results)

View File

@@ -0,0 +1,44 @@
[project]
name = "sgl-model-gateway-e2e-tests"
version = "0.1.0"
description = "E2E tests for sgl-model-gateway"
requires-python = ">=3.9"
dependencies = [
"grpcio",
"grpcio-health-checking",
"httpx",
"openai",
"py", # Required for pytest-parallel with newer pytest versions
"pytest",
"pytest-parallel",
"pytest-rerunfailures",
]
[project.optional-dependencies]
dev = [
"ruff",
]
[tool.pytest.ini_options]
testpaths = ["."]
markers = [
"e2e: mark test as end-to-end test requiring GPU workers",
"slow: mark test as slow-running",
"thread_unsafe: mark test as incompatible with parallel thread execution",
]
addopts = "-v -s"
# Explicitly disable live log to avoid "---- live log ----" dividers
# We configure logging manually in conftest.py
log_cli = false
# Parallel execution configuration:
# Use --workers 1 --tests-per-worker N to run N tests concurrently as threads
# within a single process. This enables true shared-worker parallelism where
# the session-scoped model_pool fixture is shared across all threads.
#
# Example usage:
# pytest --workers 1 --tests-per-worker 4 e2e_test/router/
#
# The thread-safe ModelPool and GPUAllocator classes enable safe concurrent
# access from multiple test threads.

View File

@@ -0,0 +1,9 @@
"""Response API E2E tests.
Tests for the Response API endpoints including:
- Basic CRUD operations for responses and conversations
- State management (previous_response_id, conversation-based)
- Streaming events and output validation
- Structured output (json_schema)
- Tool calling (function tools and MCP)
"""

View File

@@ -0,0 +1,441 @@
"""Basic CRUD tests for Response API.
Tests for Response and Conversation CRUD operations against cloud backends.
Source: Migrated from e2e_response_api/features/test_basic_crud.py
"""
from __future__ import annotations
import logging
import time
import openai
import pytest
from openai import OpenAI
from openai.types import responses
logger = logging.getLogger(__name__)
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
if resp.status == "completed":
return resp
elif resp.status == "failed":
raise AssertionError(f"Background task failed: {resp.error}")
elif resp.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"
)
# =============================================================================
# Response CRUD Tests (Memory Storage)
# =============================================================================
@pytest.mark.parametrize("setup_backend", ["openai"], indirect=True)
class TestResponseCRUD:
"""Tests for Response API CRUD operations."""
def test_create_and_get_response(self, setup_backend):
"""Test creating response and retrieving it."""
_, model, client, gateway = 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, gateway = 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, gateway = 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"
# =============================================================================
# Response CRUD Tests (Oracle Storage)
# =============================================================================
@pytest.mark.storage("oracle")
@pytest.mark.parametrize("setup_backend", ["openai"], indirect=True)
class TestResponseCRUDOracleStorage:
"""Tests for Response API CRUD operations with Oracle history backend."""
def test_create_and_get_response(self, setup_backend):
"""Test creating response and retrieving it."""
_, model, client, gateway = 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, gateway = 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, gateway = 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"
# =============================================================================
# Conversation CRUD Tests (Memory Storage)
# =============================================================================
@pytest.mark.parametrize("setup_backend", ["openai"], indirect=True)
class TestConversationCRUD:
"""Tests for Conversation API CRUD operations."""
def test_create_and_get_conversation(self, setup_backend):
"""Test creating and retrieving conversation."""
_, model, client, gateway = 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, gateway = 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 get_data["key1"] == "value1"
assert get_data["key2"] == "value2"
def test_delete_conversation(self, setup_backend):
"""Test deleting conversation."""
_, model, client, gateway = 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, gateway = 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
# =============================================================================
# Conversation CRUD Tests (Oracle Storage)
# =============================================================================
@pytest.mark.storage("oracle")
@pytest.mark.parametrize("setup_backend", ["openai"], indirect=True)
class TestConversationCRUDOracleStorage:
"""Tests for Conversation API CRUD operations with Oracle history backend."""
def test_create_and_get_conversation(self, setup_backend):
"""Test creating and retrieving conversation."""
_, model, client, gateway = 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, gateway = 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 get_data["key1"] == "value1"
assert get_data["key2"] == "value2"
def test_delete_conversation(self, setup_backend):
"""Test deleting conversation."""
_, model, client, gateway = 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, gateway = 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

View File

@@ -0,0 +1,349 @@
"""State management tests for Response API.
Tests both previous_response_id and conversation-based state management.
These tests work across local (gRPC) and cloud (OpenAI, xAI) backends.
Source: Migrated from e2e_response_api/features/test_state_management.py
"""
from __future__ import annotations
import logging
import openai
import pytest
logger = logging.getLogger(__name__)
# =============================================================================
# Cloud Backend Tests (OpenAI, xAI)
# =============================================================================
@pytest.mark.parametrize("setup_backend", ["openai", "xai"], indirect=True)
class TestStateManagementCloud:
"""State management tests against cloud APIs."""
def test_basic_response_creation(self, setup_backend):
"""Test basic response creation without state."""
_, model, client, gateway = 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, gateway = setup_backend
resp = client.responses.create(
model=model, input="Count to 5", stream=True, max_output_tokens=50
)
events = list(resp)
created_events = [e for e in events if e.type == "response.created"]
assert len(created_events) > 0
assert any(
e.type in ["response.completed", "response.in_progress"] for e in events
)
def test_previous_response_id_chaining(self, setup_backend):
"""Test chaining responses using previous_response_id."""
_, model, client, gateway = 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"
# Second response referencing first
resp2 = client.responses.create(
model=model, input="What is my name", previous_response_id=resp1.id
)
assert resp2.error is None
assert resp2.status == "completed"
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
def test_conversation_with_multiple_turns(self, setup_backend):
"""Test state management using conversation ID."""
_, model, client, gateway = setup_backend
# 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"
assert "5" in resp2.output_text or "five" in resp2.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"
assert "8" in resp3.output_text or "eight" in resp3.output_text.lower()
items = client.conversations.items.list(conversation_id)
assert items.data is not None
assert len(items.data) >= 6 # 3 inputs + 3 outputs
@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, gateway = 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_mutually_exclusive_parameters(self, setup_backend):
"""Test that previous_response_id and conversation are mutually exclusive."""
_, model, client, gateway = setup_backend
conversation_id = "conv_123"
resp1 = client.responses.create(model=model, input="Test")
with pytest.raises(openai.BadRequestError):
client.responses.create(
model=model,
input="This should fail",
previous_response_id=resp1.id,
conversation=conversation_id,
)
# =============================================================================
# Local Backend Tests (gRPC with Qwen model)
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("qwen-14b")
@pytest.mark.gateway(
extra_args=["--tool-call-parser", "qwen", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestStateManagementLocal:
"""State management tests against local gRPC backend."""
@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, gateway = 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_basic_response_creation(self, setup_backend):
"""Test basic response creation without state."""
_, model, client, gateway = 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, gateway = setup_backend
resp = client.responses.create(
model=model, input="Count to 5", stream=True, max_output_tokens=50
)
events = list(resp)
created_events = [e for e in events if e.type == "response.created"]
assert len(created_events) > 0
assert any(
e.type in ["response.completed", "response.in_progress"] for e in events
)
def test_previous_response_id_chaining(self, setup_backend):
"""Test chaining responses using previous_response_id."""
_, model, client, gateway = 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"
# Second response referencing first
resp2 = client.responses.create(
model=model, input="What is my name", previous_response_id=resp1.id
)
assert resp2.error is None
assert resp2.status == "completed"
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
def test_mutually_exclusive_parameters(self, setup_backend):
"""Test that previous_response_id and conversation are mutually exclusive."""
_, model, client, gateway = setup_backend
conversation_id = "conv_123"
resp1 = client.responses.create(model=model, input="Test")
with pytest.raises(openai.BadRequestError):
client.responses.create(
model=model,
input="This should fail",
previous_response_id=resp1.id,
conversation=conversation_id,
)
# =============================================================================
# Local Backend Tests (gRPC with Harmony/Reasoning model)
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("gpt-oss")
@pytest.mark.gateway(
extra_args=["--reasoning-parser=gpt-oss", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestStateManagementHarmony:
"""State management tests against local gRPC backend with Harmony model."""
@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, gateway = 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_basic_response_creation(self, setup_backend):
"""Test basic response creation without state."""
_, model, client, gateway = 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, gateway = setup_backend
resp = client.responses.create(
model=model, input="Count to 5", stream=True, max_output_tokens=50
)
events = list(resp)
created_events = [e for e in events if e.type == "response.created"]
assert len(created_events) > 0
assert any(
e.type in ["response.completed", "response.in_progress"] for e in events
)
def test_previous_response_id_chaining(self, setup_backend):
"""Test chaining responses using previous_response_id."""
_, model, client, gateway = 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"
# Second response referencing first
resp2 = client.responses.create(
model=model, input="What is my name", previous_response_id=resp1.id
)
assert resp2.error is None
assert resp2.status == "completed"
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
def test_mutually_exclusive_parameters(self, setup_backend):
"""Test that previous_response_id and conversation are mutually exclusive."""
_, model, client, gateway = setup_backend
conversation_id = "conv_123"
resp1 = client.responses.create(model=model, input="Test")
with pytest.raises(openai.BadRequestError):
client.responses.create(
model=model,
input="This should fail",
previous_response_id=resp1.id,
conversation=conversation_id,
)

View File

@@ -0,0 +1,256 @@
"""Streaming events tests for Response API.
Tests for streaming event validation including:
- Zero-based output_index for content
- OutputItemDone event emission and output array construction
- Reasoning content with proper output_index
Source: Migrated from e2e_response_api/features/test_streaming_events.py
"""
from __future__ import annotations
import logging
import pytest
logger = logging.getLogger(__name__)
# =============================================================================
# Local Backend Tests (gRPC with Qwen model)
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("qwen-14b")
@pytest.mark.gateway(
extra_args=["--tool-call-parser", "qwen", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestStreamingEventsLocal:
"""Streaming event tests against local gRPC backend."""
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, gateway = setup_backend
resp = client.responses.create(
model=model,
input="Count from 1 to 3",
stream=True,
max_output_tokens=50,
)
events = list(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 item in 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"
# =============================================================================
# Local Backend Tests (gRPC with Harmony/Reasoning model)
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("gpt-oss")
@pytest.mark.gateway(
extra_args=["--reasoning-parser=gpt-oss", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestStreamingEventsHarmony:
"""Streaming event tests against local gRPC backend with Harmony model."""
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, gateway = setup_backend
resp = client.responses.create(
model=model,
input="Count from 1 to 3",
stream=True,
max_output_tokens=50,
)
events = list(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 item in 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.
"""
_, model, client, gateway = setup_backend
resp = client.responses.create(
model=model,
input="What is the capital of France? Think step by step.",
stream=True,
max_output_tokens=200,
)
events = list(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

View File

@@ -0,0 +1,289 @@
"""Structured output tests for Response API.
Tests for text.format field with json_object and json_schema formats.
Source: Migrated from e2e_response_api/features/test_structured_output.py
"""
from __future__ import annotations
import json
import logging
import pytest
logger = logging.getLogger(__name__)
# =============================================================================
# Cloud Backend Tests (OpenAI)
# =============================================================================
@pytest.mark.parametrize("setup_backend", ["openai"], indirect=True)
class TestStructuredOutputCloud:
"""Structured output tests against cloud APIs."""
def test_structured_output_json_schema(self, setup_backend):
"""Test structured output with json_schema format."""
_, model, client, gateway = setup_backend
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_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
# =============================================================================
# Local Backend Tests (gRPC with Harmony model - complex schema)
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("gpt-oss")
@pytest.mark.gateway(
extra_args=["--reasoning-parser=gpt-oss", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestStructuredOutputHarmony:
"""Structured output tests against local gRPC backend with Harmony model."""
def test_structured_output_json_schema(self, setup_backend):
"""Test structured output with json_schema format."""
_, model, client, gateway = setup_backend
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
# =============================================================================
# Local Backend Tests (gRPC with Qwen model - simple schema)
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("qwen-14b")
@pytest.mark.gateway(
extra_args=["--tool-call-parser", "qwen", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestSimpleSchemaStructuredOutput:
"""Structured output tests with simpler schema for models that don't
handle complex schemas well.
"""
def test_structured_output_json_schema(self, setup_backend):
"""Test structured output with simple json_schema format."""
_, model, client, gateway = 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"

View File

@@ -0,0 +1,902 @@
"""Tool calling tests for Response API.
Tests for function calling functionality, tool choices and MCP calling
functionality across different backends.
Source: Migrated from e2e_response_api/features/test_tools_call.py
"""
from __future__ import annotations
import json
import logging
import time
import pytest
logger = logging.getLogger(__name__)
# =============================================================================
# Shared 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"],
},
}
GET_HOROSCOPE_FUNCTION = {
"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"],
},
}
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"
)
# =============================================================================
# Cloud Backend Tests (OpenAI) - Basic Function Calling
# =============================================================================
@pytest.mark.parametrize("setup_backend", ["openai"], indirect=True)
class TestToolCallingCloud:
"""Tool calling tests against cloud APIs."""
def test_basic_function_call(self, setup_backend):
"""Test basic function calling workflow."""
_, model, client, gateway = setup_backend
tools = [GET_HOROSCOPE_FUNCTION]
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. "
)
input_list = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": "What is my horoscope? I am an Aquarius."},
]
resp = client.responses.create(model=model, input=input_list, tools=tools)
assert resp.error is None
assert resp.id is not None
assert resp.status == "completed"
assert resp.output is not None
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 == "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"
# Provide function call output
input_list.append(function_call)
horoscope = f"{args['sign']}: Next Tuesday you will befriend a baby otter."
input_list.append(
{
"type": "function_call_output",
"call_id": function_call.call_id,
"output": json.dumps({"horoscope": horoscope}),
}
)
# 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"
output2 = resp2.output
assert len(output2) > 0
messages = [item for item in output2 if item.type == "message"]
assert len(messages) > 0
message = messages[0]
assert message.content is not None
text_parts = [
part.text for part in message.content if part.type == "output_text"
]
full_text = " ".join(text_parts).lower()
assert "baby otter" in full_text or "aquarius" in full_text
def test_mcp_basic_tool_call(self, setup_backend):
"""Test basic MCP tool call (non-streaming)."""
_, model, client, gateway = setup_backend
time.sleep(2) # Avoid rate limiting
resp = client.responses.create(
model=model,
input=MCP_TEST_PROMPT,
tools=[BRAVE_MCP_TOOL],
stream=False,
reasoning={"effort": "low"},
)
assert resp.error is None
assert resp.id is not None
assert resp.status == "completed"
assert resp.model is not None
assert resp.output is not None
assert len(resp.output_text) > 0
output_types = [item.type for item in resp.output]
assert "mcp_list_tools" in output_types
mcp_calls = [item for item in resp.output if item.type == "mcp_call"]
assert len(mcp_calls) > 0
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 validation for cloud backends
messages = [item for item in resp.output if item.type == "message"]
assert len(messages) > 0, "Response should contain at least one message"
for msg in messages:
assert msg.content is not None
assert isinstance(msg.content, list)
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)."""
_, model, client, gateway = setup_backend
time.sleep(2) # Avoid rate limiting
resp = client.responses.create(
model=model,
input=MCP_TEST_PROMPT,
tools=[BRAVE_MCP_TOOL],
stream=True,
reasoning={"effort": "low"},
)
events = list(resp)
assert len(events) > 0
event_types = [event.type for event in events]
assert "response.created" in event_types, "Should have response.created event"
assert (
"response.completed" in event_types
), "Should have response.completed event"
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"
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"
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
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 validation for cloud backends - 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"
# =============================================================================
# Local Backend Tests (gRPC with Harmony model) - Tool Choice
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("gpt-oss")
@pytest.mark.gateway(
extra_args=["--reasoning-parser=gpt-oss", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestToolChoiceHarmony:
"""Tool choice tests against local gRPC backend with Harmony model."""
def test_tool_choice_auto(self, setup_backend):
"""Test tool_choice="auto" allows model to decide whether to use tools."""
_, model, client, gateway = setup_backend
tools = [GET_WEATHER_FUNCTION]
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
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."""
_, model, client, gateway = setup_backend
tools = [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
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."""
_, model, client, gateway = setup_backend
tools = [SEARCH_WEB_FUNCTION, GET_WEATHER_FUNCTION]
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
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0, "Must call the specified function"
assert (
function_calls[0].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."""
_, model, client, gateway = setup_backend
tools = [CALCULATE_FUNCTION]
resp = client.responses.create(
model=model,
input="Calculate 42 * 17",
tools=tools,
tool_choice="required",
stream=True,
)
events = list(resp)
assert len(events) > 0
event_types = [e.type for e in events]
assert "response.function_call_arguments.delta" in event_types
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
def test_tool_choice_with_mcp_tools(self, setup_backend):
"""Test tool_choice parameter works with MCP tools."""
_, model, client, gateway = setup_backend
tools = [DEEPWIKI_MCP_TOOL]
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
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."""
_, model, client, gateway = setup_backend
tools = [DEEPWIKI_MCP_TOOL, LOCAL_SEARCH_FUNCTION]
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
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0
assert function_calls[0].name == "local_search"
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"
def test_basic_function_call(self, setup_backend):
"""Test basic function calling workflow."""
_, model, client, gateway = setup_backend
tools = [GET_HOROSCOPE_FUNCTION]
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. "
)
input_list = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": "What is my horoscope? I am an Aquarius."},
]
resp = client.responses.create(model=model, input=input_list, tools=tools)
assert resp.error is None
assert resp.id is not None
assert resp.status == "completed"
output = resp.output
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0
function_call = function_calls[0]
assert function_call.name == "get_horoscope"
args = json.loads(function_call.arguments)
assert "sign" in args
assert args["sign"].lower() == "aquarius"
def test_mcp_basic_tool_call(self, setup_backend):
"""Test basic MCP tool call (non-streaming)."""
_, model, client, gateway = setup_backend
time.sleep(2)
resp = client.responses.create(
model=model,
input=MCP_TEST_PROMPT,
tools=[BRAVE_MCP_TOOL],
stream=False,
reasoning={"effort": "low"},
)
assert resp.error is None
assert resp.id is not None
assert resp.status == "completed"
assert len(resp.output_text) > 0
output_types = [item.type for item in resp.output]
assert "mcp_list_tools" in output_types
mcp_calls = [item for item in resp.output if item.type == "mcp_call"]
assert len(mcp_calls) > 0
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"
def test_mcp_basic_tool_call_streaming(self, setup_backend):
"""Test basic MCP tool call (streaming)."""
_, model, client, gateway = setup_backend
time.sleep(2)
resp = client.responses.create(
model=model,
input=MCP_TEST_PROMPT,
tools=[BRAVE_MCP_TOOL],
stream=True,
reasoning={"effort": "low"},
)
events = list(resp)
assert len(events) > 0
event_types = [event.type for event in events]
assert "response.created" in event_types
assert "response.completed" in event_types
assert "response.mcp_list_tools.completed" in event_types
assert "response.mcp_call.completed" in event_types
def test_mixed_mcp_and_function_tools(self, setup_backend):
"""Test mixed MCP and function tools (non-streaming)."""
_, model, client, gateway = setup_backend
resp = client.responses.create(
model=model,
input="Give me diagnostics for the Astra-7 Core Reactor.",
tools=[BRAVE_MCP_TOOL, SYSTEM_DIAGNOSTICS_FUNCTION],
stream=False,
tool_choice="auto",
)
assert resp.error is None
assert resp.id is not None
assert resp.output is not None
output = resp.output
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0
system_diagnostics_call = function_calls[0]
assert system_diagnostics_call.name == "get_system_diagnostics"
assert system_diagnostics_call.call_id is not None
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)."""
_, model, client, gateway = setup_backend
resp = client.responses.create(
model=model,
input="Give me diagnostics for the Astra-7 Core Reactor.",
tools=[BRAVE_MCP_TOOL, SYSTEM_DIAGNOSTICS_FUNCTION],
stream=True,
tool_choice="auto",
)
events = list(resp)
assert len(events) > 0
event_types = [e.type for e in events]
assert "response.created" in event_types
assert "response.mcp_list_tools.completed" in event_types
assert "response.function_call_arguments.delta" in event_types
assert "response.function_call_arguments.done" in event_types
func_arg_deltas = [
e for e in events if e.type == "response.function_call_arguments.delta"
]
assert len(func_arg_deltas) > 0
full_delta_event = "".join(e.delta for e in func_arg_deltas)
assert (
"system_name" in full_delta_event.lower()
and "astra-7" in full_delta_event.lower()
)
# =============================================================================
# Local Backend Tests (gRPC with Qwen model) - Tool Choice
# =============================================================================
@pytest.mark.e2e
@pytest.mark.model("qwen-14b")
@pytest.mark.gateway(
extra_args=["--tool-call-parser", "qwen", "--history-backend", "memory"]
)
@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True)
class TestToolChoiceLocal:
"""Tool choice tests against local gRPC backend with Qwen model."""
def test_tool_choice_auto(self, setup_backend):
"""Test tool_choice="auto" allows model to decide whether to use tools."""
_, model, client, gateway = setup_backend
tools = [GET_WEATHER_FUNCTION]
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
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0
def test_tool_choice_required(self, setup_backend):
"""Test tool_choice="required" forces the model to call at least one tool."""
_, model, client, gateway = setup_backend
tools = [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
function_calls = [item for item in resp.output if item.type == "function_call"]
assert len(function_calls) > 0
def test_tool_choice_specific_function(self, setup_backend):
"""Test tool_choice with specific function name forces that function to be called."""
_, model, client, gateway = setup_backend
tools = [SEARCH_WEB_FUNCTION, GET_WEATHER_FUNCTION]
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
function_calls = [item for item in resp.output if item.type == "function_call"]
assert len(function_calls) > 0
assert function_calls[0].name == "search_web"
def test_mcp_basic_tool_call(self, setup_backend):
"""Test basic MCP tool call (non-streaming)."""
_, model, client, gateway = setup_backend
time.sleep(2)
resp = client.responses.create(
model=model,
input=MCP_TEST_PROMPT,
tools=[BRAVE_MCP_TOOL],
stream=False,
reasoning={"effort": "low"},
)
assert resp.error is None
assert resp.id is not None
assert resp.status == "completed"
output_types = [item.type for item in resp.output]
assert "mcp_list_tools" in output_types
mcp_calls = [item for item in resp.output if item.type == "mcp_call"]
assert len(mcp_calls) > 0
def test_mcp_basic_tool_call_streaming(self, setup_backend):
"""Test basic MCP tool call (streaming)."""
_, model, client, gateway = setup_backend
time.sleep(2)
resp = client.responses.create(
model=model,
input=MCP_TEST_PROMPT,
tools=[BRAVE_MCP_TOOL],
stream=True,
reasoning={"effort": "low"},
)
events = list(resp)
assert len(events) > 0
event_types = [event.type for event in events]
assert "response.created" in event_types
assert "response.completed" in event_types
def test_tool_choice_with_mcp_tools(self, setup_backend):
"""Test tool_choice parameter works with MCP tools."""
_, model, client, gateway = setup_backend
tools = [DEEPWIKI_MCP_TOOL]
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
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."""
_, model, client, gateway = setup_backend
tools = [DEEPWIKI_MCP_TOOL, LOCAL_SEARCH_FUNCTION]
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
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0
assert function_calls[0].name == "local_search"
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"
def test_mixed_mcp_and_function_tools(self, setup_backend):
"""Test mixed MCP and function tools (non-streaming)."""
_, model, client, gateway = setup_backend
resp = client.responses.create(
model=model,
input="Give me diagnostics for the Astra-7 Core Reactor.",
tools=[BRAVE_MCP_TOOL, SYSTEM_DIAGNOSTICS_FUNCTION],
stream=False,
tool_choice="auto",
)
assert resp.error is None
assert resp.id is not None
assert resp.output is not None
output = resp.output
function_calls = [item for item in output if item.type == "function_call"]
assert len(function_calls) > 0
system_diagnostics_call = function_calls[0]
assert system_diagnostics_call.name == "get_system_diagnostics"
assert system_diagnostics_call.call_id is not None
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)."""
_, model, client, gateway = setup_backend
resp = client.responses.create(
model=model,
input="Give me diagnostics for the Astra-7 Core Reactor.",
tools=[BRAVE_MCP_TOOL, SYSTEM_DIAGNOSTICS_FUNCTION],
stream=True,
tool_choice="auto",
)
events = list(resp)
assert len(events) > 0
event_types = [e.type for e in events]
assert "response.created" in event_types
assert "response.mcp_list_tools.completed" in event_types
assert "response.function_call_arguments.delta" in event_types
assert "response.function_call_arguments.done" in event_types
func_arg_deltas = [
e for e in events if e.type == "response.function_call_arguments.delta"
]
assert len(func_arg_deltas) > 0
full_delta_event = "".join(e.delta for e in func_arg_deltas)
assert (
"system_name" in full_delta_event.lower()
and "astra-7" in full_delta_event.lower()
)

View File

@@ -0,0 +1,76 @@
"""MMLU evaluation tests for router functionality.
Tests the router's ability to handle MMLU benchmark evaluations across
different backend configurations (gRPC and HTTP workers).
Usage:
# Run with gRPC backend only
pytest e2e_test/router/test_mmlu.py -v
# Run with specific backend
pytest e2e_test/router/test_mmlu.py -v -k "grpc"
"""
from __future__ import annotations
import logging
from types import SimpleNamespace
import pytest
from infra import run_eval
logger = logging.getLogger(__name__)
@pytest.mark.e2e
@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True)
class TestMMLU:
"""MMLU evaluation tests using local workers (gRPC and HTTP)."""
def test_mmlu_basic(self, setup_backend):
"""Basic MMLU evaluation with score threshold.
Runs MMLU evaluation with 64 examples and validates that
accuracy meets minimum threshold (>= 0.65).
Note: setup_backend fixture already waits for workers to be ready.
"""
backend, model, client, *_ = setup_backend
args = SimpleNamespace(
base_url=str(client.base_url),
model=model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert (
metrics["score"] >= 0.65
), f"MMLU score {metrics['score']:.2f} below threshold 0.65"
logger.info("MMLU score: %.2f (threshold: 0.65)", metrics["score"])
def test_mmlu_extended(self, setup_backend):
"""Extended MMLU evaluation with more examples.
Runs MMLU with 128 examples for more statistically
significant results.
"""
backend, model, client, *_ = setup_backend
args = SimpleNamespace(
base_url=str(client.base_url),
model=model,
eval_name="mmlu",
num_examples=128,
num_threads=64,
temperature=0.1,
)
metrics = run_eval(args)
assert (
metrics["score"] >= 0.65
), f"MMLU score {metrics['score']:.2f} below threshold 0.65"
logger.info("MMLU extended score: %.2f (threshold: 0.65)", metrics["score"])

View File

@@ -0,0 +1,61 @@
"""MMLU evaluation tests for PD (Prefill-Decode) disaggregated routing.
PD disaggregation separates prefill and decode phases across different
workers for improved throughput and resource utilization.
Requirements:
- sgl_kernel package
- GPUs: num_prefill + num_decode (default: 2 GPUs for 1+1)
- Optional: InfiniBand for high-performance transfers
Configuration via markers:
@pytest.mark.model("model-id") # Override default model
@pytest.mark.workers(prefill=2, decode=2) # Custom worker counts
@pytest.mark.gateway(policy="round_robin") # Gateway configuration
Usage:
# Basic (1 prefill + 1 decode)
pytest e2e_test/router/test_pd_mmlu.py -v
# Run specific test
pytest e2e_test/router/test_pd_mmlu.py::TestPDMMLU::test_pd_mmlu_basic -v
"""
from __future__ import annotations
import logging
from types import SimpleNamespace
import pytest
from infra import run_eval
logger = logging.getLogger(__name__)
@pytest.mark.e2e
@pytest.mark.parametrize("setup_backend", ["pd"], indirect=True)
class TestPDMMLU:
"""MMLU evaluation tests using PD disaggregated routing."""
def test_pd_mmlu_basic(self, setup_backend):
"""Basic MMLU evaluation with PD disaggregation.
Runs MMLU with 1 prefill + 1 decode worker and validates
accuracy meets threshold (>= 0.65).
"""
backend, model, client, *_ = setup_backend
args = SimpleNamespace(
base_url=str(client.base_url),
model=model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
assert (
metrics["score"] >= 0.65
), f"PD MMLU score {metrics['score']:.2f} below threshold 0.65"
logger.info("PD MMLU score: %.2f (threshold: 0.65)", metrics["score"])

View File

@@ -0,0 +1,220 @@
"""Tests for gateway worker management APIs.
Tests the gateway's worker management endpoints:
- GET /workers - List all workers
- POST /add_worker - Add a worker dynamically
- POST /remove_worker - Remove a worker dynamically
- GET /v1/models - List available models
Usage:
pytest e2e_test/router/test_worker_api.py -v
"""
from __future__ import annotations
import logging
import pytest
from infra import ConnectionMode, Gateway, ModelPool
logger = logging.getLogger(__name__)
@pytest.mark.e2e
@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True)
class TestWorkerAPI:
"""Tests for worker management APIs using setup_backend fixture."""
def test_list_workers(self, setup_backend):
"""Test listing workers via /workers endpoint."""
backend, model, client, gateway = setup_backend
workers = gateway.list_workers()
assert len(workers) >= 1, "Expected at least one worker"
logger.info("Found %d workers", len(workers))
for worker in workers:
logger.info(
"Worker: id=%s, url=%s, status=%s",
worker.id,
worker.url,
worker.status,
)
assert worker.url, "Worker should have a URL"
def test_list_models(self, setup_backend):
"""Test listing models via /v1/models endpoint."""
backend, model, client, gateway = setup_backend
models = gateway.list_models()
assert len(models) >= 1, "Expected at least one model"
logger.info("Found %d models", len(models))
for m in models:
logger.info("Model: %s", m.get("id", "unknown"))
assert "id" in m, "Model should have an id"
def test_health_endpoint(self, setup_backend):
"""Test health check endpoint."""
backend, model, client, gateway = setup_backend
assert gateway.health(), "Gateway should be healthy"
logger.info("Gateway health check passed")
@pytest.mark.e2e
class TestIGWMode:
"""Tests for IGW mode - start gateway empty, add workers via API.
Workers are launched on-demand via model_pool.get().
"""
def test_igw_start_empty(self, model_pool: ModelPool):
"""Test starting gateway in IGW mode with no workers."""
gateway = Gateway()
gateway.start(igw_mode=True)
try:
assert gateway.health(), "Gateway should be healthy"
assert gateway.igw_mode, "Gateway should be in IGW mode"
workers = gateway.list_workers()
logger.info("IGW gateway started with %d workers", len(workers))
finally:
gateway.shutdown()
def test_igw_add_worker(self, model_pool: ModelPool):
"""Test adding a worker to IGW gateway."""
http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP)
gateway = Gateway()
gateway.start(igw_mode=True)
try:
# Add worker
success, result = gateway.add_worker(http_instance.worker_url)
assert success, f"Failed to add worker: {result}"
logger.info("Added worker: %s", result)
# Verify worker was added
workers = gateway.list_workers()
assert len(workers) >= 1, "Expected at least one worker"
logger.info("Worker count: %d", len(workers))
# Verify models are available
models = gateway.list_models()
logger.info("Models available: %d", len(models))
finally:
gateway.shutdown()
def test_igw_add_and_remove_worker(self, model_pool: ModelPool):
"""Test adding and removing workers dynamically."""
http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP)
gateway = Gateway()
gateway.start(igw_mode=True)
try:
# Add worker
success, _ = gateway.add_worker(http_instance.worker_url)
assert success, "Failed to add worker"
initial_count = len(gateway.list_workers())
logger.info("Worker count after add: %d", initial_count)
# Remove worker
success, msg = gateway.remove_worker(http_instance.worker_url)
if success:
logger.info("Removed worker: %s", msg)
final_count = len(gateway.list_workers())
logger.info("Worker count after remove: %d", final_count)
else:
logger.warning("Remove worker not supported: %s", msg)
finally:
gateway.shutdown()
def test_igw_multiple_workers(self, model_pool: ModelPool):
"""Test adding multiple workers (HTTP + gRPC) to IGW gateway."""
http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP)
grpc_instance = model_pool.get("llama-8b", ConnectionMode.GRPC)
gateway = Gateway()
gateway.start(igw_mode=True)
try:
# Add both workers
success1, _ = gateway.add_worker(http_instance.worker_url)
success2, _ = gateway.add_worker(grpc_instance.worker_url)
if not success1 or not success2:
pytest.skip("Dynamic worker management not fully supported")
workers = gateway.list_workers()
logger.info("Worker count: %d", len(workers))
assert len(workers) >= 2, "Expected at least 2 workers"
for w in workers:
logger.info("Worker: id=%s, url=%s", w.id, w.url)
finally:
gateway.shutdown()
@pytest.mark.e2e
class TestDisableHealthCheck:
"""Tests for --disable-health-check CLI option."""
def test_disable_health_check_workers_immediately_healthy(
self, model_pool: ModelPool
):
"""Test that workers are immediately healthy when health checks are disabled."""
http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP)
gateway = Gateway()
gateway.start(
igw_mode=True,
extra_args=["--disable-health-check"],
)
try:
# Add worker - should be immediately healthy since health checks are disabled
success, worker_id = gateway.add_worker(
http_instance.worker_url,
wait_ready=True,
ready_timeout=10, # Short timeout since it should be immediate
)
assert success, f"Failed to add worker: {worker_id}"
logger.info("Added worker with health checks disabled: %s", worker_id)
# Verify worker is healthy
workers = gateway.list_workers()
assert len(workers) >= 1, "Expected at least one worker"
for worker in workers:
logger.info(
"Worker: id=%s, status=%s, disable_health_check=%s",
worker.id,
worker.status,
worker.metadata.get("disable_health_check"),
)
# Worker should be healthy immediately
assert (
worker.status == "healthy"
), "Worker should be healthy when health checks disabled"
finally:
gateway.shutdown()
def test_disable_health_check_gateway_starts_without_health_checker(
self, model_pool: ModelPool
):
"""Test that gateway starts successfully with health checks disabled."""
gateway = Gateway()
gateway.start(
igw_mode=True,
extra_args=["--disable-health-check"],
)
try:
assert gateway.health(), "Gateway should be healthy"
logger.info("Gateway started with health checks disabled")
finally:
gateway.shutdown()