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