chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[model-gateway] refactor e2e test infrastructure and add router CI (#16513)

Browse Source
This commit is contained in:
Simo Lin
2026-01-05 13:50:59 -08:00
committed by GitHub
parent c371df2f25
commit 4221b7c573
17 changed files with 2479 additions and 1151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-rust.yml
+55 -1
View File
@@ -323,6 +323,60 @@ jobs:
docker rm oracle-db || true
router-e2e-tests:
if: |
github.event_name != 'pull_request' ||
(github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) ||
(github.event.action == 'labeled' && github.event.label.name == 'run-ci')
runs-on: 4-gpu-a10
timeout-minutes: 45
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install rust dependencies
run: |
bash scripts/ci/ci_install_rust.sh
- name: Configure sccache
uses: mozilla-actions/sccache-action@v0.0.9
with:
version: "v0.12.0"
- name: Rust cache
uses: Swatinem/rust-cache@v2
with:
workspaces: sgl-model-gateway
shared-key: "rust-cache"
cache-all-crates: true
cache-on-failure: true
save-if: true
- name: Install SGLang dependencies
run: |
sudo --preserve-env=PATH bash scripts/ci/ci_install_dependency.sh
- name: Build python binding
run: |
source "$HOME/.cargo/env"
export RUSTC_WRAPPER=sccache
cd sgl-model-gateway/bindings/python
python3 -m pip install --upgrade pip maturin
pip uninstall -y sglang-router
maturin build --profile ci --features vendored-openssl --out dist
pip install dist/*.whl
- name: Install e2e test dependencies
run: |
python3 -m pip install pytest pytest-rerunfailures httpx openai grpcio grpcio-health-checking
- name: Run router e2e tests
run: |
bash scripts/killall_sglang.sh "nuk_gpus"
cd sgl-model-gateway
source "$HOME/.cargo/env"
ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" SHOW_WORKER_LOGS=0 SHOW_ROUTER_LOGS=1 pytest --reruns 2 --reruns-delay 5 e2e_test/router -s -vv -o log_cli=true --log-cli-level=INFO
docker-build-test:
if: |
github.event_name != 'pull_request' ||
@@ -347,7 +401,7 @@ jobs:
cache-to: type=gha,mode=max
finish:
needs: [maturin-build-test, router-unit-tests, router-http-tests, router-grpc-response-api-tests, docker-build-test]
needs: [maturin-build-test, router-unit-tests, router-http-tests, router-grpc-response-api-tests, router-e2e-tests, docker-build-test]
runs-on: ubuntu-latest
steps:
- name: Finish
sgl-model-gateway/e2e_test/backends.py
+47 -365
View File
@@ -1,345 +1,65 @@
"""Backend configurations for E2E tests.
"""Cloud backend configurations for E2E tests.
This module defines the available backends for E2E testing:
- grpc: Local gRPC workers with SGLang router
- http: Local HTTP workers with SGLang router
- openai: OpenAI API backend
- xai: xAI API backend
Each backend configuration specifies:
- model: Model path or name
- launcher: Function to launch the backend
- launcher_kwargs: Arguments for the launcher
- needs_workers: Whether local GPU workers are needed
- api_key_env: Environment variable for API key (if needed)
This module handles cloud API backends (OpenAI, xAI) that don't need local GPU workers.
For local backends (gRPC, HTTP), use ModelPool from infra/ to launch workers,
then launch the router separately pointing to those workers.
"""
from __future__ import annotations
import logging
import os
import signal
import socket
import subprocess
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Callable
from typing import Any
import requests
if TYPE_CHECKING:
import openai
from infra.model_specs import _resolve_model_path
from infra import get_open_port, kill_process_tree, wait_for_health
logger = logging.getLogger(__name__)
# Default ports for each backend type (can be overridden)
DEFAULT_PORTS = {
"grpc": 30030,
"grpc_harmony": 30031,
"http": 30020,
"openai": 30010,
"xai": 30011,
"oracle_store": 30040,
}
# Prometheus port offset from main port
PROMETHEUS_PORT_OFFSET = 1000
def get_open_port() -> int:
"""Get an available port by binding to port 0."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
s.listen(1)
return s.getsockname()[1]
def kill_process_tree(pid: int, sig: int = signal.SIGTERM) -> None:
"""Kill a process and all its children."""
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 wait_for_health(
url: str,
timeout: float = 60,
api_key: str | None = None,
) -> None:
"""Wait for a server's /health endpoint to return 200."""
start = time.time()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
while time.time() - start < timeout:
try:
resp = requests.get(f"{url}/health", headers=headers, timeout=5)
if resp.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(1)
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."""
start = time.time()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
while time.time() - start < timeout:
try:
resp = requests.get(f"{router_url}/workers", headers=headers, timeout=5)
if resp.status_code == 200:
data = resp.json()
if data.get("total", 0) >= expected_workers:
logger.info(
"All %d workers connected after %.1fs",
expected_workers,
time.time() - 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"
)
@dataclass
class ClusterInfo:
"""Information about a running cluster."""
class RouterInstance:
"""A running router instance (for cloud backends)."""
base_url: str
router_process: subprocess.Popen
worker_processes: list[subprocess.Popen]
model: str
backend: str
def shutdown(self) -> None:
"""Shutdown the cluster."""
# Kill router first
"""Shutdown the router."""
if self.router_process.poll() is None:
kill_process_tree(self.router_process.pid)
# Kill workers
for proc in self.worker_processes:
if proc.poll() is None:
kill_process_tree(proc.pid)
def launch_grpc_cluster(
model: str,
base_url: str | None = None,
*,
num_workers: int = 1,
tp_size: int = 1,
policy: str = "round_robin",
api_key: str | None = None,
worker_args: list[str] | None = None,
router_args: list[str] | None = None,
timeout: float = 300,
show_output: bool | None = None,
) -> ClusterInfo:
"""Launch gRPC workers and router.
Args:
model: Model path
base_url: Base URL for router (auto-assigns port if None)
num_workers: Number of workers to launch
tp_size: Tensor parallelism size
policy: Routing policy
api_key: Optional API key for router auth
worker_args: Additional worker arguments
router_args: Additional router arguments
timeout: Startup timeout in seconds
show_output: Show subprocess output (default: SHOW_ROUTER_LOGS env var)
Returns:
ClusterInfo with running processes
"""
if show_output is None:
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
# Determine router port
if base_url:
router_port = int(base_url.split(":")[-1])
else:
router_port = get_open_port()
base_url = f"http://127.0.0.1:{router_port}"
logger.info("Launching gRPC cluster: %d workers, tp=%d", num_workers, tp_size)
# Launch workers
workers = []
worker_urls = []
for i in range(num_workers):
worker_port = get_open_port()
worker_url = f"grpc://127.0.0.1:{worker_port}"
worker_urls.append(worker_url)
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
"127.0.0.1",
"--port",
str(worker_port),
"--grpc-mode",
"--mem-fraction-static",
"0.8",
"--log-level",
"warning",
]
if tp_size > 1:
cmd.extend(["--tp-size", str(tp_size)])
if worker_args:
cmd.extend(worker_args)
logger.info("Starting worker %d on port %d", i + 1, worker_port)
proc = subprocess.Popen(
cmd,
stdout=None if show_output else subprocess.PIPE,
stderr=None if show_output else subprocess.PIPE,
start_new_session=True,
)
workers.append(proc)
# Wait for workers to initialize
logger.info("Waiting for workers to initialize (20s)...")
time.sleep(20)
# Verify workers are alive
for i, worker in enumerate(workers):
if worker.poll() is not None:
# Cleanup
for w in workers:
try:
kill_process_tree(w.pid)
except Exception:
pass
raise RuntimeError(f"Worker {i + 1} died during startup")
# Launch router
router_cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(router_port),
"--prometheus-port",
str(router_port + PROMETHEUS_PORT_OFFSET),
"--policy",
policy,
"--model-path",
model,
"--log-level",
"warn",
"--worker-urls",
*worker_urls,
]
if api_key:
router_cmd.extend(["--api-key", api_key])
if router_args:
router_cmd.extend(router_args)
logger.info("Starting router on port %d", router_port)
router_proc = subprocess.Popen(
router_cmd,
stdout=None if show_output else subprocess.PIPE,
stderr=None if show_output else subprocess.PIPE,
start_new_session=True,
)
# Wait for router to be ready with all workers
try:
wait_for_workers_ready(base_url, num_workers, timeout=timeout, api_key=api_key)
except TimeoutError:
# Cleanup on failure
kill_process_tree(router_proc.pid)
for w in workers:
kill_process_tree(w.pid)
raise
logger.info("gRPC cluster ready at %s with %d workers", base_url, num_workers)
return ClusterInfo(
base_url=base_url,
router_process=router_proc,
worker_processes=workers,
model=model,
backend="grpc",
)
def launch_openai_router(
def launch_cloud_router(
backend: str, # "openai" or "xai"
base_url: str | None = None,
*,
history_backend: str = "memory",
router_args: list[str] | None = None,
timeout: float = 60,
show_output: bool | None = None,
) -> ClusterInfo:
"""Launch router with OpenAI/xAI backend.
) -> RouterInstance:
"""Launch router with cloud API backend (OpenAI/xAI).
Args:
backend: "openai" or "xai"
base_url: Base URL for router (auto-assigns port if None)
history_backend: "memory" or "oracle"
router_args: Additional router arguments
timeout: Startup timeout in seconds
show_output: Show subprocess output
Returns:
ClusterInfo with running router
RouterInstance with running router
"""
if show_output is None:
show_output = os.environ.get("SHOW_ROUTER_LOGS", "0") == "1"
# Determine port
if base_url:
router_port = int(base_url.split(":")[-1])
else:
router_port = get_open_port()
base_url = f"http://127.0.0.1:{router_port}"
router_port = get_open_port()
prometheus_port = get_open_port()
base_url = f"http://127.0.0.1:{router_port}"
# Get API key
# Get API key and worker URL
if backend == "openai":
worker_url = "https://api.openai.com"
api_key = os.environ.get("OPENAI_API_KEY")
@@ -351,7 +71,7 @@ def launch_openai_router(
if not api_key:
raise ValueError("XAI_API_KEY environment variable required")
else:
raise ValueError(f"Unsupported backend: {backend}")
raise ValueError(f"Unsupported cloud backend: {backend}")
logger.info("Launching %s router on port %d", backend, router_port)
@@ -364,7 +84,7 @@ def launch_openai_router(
"--port",
str(router_port),
"--prometheus-port",
str(router_port + PROMETHEUS_PORT_OFFSET),
str(prometheus_port),
"--backend",
"openai",
"--worker-urls",
@@ -400,105 +120,67 @@ def launch_openai_router(
logger.info("%s router ready at %s", backend, base_url)
return ClusterInfo(
return RouterInstance(
base_url=base_url,
router_process=router_proc,
worker_processes=[],
model="", # Cloud API - model specified per request
backend=backend,
)
# Backend configuration registry
BACKENDS: dict[str, dict[str, Any]] = {
"grpc": {
"description": "Local gRPC workers with SGLang router",
"model": _resolve_model_path("meta-llama/Llama-3.1-8B-Instruct"),
"launcher": launch_grpc_cluster,
"launcher_kwargs": {
"num_workers": 1,
"tp_size": 1,
"policy": "round_robin",
},
"needs_workers": True,
"api_key_env": None,
},
"grpc_harmony": {
"description": "Local gRPC workers with Harmony model",
"model": _resolve_model_path("openai/gpt-oss-20b"),
"launcher": launch_grpc_cluster,
"launcher_kwargs": {
"num_workers": 1,
"tp_size": 2,
"policy": "round_robin",
"worker_args": ["--reasoning-parser=gpt-oss"],
"router_args": ["--history-backend", "memory"],
},
"needs_workers": True,
"api_key_env": None,
},
# Cloud backend configurations
CLOUD_BACKENDS: dict[str, dict[str, Any]] = {
"openai": {
"description": "OpenAI API backend",
"model": "gpt-4o-mini",
"launcher": launch_openai_router,
"launcher_kwargs": {
"backend": "openai",
"history_backend": "memory",
},
"needs_workers": False,
"api_key_env": "OPENAI_API_KEY",
"history_backend": "memory",
},
"xai": {
"description": "xAI API backend",
"model": "grok-2-latest",
"launcher": launch_openai_router,
"launcher_kwargs": {
"backend": "xai",
"history_backend": "memory",
},
"needs_workers": False,
"api_key_env": "XAI_API_KEY",
"history_backend": "memory",
},
"oracle_store": {
"description": "OpenAI API with Oracle history backend",
"model": "gpt-4o-mini",
"launcher": launch_openai_router,
"launcher_kwargs": {
"backend": "openai",
"history_backend": "oracle",
},
"needs_workers": False,
"api_key_env": "OPENAI_API_KEY",
"history_backend": "oracle",
},
}
def get_backend_config(backend: str) -> dict[str, Any]:
"""Get configuration for a backend."""
if backend not in BACKENDS:
def get_cloud_backend_config(backend: str) -> dict[str, Any]:
"""Get configuration for a cloud backend."""
if backend not in CLOUD_BACKENDS:
raise KeyError(
f"Unknown backend: {backend}. Available: {list(BACKENDS.keys())}"
f"Unknown cloud backend: {backend}. Available: {list(CLOUD_BACKENDS.keys())}"
)
return BACKENDS[backend]
return CLOUD_BACKENDS[backend]
def launch_backend(backend: str, **kwargs: Any) -> ClusterInfo:
"""Launch a backend cluster.
def launch_cloud_backend(backend: str, **kwargs: Any) -> RouterInstance:
"""Launch a cloud backend router.
Args:
backend: Backend name from BACKENDS
backend: Backend name from CLOUD_BACKENDS
**kwargs: Override launcher kwargs
Returns:
ClusterInfo with running cluster
RouterInstance with running router
"""
cfg = get_backend_config(backend)
cfg = get_cloud_backend_config(backend)
# Merge kwargs with defaults
launcher_kwargs = {**cfg["launcher_kwargs"], **kwargs}
# Add model for grpc backends
if cfg["needs_workers"]:
return cfg["launcher"](cfg["model"], **launcher_kwargs)
# Determine actual backend type (openai or xai)
if backend == "oracle_store":
actual_backend = "openai"
else:
return cfg["launcher"](**launcher_kwargs)
actual_backend = backend
history_backend = kwargs.pop("history_backend", cfg["history_backend"])
return launch_cloud_router(
actual_backend,
history_backend=history_backend,
**kwargs,
)
sgl-model-gateway/e2e_test/conftest.py
+674 -113
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import logging
import os
import subprocess
import sys
from importlib.util import find_spec
from pathlib import Path
@@ -14,25 +15,164 @@ import pytest
if TYPE_CHECKING:
from infra import ModelPool
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)8s] %(name)s: %(message)s",
)
# ---------------------------------------------------------------------------
# Logging setup (clean output without pytest's "---- live log ----" dividers)
# ---------------------------------------------------------------------------
def _setup_logging() -> None:
"""Configure clean logging to stdout with timestamps."""
# Custom format: timestamp [logger] message
fmt = "%(asctime)s.%(msecs)03d [%(name)s] %(message)s"
datefmt = "%H:%M:%S"
# Create handler for stdout
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(logging.Formatter(fmt, datefmt))
# Configure our e2e_test and infra modules for INFO level
for logger_name in ("e2e_test", "infra"):
log = logging.getLogger(logger_name)
log.setLevel(logging.INFO)
log.addHandler(handler)
log.propagate = False # Don't double-log
# Suppress noisy third-party loggers
for logger_name in ("openai", "httpx", "httpcore", "numexpr"):
logging.getLogger(logger_name).setLevel(logging.WARNING)
_setup_logging()
logger = logging.getLogger(__name__)
# Only add bindings/python to path if the wheel is not installed (for local development)
# This ensures CI tests use the installed wheel which contains the Rust extension
_ROOT = Path(__file__).resolve().parents[1]
# ---------------------------------------------------------------------------
# Test visibility hooks
# ---------------------------------------------------------------------------
def pytest_runtest_logstart(nodeid: str, location: tuple) -> None:
"""Print clear test header at start of each test."""
# Extract test name from nodeid (e.g., "test_mmlu.py::TestMMLU::test_mmlu_basic[grpc]")
test_name = nodeid.split("::")[-1] if "::" in nodeid else nodeid
print(f"\n{'='*60}")
print(f"TEST: {test_name}")
print(f"{'='*60}")
# Path setup for imports
_ROOT = Path(__file__).resolve().parents[1] # sgl-model-gateway/
_E2E_TEST = Path(__file__).resolve().parent # e2e_test/
_SRC = _ROOT / "bindings" / "python"
# Check if sglang_router is already installed with the Rust extension
# 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
# Only add bindings/python if wheel is not installed (development mode)
if not _wheel_installed and str(_SRC) not in sys.path:
sys.path.insert(0, str(_SRC))
# Import constants after path setup
from infra import (
DEFAULT_MODEL,
DEFAULT_ROUTER_TIMEOUT,
ENV_BACKENDS,
ENV_MODEL,
ENV_MODELS,
ENV_SHOW_ROUTER_LOGS,
ENV_SKIP_BACKEND_SETUP,
ENV_SKIP_MODEL_POOL,
ENV_STARTUP_TIMEOUT,
LOCAL_MODES,
PARAM_MODEL,
PARAM_SETUP_BACKEND,
ConnectionMode,
)
# ---------------------------------------------------------------------------
# Test collection: scan for required backends
# ---------------------------------------------------------------------------
# Global storage for scanned requirements
_scanned_backends: set[str] = set() # {"grpc", "http", "openai", ...}
_scanned_models: set[str] = set() # {"llama-8b", "qwen-7b", ...}
def pytest_collection_modifyitems(
session: pytest.Session,
config: pytest.Config,
items: list[pytest.Item],
) -> None:
"""Scan collected tests to determine required backends and models.
This runs after test collection but before tests execute.
It extracts backend requirements from @pytest.mark.parametrize markers.
"""
global _scanned_backends, _scanned_models
for item in items:
# Scan parametrize markers for setup_backend
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:
# Extract backend names from parametrize values
if isinstance(param_values, (list, tuple)):
_scanned_backends.update(param_values)
elif param_name == PARAM_MODEL or PARAM_MODEL in param_name:
# Extract model names
if isinstance(param_values, (list, tuple)):
_scanned_models.update(param_values)
# Also check for @pytest.mark.model("name") markers
model_marker = item.get_closest_marker(PARAM_MODEL)
if model_marker and model_marker.args:
_scanned_models.add(model_marker.args[0])
logger.info(
"Scanned test requirements - backends: %s, models: %s",
_scanned_backends or {"(none)"},
_scanned_models or {"(none)"},
)
def get_pool_requirements() -> list[tuple[str, ConnectionMode]]:
"""Build pool requirements from scanned test markers.
Returns:
List of (model_id, ConnectionMode) tuples needed by tests.
"""
# Default model if none specified
models = _scanned_models or {DEFAULT_MODEL}
# Convert scanned string backends to ConnectionMode enums
# Filter to local backends only (grpc, http) - cloud backends don't need workers
local_modes: set[ConnectionMode] = set()
for backend in _scanned_backends:
try:
mode = ConnectionMode(backend)
if mode in LOCAL_MODES:
local_modes.add(mode)
except ValueError:
# Not a ConnectionMode (e.g., "openai", "xai") - skip
pass
# Default to HTTP if no local backends specified
if not local_modes:
local_modes = {ConnectionMode.HTTP}
# Build requirements: each model needs each mode
requirements: list[tuple[str, ConnectionMode]] = []
for model in models:
for mode in local_modes:
requirements.append((model, mode))
return requirements
# ---------------------------------------------------------------------------
# Custom pytest markers
@@ -47,7 +187,15 @@ def pytest_configure(config: pytest.Config) -> None:
)
config.addinivalue_line(
"markers",
"backend(name): mark test to use a specific backend (grpc, openai, etc.)",
"backend(name): mark test to use a specific backend (grpc, http, openai, etc.)",
)
config.addinivalue_line(
"markers",
"workers(n): number of workers to launch behind the router (default: 1)",
)
config.addinivalue_line(
"markers",
"pd(num_prefill=1, num_decode=1): PD disaggregation worker configuration",
)
config.addinivalue_line(
"markers",
@@ -67,81 +215,111 @@ def pytest_configure(config: pytest.Config) -> None:
_model_pool: "ModelPool | None" = None
def _get_requested_models(config: pytest.Config) -> list[str]:
"""Determine which models are needed based on collected tests.
This scans all test items for @pytest.mark.model() markers and returns
the unique set of models requested.
"""
models = set()
# This is called during collection, so we need to iterate items
for item in config.pluginmanager.get_plugin("main").session.items:
marker = item.get_closest_marker("model")
if marker and marker.args:
models.add(marker.args[0])
return list(models)
@pytest.fixture(scope="session")
def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
"""Session-scoped fixture providing the model pool.
"""Session-scoped fixture that manages SGLang worker processes.
The model pool pre-loads all models needed by tests in this session,
running them in parallel across available GPUs.
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.
Usage:
@pytest.mark.model("llama-8b")
def test_chat(model_pool):
client = model_pool.get_client("llama-8b")
...
Startup behavior:
- Scans test markers to determine required (model, mode) combinations
- Launches workers sequentially, but they boot up concurrently
- Waits for all workers to become healthy before returning
Test requirements are auto-detected from:
- @pytest.mark.parametrize("setup_backend", ["grpc", "http"])
- @pytest.mark.model("model-name")
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
# Import here to avoid import errors when infra is not set up
from infra import MODEL_SPECS, GPUAllocator, ModelPool
if _model_pool is not None:
return _model_pool
# Check if we should skip model startup (e.g., for unit tests)
if os.environ.get("SKIP_MODEL_POOL", "").lower() in ("1", "true", "yes"):
logger.info("SKIP_MODEL_POOL is set, skipping model pool startup")
# 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 which models to start
# For now, start models based on environment or a default set
models_env = os.environ.get("E2E_MODELS", "")
if models_env:
model_ids = [m.strip() for m in models_env.split(",") if m.strip()]
# 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}
requirements = [(m, b) for m in models for b in backend_modes]
logger.info("Using env var requirements: %s", requirements)
else:
# Default: start commonly needed models
model_ids = ["llama-8b", "qwen-7b"]
# Use scanned requirements from test markers
requirements = get_pool_requirements()
logger.info("Using scanned requirements: %s", requirements)
# Filter to available specs
model_ids = [m for m in model_ids if m in MODEL_SPECS]
# Filter to valid models
requirements = [(m, b) for m, b in requirements if m in MODEL_SPECS]
if not model_ids:
logger.warning("No models specified, model pool will be empty")
if not requirements:
logger.warning("No valid requirements, model pool will be empty")
_model_pool = ModelPool(GPUAllocator(gpus=[]))
return _model_pool
logger.info("Starting model pool with models: %s", model_ids)
# Create and start the pool
allocator = GPUAllocator()
_model_pool = ModelPool(allocator)
grpc_mode = os.environ.get("E2E_GRPC_MODE", "").lower() in ("1", "true", "yes")
startup_timeout = int(os.environ.get("E2E_STARTUP_TIMEOUT", "300"))
startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300"))
_model_pool.startup(requirements=requirements, startup_timeout=startup_timeout)
_model_pool.startup(
model_ids=model_ids,
grpc_mode=grpc_mode,
startup_timeout=startup_timeout,
)
# Pre-launch PD workers if 'pd' backend is detected
if "pd" in _scanned_backends:
logger.info("PD backend detected, pre-launching PD workers")
# Use default model for PD workers
pd_model = next(iter(_scanned_models), DEFAULT_MODEL)
if pd_model in MODEL_SPECS:
try:
_model_pool.launch_pd_workers(
model_id=pd_model,
num_prefill=1,
num_decode=1,
startup_timeout=startup_timeout,
)
except Exception as e:
logger.warning("Failed to pre-launch PD workers: %s", e)
# Log final GPU allocation summary
logger.info(_model_pool.allocator.summary())
# Register cleanup
request.addfinalizer(_model_pool.shutdown)
@@ -158,10 +336,10 @@ def model_client(request: pytest.FixtureRequest, model_pool: "ModelPool"):
def test_chat(model_client):
response = model_client.chat.completions.create(...)
"""
marker = request.node.get_closest_marker("model")
marker = request.node.get_closest_marker(PARAM_MODEL)
if marker is None:
pytest.fail(
"Test must be marked with @pytest.mark.model('model-id') to use model_client fixture"
f"Test must be marked with @pytest.mark.{PARAM_MODEL}('model-id') to use model_client fixture"
)
model_id = marker.args[0]
@@ -181,10 +359,10 @@ def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> s
def test_direct_http(model_base_url):
response = httpx.get(f"{model_base_url}/health")
"""
marker = request.node.get_closest_marker("model")
marker = request.node.get_closest_marker(PARAM_MODEL)
if marker is None:
pytest.fail(
"Test must be marked with @pytest.mark.model('model-id') to use model_base_url fixture"
f"Test must be marked with @pytest.mark.{PARAM_MODEL}('model-id') to use model_base_url fixture"
)
model_id = marker.args[0]
@@ -196,94 +374,477 @@ def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> s
# ---------------------------------------------------------------------------
# Backend fixtures (class-scoped)
# Router launching helpers
# ---------------------------------------------------------------------------
@pytest.fixture(scope="class")
def setup_backend(request: pytest.FixtureRequest):
"""Class-scoped fixture for launching backend clusters.
def launch_local_router(
worker_urls: list[str],
model_path: str,
*,
policy: str = "round_robin",
router_args: list[str] | None = None,
timeout: float = DEFAULT_ROUTER_TIMEOUT,
show_output: bool | None = None,
) -> tuple[str, subprocess.Popen]:
"""Launch a router pointing to pre-started workers.
This fixture is used with pytest.mark.parametrize to run tests
against multiple backends.
Args:
worker_urls: List of worker URLs (e.g., ["http://127.0.0.1:30000"])
model_path: Model path for the router
policy: Routing policy
router_args: Additional router arguments
timeout: Startup timeout in seconds
show_output: Show subprocess output
Returns:
Tuple of (base_url, router_process)
"""
from infra import get_open_port, wait_for_workers_ready
if show_output is None:
show_output = os.environ.get(ENV_SHOW_ROUTER_LOGS, "0") == "1"
router_port = get_open_port()
prometheus_port = get_open_port()
base_url = f"http://127.0.0.1:{router_port}"
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(router_port),
"--prometheus-port",
str(prometheus_port),
"--policy",
policy,
"--model-path",
model_path,
"--log-level",
"warn",
"--worker-urls",
*worker_urls,
]
if router_args:
cmd.extend(router_args)
logger.info("Starting router on port %d with workers: %s", router_port, worker_urls)
router_proc = subprocess.Popen(
cmd,
stdout=None if show_output else subprocess.PIPE,
stderr=None if show_output else subprocess.PIPE,
start_new_session=True,
)
try:
wait_for_workers_ready(base_url, len(worker_urls), timeout=timeout)
except TimeoutError:
from infra import kill_process_tree
kill_process_tree(router_proc.pid)
raise
logger.info("Router ready at %s", base_url)
return base_url, router_proc
def launch_pd_router(
prefills: list,
decodes: list,
*,
policy: str = "round_robin",
router_args: list[str] | None = None,
timeout: float = DEFAULT_ROUTER_TIMEOUT,
show_output: bool | None = None,
) -> tuple[str, subprocess.Popen]:
"""Launch a PD disaggregation router.
Args:
prefills: List of prefill ModelInstance objects.
decodes: List of decode ModelInstance objects.
policy: Routing policy.
router_args: Additional router arguments.
timeout: Startup timeout in seconds.
show_output: Show subprocess output.
Returns:
Tuple of (base_url, router_process)
"""
from infra import get_open_port, wait_for_health
if show_output is None:
show_output = os.environ.get(ENV_SHOW_ROUTER_LOGS, "0") == "1"
router_port = get_open_port()
prometheus_port = get_open_port()
base_url = f"http://127.0.0.1:{router_port}"
cmd = [
"python3",
"-m",
"sglang_router.launch_router",
"--host",
"127.0.0.1",
"--port",
str(router_port),
"--prometheus-port",
str(prometheus_port),
"--prometheus-host",
"127.0.0.1",
"--policy",
policy,
"--pd-disaggregation",
"--log-level",
"warn",
]
# Add prefill workers with bootstrap ports
for pf in prefills:
cmd += ["--prefill", pf.base_url, str(pf.bootstrap_port)]
# Add decode workers
for dc in decodes:
cmd += ["--decode", dc.base_url]
if router_args:
cmd.extend(router_args)
logger.info(
"Starting PD router on port %d with %d prefill, %d decode workers",
router_port,
len(prefills),
len(decodes),
)
router_proc = subprocess.Popen(
cmd,
stdout=None if show_output else subprocess.PIPE,
stderr=None if show_output else subprocess.PIPE,
start_new_session=True,
)
try:
wait_for_health(base_url, timeout=timeout)
except TimeoutError:
from infra import kill_process_tree
kill_process_tree(router_proc.pid)
raise
logger.info("PD Router ready at %s", base_url)
return base_url, router_proc
# ---------------------------------------------------------------------------
# Backend fixtures
# ---------------------------------------------------------------------------
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
@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(n): Number of workers behind router (default: 1)
- @pytest.mark.pd(num_prefill=1, num_decode=1): PD worker configuration
Returns:
Tuple of (backend_name, model_path, openai_client)
Usage:
@pytest.mark.parametrize("setup_backend", ["grpc", "openai"], indirect=True)
class TestChatCompletions:
def test_basic(self, setup_backend):
backend, model, client = setup_backend
response = client.chat.completions.create(...)
# Simple - uses defaults
@pytest.mark.parametrize("setup_backend", ["http"], indirect=True)
class TestBasic:
...
Environment variables:
- SKIP_BACKEND_SETUP: Skip backend startup (for dry runs)
- SHOW_ROUTER_LOGS: Show subprocess output
# With model override
@pytest.mark.model("qwen-7b")
@pytest.mark.parametrize("setup_backend", ["http"], indirect=True)
class TestWithModel:
...
# Load balancing with multiple workers
@pytest.mark.workers(3)
@pytest.mark.parametrize("setup_backend", ["http"], indirect=True)
class TestLoadBalancing:
...
# PD with custom configuration
@pytest.mark.pd(num_prefill=2, num_decode=2)
@pytest.mark.parametrize("setup_backend", ["pd"], indirect=True)
class TestPDScaling:
...
"""
import openai
from backends import BACKENDS, ClusterInfo, launch_backend
from infra import kill_process_tree
backend_name = request.param
# Skip if requested
if os.environ.get("SKIP_BACKEND_SETUP", "").lower() in ("1", "true", "yes"):
pytest.skip("SKIP_BACKEND_SETUP is set")
if os.environ.get(ENV_SKIP_BACKEND_SETUP, "").lower() in ("1", "true", "yes"):
pytest.skip(f"{ENV_SKIP_BACKEND_SETUP} is set")
# Check if backend requires API key
cfg = BACKENDS.get(backend_name)
if cfg is None:
# 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)
# PD disaggregation backend
if backend_name == "pd":
# Check PD requirements
try:
import sgl_kernel # noqa: F401
except ImportError:
pytest.skip("sgl_kernel not available, required for PD disaggregation")
try:
import torch
except ImportError:
pytest.skip("torch not available")
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
# Get PD configuration from marker
pd_config = _get_marker_kwargs(
request, "pd", defaults={"num_prefill": 1, "num_decode": 1}
)
num_prefill = pd_config["num_prefill"]
num_decode = pd_config["num_decode"]
# Check GPU requirements
required_gpus = num_prefill + num_decode
gpu_count = torch.cuda.device_count()
if gpu_count < required_gpus:
pytest.skip(
f"PD tests require {required_gpus} GPUs "
f"({num_prefill} prefill + {num_decode} decode), found {gpu_count}"
)
# Try to use pre-launched PD workers, or launch new ones if needed
from infra import WorkerType
existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL)
existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE)
if (
len(existing_prefills) >= num_prefill
and len(existing_decodes) >= num_decode
):
# Use pre-launched workers
prefills = existing_prefills[:num_prefill]
decodes = existing_decodes[:num_decode]
logger.info(
"Using pre-launched PD workers: %d prefill, %d decode",
len(prefills),
len(decodes),
)
else:
# Launch new PD workers (custom config or not pre-launched)
prefills, decodes = model_pool.launch_pd_workers(
model_id=model_id,
num_prefill=num_prefill,
num_decode=num_decode,
startup_timeout=300,
)
model_path = prefills[0].model_path if prefills else None
# Launch PD router
base_url, router_proc = launch_pd_router(prefills, decodes)
client = openai.OpenAI(
base_url=f"{base_url}/v1",
api_key="not-used",
)
logger.info(
"Setup PD backend: model=%s, %d prefill + %d decode workers, router=%s",
model_id,
len(prefills),
len(decodes),
base_url,
)
try:
yield backend_name, model_path, client
finally:
logger.info("Tearing down PD router")
kill_process_tree(router_proc.pid)
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 router
if is_local:
# Get number of workers from marker
num_workers = _get_marker_value(request, "workers", default=1)
try:
instance = model_pool.get(model_id, connection_mode)
except KeyError:
pytest.skip(f"Model {model_id}:{backend_name} not available in pool")
# Build worker URLs list
# For num_workers > 1, we need multiple workers from the pool
# For now, we reuse the same worker URL (router will load balance)
# TODO: Support launching multiple distinct workers for true LB testing
worker_urls = [instance.worker_url] * num_workers
# Launch router pointing to the worker(s)
base_url, router_proc = launch_local_router(
worker_urls=worker_urls,
model_path=instance.model_path,
)
client = openai.OpenAI(
base_url=f"{base_url}/v1",
api_key="not-used",
)
logger.info(
"Setup %s backend: model=%s, workers=%d, router=%s",
backend_name,
model_id,
num_workers,
base_url,
)
try:
yield backend_name, instance.model_path, client
finally:
logger.info("Tearing down router for %s backend", backend_name)
kill_process_tree(router_proc.pid)
return
# Cloud backends: launch cloud router
from backends import CLOUD_BACKENDS, launch_cloud_backend
if backend_name not in CLOUD_BACKENDS:
pytest.fail(f"Unknown backend: {backend_name}")
cfg = CLOUD_BACKENDS[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")
logger.info("Setting up backend: %s", backend_name)
logger.info("Launching cloud backend: %s", backend_name)
router = launch_cloud_backend(backend_name)
# Launch the backend
cluster: ClusterInfo = launch_backend(backend_name)
# Create OpenAI client
api_key = os.environ.get(api_key_env) if api_key_env else "not-used"
client = openai.OpenAI(
base_url=f"{cluster.base_url}/v1",
base_url=f"{router.base_url}/v1",
api_key=api_key,
)
# Yield to test
try:
yield backend_name, cfg["model"], client
finally:
logger.info("Tearing down backend: %s", backend_name)
cluster.shutdown()
logger.info("Tearing down cloud backend: %s", backend_name)
router.shutdown()
@pytest.fixture
def backend_cluster(request: pytest.FixtureRequest):
"""Function-scoped fixture for launching a fresh backend per test.
def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
"""Function-scoped fixture for launching a fresh router per test.
Unlike setup_backend (class-scoped), this creates a new cluster
for each test function. Use for tests that modify cluster state.
This launches a new router for each test, pointing to workers from the pool.
Use for tests that need isolated router state.
Usage:
@pytest.mark.parametrize("backend_cluster", ["grpc"], indirect=True)
def test_add_worker(backend_cluster):
cluster = backend_cluster
# cluster.base_url, cluster.router_process, etc.
@pytest.mark.parametrize("backend_router", ["grpc", "http"], indirect=True)
def test_router_state(backend_router):
base_url, router_proc = backend_router
# Test router-specific behavior
"""
from backends import BACKENDS, launch_backend
from infra import kill_process_tree
backend_name = request.param
model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL)
cfg = BACKENDS.get(backend_name)
if cfg is None:
pytest.fail(f"Unknown backend: {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")
cluster = launch_backend(backend_name)
# Convert string to ConnectionMode
connection_mode = ConnectionMode(backend_name)
try:
yield cluster
instance = model_pool.get(model_id, connection_mode)
except KeyError:
pytest.skip(f"Model {model_id}:{backend_name} not available in pool")
base_url, router_proc = launch_local_router(
worker_urls=[instance.worker_url],
model_path=instance.model_path,
)
try:
yield base_url, router_proc
finally:
cluster.shutdown()
kill_process_tree(router_proc.pid)
sgl-model-gateway/e2e_test/fixtures/generate_test_certs.py
-236
View File
@@ -1,236 +0,0 @@
"""
Generate self-signed certificates for mTLS integration testing.
Creates a Certificate Authority (CA), server certificates, and client certificates.
"""
import datetime
import ipaddress
from pathlib import Path
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID
def generate_private_key():
"""Generate an RSA private key."""
return rsa.generate_private_key(
public_exponent=65537,
key_size=2048,
)
def generate_ca_certificate():
"""Generate a self-signed CA certificate."""
private_key = generate_private_key()
subject = issuer = x509.Name(
[
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
x509.NameAttribute(NameOID.COMMON_NAME, "Test CA"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(issuer)
.public_key(private_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
.add_extension(
x509.BasicConstraints(ca=True, path_length=None),
critical=True,
)
.add_extension(
x509.KeyUsage(
digital_signature=True,
key_cert_sign=True,
crl_sign=True,
key_encipherment=False,
content_commitment=False,
data_encipherment=False,
key_agreement=False,
encipher_only=False,
decipher_only=False,
),
critical=True,
)
.sign(private_key, hashes.SHA256())
)
return private_key, cert
def generate_server_certificate(ca_key, ca_cert):
"""Generate a server certificate signed by the CA."""
private_key = generate_private_key()
subject = x509.Name(
[
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
x509.NameAttribute(NameOID.COMMON_NAME, "localhost"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(ca_cert.subject)
.public_key(private_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
.add_extension(
x509.SubjectAlternativeName(
[
x509.DNSName("localhost"),
x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
]
),
critical=False,
)
.add_extension(
x509.KeyUsage(
digital_signature=True,
key_encipherment=True,
key_cert_sign=False,
crl_sign=False,
content_commitment=False,
data_encipherment=False,
key_agreement=False,
encipher_only=False,
decipher_only=False,
),
critical=True,
)
.add_extension(
x509.ExtendedKeyUsage(
[
x509.oid.ExtendedKeyUsageOID.SERVER_AUTH,
]
),
critical=False,
)
.sign(ca_key, hashes.SHA256())
)
return private_key, cert
def generate_client_certificate(ca_key, ca_cert):
"""Generate a client certificate signed by the CA."""
private_key = generate_private_key()
subject = x509.Name(
[
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
x509.NameAttribute(NameOID.COMMON_NAME, "test-client"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(ca_cert.subject)
.public_key(private_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
.add_extension(
x509.KeyUsage(
digital_signature=True,
key_encipherment=True,
key_cert_sign=False,
crl_sign=False,
content_commitment=False,
data_encipherment=False,
key_agreement=False,
encipher_only=False,
decipher_only=False,
),
critical=True,
)
.add_extension(
x509.ExtendedKeyUsage(
[
x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH,
]
),
critical=False,
)
.sign(ca_key, hashes.SHA256())
)
return private_key, cert
def save_key(key, path: Path):
"""Save private key to PEM file."""
with open(path, "wb") as f:
f.write(
key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.TraditionalOpenSSL,
encryption_algorithm=serialization.NoEncryption(),
)
)
def save_cert(cert, path: Path):
"""Save certificate to PEM file."""
with open(path, "wb") as f:
f.write(cert.public_bytes(serialization.Encoding.PEM))
def generate_all_certificates(output_dir: Path):
"""Generate all certificates and keys for mTLS testing."""
output_dir.mkdir(parents=True, exist_ok=True)
print("==> Generating CA certificate...")
ca_key, ca_cert = generate_ca_certificate()
save_key(ca_key, output_dir / "ca-key.pem")
save_cert(ca_cert, output_dir / "ca-cert.pem")
print("==> Generating server certificate...")
server_key, server_cert = generate_server_certificate(ca_key, ca_cert)
save_key(server_key, output_dir / "server-key.pem")
save_cert(server_cert, output_dir / "server-cert.pem")
print("==> Generating client certificate...")
client_key, client_cert = generate_client_certificate(ca_key, ca_cert)
save_key(client_key, output_dir / "client-key.pem")
save_cert(client_cert, output_dir / "client-cert.pem")
print(f"==> Certificates generated successfully in {output_dir}")
print()
print("Files created:")
print(" - ca-cert.pem : CA certificate (for verifying server/client certs)")
print(" - ca-key.pem : CA private key")
print(" - server-cert.pem : Server certificate")
print(" - server-key.pem : Server private key")
print(" - client-cert.pem : Client certificate")
print(" - client-key.pem : Client private key")
print()
print("Test server can use: server-cert.pem + server-key.pem")
print("Test router can use: client-cert.pem + client-key.pem + ca-cert.pem")
if __name__ == "__main__":
script_dir = Path(__file__).parent
certs_dir = script_dir / "test_certs"
generate_all_certificates(certs_dir)
sgl-model-gateway/e2e_test/fixtures/mock_worker.py
-288
View File
@@ -1,288 +0,0 @@
"""
Lightweight mock worker HTTP server for router integration tests.
Implements minimal endpoints used by the router:
- GET /health, /health_generate
- POST /generate, /v1/completions, /v1/chat/completions
- POST /flush_cache
- GET /get_server_info, /get_model_info, /v1/models
Behavior knobs are controlled via CLI flags to simulate failures, latency, and load.
"""
import argparse
import asyncio
import json
import os
import random
import signal
import sys
import time
from contextlib import asynccontextmanager
from typing import Optional
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
# Global state (per-process)
_inflight = 0
_failures_seen = 0
def _parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--host", default="127.0.0.1")
p.add_argument("--port", type=int, required=True)
p.add_argument("--worker-id", default=None)
p.add_argument("--latency-ms", type=int, default=0)
p.add_argument("--timeout", action="store_true")
p.add_argument("--status-code", type=int, default=200)
p.add_argument("--fail-first-n", type=int, default=0)
p.add_argument("--random-fail-rate", type=float, default=0.0)
p.add_argument("--require-api-key", action="store_true")
p.add_argument("--api-key", default=None)
p.add_argument("--max-payload-bytes", type=int, default=10 * 1024 * 1024)
p.add_argument("--stream", action="store_true")
p.add_argument("--dp-size", type=int, default=1)
p.add_argument("--crash-on-request", action="store_true")
p.add_argument("--health-fail-after-ms", type=int, default=0)
# TLS/mTLS configuration
p.add_argument(
"--ssl-certfile", type=str, default=None, help="Path to SSL certificate file"
)
p.add_argument("--ssl-keyfile", type=str, default=None, help="Path to SSL key file")
p.add_argument(
"--ssl-ca-certs",
type=str,
default=None,
help="Path to CA certificates for client verification",
)
return p.parse_args()
def _extract_worker_id(args: argparse.Namespace) -> str:
if args.worker_id:
return str(args.worker_id)
# default to port (unique enough for tests)
return f"worker-{args.port}"
def create_app(args: argparse.Namespace) -> FastAPI:
app = FastAPI()
worker_id = _extract_worker_id(args)
start_ts = time.time()
crashed = {"done": False}
async def maybe_delay():
if args.latency_ms > 0:
await asyncio.sleep(args.latency_ms / 1000.0)
def should_fail() -> Optional[int]:
global _failures_seen
# Fail first N requests (500)
if args.fail_first_n > 0 and _failures_seen < args.fail_first_n:
_failures_seen += 1
return 500
# Random failure probability (500)
if args.random_fail_rate > 0.0 and random.random() < args.random_fail_rate:
return 500
# Forced status code override (non-200) for all responses
if args.status_code != 200:
return int(args.status_code)
return None
def check_api_key(request: Request):
if not args.require_api_key:
return
auth = request.headers.get("Authorization")
if not auth or not auth.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Unauthorized")
key = auth.split(" ", 1)[1]
if args.api_key and key != args.api_key:
raise HTTPException(status_code=401, detail="Unauthorized")
@asynccontextmanager
async def track_inflight():
global _inflight
_inflight += 1
try:
yield
finally:
_inflight -= 1
@app.get("/health")
async def health():
if (
args.health_fail_after_ms
and (time.time() - start_ts) * 1000.0 >= args.health_fail_after_ms
):
return PlainTextResponse("bad", status_code=500)
return PlainTextResponse("ok", status_code=200)
@app.get("/health_generate")
async def health_generate():
return PlainTextResponse("ok", status_code=200)
@app.post("/flush_cache")
async def flush_cache():
return PlainTextResponse("ok", status_code=200)
@app.get("/get_model_info")
async def get_model_info():
return JSONResponse({"model": "mock", "vocab_size": 32000})
@app.get("/v1/models")
async def list_models():
return JSONResponse({"data": [{"id": "mock", "object": "model"}]})
@app.get("/get_server_info")
async def get_server_info(request: Request):
# Enforce API key on server info when required (used by dp_aware probing)
check_api_key(request)
return JSONResponse(
{
"worker_id": worker_id,
"load_in_flight": _inflight,
"cache": {"size": 0, "hit_rate": 0.0},
"dp_size": int(args.dp_size),
}
)
@app.get("/get_load")
async def get_load(request: Request):
check_api_key(request)
# Return format matching real workers: array of load info per DP rank
return JSONResponse(
[
{
"dp_rank": 0,
"num_reqs": _inflight,
"num_waiting_reqs": 0,
"num_tokens": _inflight,
}
]
)
def make_json_response(obj: dict, status_code: int = 200) -> JSONResponse:
resp = JSONResponse(obj, status_code=status_code)
resp.headers["X-Worker-Id"] = worker_id
return resp
async def handle_text_request(request: Request):
# Authorization
check_api_key(request)
# Payload limit
body = await request.body()
if len(body) > args.max_payload_bytes:
return make_json_response({"error": "payload too large"}, status_code=413)
# Simulate crash on first request
if args.crash_on_request and not crashed["done"]:
crashed["done"] = True
os._exit(1)
# Optional timeout (simulate hang)
if args.timeout:
await asyncio.sleep(3600)
# Optional latency
await maybe_delay()
# Optional failures
fail_code = should_fail()
if fail_code is not None and fail_code != 200:
return make_json_response(
{"error": f"mock failure {fail_code}"}, status_code=fail_code
)
# Build response echoing minimal shape
try:
data = await request.json()
except (json.JSONDecodeError, ValueError):
data = {}
received_headers = {k.lower(): v for k, v in request.headers.items()}
now = time.time()
ret = {
"id": f"cmpl-{int(now*1000)}",
"object": "text_completion",
"created": int(now),
"model": "mock",
"choices": [
{
"text": "ok",
"index": 0,
"finish_reason": "stop",
}
],
"worker_id": worker_id,
"echo": data,
"received_headers": received_headers,
}
return make_json_response(ret, status_code=200)
async def handle_stream_request(request: Request):
check_api_key(request)
async def gen():
# minimal 2-chunk stream then [DONE]
for i in range(2):
await asyncio.sleep(0.01)
chunk = {
"choices": [{"delta": {"content": "x"}}],
"worker_id": worker_id,
}
yield f"data: {json.dumps(chunk)}\n\n"
yield "data: [DONE]\n\n"
headers = {"X-Worker-Id": worker_id}
return StreamingResponse(gen(), media_type="text/event-stream", headers=headers)
@app.post("/generate")
async def generate(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
@app.post("/v1/completions")
async def completions(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
async with track_inflight():
if args.stream:
return await handle_stream_request(request)
return await handle_text_request(request)
return app
def main() -> None:
args = _parse_args()
app = create_app(args)
# Handle SIGTERM gracefully for fast test teardown
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
# Configure SSL if certificates are provided
ssl_config = {}
if args.ssl_certfile and args.ssl_keyfile:
ssl_config["ssl_certfile"] = args.ssl_certfile
ssl_config["ssl_keyfile"] = args.ssl_keyfile
# If CA certs provided, require client certificates (mTLS)
if args.ssl_ca_certs:
ssl_config["ssl_ca_certs"] = args.ssl_ca_certs
ssl_config["ssl_cert_reqs"] = 2 # ssl.CERT_REQUIRED
uvicorn.run(app, host=args.host, port=args.port, log_level="warning", **ssl_config)
if __name__ == "__main__":
main()
sgl-model-gateway/e2e_test/infra/__init__.py
+67
View File
@@ -1,5 +1,29 @@
"""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,
PARAM_BACKEND_ROUTER,
PARAM_MODEL,
PARAM_SETUP_BACKEND,
ConnectionMode,
Runtime,
WorkerType,
)
from .gpu_allocator import (
GPUAllocator,
GPUInfo,
@@ -26,8 +50,43 @@ from .model_specs import ( # Default model paths; Model groups
MODEL_SPECS,
REASONING_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
"ConnectionMode",
"WorkerType",
"Runtime",
# 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",
# 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",
@@ -38,6 +97,12 @@ __all__ = [
"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",
# Model management
"ModelInstance",
"ModelPool",
@@ -56,4 +121,6 @@ __all__ = [
"EMBEDDING_MODELS",
"REASONING_MODELS",
"FUNCTION_CALLING_MODELS",
# Evaluation
"run_eval",
]
sgl-model-gateway/e2e_test/infra/constants.py
+60
View File
@@ -0,0 +1,60 @@
"""Constants and enums for E2E test infrastructure."""
from enum import Enum, auto
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 = 5
sgl-model-gateway/e2e_test/infra/gpu_allocator.py
+35 -14
View File
@@ -199,6 +199,7 @@ class GPUAllocator:
"""
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
def _detect_gpus(self) -> list[GPUInfo]:
"""Auto-detect available GPUs via nvidia-ml-py (NVML)."""
@@ -251,11 +252,14 @@ class GPUAllocator:
2. For each model, find the first GPU(s) that can fit it
3. For multi-GPU models, find consecutive GPUs
Note: This method tracks used GPUs across multiple calls, so subsequent
allocations will use different GPUs than previous ones.
Args:
model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys
Returns:
List of GPUSlots with assigned models
List of GPUSlots with assigned models (only the newly allocated slots)
"""
if not self.gpus:
logger.warning("No GPUs available for allocation")
@@ -268,16 +272,15 @@ class GPUAllocator:
reverse=True,
)
# Track which GPUs are used
used_gpus: set[int] = set()
slots: list[GPUSlot] = []
# Track new slots allocated in this call
new_slots: list[GPUSlot] = []
for model_id, spec in sorted_models:
memory_gb = spec.get("memory_gb", 16)
tp_size = spec.get("tp", 1)
# Find available GPUs
available = [g for g in self.gpus if g.id not in used_gpus]
# 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
@@ -289,8 +292,8 @@ class GPUAllocator:
assigned_model=model_id,
port=get_open_port(),
)
slots.append(slot)
used_gpus.add(gpu.id)
new_slots.append(slot)
self._used_gpus.add(gpu.id)
logger.info(
"Allocated GPU %d (%s, %.1fGB) for %s",
gpu.id,
@@ -301,7 +304,10 @@ class GPUAllocator:
break
else:
logger.warning(
"No GPU with %.1fGB available for %s", memory_gb, model_id
"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
@@ -320,8 +326,8 @@ class GPUAllocator:
assigned_model=model_id,
port=get_open_port(),
)
slots.append(slot)
used_gpus.update(gpu_ids)
new_slots.append(slot)
self._used_gpus.update(gpu_ids)
logger.info(
"Allocated GPUs %s (%.1fGB total) for %s (tp=%d)",
gpu_ids,
@@ -332,14 +338,16 @@ class GPUAllocator:
break
else:
logger.warning(
"No %d consecutive GPUs with %.1fGB available for %s",
"No %d consecutive GPUs with %.1fGB available for %s (used: %s)",
tp_size,
memory_gb,
model_id,
self._used_gpus,
)
self.slots = slots
return slots
# 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."""
@@ -348,10 +356,23 @@ class GPUAllocator:
return slot
return None
def release_gpus(self, gpu_ids: list[int]) -> None:
"""Release GPUs back to the available pool.
Args:
gpu_ids: List of GPU IDs to release.
"""
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 summary(self) -> str:
"""Return a summary of GPU allocations."""
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(
sgl-model-gateway/e2e_test/infra/model_pool.py
+461 -91
View File
@@ -4,10 +4,9 @@ from __future__ import annotations
import logging
import os
import signal
import subprocess
import time
from dataclasses import dataclass, field
from dataclasses import dataclass
from typing import TYPE_CHECKING
import httpx
@@ -15,61 +14,179 @@ import httpx
if TYPE_CHECKING:
import openai
from .gpu_allocator import GPUAllocator, GPUSlot
from .constants import (
DEFAULT_HOST,
DEFAULT_MODEL,
DEFAULT_STARTUP_TIMEOUT,
ENV_SHOW_WORKER_LOGS,
HEALTH_CHECK_INTERVAL,
LOCAL_MODES,
ConnectionMode,
WorkerType,
)
from .gpu_allocator import GPUAllocator, GPUSlot, get_open_port
from .model_specs import MODEL_SPECS, get_model_spec
from .process_utils import detect_ib_device
logger = logging.getLogger(__name__)
# Default timeout for model startup (seconds)
DEFAULT_STARTUP_TIMEOUT = 300
# Health check interval (seconds)
HEALTH_CHECK_INTERVAL = 5
# Host for model servers
DEFAULT_HOST = "127.0.0.1"
@dataclass
class ModelInstance:
"""A running model instance."""
model_id: str
mode: ConnectionMode
model_path: str
base_url: str
port: int
process: subprocess.Popen
gpu_slot: GPUSlot
grpc_mode: bool = False
gpu_slot: GPUSlot | None
worker_type: WorkerType = WorkerType.REGULAR
bootstrap_port: int | None = None # For prefill workers in PD mode
@property
def key(self) -> str:
"""Unique key for this instance.
Regular: 'model_id:mode' (e.g., 'llama-8b:http')
PD workers: 'model_id:mode:worker_type' (e.g., 'llama-8b:http:prefill')
"""
if self.worker_type == WorkerType.REGULAR:
return f"{self.model_id}:{self.mode.value}"
return f"{self.model_id}:{self.mode.value}:{self.worker_type.value}"
@property
def worker_url(self) -> str:
"""URL to use when connecting router to this worker."""
if self.mode == ConnectionMode.GRPC:
return f"grpc://{DEFAULT_HOST}:{self.port}"
return self.base_url
def is_alive(self) -> bool:
"""Check if the process is still running."""
return self.process.poll() is None
def health_check(self, timeout: float = 5.0) -> bool:
"""Check if the model server is healthy via HTTP."""
"""Check if the model server is healthy.
Uses HTTP /health endpoint for HTTP workers, gRPC health check for gRPC workers.
"""
if self.mode == ConnectionMode.GRPC:
return self._grpc_health_check(timeout)
return self._http_health_check(timeout)
def _http_health_check(self, timeout: float = 5.0) -> bool:
"""Check health via HTTP /health endpoint."""
try:
resp = httpx.get(f"{self.base_url}/health", timeout=timeout)
return resp.status_code == 200
except (httpx.RequestError, httpx.TimeoutException):
return False
def deep_health_check(self, timeout: float = 30.0) -> bool:
"""Deep health check that verifies the model can actually generate.
Uses /health_generate for HTTP workers (runs actual inference).
For gRPC workers, falls back to standard health check.
"""
if self.mode == ConnectionMode.GRPC:
# For gRPC, use standard health check (no /health_generate equivalent)
return self._grpc_health_check(timeout)
try:
resp = httpx.get(f"{self.base_url}/health_generate", timeout=timeout)
return resp.status_code == 200
except (httpx.RequestError, httpx.TimeoutException):
return False
def _grpc_health_check(self, timeout: float = 5.0) -> bool:
"""Check health via gRPC health check protocol."""
try:
import grpc
from grpc_health.v1 import health_pb2, health_pb2_grpc
except ImportError as e:
logger.debug("gRPC libraries not available: %s", e)
return False
try:
channel = grpc.insecure_channel(f"{DEFAULT_HOST}:{self.port}")
try:
stub = health_pb2_grpc.HealthStub(channel)
request = health_pb2.HealthCheckRequest(service="")
response = stub.Check(request, timeout=timeout)
is_serving = response.status == health_pb2.HealthCheckResponse.SERVING
if is_serving:
logger.debug(
"gRPC health check passed for port %d (status: SERVING)",
self.port,
)
return is_serving
finally:
channel.close()
except grpc.RpcError as e:
# gRPC-specific errors (connection refused, deadline exceeded, etc.)
logger.debug(
"gRPC health check failed for port %d: %s",
self.port,
e.code() if hasattr(e, "code") else str(e),
)
return False
except Exception as e:
# Other errors
logger.debug(
"gRPC health check error for port %d: %s",
self.port,
str(e),
)
return False
def terminate(self, timeout: float = 10.0) -> None:
"""Terminate the model server process."""
if self.process.poll() is not None:
return # Already terminated
logger.info("Terminating model %s (PID %d)", self.model_id, self.process.pid)
logger.info("Terminating %s (PID %d)", self.key, self.process.pid)
# Try graceful shutdown first
self.process.terminate()
try:
self.process.wait(timeout=timeout)
except subprocess.TimeoutExpired:
logger.warning("Model %s did not terminate, killing", self.model_id)
logger.warning("%s did not terminate, killing", self.key)
self.process.kill()
self.process.wait()
class ModelPool:
"""Manages a pool of pre-loaded models across GPUs."""
"""Manages long-running SGLang worker processes across GPUs.
Workers are expensive to start (~30-60s due to model loading), so this pool
keeps them running and allows reuse across multiple tests. Routers can then
be launched cheaply (~1-2s) pointing to these workers.
Startup behavior:
- Workers are launched sequentially (one subprocess.Popen at a time)
- But they boot up concurrently (overlapping model loading)
- _wait_all_healthy() blocks until all workers respond to health checks
Instance keys:
- Regular workers: "model_id:mode" (e.g., "llama-8b:http")
- PD workers: "model_id:mode:worker_type" (e.g., "llama-8b:http:prefill")
Limitations:
- Currently one worker instance per (model_id, mode) combination
- @pytest.mark.workers(n) duplicates URLs to router, not distinct workers
- For true multi-worker LB testing, extend to support multiple instances
Usage:
pool = ModelPool()
pool.startup(requirements=[("llama-8b", ConnectionMode.HTTP)])
instance = pool.get("llama-8b", "http")
# instance.base_url -> "http://127.0.0.1:30000"
# instance.worker_url -> URL for router to connect to
"""
def __init__(self, allocator: GPUAllocator | None = None):
"""Initialize the model pool.
@@ -78,68 +195,114 @@ class ModelPool:
allocator: GPU allocator to use. If None, creates a new one.
"""
self.allocator = allocator or GPUAllocator()
self.instances: dict[str, ModelInstance] = {}
self.instances: dict[str, ModelInstance] = {} # key = "model_id:mode"
self._startup_timeout = DEFAULT_STARTUP_TIMEOUT
def startup(
self,
model_ids: list[str] | None = None,
grpc_mode: bool = False,
requirements: list[tuple[str, ConnectionMode]] | None = None,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
) -> None:
"""Spin up models in parallel on assigned GPU slots.
"""Start worker processes for the required models.
Workers are launched sequentially (one Popen at a time) but boot up
concurrently since model loading happens in parallel across processes.
This method blocks until all workers pass health checks.
Args:
model_ids: List of model IDs to start. If None, starts all in MODEL_SPECS.
grpc_mode: If True, launch workers in gRPC mode.
startup_timeout: Timeout in seconds for each model to become healthy.
requirements: List of (model_id, mode) tuples specifying what to start.
mode is ConnectionMode.HTTP or ConnectionMode.GRPC.
If None, starts default model in HTTP mode.
startup_timeout: Timeout in seconds for all models to become healthy.
"""
self._startup_timeout = startup_timeout
# Determine which models to start
if model_ids is None:
model_ids = list(MODEL_SPECS.keys())
if requirements is None:
requirements = [(DEFAULT_MODEL, ConnectionMode.HTTP)]
# Filter to models we have specs for
specs_to_start = {
mid: MODEL_SPECS[mid] for mid in model_ids if mid in MODEL_SPECS
}
# Deduplicate and validate
requirements = list(set(requirements))
valid_requirements = []
for model_id, mode in requirements:
if model_id not in MODEL_SPECS:
logger.warning("Unknown model %s, skipping", model_id)
continue
if mode not in LOCAL_MODES:
logger.warning("Invalid mode %s for %s, skipping", mode, model_id)
continue
valid_requirements.append((model_id, mode))
if not specs_to_start:
logger.warning("No valid model specs to start")
if not valid_requirements:
logger.warning("No valid requirements to start")
return
logger.info("Starting model pool with: %s", valid_requirements)
# Build allocation specs - each (model, mode) combo needs its own slot
# Use "model_id:mode" as the allocation key
allocation_specs = {}
for model_id, mode in valid_requirements:
spec = MODEL_SPECS[model_id]
key = f"{model_id}:{mode.value}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": spec.get("tp", 1),
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(specs_to_start)
slots = self.allocator.allocate_slots(allocation_specs)
if not slots:
logger.warning("No GPU slots allocated")
return
logger.info(self.allocator.summary())
# Launch all models in parallel
for slot in slots:
if slot.assigned_model:
self._launch_model(slot, grpc_mode=grpc_mode)
logger.warning("No GPU slots allocated, launching without GPU assignment")
# Fallback: launch without specific GPU assignment
for model_id, mode in valid_requirements:
self._launch_model(model_id, mode, gpu_slot=None)
else:
# Launch on allocated slots
for slot in slots:
if slot.assigned_model:
# Parse "model_id:mode" back
model_id, mode_str = slot.assigned_model.rsplit(":", 1)
mode = ConnectionMode(mode_str)
self._launch_model(model_id, mode, gpu_slot=slot)
# Wait for all to be healthy
self._wait_all_healthy()
def _launch_model(self, slot: GPUSlot, grpc_mode: bool = False) -> None:
"""Launch a model on the given GPU slot."""
model_id = slot.assigned_model
if not model_id:
return
def _launch_model(
self,
model_id: str,
mode: ConnectionMode,
gpu_slot: GPUSlot | None = None,
worker_type: WorkerType = WorkerType.REGULAR,
bootstrap_port: int | None = None,
ib_device: str | None = None,
) -> ModelInstance:
"""Launch a model instance.
Args:
model_id: Model identifier from MODEL_SPECS.
mode: Connection mode (HTTP or GRPC).
gpu_slot: GPU slot assignment, or None for auto.
worker_type: Worker type (REGULAR, PREFILL, or DECODE).
bootstrap_port: Bootstrap port for prefill workers in PD mode.
ib_device: InfiniBand device for PD disaggregation.
Returns:
The launched ModelInstance.
"""
spec = get_model_spec(model_id)
model_path = spec["model"]
tp_size = spec.get("tp", 1)
port = slot.port
# Build environment with CUDA_VISIBLE_DEVICES
# Get port - use slot's port if available, otherwise find open port
port = gpu_slot.port if gpu_slot else get_open_port()
# Build environment
env = os.environ.copy()
env["CUDA_VISIBLE_DEVICES"] = slot.cuda_visible_devices()
if gpu_slot:
env["CUDA_VISIBLE_DEVICES"] = gpu_slot.cuda_visible_devices()
# Build command
cmd = [
@@ -148,6 +311,8 @@ class ModelPool:
"sglang.launch_server",
"--model-path",
model_path,
"--host",
DEFAULT_HOST,
"--port",
str(port),
"--tp-size",
@@ -156,52 +321,81 @@ class ModelPool:
"warning",
]
if grpc_mode:
if mode == ConnectionMode.GRPC:
cmd.append("--grpc-mode")
logger.info(
"Launching %s on GPUs %s port %d: %s",
model_id,
slot.gpu_ids,
port,
" ".join(cmd),
)
# PD disaggregation arguments
if worker_type == WorkerType.PREFILL:
cmd.extend(["--disaggregation-mode", "prefill"])
if bootstrap_port:
cmd.extend(["--disaggregation-bootstrap-port", str(bootstrap_port)])
if ib_device:
cmd.extend(["--disaggregation-ib-device", ib_device])
elif worker_type == WorkerType.DECODE:
cmd.extend(["--disaggregation-mode", "decode"])
if ib_device:
cmd.extend(["--disaggregation-ib-device", ib_device])
# Build key based on worker type
if worker_type == WorkerType.REGULAR:
key = f"{model_id}:{mode.value}"
else:
key = f"{model_id}:{mode.value}:{worker_type.value}"
gpu_info = gpu_slot.gpu_ids if gpu_slot else "auto"
logger.info("Launching %s on GPUs %s port %d", key, gpu_info, port)
show_output = os.environ.get(ENV_SHOW_WORKER_LOGS, "0") == "1"
# Start the process
proc = subprocess.Popen(
cmd,
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
# Use process group for clean shutdown
stdout=None if show_output else subprocess.PIPE,
stderr=None if show_output else subprocess.PIPE,
start_new_session=True,
)
base_url = f"http://{DEFAULT_HOST}:{port}"
instance = ModelInstance(
model_id=model_id,
mode=mode,
model_path=model_path,
base_url=base_url,
port=port,
process=proc,
gpu_slot=slot,
grpc_mode=grpc_mode,
gpu_slot=gpu_slot,
worker_type=worker_type,
bootstrap_port=bootstrap_port,
)
self.instances[model_id] = instance
self.instances[key] = instance
return instance
def _wait_all_healthy(self) -> None:
"""Wait for all model instances to become healthy."""
start_time = time.time()
pending = set(self.instances.keys())
check_count = 0
logger.info(
"Waiting for %d workers to become healthy (timeout: %ds)...",
len(pending),
self._startup_timeout,
)
while pending and (time.time() - start_time) < self._startup_timeout:
for model_id in list(pending):
instance = self.instances[model_id]
check_count += 1
elapsed = time.time() - start_time
for key in list(pending):
instance = self.instances[key]
# Check if process died
if not instance.is_alive():
logger.error(
"Model %s (PID %d) died during startup",
model_id,
"[%.1fs] %s (PID %d) died during startup",
elapsed,
key,
instance.process.pid,
)
# Read stderr for debugging
@@ -209,62 +403,238 @@ class ModelPool:
stderr = instance.process.stderr.read()
if stderr:
logger.error("Stderr: %s", stderr.decode()[-2000:])
pending.discard(model_id)
pending.discard(key)
continue
# Check health
if instance.health_check():
logger.info(
"Model %s is healthy at %s", model_id, instance.base_url
"[%.1fs] %s is healthy at %s (check #%d)",
elapsed,
key,
instance.base_url,
check_count,
)
pending.discard(model_id)
pending.discard(key)
if pending:
# Log progress every 30 seconds
if check_count % 15 == 0: # ~30s at 2s interval
logger.info(
"[%.1fs] Still waiting for %d workers: %s",
elapsed,
len(pending),
list(pending),
)
time.sleep(HEALTH_CHECK_INTERVAL)
if pending:
elapsed = time.time() - start_time
logger.error(
"Models failed to start within %ds: %s",
"[%.1fs] Models failed to start within %ds: %s",
elapsed,
self._startup_timeout,
pending,
)
# Terminate failed instances
for model_id in pending:
self.instances[model_id].terminate()
del self.instances[model_id]
for key in pending:
self.instances[key].terminate()
del self.instances[key]
else:
elapsed = time.time() - start_time
logger.info(
"[%.1fs] All %d workers healthy after %d health checks",
elapsed,
len(self.instances),
check_count,
)
def get_client(self, model_id: str) -> "openai.OpenAI":
def get(
self,
model_id: str,
mode: ConnectionMode | str,
worker_type: WorkerType | str = WorkerType.REGULAR,
) -> ModelInstance:
"""Get a model instance by model_id, mode, and worker_type.
Args:
model_id: The model ID (e.g., "llama-8b")
mode: The mode (ConnectionMode.HTTP or ConnectionMode.GRPC, or string)
worker_type: The worker type (REGULAR, PREFILL, DECODE). Defaults to REGULAR.
Returns:
ModelInstance for the requested model/mode/worker_type.
Raises:
KeyError: If model/mode/worker_type combination is not running.
"""
# Accept both enum and string for convenience
if isinstance(mode, str):
mode = ConnectionMode(mode)
if isinstance(worker_type, str):
worker_type = WorkerType(worker_type)
if worker_type == WorkerType.REGULAR:
key = f"{model_id}:{mode.value}"
else:
key = f"{model_id}:{mode.value}:{worker_type.value}"
if key not in self.instances:
raise KeyError(
f"{key} not running. Available: {list(self.instances.keys())}"
)
instance = self.instances[key]
# Verify worker is still alive and healthy
if not instance.is_alive():
raise RuntimeError(f"Worker {key} process died (was healthy at startup)")
if not instance.deep_health_check(timeout=30.0):
raise RuntimeError(
f"Worker {key} failed deep health check (health_generate) - "
"model may be stuck or crashed"
)
logger.info("Worker %s passed deep health check", key)
return instance
def get_workers_by_type(
self, model_id: str, worker_type: WorkerType
) -> list[ModelInstance]:
"""Get all workers of a specific type for a model.
Args:
model_id: The model ID.
worker_type: The worker type to filter by.
Returns:
List of matching ModelInstance objects.
"""
return [
inst
for inst in self.instances.values()
if inst.model_id == model_id and inst.worker_type == worker_type
]
def launch_pd_workers(
self,
model_id: str,
num_prefill: int = 1,
num_decode: int = 1,
mode: ConnectionMode = ConnectionMode.HTTP,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
) -> tuple[list[ModelInstance], list[ModelInstance]]:
"""Launch prefill and decode workers for PD disaggregation.
Args:
model_id: Model identifier from MODEL_SPECS.
num_prefill: Number of prefill workers to launch. Defaults to 1.
num_decode: Number of decode workers to launch. Defaults to 1.
mode: Connection mode (HTTP or GRPC).
startup_timeout: Timeout for workers to become healthy.
Returns:
Tuple of (prefill_instances, decode_instances).
"""
self._startup_timeout = startup_timeout
if model_id not in MODEL_SPECS:
raise ValueError(f"Unknown model: {model_id}")
spec = get_model_spec(model_id)
ib_device = detect_ib_device()
if ib_device:
logger.info("Detected InfiniBand device: %s", ib_device)
# Build allocation specs for all PD workers
# Each worker needs its own GPU slot
allocation_specs = {}
for i in range(num_prefill):
key = f"{model_id}:{mode.value}:prefill_{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": spec.get("tp", 1),
}
for i in range(num_decode):
key = f"{model_id}:{mode.value}:decode_{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": spec.get("tp", 1),
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
slot_map = {slot.assigned_model: slot for slot in slots}
if not slots:
logger.warning(
"No GPU slots allocated for PD workers, launching without GPU assignment"
)
prefill_instances: list[ModelInstance] = []
decode_instances: list[ModelInstance] = []
# Launch prefill workers
for i in range(num_prefill):
key = f"{model_id}:{mode.value}:prefill_{i}"
gpu_slot = slot_map.get(key)
bootstrap_port = get_open_port()
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.PREFILL,
bootstrap_port=bootstrap_port,
ib_device=ib_device,
)
prefill_instances.append(instance)
# Launch decode workers
for i in range(num_decode):
key = f"{model_id}:{mode.value}:decode_{i}"
gpu_slot = slot_map.get(key)
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.DECODE,
ib_device=ib_device,
)
decode_instances.append(instance)
# Wait for all to be healthy
self._wait_all_healthy()
return prefill_instances, decode_instances
def get_client(
self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP
) -> "openai.OpenAI":
"""Get OpenAI client for a specific model.
Args:
model_id: The model ID to get a client for.
mode: The mode (ConnectionMode.HTTP or ConnectionMode.GRPC). Defaults to HTTP.
Returns:
OpenAI client configured for this model.
Raises:
KeyError: If model is not running.
"""
import openai
if model_id not in self.instances:
raise KeyError(
f"Model {model_id} not running. Available: {list(self.instances.keys())}"
)
instance = self.instances[model_id]
instance = self.get(model_id, mode)
return openai.OpenAI(
base_url=f"{instance.base_url}/v1",
api_key="not-used",
)
def get_base_url(self, model_id: str) -> str:
def get_base_url(
self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP
) -> str:
"""Get the base URL for a specific model."""
if model_id not in self.instances:
raise KeyError(
f"Model {model_id} not running. Available: {list(self.instances.keys())}"
)
return self.instances[model_id].base_url
return self.get(model_id, mode).base_url
def shutdown(self) -> None:
"""Tear down all models."""
sgl-model-gateway/e2e_test/infra/process_utils.py
+161
View File
@@ -0,0 +1,161 @@
"""Process management utilities for E2E tests."""
from __future__ import annotations
import logging
import os
import signal
import subprocess
import time
import requests
logger = logging.getLogger(__name__)
def kill_process_tree(pid: int, sig: int = signal.SIGTERM) -> None:
"""Kill a process and all its children.
Args:
pid: Process ID to kill
sig: Signal to send (default: SIGTERM)
"""
try:
import psutil
parent = psutil.Process(pid)
children = parent.children(recursive=True)
for child in children:
try:
child.send_signal(sig)
except psutil.NoSuchProcess:
pass
parent.send_signal(sig)
except ImportError:
# Fallback if psutil not available
os.kill(pid, sig)
except Exception as e:
logger.warning("Failed to kill process tree for PID %d: %s", pid, e)
def terminate_process(proc: subprocess.Popen, timeout: float = 30) -> None:
"""Gracefully terminate a process, kill if needed.
Args:
proc: Process to terminate
timeout: Seconds to wait before force-killing
"""
if proc is None or proc.poll() is not None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
def wait_for_health(
url: str,
timeout: float = 60,
api_key: str | None = None,
check_interval: float = 1.0,
) -> None:
"""Wait for a server's /health endpoint to return 200.
Args:
url: Base URL of the server
timeout: Seconds to wait before timing out
api_key: Optional API key for auth header
check_interval: Seconds between health checks
"""
start = time.perf_counter()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
resp = session.get(f"{url}/health", headers=headers, timeout=5)
if resp.status_code == 200:
logger.info("Service healthy at %s", url)
return
except requests.RequestException:
pass
time.sleep(check_interval)
raise TimeoutError(f"Server at {url} did not become healthy within {timeout}s")
def wait_for_workers_ready(
router_url: str,
expected_workers: int,
timeout: float = 300,
api_key: str | None = None,
) -> None:
"""Wait for router to have all workers connected.
Args:
router_url: Base URL of the router
expected_workers: Number of workers to wait for
timeout: Seconds to wait before timing out
api_key: Optional API key for auth header
"""
start = time.perf_counter()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
while time.perf_counter() - start < timeout:
try:
resp = requests.get(f"{router_url}/workers", headers=headers, timeout=5)
if resp.status_code == 200:
data = resp.json()
total = data.get("total", len(data.get("workers", [])))
if total >= expected_workers:
logger.info(
"All %d workers connected after %.1fs",
expected_workers,
time.perf_counter() - start,
)
return
except requests.RequestException:
pass
time.sleep(2)
raise TimeoutError(
f"Router at {router_url} did not get {expected_workers} workers within {timeout}s"
)
def detect_ib_device() -> str | None:
"""Detect first active InfiniBand device (e.g., mlx5_0).
Returns:
Device name if found (e.g., "mlx5_0"), None otherwise.
"""
try:
subprocess.run(
["ibv_devinfo", "-l"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=1,
)
except (FileNotFoundError, subprocess.TimeoutExpired):
return None
for i in range(12):
dev = f"mlx5_{i}"
try:
res = subprocess.run(
["ibv_devinfo", dev],
capture_output=True,
text=True,
timeout=2,
)
if res.returncode == 0 and "state:" in res.stdout:
for line in res.stdout.splitlines():
if "state:" in line and "PORT_ACTIVE" in line:
logger.info("Detected IB device: %s", dev)
return dev
except Exception:
pass
return None
sgl-model-gateway/e2e_test/infra/run_eval.py
+139
View File
@@ -0,0 +1,139 @@
"""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
logger = logging.getLogger(__name__)
# MMLU dataset URL
MMLU_DATASET_URL = "https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv"
@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:
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
sgl-model-gateway/e2e_test/infra/simple_eval_common.py
+485
View File
@@ -0,0 +1,485 @@
# 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
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",
fovea: int = 768,
):
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 < 6: # 126 seconds in total
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
logger.debug(
"Rate limit, retry %d after %ds: %s",
trial,
exception_backoff,
e,
)
time.sleep(exception_backoff)
trial += 1
logger.warning("All retry attempts exhausted, returning empty response")
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)
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)
sgl-model-gateway/e2e_test/infra/simple_eval_mmlu.py
+126
View File
@@ -0,0 +1,126 @@
# 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
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):
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)
sgl-model-gateway/e2e_test/pyproject.toml
+30
View File
@@ -0,0 +1,30 @@
[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",
"pytest",
"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",
]
addopts = "-v -s"
# Explicitly disable live log to avoid "---- live log ----" dividers
# We configure logging manually in conftest.py
log_cli = false
sgl-model-gateway/e2e_test/router/test_mmlu.py
+78
View File
@@ -0,0 +1,78 @@
"""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
base_url = str(client.base_url).rstrip("/v1")
args = SimpleNamespace(
base_url=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
base_url = str(client.base_url).rstrip("/v1")
args = SimpleNamespace(
base_url=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"])
sgl-model-gateway/e2e_test/router/test_pd_mmlu.py
+61
View File
@@ -0,0 +1,61 @@
"""MMLU evaluation tests for PD (Prefill-Decode) disaggregated routing.
PD disaggregation separates prefill and decode phases across different
workers for improved throughput and resource utilization.
Requirements:
- sgl_kernel package
- GPUs: num_prefill + num_decode (default: 2 GPUs for 1+1)
- Optional: InfiniBand for high-performance transfers
Configuration via markers:
@pytest.mark.model("model-id") # Override default model
@pytest.mark.pd(num_prefill=2, num_decode=2) # Custom worker counts
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
base_url = str(client.base_url).rstrip("/v1")
args = SimpleNamespace(
base_url=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"])
sgl-model-gateway/e2e_test/utils.py
-43
View File
@@ -182,49 +182,6 @@ def get_tokenizer_from_processor(processor):
return processor.tokenizer
# =============================================================================
# Pytest Utilities
# =============================================================================
def pytest_retry(max_retries: int = 3):
"""Decorator for pytest test functions with retry support.
Args:
max_retries: Maximum number of retry attempts
Example:
@pytest_retry(max_retries=3)
def test_flaky_operation():
# Test that might occasionally fail
pass
"""
import functools
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
last_exception = None
for attempt in range(max_retries + 1):
try:
return func(*args, **kwargs)
except Exception as e:
last_exception = e
if attempt < max_retries:
logger.info(
"Test %s failed on attempt %d/%d, retrying...",
func.__name__,
attempt + 1,
max_retries + 1,
)
continue
raise last_exception
return wrapper
return decorator
# =============================================================================
# Environment Utilities
# =============================================================================