596 lines
20 KiB
Python
596 lines
20 KiB
Python
"""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
|