chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
1
third_party/sglang/sgl-model-gateway/e2e_test/__init__.py
vendored
Normal file
1
third_party/sglang/sgl-model-gateway/e2e_test/__init__.py
vendored
Normal file
@@ -0,0 +1 @@
|
||||
"""Test package root for router Python tests."""
|
||||
0
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/__init__.py
vendored
Normal file
0
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/__init__.py
vendored
Normal file
222
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/conftest.py
vendored
Normal file
222
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/conftest.py
vendored
Normal 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
|
||||
98
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/results.py
vendored
Normal file
98
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/results.py
vendored
Normal 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
|
||||
119
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/summarize.py
vendored
Normal file
119
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/summarize.py
vendored
Normal 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()
|
||||
26
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/test_pd_perf.py
vendored
Normal file
26
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/test_pd_perf.py
vendored
Normal 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,
|
||||
},
|
||||
)
|
||||
27
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/test_regular_perf.py
vendored
Normal file
27
third_party/sglang/sgl-model-gateway/e2e_test/benchmarks/test_regular_perf.py
vendored
Normal 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,
|
||||
},
|
||||
)
|
||||
0
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/__init__.py
vendored
Normal file
0
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/__init__.py
vendored
Normal file
168
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_enable_thinking.py
vendored
Normal file
168
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_enable_thinking.py
vendored
Normal 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"
|
||||
1529
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_function_calling.py
vendored
Normal file
1529
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_function_calling.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
316
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_openai_server.py
vendored
Normal file
316
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_openai_server.py
vendored
Normal 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
|
||||
165
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_reasoning_content.py
vendored
Normal file
165
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_reasoning_content.py
vendored
Normal 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
|
||||
167
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_validation.py
vendored
Normal file
167
third_party/sglang/sgl-model-gateway/e2e_test/chat_completions/test_validation.py
vendored
Normal 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}"
|
||||
)
|
||||
225
third_party/sglang/sgl-model-gateway/e2e_test/conftest.py
vendored
Normal file
225
third_party/sglang/sgl-model-gateway/e2e_test/conftest.py
vendored
Normal 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",
|
||||
]
|
||||
0
third_party/sglang/sgl-model-gateway/e2e_test/embeddings/__init__.py
vendored
Normal file
0
third_party/sglang/sgl-model-gateway/e2e_test/embeddings/__init__.py
vendored
Normal file
143
third_party/sglang/sgl-model-gateway/e2e_test/embeddings/test_basic.py
vendored
Normal file
143
third_party/sglang/sgl-model-gateway/e2e_test/embeddings/test_basic.py
vendored
Normal 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))
|
||||
262
third_party/sglang/sgl-model-gateway/e2e_test/embeddings/test_correctness.py
vendored
Normal file
262
third_party/sglang/sgl-model-gateway/e2e_test/embeddings/test_correctness.py
vendored
Normal 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")
|
||||
51
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/__init__.py
vendored
Normal file
51
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/__init__.py
vendored
Normal 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",
|
||||
]
|
||||
430
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/hooks.py
vendored
Normal file
430
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/hooks.py
vendored
Normal 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}")
|
||||
57
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/markers.py
vendored
Normal file
57
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/markers.py
vendored
Normal 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
|
||||
242
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/pool.py
vendored
Normal file
242
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/pool.py
vendored
Normal 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()
|
||||
14
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/ports.py
vendored
Normal file
14
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/ports.py
vendored
Normal 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]
|
||||
449
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/setup_backend.py
vendored
Normal file
449
third_party/sglang/sgl-model-gateway/e2e_test/fixtures/setup_backend.py
vendored
Normal 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()
|
||||
144
third_party/sglang/sgl-model-gateway/e2e_test/infra/__init__.py
vendored
Normal file
144
third_party/sglang/sgl-model-gateway/e2e_test/infra/__init__.py
vendored
Normal 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",
|
||||
]
|
||||
74
third_party/sglang/sgl-model-gateway/e2e_test/infra/constants.py
vendored
Normal file
74
third_party/sglang/sgl-model-gateway/e2e_test/infra/constants.py
vendored
Normal 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)
|
||||
595
third_party/sglang/sgl-model-gateway/e2e_test/infra/gateway.py
vendored
Normal file
595
third_party/sglang/sgl-model-gateway/e2e_test/infra/gateway.py
vendored
Normal 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
|
||||
437
third_party/sglang/sgl-model-gateway/e2e_test/infra/gpu_allocator.py
vendored
Normal file
437
third_party/sglang/sgl-model-gateway/e2e_test/infra/gpu_allocator.py
vendored
Normal 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)
|
||||
329
third_party/sglang/sgl-model-gateway/e2e_test/infra/gpu_monitor.py
vendored
Normal file
329
third_party/sglang/sgl-model-gateway/e2e_test/infra/gpu_monitor.py
vendored
Normal 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")
|
||||
1231
third_party/sglang/sgl-model-gateway/e2e_test/infra/model_pool.py
vendored
Normal file
1231
third_party/sglang/sgl-model-gateway/e2e_test/infra/model_pool.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
151
third_party/sglang/sgl-model-gateway/e2e_test/infra/model_specs.py
vendored
Normal file
151
third_party/sglang/sgl-model-gateway/e2e_test/infra/model_specs.py
vendored
Normal 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",
|
||||
},
|
||||
}
|
||||
161
third_party/sglang/sgl-model-gateway/e2e_test/infra/process_utils.py
vendored
Normal file
161
third_party/sglang/sgl-model-gateway/e2e_test/infra/process_utils.py
vendored
Normal 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
|
||||
138
third_party/sglang/sgl-model-gateway/e2e_test/infra/run_eval.py
vendored
Normal file
138
third_party/sglang/sgl-model-gateway/e2e_test/infra/run_eval.py
vendored
Normal 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
|
||||
492
third_party/sglang/sgl-model-gateway/e2e_test/infra/simple_eval_common.py
vendored
Normal file
492
third_party/sglang/sgl-model-gateway/e2e_test/infra/simple_eval_common.py
vendored
Normal 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)
|
||||
132
third_party/sglang/sgl-model-gateway/e2e_test/infra/simple_eval_mmlu.py
vendored
Normal file
132
third_party/sglang/sgl-model-gateway/e2e_test/infra/simple_eval_mmlu.py
vendored
Normal 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)
|
||||
44
third_party/sglang/sgl-model-gateway/e2e_test/pyproject.toml
vendored
Normal file
44
third_party/sglang/sgl-model-gateway/e2e_test/pyproject.toml
vendored
Normal 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.
|
||||
9
third_party/sglang/sgl-model-gateway/e2e_test/responses/__init__.py
vendored
Normal file
9
third_party/sglang/sgl-model-gateway/e2e_test/responses/__init__.py
vendored
Normal 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)
|
||||
"""
|
||||
441
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_basic_crud.py
vendored
Normal file
441
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_basic_crud.py
vendored
Normal 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
|
||||
349
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_state_management.py
vendored
Normal file
349
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_state_management.py
vendored
Normal 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,
|
||||
)
|
||||
256
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_streaming_events.py
vendored
Normal file
256
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_streaming_events.py
vendored
Normal 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
|
||||
289
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_structured_output.py
vendored
Normal file
289
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_structured_output.py
vendored
Normal 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"
|
||||
902
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_tools_call.py
vendored
Normal file
902
third_party/sglang/sgl-model-gateway/e2e_test/responses/test_tools_call.py
vendored
Normal 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()
|
||||
)
|
||||
0
third_party/sglang/sgl-model-gateway/e2e_test/router/__init__.py
vendored
Normal file
0
third_party/sglang/sgl-model-gateway/e2e_test/router/__init__.py
vendored
Normal file
76
third_party/sglang/sgl-model-gateway/e2e_test/router/test_mmlu.py
vendored
Normal file
76
third_party/sglang/sgl-model-gateway/e2e_test/router/test_mmlu.py
vendored
Normal 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"])
|
||||
61
third_party/sglang/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py
vendored
Normal file
61
third_party/sglang/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py
vendored
Normal 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"])
|
||||
220
third_party/sglang/sgl-model-gateway/e2e_test/router/test_worker_api.py
vendored
Normal file
220
third_party/sglang/sgl-model-gateway/e2e_test/router/test_worker_api.py
vendored
Normal 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()
|
||||
Reference in New Issue
Block a user