gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: vendor sglang v0.5.10 snapshot

Browse Source
This commit is contained in:
Gahow Wang
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

144
third_party/sglang/sgl-model-gateway/e2e_test/infra/__init__.py vendored Normal file
View File

@@ -0,0 +1,144 @@
"""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,
LOG_SEPARATOR_WIDTH,
MAX_RETRY_ATTEMPTS,
PARAM_BACKEND_ROUTER,
PARAM_MODEL,
PARAM_SETUP_BACKEND,
ConnectionMode,
Runtime,
WorkerType,
)
from .gateway import Gateway, WorkerInfo, launch_cloud_gateway
from .gpu_allocator import (
GPUAllocator,
GPUInfo,
GPUSlot,
get_gpu_memory_usage,
get_open_port,
get_physical_device_indices,
nvml_context,
wait_for_gpu_memory_to_clear,
)
from .gpu_monitor import GPUMonitor
from .gpu_monitor import should_monitor as should_monitor_gpu
from .model_pool import ModelInstance, ModelPool, WorkerIdentity
from .model_specs import ( # Default model paths; Model groups
CHAT_MODELS,
DEFAULT_EMBEDDING_MODEL_PATH,
DEFAULT_ENABLE_THINKING_MODEL_PATH,
DEFAULT_GPT_OSS_MODEL_PATH,
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_MODEL_PATH,
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_REASONING_MODEL_PATH,
DEFAULT_SMALL_MODEL_PATH,
EMBEDDING_MODELS,
FUNCTION_CALLING_MODELS,
MODEL_SPECS,
REASONING_MODELS,
THIRD_PARTY_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 and Identity
"ConnectionMode",
"WorkerType",
"Runtime",
"WorkerIdentity",
# 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",
"MAX_RETRY_ATTEMPTS",
"LOG_SEPARATOR_WIDTH",
# 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",
"GPUSlot",
# GPU utilities
"nvml_context",
"get_open_port",
"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",
# GPU monitoring
"GPUMonitor",
"should_monitor_gpu",
# Model management
"ModelInstance",
"ModelPool",
"MODEL_SPECS",
# Gateway
"Gateway",
"WorkerInfo",
"launch_cloud_gateway",
# Default model paths
"DEFAULT_MODEL_PATH",
"DEFAULT_SMALL_MODEL_PATH",
"DEFAULT_REASONING_MODEL_PATH",
"DEFAULT_ENABLE_THINKING_MODEL_PATH",
"DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH",
"DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH",
"DEFAULT_GPT_OSS_MODEL_PATH",
"DEFAULT_EMBEDDING_MODEL_PATH",
# Model groups
"CHAT_MODELS",
"EMBEDDING_MODELS",
"REASONING_MODELS",
"FUNCTION_CALLING_MODELS",
# Third-party models
"THIRD_PARTY_MODELS",
# Evaluation
"run_eval",
]

74
third_party/sglang/sgl-model-gateway/e2e_test/infra/constants.py vendored Normal file
View File

@@ -0,0 +1,74 @@
"""Constants and enums for E2E test infrastructure."""
from enum import Enum
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 = 2 # Check every 2s (was 5s)
# Model loading configuration
INITIAL_GRACE_PERIOD = 30 # Wait before first health check (model loading time)
LAUNCH_STAGGER_DELAY = (
10 # Delay between launching multiple workers (avoid I/O contention)
)
# Retry configuration
MAX_RETRY_ATTEMPTS = (
6 # Max retries with exponential backoff (total ~63s: 1+2+4+8+16+32)
)
# Display formatting
LOG_SEPARATOR_WIDTH = 60 # Width for log separator lines (e.g., "="*60)

595
third_party/sglang/sgl-model-gateway/e2e_test/infra/gateway.py vendored Normal file
View 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

437
third_party/sglang/sgl-model-gateway/e2e_test/infra/gpu_allocator.py vendored Normal file
View File

@@ -0,0 +1,437 @@
"""GPU detection and slot allocation for parallel test execution."""
from __future__ import annotations
import logging
import os
import socket
import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass
logger = logging.getLogger(__name__)
# Try to import nvidia-ml-py for GPU detection
try:
import pynvml
NVML_AVAILABLE = True
except ImportError:
NVML_AVAILABLE = False
logger.debug("nvidia-ml-py not available, GPU detection will be limited")
@contextmanager
def nvml_context():
"""Context manager for NVML initialization/shutdown.
Usage:
with nvml_context():
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
...
"""
if not NVML_AVAILABLE:
yield
return
try:
pynvml.nvmlInit()
yield
finally:
pynvml.nvmlShutdown()
@dataclass
class GPUInfo:
"""Information about a single GPU."""
id: int
name: str
memory_mb: int
@property
def memory_gb(self) -> float:
return self.memory_mb / 1024
@dataclass
class GPUSlot:
"""A slot representing one or more GPUs allocated for a model."""
gpu_ids: list[int]
total_memory_mb: int
assigned_model: str | None = None
port: int | None = None
@property
def total_memory_gb(self) -> float:
return self.total_memory_mb / 1024
def cuda_visible_devices(self) -> str:
"""Return CUDA_VISIBLE_DEVICES string for this slot."""
return ",".join(str(g) for g in self.gpu_ids)
def get_open_port() -> int:
"""Get an available port by binding to port 0 and reading the assigned port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
s.listen(1)
port = s.getsockname()[1]
return port
def get_physical_device_indices(devices: list[int]) -> list[int]:
"""Map logical device indices to physical indices based on CUDA_VISIBLE_DEVICES."""
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
if visible_devices is None:
return devices
visible_indices = [int(x) for x in visible_devices.split(",")]
index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
return [index_mapping[i] for i in devices if i in index_mapping]
def get_gpu_memory_usage(device_id: int) -> tuple[float, float]:
"""Get GPU memory usage in GB (used, total).
Args:
device_id: Physical GPU device ID
Returns:
Tuple of (used_gb, total_gb)
"""
if not NVML_AVAILABLE:
return (0.0, 0.0)
with nvml_context():
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
return (mem_info.used / (1024**3), mem_info.total / (1024**3))
def wait_for_gpu_memory_to_clear(
*,
devices: list[int],
threshold_bytes: int | None = None,
threshold_ratio: float | None = None,
timeout_s: float = 120,
) -> None:
"""Wait for GPU memory to be freed below a threshold.
Args:
devices: List of logical GPU device IDs to check
threshold_bytes: Memory threshold in bytes (used <= threshold)
threshold_ratio: Memory threshold as ratio (used/total <= ratio)
timeout_s: Timeout in seconds
Raises:
ValueError: If memory doesn't clear within timeout
"""
if not NVML_AVAILABLE:
logger.warning("nvidia-ml-py not available, skipping memory wait")
return
if threshold_bytes is None and threshold_ratio is None:
raise ValueError("Must specify threshold_bytes or threshold_ratio")
physical_devices = get_physical_device_indices(devices)
start_time = time.time()
with nvml_context():
while True:
output: dict[int, str] = {}
output_raw: dict[int, tuple[float, float]] = {}
for device in physical_devices:
handle = pynvml.nvmlDeviceGetHandleByIndex(device)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
gb_used = mem_info.used / (1024**3)
gb_total = mem_info.total / (1024**3)
output_raw[device] = (gb_used, gb_total)
output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
logger.debug(
"GPU memory used/total (GiB): %s",
" ".join(f"{k}={v}" for k, v in output.items()),
)
if threshold_bytes is not None:
def is_free(used: float, total: float) -> bool:
return used <= threshold_bytes / (1024**3)
threshold_desc = f"{threshold_bytes / (1024**3):.1f} GiB"
else:
def is_free(used: float, total: float) -> bool:
return used / total <= threshold_ratio # type: ignore[operator]
threshold_desc = f"{threshold_ratio:.2%}" # type: ignore[str-format]
dur_s = time.time() - start_time
if all(is_free(used, total) for used, total in output_raw.values()):
logger.info(
"GPU memory cleared on devices %s (threshold=%s) in %.1fs",
devices,
threshold_desc,
dur_s,
)
return
if dur_s >= timeout_s:
raise ValueError(
f"GPU memory on devices {devices} not freed after {dur_s:.1f}s "
f"(threshold={threshold_desc})"
)
time.sleep(5)
class GPUAllocator:
"""Detects GPUs and assigns them to model slots using bin-packing."""
def __init__(self, gpus: list[GPUInfo] | None = None):
"""Initialize the allocator.
Args:
gpus: Optional list of GPUs. If None, auto-detects via nvidia-ml-py.
"""
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
self._lock = threading.RLock() # Protects slots and _used_gpus
def _detect_gpus(self) -> list[GPUInfo]:
"""Auto-detect available GPUs via nvidia-ml-py (NVML)."""
if not NVML_AVAILABLE:
logger.warning("nvidia-ml-py not available - no GPUs detected")
return []
# Check for CUDA_VISIBLE_DEVICES restriction
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
allowed_ids: set[int] | None = None
if visible_devices:
allowed_ids = set(int(x) for x in visible_devices.split(",") if x.strip())
try:
with nvml_context():
device_count = pynvml.nvmlDeviceGetCount()
gpus = []
for idx in range(device_count):
# Skip GPUs not in CUDA_VISIBLE_DEVICES if set
if allowed_ids is not None and idx not in allowed_ids:
continue
handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
name = pynvml.nvmlDeviceGetName(handle)
# Handle bytes vs string return type (varies by pynvml version)
if isinstance(name, bytes):
name = name.decode("utf-8", errors="replace")
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
# Convert bytes to MB
memory_mb = mem_info.total // (1024 * 1024)
gpus.append(GPUInfo(idx, name, memory_mb))
logger.info("Detected %d GPUs: %s", len(gpus), [g.name for g in gpus])
return gpus
except pynvml.NVMLError as e:
logger.warning("NVML error during GPU detection: %s", e)
return []
except Exception as e:
logger.warning("Failed to detect GPUs: %s", e)
return []
def allocate_slots(
self, model_specs: dict[str, dict], preserve_order: bool = False
) -> list[GPUSlot]:
"""Allocate GPU slots based on model memory requirements.
Uses a first-fit decreasing bin-packing algorithm by default:
1. Sort models by memory requirement (largest first)
2. For each model, find the first GPU(s) that can fit it
3. For multi-GPU models, find consecutive GPUs
When preserve_order=True, processes models in dict insertion order
(test collection order) instead of sorting by memory. This ensures
models needed by earlier tests are allocated first.
Note: This method tracks used GPUs across multiple calls, so subsequent
allocations will use different GPUs than previous ones.
Thread-safe: Protected by internal lock.
Args:
model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys
preserve_order: If True, allocate in dict order (test order) instead
of sorting by memory size. Default False.
Returns:
List of GPUSlots with assigned models (only the newly allocated slots)
"""
with self._lock:
return self._allocate_slots_unlocked(model_specs, preserve_order)
def _allocate_slots_unlocked(
self, model_specs: dict[str, dict], preserve_order: bool = False
) -> list[GPUSlot]:
"""Internal allocation logic. Caller must hold _lock."""
if not self.gpus:
logger.warning("No GPUs available for allocation")
return []
if preserve_order:
# Process in dict insertion order (test collection order)
ordered_models = list(model_specs.items())
else:
# Sort models by memory requirement (largest first for better packing)
ordered_models = sorted(
model_specs.items(),
key=lambda x: x[1].get("memory_gb", 0),
reverse=True,
)
# Track new slots allocated in this call
new_slots: list[GPUSlot] = []
for model_id, spec in ordered_models:
memory_gb = spec.get("memory_gb", 16)
tp_size = spec.get("tp", 1)
# 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
for gpu in available:
if gpu.memory_gb >= memory_gb:
slot = GPUSlot(
gpu_ids=[gpu.id],
total_memory_mb=gpu.memory_mb,
assigned_model=model_id,
port=get_open_port(),
)
new_slots.append(slot)
self._used_gpus.add(gpu.id)
logger.info(
"Allocated GPU %d (%s, %.1fGB) for %s",
gpu.id,
gpu.name,
gpu.memory_gb,
model_id,
)
break
else:
logger.warning(
"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
# Sort available by ID for consecutive allocation
available_sorted = sorted(available, key=lambda g: g.id)
for i in range(len(available_sorted) - tp_size + 1):
candidate_gpus = available_sorted[i : i + tp_size]
total_mem = sum(g.memory_mb for g in candidate_gpus)
if total_mem >= memory_gb * 1024:
gpu_ids = [g.id for g in candidate_gpus]
slot = GPUSlot(
gpu_ids=gpu_ids,
total_memory_mb=total_mem,
assigned_model=model_id,
port=get_open_port(),
)
new_slots.append(slot)
self._used_gpus.update(gpu_ids)
logger.info(
"Allocated GPUs %s (%.1fGB total) for %s (tp=%d)",
gpu_ids,
total_mem / 1024,
model_id,
tp_size,
)
break
else:
logger.warning(
"No %d consecutive GPUs with %.1fGB available for %s (used: %s)",
tp_size,
memory_gb,
model_id,
self._used_gpus,
)
# 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.
Thread-safe: Protected by internal lock.
"""
with self._lock:
for slot in self.slots:
if slot.assigned_model == model_id:
return slot
return None
def release_gpus(self, gpu_ids: list[int]) -> None:
"""Release GPUs back to the available pool.
Thread-safe: Protected by internal lock.
Args:
gpu_ids: List of GPU IDs to release.
"""
with self._lock:
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 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.
Thread-safe: Protected by internal lock.
Returns:
List of GPU IDs that are not currently allocated.
"""
with self._lock:
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.
Thread-safe: Protected by internal lock.
"""
with self._lock:
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(
f" - {slot.assigned_model}: GPUs {slot.gpu_ids} "
f"({slot.total_memory_gb:.1f}GB) port={slot.port}"
)
return "\n".join(lines)

329
third_party/sglang/sgl-model-gateway/e2e_test/infra/gpu_monitor.py vendored Normal file
View File

@@ -0,0 +1,329 @@
"""GPU utilization monitoring for benchmarks.
This module provides a low-impact GPU monitor that runs in a separate process
and collects utilization samples using NVML.
"""
from __future__ import annotations
import json
import logging
import os
import time
from multiprocessing import Process
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
logger = logging.getLogger(__name__)
def _percentile(samples: list[float], p: float) -> float:
"""Calculate percentile from sorted samples."""
if not samples:
return 0.0
sorted_samples = sorted(samples)
idx = max(
0,
min(
len(sorted_samples) - 1, int(round((p / 100.0) * (len(sorted_samples) - 1)))
),
)
return float(sorted_samples[idx])
def _compute_stats(samples: list[float]) -> dict[str, float]:
"""Compute statistics for a list of samples."""
if not samples:
return {
"mean": 0.0,
"min": 0.0,
"max": 0.0,
"p5": 0.0,
"p10": 0.0,
"p25": 0.0,
"p50": 0.0,
"p75": 0.0,
"p90": 0.0,
"p95": 0.0,
"count": 0,
}
return {
"mean": sum(samples) / len(samples),
"min": min(samples),
"max": max(samples),
"p5": _percentile(samples, 5),
"p10": _percentile(samples, 10),
"p25": _percentile(samples, 25),
"p50": _percentile(samples, 50),
"p75": _percentile(samples, 75),
"p90": _percentile(samples, 90),
"p95": _percentile(samples, 95),
"count": len(samples),
}
def _monitor_loop(pid: int, output_path: str, interval: float) -> None:
"""Main monitoring loop - runs in separate process.
Monitors GPU utilization until the target process exits, then writes
results to output_path as JSON.
"""
# Lower process priority to minimize impact on benchmark
try:
os.nice(10)
except Exception:
pass
# Initialize NVML
try:
import pynvml
pynvml.nvmlInit()
except Exception as e:
logger.warning("Failed to initialize NVML: %s", e)
_write_empty_result(output_path)
return
# Get GPU handles
try:
device_count = pynvml.nvmlDeviceGetCount()
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(device_count)]
except Exception as e:
logger.warning("Failed to get GPU handles: %s", e)
_write_empty_result(output_path)
_shutdown_nvml()
return
# Collect samples
per_gpu_samples: dict[str, list[float]] = {str(i): [] for i in range(device_count)}
overall_samples: list[float] = []
try:
while _process_alive(pid):
try:
gpu_utils = []
for idx, handle in enumerate(handles):
try:
util = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
gpu_utils.append(float(util))
per_gpu_samples[str(idx)].append(float(util))
except Exception:
continue
if gpu_utils:
avg = sum(gpu_utils) / len(gpu_utils)
overall_samples.append(avg)
except Exception:
pass
time.sleep(interval)
finally:
# Write results
_write_result(output_path, pid, interval, overall_samples, per_gpu_samples)
_shutdown_nvml()
def _process_alive(pid: int) -> bool:
"""Check if process is still running."""
try:
os.kill(pid, 0)
return True
except (OSError, ProcessLookupError):
return False
def _write_empty_result(path: str) -> None:
"""Write empty result file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {"overall": [], "per_gpu": {}},
},
f,
)
except Exception:
pass
def _write_result(
path: str,
pid: int,
interval: float,
overall_samples: list[float],
per_gpu_samples: dict[str, list[float]],
) -> None:
"""Write monitoring results to JSON file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"bench_pid": pid,
"interval_sec": interval,
"count": len(overall_samples),
"overall": _compute_stats(overall_samples),
"per_gpu": {
k: _compute_stats(v) for k, v in per_gpu_samples.items()
},
"raw": {
"overall": overall_samples,
"per_gpu": per_gpu_samples,
},
},
f,
)
except Exception as e:
logger.warning("Failed to write GPU monitor results: %s", e)
def _shutdown_nvml() -> None:
"""Shutdown NVML."""
try:
import pynvml
pynvml.nvmlShutdown()
except Exception:
pass
class GPUMonitor:
"""GPU utilization monitor for benchmarks.
Usage:
monitor = GPUMonitor(output_dir="benchmark_results")
monitor.start(target_pid=12345)
# ... run benchmark ...
result = monitor.stop()
monitor.assert_thresholds({"gpu_util_p50_min": 99})
"""
def __init__(
self,
output_dir: str | Path = ".",
interval: float = 2.0,
):
self.output_dir = Path(output_dir)
self.interval = interval
self._process: Process | None = None
self._output_path: str | None = None
self._result: dict[str, Any] | None = None
@property
def output_path(self) -> str | None:
"""Path to the GPU utilization JSON file."""
return self._output_path
def start(self, target_pid: int) -> None:
"""Start monitoring GPU utilization for the target process."""
self._output_path = str(self.output_dir / "gpu_utilization.json")
self._result = None
self._process = Process(
target=_monitor_loop,
args=(target_pid, self._output_path, self.interval),
daemon=True,
)
self._process.start()
logger.debug("Started GPU monitor for PID %d", target_pid)
def stop(self, timeout: float = 5.0) -> dict[str, Any] | None:
"""Stop monitoring and return results."""
if self._process is None:
return None
try:
self._process.join(timeout=timeout)
except Exception:
pass
if self._process.is_alive():
try:
self._process.terminate()
except Exception:
pass
self._process = None
self._result = self._read_result()
return self._result
def _read_result(self) -> dict[str, Any] | None:
"""Read results from output file."""
if not self._output_path or not os.path.exists(self._output_path):
return None
try:
with open(self._output_path) as f:
return json.load(f)
except Exception as e:
logger.warning("Failed to read GPU monitor result: %s", e)
return None
def log_summary(self) -> None:
"""Log a summary of GPU utilization."""
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
logger.info(
"GPU utilization: mean=%.2f%% p50=%.2f%% (samples=%d)",
overall.get("mean", 0.0),
overall.get("p50", 0.0),
result.get("count", 0),
)
def assert_thresholds(self, thresholds: dict[str, float] | None) -> None:
"""Assert GPU utilization meets thresholds.
Supported thresholds:
- gpu_util_mean_min: Minimum mean GPU utilization %
- gpu_util_p50_min: Minimum p50 GPU utilization %
"""
if not thresholds:
return
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
mean_threshold = thresholds.get("gpu_util_mean_min")
if mean_threshold is not None:
mean_value = overall.get("mean", 0.0)
assert (
mean_value >= mean_threshold
), f"GPU utilization mean below threshold: {mean_value:.2f}% < {mean_threshold}%"
p50_threshold = thresholds.get("gpu_util_p50_min")
if p50_threshold is not None:
p50_value = overall.get("p50")
if p50_value is not None:
assert (
p50_value >= p50_threshold
), f"GPU utilization p50 below threshold: {p50_value:.2f}% < {p50_threshold}%"
def should_monitor(thresholds: dict[str, Any] | None) -> bool:
"""Check if GPU monitoring should be enabled.
Returns True if:
- thresholds contains gpu_util_mean_min or gpu_util_p50_min, OR
- GPU_UTIL_LOG environment variable is truthy
"""
if thresholds:
if thresholds.get("gpu_util_mean_min") is not None:
return True
if thresholds.get("gpu_util_p50_min") is not None:
return True
return os.environ.get("GPU_UTIL_LOG", "").lower() in ("1", "true", "yes")

1231
third_party/sglang/sgl-model-gateway/e2e_test/infra/model_pool.py vendored Normal file
View File

File diff suppressed because it is too large Load Diff

151
third_party/sglang/sgl-model-gateway/e2e_test/infra/model_specs.py vendored Normal file
View File

@@ -0,0 +1,151 @@
"""Model specifications for E2E tests.
Each model spec defines:
- model: HuggingFace model path or local path
- memory_gb: Estimated GPU memory required
- tp: Tensor parallelism size (number of GPUs needed)
- features: List of features this model supports (for test filtering)
"""
from __future__ import annotations
import os
# Environment variable for local model paths (CI uses local copies for speed)
ROUTER_LOCAL_MODEL_PATH = os.environ.get("ROUTER_LOCAL_MODEL_PATH", "")
def _resolve_model_path(hf_path: str) -> str:
"""Resolve model path, preferring local path if available."""
if ROUTER_LOCAL_MODEL_PATH:
local_path = os.path.join(ROUTER_LOCAL_MODEL_PATH, hf_path)
if os.path.exists(local_path):
return local_path
return hf_path
MODEL_SPECS: dict[str, dict] = {
# Primary chat model - used for most tests
"llama-8b": {
"model": _resolve_model_path("meta-llama/Llama-3.1-8B-Instruct"),
"memory_gb": 16,
"tp": 1,
"features": ["chat", "streaming", "function_calling"],
},
# Small model for quick tests
"llama-1b": {
"model": _resolve_model_path("meta-llama/Llama-3.2-1B-Instruct"),
"memory_gb": 4,
"tp": 1,
"features": ["chat", "streaming", "tool_choice"],
},
# Function calling specialist
"qwen-7b": {
"model": _resolve_model_path("Qwen/Qwen2.5-7B-Instruct"),
"memory_gb": 14,
"tp": 1,
"features": ["chat", "streaming", "function_calling", "pythonic_tools"],
},
# Function calling specialist (larger, for Response API tests)
"qwen-14b": {
"model": _resolve_model_path("Qwen/Qwen2.5-14B-Instruct"),
"memory_gb": 28,
"tp": 2,
"features": ["chat", "streaming", "function_calling", "pythonic_tools"],
"worker_args": [
"--context-length=1000"
], # Faster startup, prevents memory issues
},
# Reasoning model
"deepseek-7b": {
"model": _resolve_model_path("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"),
"memory_gb": 14,
"tp": 1,
"features": ["chat", "streaming", "reasoning"],
},
# Thinking/reasoning model (larger)
"qwen-30b": {
"model": _resolve_model_path("Qwen/Qwen3-30B-A3B"),
"memory_gb": 60,
"tp": 4,
"features": ["chat", "streaming", "thinking", "reasoning"],
},
# Mistral for function calling
"mistral-7b": {
"model": _resolve_model_path("mistralai/Mistral-7B-Instruct-v0.3"),
"memory_gb": 14,
"tp": 1,
"features": ["chat", "streaming", "function_calling"],
},
# Embedding model
"embedding": {
"model": _resolve_model_path("intfloat/e5-mistral-7b-instruct"),
"memory_gb": 14,
"tp": 1,
"features": ["embedding"],
},
# GPT-OSS model (Harmony)
"gpt-oss": {
"model": _resolve_model_path("openai/gpt-oss-20b"),
"memory_gb": 40,
"tp": 2,
"features": ["chat", "streaming", "reasoning", "harmony"],
},
}
def get_models_with_feature(feature: str) -> list[str]:
"""Get list of model IDs that support a specific feature."""
return [
model_id
for model_id, spec in MODEL_SPECS.items()
if feature in spec.get("features", [])
]
def get_model_spec(model_id: str) -> dict:
"""Get spec for a specific model, raising KeyError if not found."""
if model_id not in MODEL_SPECS:
raise KeyError(
f"Unknown model: {model_id}. Available: {list(MODEL_SPECS.keys())}"
)
return MODEL_SPECS[model_id]
# Convenience groupings for test parametrization
CHAT_MODELS = get_models_with_feature("chat")
EMBEDDING_MODELS = get_models_with_feature("embedding")
REASONING_MODELS = get_models_with_feature("reasoning")
FUNCTION_CALLING_MODELS = get_models_with_feature("function_calling")
# =============================================================================
# Default model path constants (for backward compatibility with existing tests)
# =============================================================================
DEFAULT_MODEL_PATH = MODEL_SPECS["llama-8b"]["model"]
DEFAULT_SMALL_MODEL_PATH = MODEL_SPECS["llama-1b"]["model"]
DEFAULT_REASONING_MODEL_PATH = MODEL_SPECS["deepseek-7b"]["model"]
DEFAULT_ENABLE_THINKING_MODEL_PATH = MODEL_SPECS["qwen-30b"]["model"]
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH = MODEL_SPECS["qwen-7b"]["model"]
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH = MODEL_SPECS["mistral-7b"]["model"]
DEFAULT_GPT_OSS_MODEL_PATH = MODEL_SPECS["gpt-oss"]["model"]
DEFAULT_EMBEDDING_MODEL_PATH = MODEL_SPECS["embedding"]["model"]
# =============================================================================
# Third-party model configurations (cloud APIs)
# =============================================================================
THIRD_PARTY_MODELS: dict[str, dict] = {
"openai": {
"description": "OpenAI API",
"model": "gpt-5-nano",
"api_key_env": "OPENAI_API_KEY",
},
"xai": {
"description": "xAI API",
"model": "grok-4-fast",
"api_key_env": "XAI_API_KEY",
},
}

161
third_party/sglang/sgl-model-gateway/e2e_test/infra/process_utils.py vendored Normal file
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

138
third_party/sglang/sgl-model-gateway/e2e_test/infra/run_eval.py vendored Normal file
View File

@@ -0,0 +1,138 @@
"""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
from .simple_eval_mmlu import MMLU_DATASET_URL
logger = logging.getLogger(__name__)
@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:
base_url = base_url.rstrip("/") # Remove trailing slashes
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

492
third_party/sglang/sgl-model-gateway/e2e_test/infra/simple_eval_common.py vendored Normal file
View File

@@ -0,0 +1,492 @@
# 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
from .constants import MAX_RETRY_ATTEMPTS
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",
):
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 < MAX_RETRY_ATTEMPTS:
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
# 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,
MAX_RETRY_ATTEMPTS,
exception_backoff,
e,
)
time.sleep(exception_backoff)
trial += 1
logger.warning(
"All retry attempts exhausted after %d retries, returning empty response",
MAX_RETRY_ATTEMPTS,
)
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, timeout=30)
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)

132
third_party/sglang/sgl-model-gateway/e2e_test/infra/simple_eval_mmlu.py vendored Normal file
View File

@@ -0,0 +1,132 @@
# 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
# MMLU dataset URL (hosted by OpenAI)
MMLU_DATASET_URL = "https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv"
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):
if "://" in filename:
df = pandas.read_csv(filename, storage_options={"timeout": 30})
else:
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)