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

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

Browse Source
This commit is contained in:
Simo Lin
2026-01-05 13:50:59 -08:00
committed by GitHub
parent c371df2f25
commit 4221b7c573
17 changed files with 2479 additions and 1151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/e2e_test/infra/__init__.py
+67
View File
@@ -1,5 +1,29 @@
"""Infrastructure for parallel GPU test execution."""
from .constants import ( # Enums; Convenience sets; Fixture parameters; Defaults; Environment variables
CLOUD_RUNTIMES,
DEFAULT_HOST,
DEFAULT_MODEL,
DEFAULT_ROUTER_TIMEOUT,
DEFAULT_STARTUP_TIMEOUT,
ENV_BACKENDS,
ENV_MODEL,
ENV_MODELS,
ENV_SHOW_ROUTER_LOGS,
ENV_SHOW_WORKER_LOGS,
ENV_SKIP_BACKEND_SETUP,
ENV_SKIP_MODEL_POOL,
ENV_STARTUP_TIMEOUT,
HEALTH_CHECK_INTERVAL,
LOCAL_MODES,
LOCAL_RUNTIMES,
PARAM_BACKEND_ROUTER,
PARAM_MODEL,
PARAM_SETUP_BACKEND,
ConnectionMode,
Runtime,
WorkerType,
)
from .gpu_allocator import (
GPUAllocator,
GPUInfo,
@@ -26,8 +50,43 @@ from .model_specs import ( # Default model paths; Model groups
MODEL_SPECS,
REASONING_MODELS,
)
from .process_utils import (
detect_ib_device,
kill_process_tree,
terminate_process,
wait_for_health,
wait_for_workers_ready,
)
from .run_eval import run_eval
__all__ = [
# Enums
"ConnectionMode",
"WorkerType",
"Runtime",
# Convenience sets
"LOCAL_MODES",
"LOCAL_RUNTIMES",
"CLOUD_RUNTIMES",
# Fixture params
"PARAM_SETUP_BACKEND",
"PARAM_BACKEND_ROUTER",
"PARAM_MODEL",
# Defaults
"DEFAULT_MODEL",
"DEFAULT_HOST",
"DEFAULT_STARTUP_TIMEOUT",
"DEFAULT_ROUTER_TIMEOUT",
"HEALTH_CHECK_INTERVAL",
# Env vars
"ENV_MODELS",
"ENV_BACKENDS",
"ENV_MODEL",
"ENV_STARTUP_TIMEOUT",
"ENV_SKIP_MODEL_POOL",
"ENV_SKIP_BACKEND_SETUP",
"ENV_SHOW_ROUTER_LOGS",
"ENV_SHOW_WORKER_LOGS",
# GPU allocation
"GPUAllocator",
"GPUInfo",
@@ -38,6 +97,12 @@ __all__ = [
"get_physical_device_indices",
"get_gpu_memory_usage",
"wait_for_gpu_memory_to_clear",
# Process utilities
"kill_process_tree",
"terminate_process",
"wait_for_health",
"wait_for_workers_ready",
"detect_ib_device",
# Model management
"ModelInstance",
"ModelPool",
@@ -56,4 +121,6 @@ __all__ = [
"EMBEDDING_MODELS",
"REASONING_MODELS",
"FUNCTION_CALLING_MODELS",
# Evaluation
"run_eval",
]
sgl-model-gateway/e2e_test/infra/constants.py
+60
View File
@@ -0,0 +1,60 @@
"""Constants and enums for E2E test infrastructure."""
from enum import Enum, auto
class ConnectionMode(str, Enum):
"""Worker connection protocol."""
HTTP = "http"
GRPC = "grpc"
class WorkerType(str, Enum):
"""Worker specialization type."""
REGULAR = "regular"
PREFILL = "prefill"
DECODE = "decode"
class Runtime(str, Enum):
"""Inference runtime/backend."""
SGLANG = "sglang"
VLLM = "vllm"
OPENAI = "openai"
XAI = "xai"
GEMINI = "gemini"
# Convenience sets
LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC})
LOCAL_RUNTIMES = frozenset({Runtime.SGLANG, Runtime.VLLM})
CLOUD_RUNTIMES = frozenset({Runtime.OPENAI, Runtime.XAI, Runtime.GEMINI})
# Fixture parameter names (used in @pytest.mark.parametrize)
PARAM_SETUP_BACKEND = "setup_backend"
PARAM_BACKEND_ROUTER = "backend_router"
PARAM_MODEL = "model"
# Default model
DEFAULT_MODEL = "llama-8b"
# Environment variable names
ENV_MODELS = "E2E_MODELS"
ENV_BACKENDS = "E2E_BACKENDS"
ENV_MODEL = "E2E_MODEL"
ENV_STARTUP_TIMEOUT = "E2E_STARTUP_TIMEOUT"
ENV_SKIP_MODEL_POOL = "SKIP_MODEL_POOL"
ENV_SKIP_BACKEND_SETUP = "SKIP_BACKEND_SETUP"
ENV_SHOW_ROUTER_LOGS = "SHOW_ROUTER_LOGS"
ENV_SHOW_WORKER_LOGS = "SHOW_WORKER_LOGS"
# Network
DEFAULT_HOST = "127.0.0.1"
# Timeouts (seconds)
DEFAULT_STARTUP_TIMEOUT = 300
DEFAULT_ROUTER_TIMEOUT = 60
HEALTH_CHECK_INTERVAL = 5
sgl-model-gateway/e2e_test/infra/gpu_allocator.py
+35 -14
View File
@@ -199,6 +199,7 @@ class GPUAllocator:
"""
self.gpus = gpus if gpus is not None else self._detect_gpus()
self.slots: list[GPUSlot] = []
self._used_gpus: set[int] = set() # Track GPUs used across all allocations
def _detect_gpus(self) -> list[GPUInfo]:
"""Auto-detect available GPUs via nvidia-ml-py (NVML)."""
@@ -251,11 +252,14 @@ class GPUAllocator:
2. For each model, find the first GPU(s) that can fit it
3. For multi-GPU models, find consecutive GPUs
Note: This method tracks used GPUs across multiple calls, so subsequent
allocations will use different GPUs than previous ones.
Args:
model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys
Returns:
List of GPUSlots with assigned models
List of GPUSlots with assigned models (only the newly allocated slots)
"""
if not self.gpus:
logger.warning("No GPUs available for allocation")
@@ -268,16 +272,15 @@ class GPUAllocator:
reverse=True,
)
# Track which GPUs are used
used_gpus: set[int] = set()
slots: list[GPUSlot] = []
# Track new slots allocated in this call
new_slots: list[GPUSlot] = []
for model_id, spec in sorted_models:
memory_gb = spec.get("memory_gb", 16)
tp_size = spec.get("tp", 1)
# Find available GPUs
available = [g for g in self.gpus if g.id not in used_gpus]
# Find available GPUs (not used by any previous allocation)
available = [g for g in self.gpus if g.id not in self._used_gpus]
if tp_size == 1:
# Single GPU - find one with enough memory
@@ -289,8 +292,8 @@ class GPUAllocator:
assigned_model=model_id,
port=get_open_port(),
)
slots.append(slot)
used_gpus.add(gpu.id)
new_slots.append(slot)
self._used_gpus.add(gpu.id)
logger.info(
"Allocated GPU %d (%s, %.1fGB) for %s",
gpu.id,
@@ -301,7 +304,10 @@ class GPUAllocator:
break
else:
logger.warning(
"No GPU with %.1fGB available for %s", memory_gb, model_id
"No GPU with %.1fGB available for %s (used: %s)",
memory_gb,
model_id,
self._used_gpus,
)
else:
# Multi-GPU - find consecutive GPUs with enough total memory
@@ -320,8 +326,8 @@ class GPUAllocator:
assigned_model=model_id,
port=get_open_port(),
)
slots.append(slot)
used_gpus.update(gpu_ids)
new_slots.append(slot)
self._used_gpus.update(gpu_ids)
logger.info(
"Allocated GPUs %s (%.1fGB total) for %s (tp=%d)",
gpu_ids,
@@ -332,14 +338,16 @@ class GPUAllocator:
break
else:
logger.warning(
"No %d consecutive GPUs with %.1fGB available for %s",
"No %d consecutive GPUs with %.1fGB available for %s (used: %s)",
tp_size,
memory_gb,
model_id,
self._used_gpus,
)
self.slots = slots
return slots
# Add new slots to existing slots list
self.slots.extend(new_slots)
return new_slots
def get_slot_for_model(self, model_id: str) -> GPUSlot | None:
"""Get the slot assigned to a specific model."""
@@ -348,10 +356,23 @@ class GPUAllocator:
return slot
return None
def release_gpus(self, gpu_ids: list[int]) -> None:
"""Release GPUs back to the available pool.
Args:
gpu_ids: List of GPU IDs to release.
"""
for gpu_id in gpu_ids:
self._used_gpus.discard(gpu_id)
# Remove slots that used these GPUs
self.slots = [s for s in self.slots if not any(g in gpu_ids for g in s.gpu_ids)]
logger.info("Released GPUs %s, now used: %s", gpu_ids, self._used_gpus)
def summary(self) -> str:
"""Return a summary of GPU allocations."""
lines = ["GPU Allocation Summary:"]
lines.append(f" Total GPUs: {len(self.gpus)}")
lines.append(f" Used GPUs: {sorted(self._used_gpus)}")
lines.append(f" Allocated Slots: {len(self.slots)}")
for slot in self.slots:
lines.append(
sgl-model-gateway/e2e_test/infra/model_pool.py
+461 -91
View File
@@ -4,10 +4,9 @@ from __future__ import annotations
import logging
import os
import signal
import subprocess
import time
from dataclasses import dataclass, field
from dataclasses import dataclass
from typing import TYPE_CHECKING
import httpx
@@ -15,61 +14,179 @@ import httpx
if TYPE_CHECKING:
import openai
from .gpu_allocator import GPUAllocator, GPUSlot
from .constants import (
DEFAULT_HOST,
DEFAULT_MODEL,
DEFAULT_STARTUP_TIMEOUT,
ENV_SHOW_WORKER_LOGS,
HEALTH_CHECK_INTERVAL,
LOCAL_MODES,
ConnectionMode,
WorkerType,
)
from .gpu_allocator import GPUAllocator, GPUSlot, get_open_port
from .model_specs import MODEL_SPECS, get_model_spec
from .process_utils import detect_ib_device
logger = logging.getLogger(__name__)
# Default timeout for model startup (seconds)
DEFAULT_STARTUP_TIMEOUT = 300
# Health check interval (seconds)
HEALTH_CHECK_INTERVAL = 5
# Host for model servers
DEFAULT_HOST = "127.0.0.1"
@dataclass
class ModelInstance:
"""A running model instance."""
model_id: str
mode: ConnectionMode
model_path: str
base_url: str
port: int
process: subprocess.Popen
gpu_slot: GPUSlot
grpc_mode: bool = False
gpu_slot: GPUSlot | None
worker_type: WorkerType = WorkerType.REGULAR
bootstrap_port: int | None = None # For prefill workers in PD mode
@property
def key(self) -> str:
"""Unique key for this instance.
Regular: 'model_id:mode' (e.g., 'llama-8b:http')
PD workers: 'model_id:mode:worker_type' (e.g., 'llama-8b:http:prefill')
"""
if self.worker_type == WorkerType.REGULAR:
return f"{self.model_id}:{self.mode.value}"
return f"{self.model_id}:{self.mode.value}:{self.worker_type.value}"
@property
def worker_url(self) -> str:
"""URL to use when connecting router to this worker."""
if self.mode == ConnectionMode.GRPC:
return f"grpc://{DEFAULT_HOST}:{self.port}"
return self.base_url
def is_alive(self) -> bool:
"""Check if the process is still running."""
return self.process.poll() is None
def health_check(self, timeout: float = 5.0) -> bool:
"""Check if the model server is healthy via HTTP."""
"""Check if the model server is healthy.
Uses HTTP /health endpoint for HTTP workers, gRPC health check for gRPC workers.
"""
if self.mode == ConnectionMode.GRPC:
return self._grpc_health_check(timeout)
return self._http_health_check(timeout)
def _http_health_check(self, timeout: float = 5.0) -> bool:
"""Check health via HTTP /health endpoint."""
try:
resp = httpx.get(f"{self.base_url}/health", timeout=timeout)
return resp.status_code == 200
except (httpx.RequestError, httpx.TimeoutException):
return False
def deep_health_check(self, timeout: float = 30.0) -> bool:
"""Deep health check that verifies the model can actually generate.
Uses /health_generate for HTTP workers (runs actual inference).
For gRPC workers, falls back to standard health check.
"""
if self.mode == ConnectionMode.GRPC:
# For gRPC, use standard health check (no /health_generate equivalent)
return self._grpc_health_check(timeout)
try:
resp = httpx.get(f"{self.base_url}/health_generate", timeout=timeout)
return resp.status_code == 200
except (httpx.RequestError, httpx.TimeoutException):
return False
def _grpc_health_check(self, timeout: float = 5.0) -> bool:
"""Check health via gRPC health check protocol."""
try:
import grpc
from grpc_health.v1 import health_pb2, health_pb2_grpc
except ImportError as e:
logger.debug("gRPC libraries not available: %s", e)
return False
try:
channel = grpc.insecure_channel(f"{DEFAULT_HOST}:{self.port}")
try:
stub = health_pb2_grpc.HealthStub(channel)
request = health_pb2.HealthCheckRequest(service="")
response = stub.Check(request, timeout=timeout)
is_serving = response.status == health_pb2.HealthCheckResponse.SERVING
if is_serving:
logger.debug(
"gRPC health check passed for port %d (status: SERVING)",
self.port,
)
return is_serving
finally:
channel.close()
except grpc.RpcError as e:
# gRPC-specific errors (connection refused, deadline exceeded, etc.)
logger.debug(
"gRPC health check failed for port %d: %s",
self.port,
e.code() if hasattr(e, "code") else str(e),
)
return False
except Exception as e:
# Other errors
logger.debug(
"gRPC health check error for port %d: %s",
self.port,
str(e),
)
return False
def terminate(self, timeout: float = 10.0) -> None:
"""Terminate the model server process."""
if self.process.poll() is not None:
return # Already terminated
logger.info("Terminating model %s (PID %d)", self.model_id, self.process.pid)
logger.info("Terminating %s (PID %d)", self.key, self.process.pid)
# Try graceful shutdown first
self.process.terminate()
try:
self.process.wait(timeout=timeout)
except subprocess.TimeoutExpired:
logger.warning("Model %s did not terminate, killing", self.model_id)
logger.warning("%s did not terminate, killing", self.key)
self.process.kill()
self.process.wait()
class ModelPool:
"""Manages a pool of pre-loaded models across GPUs."""
"""Manages long-running SGLang worker processes across GPUs.
Workers are expensive to start (~30-60s due to model loading), so this pool
keeps them running and allows reuse across multiple tests. Routers can then
be launched cheaply (~1-2s) pointing to these workers.
Startup behavior:
- Workers are launched sequentially (one subprocess.Popen at a time)
- But they boot up concurrently (overlapping model loading)
- _wait_all_healthy() blocks until all workers respond to health checks
Instance keys:
- Regular workers: "model_id:mode" (e.g., "llama-8b:http")
- PD workers: "model_id:mode:worker_type" (e.g., "llama-8b:http:prefill")
Limitations:
- Currently one worker instance per (model_id, mode) combination
- @pytest.mark.workers(n) duplicates URLs to router, not distinct workers
- For true multi-worker LB testing, extend to support multiple instances
Usage:
pool = ModelPool()
pool.startup(requirements=[("llama-8b", ConnectionMode.HTTP)])
instance = pool.get("llama-8b", "http")
# instance.base_url -> "http://127.0.0.1:30000"
# instance.worker_url -> URL for router to connect to
"""
def __init__(self, allocator: GPUAllocator | None = None):
"""Initialize the model pool.
@@ -78,68 +195,114 @@ class ModelPool:
allocator: GPU allocator to use. If None, creates a new one.
"""
self.allocator = allocator or GPUAllocator()
self.instances: dict[str, ModelInstance] = {}
self.instances: dict[str, ModelInstance] = {} # key = "model_id:mode"
self._startup_timeout = DEFAULT_STARTUP_TIMEOUT
def startup(
self,
model_ids: list[str] | None = None,
grpc_mode: bool = False,
requirements: list[tuple[str, ConnectionMode]] | None = None,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
) -> None:
"""Spin up models in parallel on assigned GPU slots.
"""Start worker processes for the required models.
Workers are launched sequentially (one Popen at a time) but boot up
concurrently since model loading happens in parallel across processes.
This method blocks until all workers pass health checks.
Args:
model_ids: List of model IDs to start. If None, starts all in MODEL_SPECS.
grpc_mode: If True, launch workers in gRPC mode.
startup_timeout: Timeout in seconds for each model to become healthy.
requirements: List of (model_id, mode) tuples specifying what to start.
mode is ConnectionMode.HTTP or ConnectionMode.GRPC.
If None, starts default model in HTTP mode.
startup_timeout: Timeout in seconds for all models to become healthy.
"""
self._startup_timeout = startup_timeout
# Determine which models to start
if model_ids is None:
model_ids = list(MODEL_SPECS.keys())
if requirements is None:
requirements = [(DEFAULT_MODEL, ConnectionMode.HTTP)]
# Filter to models we have specs for
specs_to_start = {
mid: MODEL_SPECS[mid] for mid in model_ids if mid in MODEL_SPECS
}
# Deduplicate and validate
requirements = list(set(requirements))
valid_requirements = []
for model_id, mode in requirements:
if model_id not in MODEL_SPECS:
logger.warning("Unknown model %s, skipping", model_id)
continue
if mode not in LOCAL_MODES:
logger.warning("Invalid mode %s for %s, skipping", mode, model_id)
continue
valid_requirements.append((model_id, mode))
if not specs_to_start:
logger.warning("No valid model specs to start")
if not valid_requirements:
logger.warning("No valid requirements to start")
return
logger.info("Starting model pool with: %s", valid_requirements)
# Build allocation specs - each (model, mode) combo needs its own slot
# Use "model_id:mode" as the allocation key
allocation_specs = {}
for model_id, mode in valid_requirements:
spec = MODEL_SPECS[model_id]
key = f"{model_id}:{mode.value}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": spec.get("tp", 1),
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(specs_to_start)
slots = self.allocator.allocate_slots(allocation_specs)
if not slots:
logger.warning("No GPU slots allocated")
return
logger.info(self.allocator.summary())
# Launch all models in parallel
for slot in slots:
if slot.assigned_model:
self._launch_model(slot, grpc_mode=grpc_mode)
logger.warning("No GPU slots allocated, launching without GPU assignment")
# Fallback: launch without specific GPU assignment
for model_id, mode in valid_requirements:
self._launch_model(model_id, mode, gpu_slot=None)
else:
# Launch on allocated slots
for slot in slots:
if slot.assigned_model:
# Parse "model_id:mode" back
model_id, mode_str = slot.assigned_model.rsplit(":", 1)
mode = ConnectionMode(mode_str)
self._launch_model(model_id, mode, gpu_slot=slot)
# Wait for all to be healthy
self._wait_all_healthy()
def _launch_model(self, slot: GPUSlot, grpc_mode: bool = False) -> None:
"""Launch a model on the given GPU slot."""
model_id = slot.assigned_model
if not model_id:
return
def _launch_model(
self,
model_id: str,
mode: ConnectionMode,
gpu_slot: GPUSlot | None = None,
worker_type: WorkerType = WorkerType.REGULAR,
bootstrap_port: int | None = None,
ib_device: str | None = None,
) -> ModelInstance:
"""Launch a model instance.
Args:
model_id: Model identifier from MODEL_SPECS.
mode: Connection mode (HTTP or GRPC).
gpu_slot: GPU slot assignment, or None for auto.
worker_type: Worker type (REGULAR, PREFILL, or DECODE).
bootstrap_port: Bootstrap port for prefill workers in PD mode.
ib_device: InfiniBand device for PD disaggregation.
Returns:
The launched ModelInstance.
"""
spec = get_model_spec(model_id)
model_path = spec["model"]
tp_size = spec.get("tp", 1)
port = slot.port
# Build environment with CUDA_VISIBLE_DEVICES
# Get port - use slot's port if available, otherwise find open port
port = gpu_slot.port if gpu_slot else get_open_port()
# Build environment
env = os.environ.copy()
env["CUDA_VISIBLE_DEVICES"] = slot.cuda_visible_devices()
if gpu_slot:
env["CUDA_VISIBLE_DEVICES"] = gpu_slot.cuda_visible_devices()
# Build command
cmd = [
@@ -148,6 +311,8 @@ class ModelPool:
"sglang.launch_server",
"--model-path",
model_path,
"--host",
DEFAULT_HOST,
"--port",
str(port),
"--tp-size",
@@ -156,52 +321,81 @@ class ModelPool:
"warning",
]
if grpc_mode:
if mode == ConnectionMode.GRPC:
cmd.append("--grpc-mode")
logger.info(
"Launching %s on GPUs %s port %d: %s",
model_id,
slot.gpu_ids,
port,
" ".join(cmd),
)
# PD disaggregation arguments
if worker_type == WorkerType.PREFILL:
cmd.extend(["--disaggregation-mode", "prefill"])
if bootstrap_port:
cmd.extend(["--disaggregation-bootstrap-port", str(bootstrap_port)])
if ib_device:
cmd.extend(["--disaggregation-ib-device", ib_device])
elif worker_type == WorkerType.DECODE:
cmd.extend(["--disaggregation-mode", "decode"])
if ib_device:
cmd.extend(["--disaggregation-ib-device", ib_device])
# Build key based on worker type
if worker_type == WorkerType.REGULAR:
key = f"{model_id}:{mode.value}"
else:
key = f"{model_id}:{mode.value}:{worker_type.value}"
gpu_info = gpu_slot.gpu_ids if gpu_slot else "auto"
logger.info("Launching %s on GPUs %s port %d", key, gpu_info, port)
show_output = os.environ.get(ENV_SHOW_WORKER_LOGS, "0") == "1"
# Start the process
proc = subprocess.Popen(
cmd,
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
# Use process group for clean shutdown
stdout=None if show_output else subprocess.PIPE,
stderr=None if show_output else subprocess.PIPE,
start_new_session=True,
)
base_url = f"http://{DEFAULT_HOST}:{port}"
instance = ModelInstance(
model_id=model_id,
mode=mode,
model_path=model_path,
base_url=base_url,
port=port,
process=proc,
gpu_slot=slot,
grpc_mode=grpc_mode,
gpu_slot=gpu_slot,
worker_type=worker_type,
bootstrap_port=bootstrap_port,
)
self.instances[model_id] = instance
self.instances[key] = instance
return instance
def _wait_all_healthy(self) -> None:
"""Wait for all model instances to become healthy."""
start_time = time.time()
pending = set(self.instances.keys())
check_count = 0
logger.info(
"Waiting for %d workers to become healthy (timeout: %ds)...",
len(pending),
self._startup_timeout,
)
while pending and (time.time() - start_time) < self._startup_timeout:
for model_id in list(pending):
instance = self.instances[model_id]
check_count += 1
elapsed = time.time() - start_time
for key in list(pending):
instance = self.instances[key]
# Check if process died
if not instance.is_alive():
logger.error(
"Model %s (PID %d) died during startup",
model_id,
"[%.1fs] %s (PID %d) died during startup",
elapsed,
key,
instance.process.pid,
)
# Read stderr for debugging
@@ -209,62 +403,238 @@ class ModelPool:
stderr = instance.process.stderr.read()
if stderr:
logger.error("Stderr: %s", stderr.decode()[-2000:])
pending.discard(model_id)
pending.discard(key)
continue
# Check health
if instance.health_check():
logger.info(
"Model %s is healthy at %s", model_id, instance.base_url
"[%.1fs] %s is healthy at %s (check #%d)",
elapsed,
key,
instance.base_url,
check_count,
)
pending.discard(model_id)
pending.discard(key)
if pending:
# Log progress every 30 seconds
if check_count % 15 == 0: # ~30s at 2s interval
logger.info(
"[%.1fs] Still waiting for %d workers: %s",
elapsed,
len(pending),
list(pending),
)
time.sleep(HEALTH_CHECK_INTERVAL)
if pending:
elapsed = time.time() - start_time
logger.error(
"Models failed to start within %ds: %s",
"[%.1fs] Models failed to start within %ds: %s",
elapsed,
self._startup_timeout,
pending,
)
# Terminate failed instances
for model_id in pending:
self.instances[model_id].terminate()
del self.instances[model_id]
for key in pending:
self.instances[key].terminate()
del self.instances[key]
else:
elapsed = time.time() - start_time
logger.info(
"[%.1fs] All %d workers healthy after %d health checks",
elapsed,
len(self.instances),
check_count,
)
def get_client(self, model_id: str) -> "openai.OpenAI":
def get(
self,
model_id: str,
mode: ConnectionMode | str,
worker_type: WorkerType | str = WorkerType.REGULAR,
) -> ModelInstance:
"""Get a model instance by model_id, mode, and worker_type.
Args:
model_id: The model ID (e.g., "llama-8b")
mode: The mode (ConnectionMode.HTTP or ConnectionMode.GRPC, or string)
worker_type: The worker type (REGULAR, PREFILL, DECODE). Defaults to REGULAR.
Returns:
ModelInstance for the requested model/mode/worker_type.
Raises:
KeyError: If model/mode/worker_type combination is not running.
"""
# Accept both enum and string for convenience
if isinstance(mode, str):
mode = ConnectionMode(mode)
if isinstance(worker_type, str):
worker_type = WorkerType(worker_type)
if worker_type == WorkerType.REGULAR:
key = f"{model_id}:{mode.value}"
else:
key = f"{model_id}:{mode.value}:{worker_type.value}"
if key not in self.instances:
raise KeyError(
f"{key} not running. Available: {list(self.instances.keys())}"
)
instance = self.instances[key]
# Verify worker is still alive and healthy
if not instance.is_alive():
raise RuntimeError(f"Worker {key} process died (was healthy at startup)")
if not instance.deep_health_check(timeout=30.0):
raise RuntimeError(
f"Worker {key} failed deep health check (health_generate) - "
"model may be stuck or crashed"
)
logger.info("Worker %s passed deep health check", key)
return instance
def get_workers_by_type(
self, model_id: str, worker_type: WorkerType
) -> list[ModelInstance]:
"""Get all workers of a specific type for a model.
Args:
model_id: The model ID.
worker_type: The worker type to filter by.
Returns:
List of matching ModelInstance objects.
"""
return [
inst
for inst in self.instances.values()
if inst.model_id == model_id and inst.worker_type == worker_type
]
def launch_pd_workers(
self,
model_id: str,
num_prefill: int = 1,
num_decode: int = 1,
mode: ConnectionMode = ConnectionMode.HTTP,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
) -> tuple[list[ModelInstance], list[ModelInstance]]:
"""Launch prefill and decode workers for PD disaggregation.
Args:
model_id: Model identifier from MODEL_SPECS.
num_prefill: Number of prefill workers to launch. Defaults to 1.
num_decode: Number of decode workers to launch. Defaults to 1.
mode: Connection mode (HTTP or GRPC).
startup_timeout: Timeout for workers to become healthy.
Returns:
Tuple of (prefill_instances, decode_instances).
"""
self._startup_timeout = startup_timeout
if model_id not in MODEL_SPECS:
raise ValueError(f"Unknown model: {model_id}")
spec = get_model_spec(model_id)
ib_device = detect_ib_device()
if ib_device:
logger.info("Detected InfiniBand device: %s", ib_device)
# Build allocation specs for all PD workers
# Each worker needs its own GPU slot
allocation_specs = {}
for i in range(num_prefill):
key = f"{model_id}:{mode.value}:prefill_{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": spec.get("tp", 1),
}
for i in range(num_decode):
key = f"{model_id}:{mode.value}:decode_{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": spec.get("tp", 1),
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
slot_map = {slot.assigned_model: slot for slot in slots}
if not slots:
logger.warning(
"No GPU slots allocated for PD workers, launching without GPU assignment"
)
prefill_instances: list[ModelInstance] = []
decode_instances: list[ModelInstance] = []
# Launch prefill workers
for i in range(num_prefill):
key = f"{model_id}:{mode.value}:prefill_{i}"
gpu_slot = slot_map.get(key)
bootstrap_port = get_open_port()
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.PREFILL,
bootstrap_port=bootstrap_port,
ib_device=ib_device,
)
prefill_instances.append(instance)
# Launch decode workers
for i in range(num_decode):
key = f"{model_id}:{mode.value}:decode_{i}"
gpu_slot = slot_map.get(key)
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.DECODE,
ib_device=ib_device,
)
decode_instances.append(instance)
# Wait for all to be healthy
self._wait_all_healthy()
return prefill_instances, decode_instances
def get_client(
self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP
) -> "openai.OpenAI":
"""Get OpenAI client for a specific model.
Args:
model_id: The model ID to get a client for.
mode: The mode (ConnectionMode.HTTP or ConnectionMode.GRPC). Defaults to HTTP.
Returns:
OpenAI client configured for this model.
Raises:
KeyError: If model is not running.
"""
import openai
if model_id not in self.instances:
raise KeyError(
f"Model {model_id} not running. Available: {list(self.instances.keys())}"
)
instance = self.instances[model_id]
instance = self.get(model_id, mode)
return openai.OpenAI(
base_url=f"{instance.base_url}/v1",
api_key="not-used",
)
def get_base_url(self, model_id: str) -> str:
def get_base_url(
self, model_id: str, mode: ConnectionMode | str = ConnectionMode.HTTP
) -> str:
"""Get the base URL for a specific model."""
if model_id not in self.instances:
raise KeyError(
f"Model {model_id} not running. Available: {list(self.instances.keys())}"
)
return self.instances[model_id].base_url
return self.get(model_id, mode).base_url
def shutdown(self) -> None:
"""Tear down all models."""
sgl-model-gateway/e2e_test/infra/process_utils.py
+161
View File
@@ -0,0 +1,161 @@
"""Process management utilities for E2E tests."""
from __future__ import annotations
import logging
import os
import signal
import subprocess
import time
import requests
logger = logging.getLogger(__name__)
def kill_process_tree(pid: int, sig: int = signal.SIGTERM) -> None:
"""Kill a process and all its children.
Args:
pid: Process ID to kill
sig: Signal to send (default: SIGTERM)
"""
try:
import psutil
parent = psutil.Process(pid)
children = parent.children(recursive=True)
for child in children:
try:
child.send_signal(sig)
except psutil.NoSuchProcess:
pass
parent.send_signal(sig)
except ImportError:
# Fallback if psutil not available
os.kill(pid, sig)
except Exception as e:
logger.warning("Failed to kill process tree for PID %d: %s", pid, e)
def terminate_process(proc: subprocess.Popen, timeout: float = 30) -> None:
"""Gracefully terminate a process, kill if needed.
Args:
proc: Process to terminate
timeout: Seconds to wait before force-killing
"""
if proc is None or proc.poll() is not None:
return
proc.terminate()
start = time.perf_counter()
while proc.poll() is None:
if time.perf_counter() - start > timeout:
proc.kill()
break
time.sleep(1)
def wait_for_health(
url: str,
timeout: float = 60,
api_key: str | None = None,
check_interval: float = 1.0,
) -> None:
"""Wait for a server's /health endpoint to return 200.
Args:
url: Base URL of the server
timeout: Seconds to wait before timing out
api_key: Optional API key for auth header
check_interval: Seconds between health checks
"""
start = time.perf_counter()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
with requests.Session() as session:
while time.perf_counter() - start < timeout:
try:
resp = session.get(f"{url}/health", headers=headers, timeout=5)
if resp.status_code == 200:
logger.info("Service healthy at %s", url)
return
except requests.RequestException:
pass
time.sleep(check_interval)
raise TimeoutError(f"Server at {url} did not become healthy within {timeout}s")
def wait_for_workers_ready(
router_url: str,
expected_workers: int,
timeout: float = 300,
api_key: str | None = None,
) -> None:
"""Wait for router to have all workers connected.
Args:
router_url: Base URL of the router
expected_workers: Number of workers to wait for
timeout: Seconds to wait before timing out
api_key: Optional API key for auth header
"""
start = time.perf_counter()
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
while time.perf_counter() - start < timeout:
try:
resp = requests.get(f"{router_url}/workers", headers=headers, timeout=5)
if resp.status_code == 200:
data = resp.json()
total = data.get("total", len(data.get("workers", [])))
if total >= expected_workers:
logger.info(
"All %d workers connected after %.1fs",
expected_workers,
time.perf_counter() - start,
)
return
except requests.RequestException:
pass
time.sleep(2)
raise TimeoutError(
f"Router at {router_url} did not get {expected_workers} workers within {timeout}s"
)
def detect_ib_device() -> str | None:
"""Detect first active InfiniBand device (e.g., mlx5_0).
Returns:
Device name if found (e.g., "mlx5_0"), None otherwise.
"""
try:
subprocess.run(
["ibv_devinfo", "-l"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=1,
)
except (FileNotFoundError, subprocess.TimeoutExpired):
return None
for i in range(12):
dev = f"mlx5_{i}"
try:
res = subprocess.run(
["ibv_devinfo", dev],
capture_output=True,
text=True,
timeout=2,
)
if res.returncode == 0 and "state:" in res.stdout:
for line in res.stdout.splitlines():
if "state:" in line and "PORT_ACTIVE" in line:
logger.info("Detected IB device: %s", dev)
return dev
except Exception:
pass
return None
sgl-model-gateway/e2e_test/infra/run_eval.py
+139
View File
@@ -0,0 +1,139 @@
"""MMLU evaluation runner for E2E tests.
Simplified evaluation runner that uses local eval implementations
with cleaner logging for CI/CD environments.
Usage:
from infra.run_eval import run_eval
from types import SimpleNamespace
args = SimpleNamespace(
base_url="http://127.0.0.1:30000",
model="meta-llama/Llama-3.1-8B-Instruct",
eval_name="mmlu",
num_examples=64,
num_threads=32,
temperature=0.1,
)
metrics = run_eval(args)
"""
from __future__ import annotations
import logging
import os
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from .simple_eval_common import Eval
from .simple_eval_common import ChatCompletionSampler, set_ulimit
logger = logging.getLogger(__name__)
# MMLU dataset URL
MMLU_DATASET_URL = "https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv"
@dataclass
class EvalConfig:
"""Configuration for running an evaluation."""
base_url: str
model: str | None = None
eval_name: str = "mmlu"
num_examples: int = 64
num_threads: int = 32
temperature: float = 0.0
max_tokens: int = 2048
host: str = "127.0.0.1"
port: int = 30000
def _get_eval(eval_name: str, num_examples: int, num_threads: int) -> "Eval":
"""Get the evaluation object by name."""
if eval_name == "mmlu":
from .simple_eval_mmlu import MMLUEval
return MMLUEval(MMLU_DATASET_URL, num_examples, num_threads)
else:
raise ValueError(f"Unknown eval: {eval_name}. Supported: mmlu")
def run_eval(args: Any) -> dict:
"""Run an evaluation and return metrics.
Args:
args: Configuration object with attributes:
- base_url: Base URL of the server (e.g., "http://127.0.0.1:30000")
- model: Model name/path (optional, will be auto-detected)
- eval_name: Evaluation name ("mmlu")
- num_examples: Number of examples to evaluate
- num_threads: Number of parallel threads
- temperature: Sampling temperature
Returns:
Dict with metrics including 'score' key.
"""
set_ulimit()
if "OPENAI_API_KEY" not in os.environ:
os.environ["OPENAI_API_KEY"] = "EMPTY"
# Build base URL
base_url = getattr(args, "base_url", None)
if base_url:
if not base_url.endswith("/v1"):
base_url = f"{base_url}/v1"
else:
host = getattr(args, "host", "127.0.0.1")
port = getattr(args, "port", 30000)
base_url = f"http://{host}:{port}/v1"
eval_name = getattr(args, "eval_name", "mmlu")
num_examples = getattr(args, "num_examples", 64)
num_threads = getattr(args, "num_threads", 32)
temperature = getattr(args, "temperature", 0.0)
max_tokens = getattr(args, "max_tokens", 2048)
model = getattr(args, "model", None)
logger.info(
"Starting %s eval: %d examples, %d threads, temp=%.2f",
eval_name,
num_examples,
num_threads,
temperature,
)
# Create sampler
sampler = ChatCompletionSampler(
model=model,
max_tokens=max_tokens,
base_url=base_url,
temperature=temperature,
)
# Get eval object
eval_obj = _get_eval(eval_name, num_examples, num_threads)
# Run evaluation
start_time = time.perf_counter()
result = eval_obj(sampler)
latency = time.perf_counter() - start_time
# Build metrics
metrics = result.metrics.copy() if result.metrics else {}
metrics["score"] = result.score
metrics["latency"] = latency
logger.info(
"%s eval complete: score=%.3f, latency=%.1fs, model=%s",
eval_name,
result.score,
latency,
sampler.model,
)
return metrics
sgl-model-gateway/e2e_test/infra/simple_eval_common.py
+485
View File
@@ -0,0 +1,485 @@
# Adapted from https://github.com/openai/simple-evals/
"""Common utilities for simple evaluations."""
from __future__ import annotations
import logging
import os
import resource
import time
from collections import defaultdict
from dataclasses import dataclass, field
from multiprocessing.pool import ThreadPool
from typing import Any
import httpx
import jinja2
import numpy as np
import openai
import requests
from openai import OpenAI
from tqdm import tqdm
logger = logging.getLogger(__name__)
OPENAI_SYSTEM_MESSAGE_API = "You are a helpful assistant."
OPENAI_SYSTEM_MESSAGE_CHATGPT = (
"You are ChatGPT, a large language model trained by OpenAI, based on the GPT-4 architecture."
+ "\nKnowledge cutoff: 2023-12\nCurrent date: 2024-04-01"
)
Message = dict[str, Any] # keys role, content
MessageList = list[Message]
class SamplerBase:
"""
Base class for defining a sampling model, which can be evaluated,
or used as part of the grading process.
"""
def __call__(self, message_list: MessageList) -> str:
raise NotImplementedError()
@dataclass
class EvalResult:
"""Result of running an evaluation (usually consisting of many samples)."""
score: float | None # top-line metric
metrics: dict[str, float] | None # other metrics
htmls: list[str] # strings of valid HTML
convos: list[MessageList] # sampled conversations
@dataclass
class SingleEvalResult:
"""Result of evaluating a single sample."""
score: float | None
metrics: dict[str, float] = field(default_factory=dict)
html: str | None = None
convo: MessageList | None = None # sampled conversation
class Eval:
"""
Base class for defining an evaluation.
"""
def __call__(self, sampler: SamplerBase) -> EvalResult:
raise NotImplementedError()
class LargerHttpxClient(httpx.Client):
def __init__(self):
timeout_config = httpx.Timeout(3600)
limits = httpx.Limits(
max_keepalive_connections=3600,
max_connections=3600,
)
super().__init__(timeout=timeout_config, limits=limits)
class ChatCompletionSampler(SamplerBase):
"""Sample from OpenAI's chat completion API."""
def __init__(
self,
base_url: str | None = None,
model: str | None = None,
system_message: str | None = None,
temperature: float = 0.0,
reasoning_effort: str | None = None,
max_tokens: int = 2048,
extra_body: dict[str, Any] | None = None,
):
self.client = OpenAI(base_url=base_url, http_client=LargerHttpxClient())
if model is None:
model = self.client.models.list().data[0].id
self.model = model
self.system_message = system_message
self.temperature = temperature
self.max_tokens = max_tokens
self.reasoning_effort = reasoning_effort
self.extra_body = extra_body
self.image_format = "url"
logger.debug(
"ChatCompletionSampler: model=%s, temp=%.2f, max_tokens=%d",
self.model,
self.temperature,
self.max_tokens,
)
def _handle_image(
self,
image: str,
encoding: str = "base64",
format: str = "png",
fovea: int = 768,
):
new_image = {
"type": "image_url",
"image_url": {
"url": f"data:image/{format};{encoding},{image}",
},
}
return new_image
def _handle_text(self, text: str):
return {"type": "text", "text": text}
def _pack_message(self, role: str, content: Any):
return {"role": str(role), "content": content}
def __call__(self, message_list: MessageList) -> str:
if self.system_message:
message_list = [
self._pack_message("system", self.system_message)
] + message_list
trial = 0
while trial < 6: # 126 seconds in total
try:
response = self.client.chat.completions.create(
model=self.model,
messages=message_list,
temperature=self.temperature,
max_tokens=self.max_tokens,
reasoning_effort=self.reasoning_effort,
extra_body=self.extra_body,
)
return response.choices[0].message.content or ""
except openai.BadRequestError as e:
logger.warning("Bad request error: %s", e)
return ""
except Exception as e:
exception_backoff = 2**trial # exponential back off
logger.debug(
"Rate limit, retry %d after %ds: %s",
trial,
exception_backoff,
e,
)
time.sleep(exception_backoff)
trial += 1
logger.warning("All retry attempts exhausted, returning empty response")
return ""
QUERY_TEMPLATE_MULTICHOICE = """
Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering.
{Question}
A) {A}
B) {B}
C) {C}
D) {D}
""".strip()
ANSWER_PATTERN_MULTICHOICE = r"(?i)Answer\s*:\s*([A-D])"
ANSWER_PATTERN = r"(?i)Answer\s*:\s*([^\n]+)"
EQUALITY_TEMPLATE = r"""
Look at the following two expressions (answers to a math problem) and judge whether they are equivalent. Only perform trivial simplifications
Examples:
Expression 1: $2x+3$
Expression 2: $3+2x$
Yes
Expression 1: 3/2
Expression 2: 1.5
Yes
Expression 1: $x^2+2x+1$
Expression 2: $y^2+2y+1$
No
Expression 1: $x^2+2x+1$
Expression 2: $(x+1)^2$
Yes
Expression 1: 3245/5
Expression 2: 649
No
(these are actually equal, don't mark them equivalent if you need to do nontrivial simplifications)
Expression 1: 2/(-3)
Expression 2: -2/3
Yes
(trivial simplifications are allowed)
Expression 1: 72 degrees
Expression 2: 72
Yes
(give benefit of the doubt to units)
Expression 1: 64
Expression 2: 64 square feet
Yes
(give benefit of the doubt to units)
---
YOUR TASK
Respond with only "Yes" or "No" (without quotes). Do not include a rationale.
Expression 1: %(expression1)s
Expression 2: %(expression2)s
""".strip()
HTML_JINJA = """
<h3>Prompt conversation</h3>
{% for message in prompt_messages %}
{{ message_to_html(message) | safe }}
{% endfor %}
<h3>Sampled message</h3>
{{ message_to_html(next_message) | safe }}
<h3>Results</h3>
<p>Correct Answer: {{ correct_answer }}</p>
<p>Extracted Answer: {{ extracted_answer }}</p>
<p>Score: {{ score }}</p>
"""
def format_multichoice_question(row):
return QUERY_TEMPLATE_MULTICHOICE.format(**row)
def check_equality(sampler: SamplerBase, expr1: str, expr2: str):
prompt = EQUALITY_TEMPLATE % {"expression1": expr1, "expression2": expr2}
response = sampler([dict(content=prompt, role="user")])
return (response or "").lower().strip() == "yes"
def _compute_stat(values: list, stat: str):
if stat == "mean":
return np.mean(values)
elif stat == "std":
return np.std(values)
elif stat == "min":
return np.min(values)
elif stat == "max":
return np.max(values)
else:
raise ValueError(f"Unknown {stat =}")
def aggregate_results(
single_eval_results: list[SingleEvalResult],
default_stats: tuple[str, ...] = ("mean", "std"),
name2stats: dict[str, tuple[str, ...]] | None = None,
) -> EvalResult:
"""
Aggregate results from multiple evaluations into a single EvalResult.
"""
name2stats = name2stats or {}
name2values = defaultdict(list)
htmls = []
convos = []
for single_eval_result in single_eval_results:
# Skip None results
if single_eval_result is None:
continue
for name, value in single_eval_result.metrics.items():
name2values[name].append(value)
if single_eval_result.score is not None:
name2values["score"].append(single_eval_result.score)
htmls.append(single_eval_result.html)
convos.append(single_eval_result.convo)
final_metrics = {}
for name, values in name2values.items():
stats = name2stats.get(name, default_stats)
for stat in stats:
key = name if stat == "mean" else f"{name}:{stat}"
final_metrics[key] = _compute_stat(values, stat)
return EvalResult(
score=final_metrics.pop("score", None),
metrics=final_metrics,
htmls=htmls,
convos=convos,
)
def map_with_progress(f: callable, xs: list[Any], num_threads: int) -> list[Any]:
"""Apply f to each element of xs, using a ThreadPool, and show progress."""
# Use quiet progress bar that doesn't pollute logs
if os.getenv("debug"):
return list(map(f, tqdm(xs, total=len(xs), leave=False)))
else:
with ThreadPool(min(num_threads, len(xs))) as pool:
return list(tqdm(pool.imap(f, xs), total=len(xs), leave=False))
jinja_env = jinja2.Environment(
loader=jinja2.BaseLoader(),
undefined=jinja2.StrictUndefined,
autoescape=jinja2.select_autoescape(["html", "xml"]),
)
_message_template = """
<div class="message {{ role }}">
<div class="role">
{{ role }}
{% if variant %}<span class="variant">({{ variant }})</span>{% endif %}
</div>
<div class="content">
<pre>{{ content }}</pre>
</div>
</div>
"""
def message_to_html(message: Message) -> str:
"""
Generate HTML snippet (inside a <div>) for a message.
"""
return jinja_env.from_string(_message_template).render(
role=message["role"],
content=message["content"],
variant=message.get("variant", None),
)
jinja_env.globals["message_to_html"] = message_to_html
_report_template = """<!DOCTYPE html>
<html>
<head>
<style>
.message {
padding: 8px 16px;
margin-bottom: 8px;
border-radius: 4px;
}
.message.user {
background-color: #B2DFDB;
color: #00695C;
}
.message.assistant {
background-color: #B39DDB;
color: #4527A0;
}
.message.system {
background-color: #EEEEEE;
color: #212121;
}
.role {
font-weight: bold;
margin-bottom: 4px;
}
.variant {
color: #795548;
}
table, th, td {
border: 1px solid black;
}
pre {
white-space: pre-wrap;
}
</style>
</head>
<body>
{% if metrics %}
<h1>Metrics</h1>
<table>
<tr>
<th>Metric</th>
<th>Value</th>
</tr>
<tr>
<td><b>Score</b></td>
<td>{{ score | float | round(3) }}</td>
</tr>
{% for name, value in metrics.items() %}
<tr>
<td>{{ name }}</td>
<td>{{ value }}</td>
</tr>
{% endfor %}
</table>
{% endif %}
<h1>Examples</h1>
{% for html in htmls %}
{{ html | safe }}
<hr>
{% endfor %}
</body>
</html>
"""
def make_report(eval_result: EvalResult) -> str:
"""
Create a standalone HTML report from an EvalResult.
"""
return jinja_env.from_string(_report_template).render(
score=eval_result.score,
metrics=eval_result.metrics,
htmls=eval_result.htmls,
)
def make_report_from_example_htmls(htmls: List[str]):
"""
Create a standalone HTML report from a list of example htmls
"""
return jinja_env.from_string(_report_template).render(
score=None, metrics={}, htmls=htmls
)
def download_dataset(path: str, url: str) -> None:
"""Download a dataset from URL to path."""
logger.info("Downloading dataset from %s", url)
try:
response = requests.get(url, stream=True)
response.raise_for_status()
total_size = int(response.headers.get("content-length", 0))
block_size = 8192
with open(path, "wb") as f, tqdm(
desc="Downloading",
total=total_size,
unit="iB",
unit_scale=True,
unit_divisor=1024,
leave=False,
) as progress_bar:
for data in response.iter_content(block_size):
size = f.write(data)
progress_bar.update(size)
logger.debug("Dataset saved to %s", path)
except requests.RequestException as e:
raise RuntimeError(f"Failed to download dataset: {e}") from e
def set_ulimit(target_soft_limit: int = 65535) -> None:
"""Set the file descriptor limit for parallel requests."""
resource_type = resource.RLIMIT_NOFILE
current_soft, current_hard = resource.getrlimit(resource_type)
if current_soft < target_soft_limit:
try:
resource.setrlimit(resource_type, (target_soft_limit, current_hard))
except ValueError as e:
logger.debug("Could not set RLIMIT_NOFILE: %s", e)
sgl-model-gateway/e2e_test/infra/simple_eval_mmlu.py
+126
View File
@@ -0,0 +1,126 @@
# Adapted from https://github.com/openai/simple-evals/
"""
MMLU Evaluation - Measuring Massive Multitask Language Understanding
Dan Hendrycks et al. https://arxiv.org/abs/2009.03300
"""
from __future__ import annotations
import random
import re
from typing import TYPE_CHECKING
import pandas
from . import simple_eval_common as common
from .simple_eval_common import (
ANSWER_PATTERN_MULTICHOICE,
HTML_JINJA,
Eval,
EvalResult,
SingleEvalResult,
format_multichoice_question,
)
if TYPE_CHECKING:
from .simple_eval_common import SamplerBase
SUBJECT_TO_CATEGORY = {
"abstract_algebra": "stem",
"anatomy": "other",
"astronomy": "stem",
"business_ethics": "other",
"clinical_knowledge": "other",
"college_biology": "stem",
"college_chemistry": "stem",
"college_computer_science": "stem",
"college_mathematics": "stem",
"college_medicine": "other",
"college_physics": "stem",
"computer_security": "stem",
"conceptual_physics": "stem",
"econometrics": "social_sciences",
"electrical_engineering": "stem",
"elementary_mathematics": "stem",
"formal_logic": "humanities",
"global_facts": "other",
"high_school_biology": "stem",
"high_school_chemistry": "stem",
"high_school_computer_science": "stem",
"high_school_european_history": "humanities",
"high_school_geography": "social_sciences",
"high_school_government_and_politics": "social_sciences",
"high_school_macroeconomics": "social_sciences",
"high_school_mathematics": "stem",
"high_school_microeconomics": "social_sciences",
"high_school_physics": "stem",
"high_school_psychology": "social_sciences",
"high_school_statistics": "stem",
"high_school_us_history": "humanities",
"high_school_world_history": "humanities",
"human_aging": "other",
"human_sexuality": "social_sciences",
"international_law": "humanities",
"jurisprudence": "humanities",
"logical_fallacies": "humanities",
"machine_learning": "stem",
"management": "other",
"marketing": "other",
"medical_genetics": "other",
"miscellaneous": "other",
"moral_disputes": "humanities",
"moral_scenarios": "humanities",
"nutrition": "other",
"philosophy": "humanities",
"prehistory": "humanities",
"professional_accounting": "other",
"professional_law": "humanities",
"professional_medicine": "other",
"professional_psychology": "social_sciences",
"public_relations": "social_sciences",
"security_studies": "social_sciences",
"sociology": "social_sciences",
"us_foreign_policy": "social_sciences",
"virology": "other",
"world_religions": "humanities",
}
class MMLUEval(Eval):
"""MMLU benchmark evaluation."""
def __init__(self, filename: str, num_examples: int | None, num_threads: int):
df = pandas.read_csv(filename)
examples = [row.to_dict() for _, row in df.iterrows()]
if num_examples:
examples = random.Random(0).sample(examples, num_examples)
self.examples = examples
self.num_threads = num_threads
def __call__(self, sampler: "SamplerBase") -> EvalResult:
def fn(row: dict) -> SingleEvalResult:
prompt_messages = [
sampler._pack_message(
content=format_multichoice_question(row), role="user"
)
]
response_text = sampler(prompt_messages)
response_text = response_text or ""
match = re.search(ANSWER_PATTERN_MULTICHOICE, response_text)
extracted_answer = match.group(1) if match else None
score = 1.0 if extracted_answer == row["Answer"] else 0.0
html = common.jinja_env.from_string(HTML_JINJA).render(
prompt_messages=prompt_messages,
next_message=dict(content=response_text, role="assistant"),
score=score,
correct_answer=row["Answer"],
extracted_answer=extracted_answer,
)
convo = prompt_messages + [dict(content=response_text, role="assistant")]
category = SUBJECT_TO_CATEGORY.get(row["Subject"], "other")
return SingleEvalResult(
html=html, score=score, metrics={category: score}, convo=convo
)
results = common.map_with_progress(fn, self.examples, self.num_threads)
return common.aggregate_results(results)