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

[model-gateway] Add model scope support and LRU eviction for GPU-constrained environments (#16525)

Browse Source
This commit is contained in:
Simo Lin
2026-01-05 18:28:07 -08:00
committed by GitHub
parent 76c71d1d34
commit 402a0bd6dc
10 changed files with 1122 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files

453
sgl-model-gateway/e2e_test/conftest.py
View File

@@ -1,4 +1,103 @@
"""Pytest configuration for E2E tests."""
"""Pytest configuration for E2E tests.
Markers
-------
This module defines several pytest markers for configuring E2E tests:
@pytest.mark.model(name, scope="session")
Specify which model to use for the test.
Args:
name: Model ID from MODEL_SPECS (e.g., "llama-8b", "qwen-7b")
scope: "session" (default) or "class"
- session: Pre-launched at test session start. Stays running.
- class: Launched on-demand when test class starts.
GPU Resource Management:
When GPUs are limited (e.g., 4 GPUs, 6 models), the model pool uses
LRU (Least Recently Used) eviction:
1. Session models are pre-launched until GPUs are full
2. Overflow models are queued for on-demand launch
3. When a queued model is needed, LRU model is evicted
4. Evicted models go back to queue and can be re-launched later
Examples:
@pytest.mark.model("llama-8b") # session scope, pre-launched
@pytest.mark.model("qwen-72b", scope="class") # on-demand only
@pytest.mark.workers(count=1, prefill=None, decode=None)
Configure worker topology for the test.
Args:
count: Number of regular workers (default: 1)
prefill: Number of prefill workers for PD disaggregation
decode: Number of decode workers for PD disaggregation
Examples:
@pytest.mark.workers(count=3) # 3 regular workers
@pytest.mark.workers(prefill=2, decode=2) # PD mode
@pytest.mark.gateway(policy="round_robin", timeout=None, extra_args=None)
Configure the gateway/router.
Args:
policy: Routing policy ("round_robin", "random", etc.)
timeout: Startup timeout in seconds
extra_args: Additional CLI arguments for the router
Examples:
@pytest.mark.gateway(policy="random")
@pytest.mark.gateway(extra_args=["--cache-routing"])
@pytest.mark.e2e
Mark test as an end-to-end test requiring GPU workers.
@pytest.mark.slow
Mark test as slow-running.
Fixtures
--------
model_pool: Session-scoped fixture managing SGLang worker processes.
setup_backend: Class-scoped fixture that launches gateway + provides client.
Usage Examples
--------------
Basic test with default model:
@pytest.mark.e2e
@pytest.mark.parametrize("setup_backend", ["http"], indirect=True)
class TestBasic:
def test_chat(self, setup_backend):
backend, model, client, gateway = setup_backend
response = client.chat.completions.create(...)
Test with specific model and multiple backends:
@pytest.mark.e2e
@pytest.mark.model("qwen-7b")
@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True)
class TestQwen:
def test_generate(self, setup_backend):
...
Large model loaded on-demand (class scope):
@pytest.mark.e2e
@pytest.mark.model("llama-70b", scope="class")
@pytest.mark.parametrize("setup_backend", ["http"], indirect=True)
class TestLargeModel:
def test_inference(self, setup_backend):
...
PD disaggregation mode:
@pytest.mark.e2e
@pytest.mark.workers(prefill=1, decode=1)
@pytest.mark.parametrize("setup_backend", ["pd"], indirect=True)
class TestPD:
def test_pd_inference(self, setup_backend):
...
"""
from __future__ import annotations
@@ -96,7 +195,8 @@ from infra import (
# Global storage for scanned requirements
_scanned_backends: set[str] = set() # {"grpc", "http", "openai", ...}
_scanned_models: set[str] = set() # {"llama-8b", "qwen-7b", ...}
_session_models: set[str] = set() # Models to pre-launch at session start
_class_models: set[str] = set() # Models to launch on-demand per class
def pytest_collection_modifyitems(
@@ -108,8 +208,12 @@ def pytest_collection_modifyitems(
This runs after test collection but before tests execute.
It extracts backend requirements from @pytest.mark.parametrize markers.
Models are categorized by scope:
- session: Pre-launched at session start (default)
- class: Launched on-demand when test class starts
"""
global _scanned_backends, _scanned_models
global _scanned_backends, _session_models, _class_models
for item in items:
# Scan parametrize markers for setup_backend
@@ -124,30 +228,45 @@ def pytest_collection_modifyitems(
_scanned_backends.update(param_values)
elif param_name == PARAM_MODEL or PARAM_MODEL in param_name:
# Extract model names
# Extract model names from parametrize - default to session scope
if isinstance(param_values, (list, tuple)):
_scanned_models.update(param_values)
_session_models.update(param_values)
# Also check for @pytest.mark.model("name") markers
# Check for @pytest.mark.model("name", scope="...") markers
model_marker = item.get_closest_marker(PARAM_MODEL)
if model_marker and model_marker.args:
_scanned_models.add(model_marker.args[0])
model_name = model_marker.args[0]
scope = model_marker.kwargs.get("scope", "session")
if scope == "class":
_class_models.add(model_name)
else:
_session_models.add(model_name)
# Remove class models from session models (class scope takes precedence if mixed)
# Actually, keep both - a model can be used by both session and class scoped tests
# The model_pool will handle this by keeping session models running
logger.info(
"Scanned test requirements - backends: %s, models: %s",
"Scanned test requirements - backends: %s, session models: %s, class models: %s",
_scanned_backends or {"(none)"},
_scanned_models or {"(none)"},
_session_models or {"(none)"},
_class_models or {"(none)"},
)
def get_pool_requirements() -> list[tuple[str, ConnectionMode]]:
"""Build pool requirements from scanned test markers.
Only returns session-scoped models for pre-launching.
Class-scoped models are launched on-demand by model_pool.get().
Returns:
List of (model_id, ConnectionMode) tuples needed by tests.
List of (model_id, ConnectionMode) tuples to pre-launch.
"""
# Only pre-launch session-scoped models
# Default model if none specified
models = _scanned_models or {DEFAULT_MODEL}
models = _session_models or {DEFAULT_MODEL}
# Convert scanned string backends to ConnectionMode enums
# Filter to local backends only (grpc, http) - cloud backends don't need workers
@@ -174,6 +293,15 @@ def get_pool_requirements() -> list[tuple[str, ConnectionMode]]:
return requirements
def get_class_scoped_models() -> set[str]:
"""Get models that are class-scoped (launched on-demand).
Returns:
Set of model IDs that should be launched on-demand.
"""
return _class_models.copy()
# ---------------------------------------------------------------------------
# Custom pytest markers
# ---------------------------------------------------------------------------
@@ -183,7 +311,8 @@ def pytest_configure(config: pytest.Config) -> None:
"""Register custom markers."""
config.addinivalue_line(
"markers",
"model(name): mark test to use a specific model from the model pool",
"model(name, scope='session'): mark test to use a specific model "
"(scope: 'session' for pre-launched, 'class' for on-demand)",
)
config.addinivalue_line(
"markers",
@@ -191,11 +320,14 @@ def pytest_configure(config: pytest.Config) -> None:
)
config.addinivalue_line(
"markers",
"workers(n): number of workers to launch behind the router (default: 1)",
"workers(count=1, prefill=None, decode=None): "
"worker configuration - use count for regular workers, "
"or prefill/decode for PD disaggregation mode",
)
config.addinivalue_line(
"markers",
"pd(num_prefill=1, num_decode=1): PD disaggregation worker configuration",
"gateway(policy='round_robin', timeout=None, extra_args=None): "
"gateway/router configuration",
)
config.addinivalue_line(
"markers",
@@ -299,6 +431,11 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
allocator = GPUAllocator()
_model_pool = ModelPool(allocator)
# Register class-scoped models for on-demand launching
class_models = get_class_scoped_models()
if class_models:
_model_pool.register_class_scoped_models(class_models)
startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300"))
_model_pool.startup(requirements=requirements, startup_timeout=startup_timeout)
@@ -306,7 +443,7 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
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)
pd_model = next(iter(_session_models), DEFAULT_MODEL)
if pd_model in MODEL_SPECS:
try:
_model_pool.launch_pd_workers(
@@ -373,173 +510,6 @@ def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> s
pytest.skip(f"Model {model_id} not available in model pool")
# ---------------------------------------------------------------------------
# Router launching helpers
# ---------------------------------------------------------------------------
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.
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
# ---------------------------------------------------------------------------
@@ -607,11 +577,12 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
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
- @pytest.mark.workers(count=1): Number of regular workers behind router
- @pytest.mark.workers(prefill=1, decode=1): PD worker configuration
- @pytest.mark.gateway(policy="round_robin", timeout=60): Gateway configuration
Returns:
Tuple of (backend_name, model_path, openai_client)
Tuple of (backend_name, model_path, openai_client, gateway)
Usage:
# Simple - uses defaults
@@ -626,19 +597,21 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
...
# Load balancing with multiple workers
@pytest.mark.workers(3)
@pytest.mark.workers(count=3)
@pytest.mark.gateway(policy="round_robin")
@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.workers(prefill=2, decode=2)
@pytest.mark.gateway(policy="round_robin")
@pytest.mark.parametrize("setup_backend", ["pd"], indirect=True)
class TestPDScaling:
...
"""
import openai
from infra import kill_process_tree
from infra import DEFAULT_ROUTER_TIMEOUT, Gateway, WorkerType
backend_name = request.param
@@ -651,6 +624,28 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
if model_id is None:
model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL)
# Get model scope from marker (session or class)
model_marker = request.node.get_closest_marker("model")
model_scope = "session"
if model_marker:
model_scope = model_marker.kwargs.get("scope", "session")
# Get worker configuration from marker
workers_config = _get_marker_kwargs(
request, "workers", defaults={"count": 1, "prefill": None, "decode": None}
)
# Get gateway configuration from marker
gateway_config = _get_marker_kwargs(
request,
"gateway",
defaults={
"policy": "round_robin",
"timeout": DEFAULT_ROUTER_TIMEOUT,
"extra_args": None,
},
)
# PD disaggregation backend
if backend_name == "pd":
# Check PD requirements
@@ -667,12 +662,9 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
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"]
# Get PD configuration from workers marker
num_prefill = workers_config.get("prefill") or 1
num_decode = workers_config.get("decode") or 1
# Check GPU requirements
required_gpus = num_prefill + num_decode
@@ -684,8 +676,6 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
)
# 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)
@@ -712,27 +702,36 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
model_path = prefills[0].model_path if prefills else None
# Launch PD router
base_url, router_proc = launch_pd_router(prefills, decodes)
# Launch PD gateway with configuration
gateway = Gateway()
gateway.start(
prefill_workers=prefills,
decode_workers=decodes,
policy=gateway_config["policy"],
timeout=gateway_config["timeout"],
extra_args=gateway_config["extra_args"],
)
client = openai.OpenAI(
base_url=f"{base_url}/v1",
base_url=f"{gateway.base_url}/v1",
api_key="not-used",
)
logger.info(
"Setup PD backend: model=%s, %d prefill + %d decode workers, router=%s",
"Setup PD backend: model=%s, %d prefill + %d decode workers, "
"gateway=%s, policy=%s",
model_id,
len(prefills),
len(decodes),
base_url,
gateway.base_url,
gateway_config["policy"],
)
try:
yield backend_name, model_path, client
yield backend_name, model_path, client, gateway
finally:
logger.info("Tearing down PD router")
kill_process_tree(router_proc.pid)
logger.info("Tearing down PD gateway")
gateway.shutdown()
return
# Check if this is a local backend (grpc, http)
@@ -743,15 +742,17 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
is_local = False
connection_mode = None
# Local backends: use worker from pool + launch router
# Local backends: use worker from pool + launch gateway
if is_local:
# Get number of workers from marker
num_workers = _get_marker_value(request, "workers", default=1)
num_workers = workers_config.get("count") or 1
try:
instance = model_pool.get(model_id, connection_mode)
instance = model_pool.get(model_id, connection_mode, scope=model_scope)
except KeyError:
pytest.skip(f"Model {model_id}:{backend_name} not available in pool")
except RuntimeError as e:
pytest.fail(str(e))
# Build worker URLs list
# For num_workers > 1, we need multiple workers from the pool
@@ -759,30 +760,35 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
# 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(
# Launch gateway with configuration
gateway = Gateway()
gateway.start(
worker_urls=worker_urls,
model_path=instance.model_path,
policy=gateway_config["policy"],
timeout=gateway_config["timeout"],
extra_args=gateway_config["extra_args"],
)
client = openai.OpenAI(
base_url=f"{base_url}/v1",
base_url=f"{gateway.base_url}/v1",
api_key="not-used",
)
logger.info(
"Setup %s backend: model=%s, workers=%d, router=%s",
"Setup %s backend: model=%s, workers=%d, gateway=%s, policy=%s",
backend_name,
model_id,
num_workers,
base_url,
gateway.base_url,
gateway_config["policy"],
)
try:
yield backend_name, instance.model_path, client
yield backend_name, instance.model_path, client, gateway
finally:
logger.info("Tearing down router for %s backend", backend_name)
kill_process_tree(router_proc.pid)
logger.info("Tearing down gateway for %s backend", backend_name)
gateway.shutdown()
return
# Cloud backends: launch cloud router
@@ -817,16 +823,16 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
"""Function-scoped fixture for launching a fresh router per test.
This launches a new router for each test, pointing to workers from the pool.
This launches a new Gateway for each test, pointing to workers from the pool.
Use for tests that need isolated router state.
Usage:
@pytest.mark.parametrize("backend_router", ["grpc", "http"], indirect=True)
def test_router_state(backend_router):
base_url, router_proc = backend_router
# Test router-specific behavior
gateway = backend_router
# Test gateway-specific behavior
"""
from infra import kill_process_tree
from infra import Gateway
backend_name = request.param
model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL)
@@ -838,13 +844,16 @@ def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
instance = model_pool.get(model_id, connection_mode)
except KeyError:
pytest.skip(f"Model {model_id}:{backend_name} not available in pool")
except RuntimeError as e:
pytest.fail(str(e))
base_url, router_proc = launch_local_router(
gateway = Gateway()
gateway.start(
worker_urls=[instance.worker_url],
model_path=instance.model_path,
)
try:
yield base_url, router_proc
yield gateway
finally:
kill_process_tree(router_proc.pid)
gateway.shutdown()

4
sgl-model-gateway/e2e_test/infra/__init__.py
View File

@@ -24,6 +24,7 @@ from .constants import ( # Enums; Convenience sets; Fixture parameters; Default
Runtime,
WorkerType,
)
from .gateway import Gateway, WorkerInfo
from .gpu_allocator import (
GPUAllocator,
GPUInfo,
@@ -107,6 +108,9 @@ __all__ = [
"ModelInstance",
"ModelPool",
"MODEL_SPECS",
# Gateway
"Gateway",
"WorkerInfo",
# Default model paths
"DEFAULT_MODEL_PATH",
"DEFAULT_SMALL_MODEL_PATH",

503
sgl-model-gateway/e2e_test/infra/gateway.py Normal file
View File

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

16
sgl-model-gateway/e2e_test/infra/gpu_allocator.py
View File

@@ -368,6 +368,22 @@ class GPUAllocator:
self.slots = [s for s in self.slots if not any(g in gpu_ids for g in s.gpu_ids)]
logger.info("Released GPUs %s, now used: %s", gpu_ids, self._used_gpus)
def release_slot(self, slot: GPUSlot) -> None:
"""Release a GPU slot back to the available pool.
Args:
slot: The GPUSlot to release.
"""
self.release_gpus(slot.gpu_ids)
def available_gpus(self) -> list[int]:
"""Get list of available (unused) GPU IDs.
Returns:
List of GPU IDs that are not currently allocated.
"""
return [g.id for g in self.gpus if g.id not in self._used_gpus]
def summary(self) -> str:
"""Return a summary of GPU allocations."""
lines = ["GPU Allocation Summary:"]

209
sgl-model-gateway/e2e_test/infra/model_pool.py
View File

@@ -44,6 +44,9 @@ class ModelInstance:
gpu_slot: GPUSlot | None
worker_type: WorkerType = WorkerType.REGULAR
bootstrap_port: int | None = None # For prefill workers in PD mode
scope: str = "session" # "session" or "class"
last_used: float = 0.0 # Timestamp for LRU eviction
_healthy: bool = False # Track if initial health check passed
@property
def key(self) -> str:
@@ -165,10 +168,14 @@ class ModelPool:
keeps them running and allows reuse across multiple tests. Routers can then
be launched cheaply (~1-2s) pointing to these workers.
Model scopes:
- session: Pre-launched at session start, never evicted
- class: Launched on-demand, can be evicted when GPUs are needed
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
- Session-scoped workers are launched at startup
- Class-scoped workers are launched on-demand via get()
- When GPUs are full, class-scoped workers are evicted (LRU)
Instance keys:
- Regular workers: "model_id:mode" (e.g., "llama-8b:http")
@@ -176,16 +183,18 @@ class ModelPool:
Limitations:
- Currently one worker instance per (model_id, mode) combination
- @pytest.mark.workers(n) duplicates URLs to router, not distinct workers
- @pytest.mark.workers(count=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)])
# Session-scoped (pre-launched)
instance = pool.get("llama-8b", "http")
# instance.base_url -> "http://127.0.0.1:30000"
# instance.worker_url -> URL for router to connect to
# Class-scoped (on-demand)
instance = pool.get("qwen-7b", "http", scope="class")
"""
def __init__(self, allocator: GPUAllocator | None = None):
@@ -197,6 +206,22 @@ class ModelPool:
self.allocator = allocator or GPUAllocator()
self.instances: dict[str, ModelInstance] = {} # key = "model_id:mode"
self._startup_timeout = DEFAULT_STARTUP_TIMEOUT
self._class_scoped_models: set[str] = (
set()
) # Models that can be launched on-demand
self._queued_models: set[str] = (
set()
) # Session models that couldn't be pre-launched
def register_class_scoped_models(self, models: set[str]) -> None:
"""Register models that may be launched on-demand.
Args:
models: Set of model IDs that are class-scoped.
"""
self._class_scoped_models = models
if models:
logger.info("Registered class-scoped models: %s", models)
def startup(
self,
@@ -253,11 +278,15 @@ class ModelPool:
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
# Track which models got slots
launched_keys = set()
if not slots:
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)
launched_keys.add(f"{model_id}:{mode.value}")
else:
# Launch on allocated slots
for slot in slots:
@@ -266,8 +295,20 @@ class ModelPool:
model_id, mode_str = slot.assigned_model.rsplit(":", 1)
mode = ConnectionMode(mode_str)
self._launch_model(model_id, mode, gpu_slot=slot)
launched_keys.add(slot.assigned_model)
# Wait for all to be healthy
# Track queued models (requested but couldn't be launched due to GPU constraints)
all_keys = set(allocation_specs.keys())
queued_keys = all_keys - launched_keys
if queued_keys:
self._queued_models.update(queued_keys)
logger.info(
"Queued %d models for on-demand launch (GPU constraints): %s",
len(queued_keys),
queued_keys,
)
# Wait for all launched models to be healthy
self._wait_all_healthy()
def _launch_model(
@@ -278,6 +319,7 @@ class ModelPool:
worker_type: WorkerType = WorkerType.REGULAR,
bootstrap_port: int | None = None,
ib_device: str | None = None,
scope: str = "session",
) -> ModelInstance:
"""Launch a model instance.
@@ -288,6 +330,7 @@ class ModelPool:
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.
scope: Model scope ("session" or "class").
Returns:
The launched ModelInstance.
@@ -367,16 +410,27 @@ class ModelPool:
gpu_slot=gpu_slot,
worker_type=worker_type,
bootstrap_port=bootstrap_port,
scope=scope,
last_used=time.time(),
)
self.instances[key] = instance
return instance
def _wait_all_healthy(self) -> None:
"""Wait for all model instances to become healthy."""
"""Wait for all model instances to become healthy.
Only checks workers that haven't been marked healthy yet,
avoiding redundant health checks on already-verified workers.
"""
start_time = time.time()
pending = set(self.instances.keys())
# Only wait for workers that haven't been verified healthy yet
pending = {key for key, inst in self.instances.items() if not inst._healthy}
check_count = 0
if not pending:
logger.info("All workers already healthy, skipping health check")
return
logger.info(
"Waiting for %d workers to become healthy (timeout: %ds)...",
len(pending),
@@ -415,6 +469,7 @@ class ModelPool:
instance.base_url,
check_count,
)
instance._healthy = True
pending.discard(key)
if pending:
@@ -454,19 +509,26 @@ class ModelPool:
model_id: str,
mode: ConnectionMode | str,
worker_type: WorkerType | str = WorkerType.REGULAR,
scope: str = "session",
) -> ModelInstance:
"""Get a model instance by model_id, mode, and worker_type.
For session-scoped models, raises KeyError if not pre-launched.
For class-scoped models, launches on-demand if not running.
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.
scope: Model scope ("session" or "class"). Class-scoped models are
launched on-demand if not running.
Returns:
ModelInstance for the requested model/mode/worker_type.
Raises:
KeyError: If model/mode/worker_type combination is not running.
KeyError: If session-scoped model is not running.
RuntimeError: If worker process died or failed health check.
"""
# Accept both enum and string for convenience
if isinstance(mode, str):
@@ -479,13 +541,36 @@ class ModelPool:
else:
key = f"{model_id}:{mode.value}:{worker_type.value}"
# Check if instance exists
if key not in self.instances:
raise KeyError(
f"{key} not running. Available: {list(self.instances.keys())}"
)
# Check if this model can be launched on-demand
is_class_scoped = scope == "class" or model_id in self._class_scoped_models
is_queued = key in self._queued_models
if is_class_scoped or is_queued:
launch_scope = "class" if is_class_scoped else "session"
logger.info(
"Launching %s model %s on-demand (queued=%s)",
launch_scope,
key,
is_queued,
)
self._ensure_gpu_available(model_id)
self._launch_model(model_id, mode, scope=launch_scope)
self._wait_for_instance(key)
# Remove from queued if it was there
self._queued_models.discard(key)
else:
raise KeyError(
f"{key} not running. Available: {list(self.instances.keys())}"
)
instance = self.instances[key]
# Update last_used timestamp
instance.last_used = time.time()
# Verify worker is still alive and healthy
if not instance.is_alive():
raise RuntimeError(f"Worker {key} process died (was healthy at startup)")
@@ -499,6 +584,104 @@ class ModelPool:
logger.info("Worker %s passed deep health check", key)
return instance
def _ensure_gpu_available(self, model_id: str) -> None:
"""Ensure GPU is available, evicting models if needed (LRU).
All models can be evicted when GPU resources are needed.
Uses LRU (least recently used) eviction strategy.
Args:
model_id: Model ID that needs GPU resources.
"""
spec = get_model_spec(model_id)
required_gpus = spec.get("tp", 1)
# Check if we have enough free GPUs
available = self.allocator.available_gpus()
if len(available) >= required_gpus:
return # Enough GPUs available
# Need to evict models to free up GPUs
# Sort by last_used (LRU eviction) - evict least recently used first
evictable = [
inst
for inst in self.instances.values()
if inst.worker_type == WorkerType.REGULAR
]
evictable.sort(key=lambda x: x.last_used)
freed_gpus = 0
for inst in evictable:
if freed_gpus >= required_gpus:
break
logger.info(
"Evicting model %s (LRU) to free GPUs for %s", inst.key, model_id
)
self._evict_instance(inst.key)
if inst.gpu_slot:
freed_gpus += len(inst.gpu_slot.gpu_ids)
# Recheck available GPUs
available = self.allocator.available_gpus()
if len(available) < required_gpus:
raise RuntimeError(
f"Cannot launch {model_id}: need {required_gpus} GPUs, "
f"only {len(available)} available after eviction"
)
def _evict_instance(self, key: str) -> None:
"""Evict a model instance and free its resources.
Evicted models are added back to the queue for potential re-launch.
Args:
key: Instance key to evict.
"""
if key not in self.instances:
return
instance = self.instances[key]
instance.terminate()
# Release GPU slot back to allocator
if instance.gpu_slot:
self.allocator.release_slot(instance.gpu_slot)
# Add to queued so it can be re-launched on-demand
self._queued_models.add(key)
del self.instances[key]
logger.info("Evicted instance %s (added to queue for re-launch)", key)
def _wait_for_instance(self, key: str, timeout: float | None = None) -> None:
"""Wait for a specific instance to become healthy.
Args:
key: Instance key to wait for.
timeout: Timeout in seconds. Defaults to _startup_timeout.
"""
if timeout is None:
timeout = self._startup_timeout
start_time = time.time()
instance = self.instances.get(key)
if not instance:
raise KeyError(f"Instance {key} not found")
while (time.time() - start_time) < timeout:
if not instance.is_alive():
raise RuntimeError(f"Worker {key} died during startup")
if instance.health_check():
logger.info("Instance %s is healthy", key)
instance._healthy = True
return
time.sleep(HEALTH_CHECK_INTERVAL)
raise TimeoutError(f"Instance {key} did not become healthy within {timeout}s")
def get_workers_by_type(
self, model_id: str, worker_type: WorkerType
) -> list[ModelInstance]:

15
sgl-model-gateway/e2e_test/infra/simple_eval_common.py
View File

@@ -141,7 +141,7 @@ class ChatCompletionSampler(SamplerBase):
self._pack_message("system", self.system_message)
] + message_list
trial = 0
while trial < 6: # 126 seconds in total
while trial < 6: # Max 63 seconds backoff (1+2+4+8+16+32)
try:
response = self.client.chat.completions.create(
model=self.model,
@@ -157,15 +157,20 @@ class ChatCompletionSampler(SamplerBase):
return ""
except Exception as e:
exception_backoff = 2**trial # exponential back off
logger.debug(
"Rate limit, retry %d after %ds: %s",
trial,
# Log first few retries at debug, later ones at warning
log_fn = logger.warning if trial >= 3 else logger.debug
log_fn(
"Request failed (retry %d/%d, backoff %ds): %s",
trial + 1,
6,
exception_backoff,
e,
)
time.sleep(exception_backoff)
trial += 1
logger.warning("All retry attempts exhausted, returning empty response")
logger.warning(
"All retry attempts exhausted after 6 retries, returning empty response"
)
return ""

4
sgl-model-gateway/e2e_test/router/test_mmlu.py
View File

@@ -35,7 +35,7 @@ class TestMMLU:
Note: setup_backend fixture already waits for workers to be ready.
"""
backend, model, client = setup_backend
backend, model, client, *_ = setup_backend
base_url = str(client.base_url).rstrip("/v1")
args = SimpleNamespace(
@@ -59,7 +59,7 @@ class TestMMLU:
Runs MMLU with 128 examples for more statistically
significant results.
"""
backend, model, client = setup_backend
backend, model, client, *_ = setup_backend
base_url = str(client.base_url).rstrip("/v1")
args = SimpleNamespace(

5
sgl-model-gateway/e2e_test/router/test_pd_mmlu.py
View File

@@ -10,7 +10,8 @@ Requirements:
Configuration via markers:
@pytest.mark.model("model-id") # Override default model
@pytest.mark.pd(num_prefill=2, num_decode=2) # Custom worker counts
@pytest.mark.workers(prefill=2, decode=2) # Custom worker counts
@pytest.mark.gateway(policy="round_robin") # Gateway configuration
Usage:
# Basic (1 prefill + 1 decode)
@@ -42,7 +43,7 @@ class TestPDMMLU:
Runs MMLU with 1 prefill + 1 decode worker and validates
accuracy meets threshold (>= 0.65).
"""
backend, model, client = setup_backend
backend, model, client, *_ = setup_backend
base_url = str(client.base_url).rstrip("/v1")
args = SimpleNamespace(

156
sgl-model-gateway/e2e_test/router/test_worker_api.py Normal file
View File

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

2
sgl-model-gateway/e2e_test/utils.py
View File

@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING
logger = logging.getLogger(__name__)
# Re-export commonly used items from submodules
from backends import kill_process_tree # noqa: F401
from infra import kill_process_tree # noqa: F401
from infra.model_specs import ( # noqa: F401; Default model paths
DEFAULT_EMBEDDING_MODEL_PATH,
DEFAULT_ENABLE_THINKING_MODEL_PATH,