minke.yu/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
sgl-model-gateway/e2e_test/infra/__init__.py
+4
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",
sgl-model-gateway/e2e_test/infra/gateway.py
+503
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()
sgl-model-gateway/e2e_test/infra/gpu_allocator.py
+16
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:"]
sgl-model-gateway/e2e_test/infra/model_pool.py
+196 -13
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]:
sgl-model-gateway/e2e_test/infra/simple_eval_common.py
+10 -5
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 ""