[model-gateway] add GPU allocator and model pool infrastructure for parallel E2E tests (#16460)

sgl-model-gateway/e2e_test/conftest.py

@@ -1,6 +1,25 @@
+"""Pytest configuration for E2E tests."""
+
+from __future__ import annotations
+
+import logging
+import os
 import sys
 from importlib.util import find_spec
 from pathlib import Path
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from infra import ModelPool
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)8s] %(name)s: %(message)s",
+)
+logger = logging.getLogger(__name__)
 
 # Only add bindings/python to path if the wheel is not installed (for local development)
 # This ensures CI tests use the installed wheel which contains the Rust extension
@@ -13,3 +32,160 @@ _wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None
 # Only add bindings/python if wheel is not installed (development mode)
 if not _wheel_installed and str(_SRC) not in sys.path:
     sys.path.insert(0, str(_SRC))
+
+
+# ---------------------------------------------------------------------------
+# Custom pytest markers
+# ---------------------------------------------------------------------------
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    """Register custom markers."""
+    config.addinivalue_line(
+        "markers",
+        "model(name): mark test to use a specific model from the model pool",
+    )
+    config.addinivalue_line(
+        "markers",
+        "e2e: mark test as an end-to-end test requiring GPU workers",
+    )
+    config.addinivalue_line(
+        "markers",
+        "slow: mark test as slow-running",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Model pool fixtures (session-scoped)
+# ---------------------------------------------------------------------------
+
+# Global model pool instance
+_model_pool: "ModelPool | None" = None
+
+
+def _get_requested_models(config: pytest.Config) -> list[str]:
+    """Determine which models are needed based on collected tests.
+
+    This scans all test items for @pytest.mark.model() markers and returns
+    the unique set of models requested.
+    """
+    models = set()
+
+    # This is called during collection, so we need to iterate items
+    for item in config.pluginmanager.get_plugin("main").session.items:
+        marker = item.get_closest_marker("model")
+        if marker and marker.args:
+            models.add(marker.args[0])
+
+    return list(models)
+
+
+@pytest.fixture(scope="session")
+def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
+    """Session-scoped fixture providing the model pool.
+
+    The model pool pre-loads all models needed by tests in this session,
+    running them in parallel across available GPUs.
+
+    Usage:
+        @pytest.mark.model("llama-8b")
+        def test_chat(model_pool):
+            client = model_pool.get_client("llama-8b")
+            ...
+    """
+    global _model_pool
+
+    # Import here to avoid import errors when infra is not set up
+    from infra import MODEL_SPECS, GPUAllocator, ModelPool
+
+    if _model_pool is not None:
+        return _model_pool
+
+    # Check if we should skip model startup (e.g., for unit tests)
+    if os.environ.get("SKIP_MODEL_POOL", "").lower() in ("1", "true", "yes"):
+        logger.info("SKIP_MODEL_POOL is set, skipping model pool startup")
+        _model_pool = ModelPool(GPUAllocator(gpus=[]))
+        return _model_pool
+
+    # Determine which models to start
+    # For now, start models based on environment or a default set
+    models_env = os.environ.get("E2E_MODELS", "")
+    if models_env:
+        model_ids = [m.strip() for m in models_env.split(",") if m.strip()]
+    else:
+        # Default: start commonly needed models
+        model_ids = ["llama-8b", "qwen-7b"]
+
+    # Filter to available specs
+    model_ids = [m for m in model_ids if m in MODEL_SPECS]
+
+    if not model_ids:
+        logger.warning("No models specified, model pool will be empty")
+        _model_pool = ModelPool(GPUAllocator(gpus=[]))
+        return _model_pool
+
+    logger.info("Starting model pool with models: %s", model_ids)
+
+    # Create and start the pool
+    allocator = GPUAllocator()
+    _model_pool = ModelPool(allocator)
+
+    grpc_mode = os.environ.get("E2E_GRPC_MODE", "").lower() in ("1", "true", "yes")
+    startup_timeout = int(os.environ.get("E2E_STARTUP_TIMEOUT", "300"))
+
+    _model_pool.startup(
+        model_ids=model_ids,
+        grpc_mode=grpc_mode,
+        startup_timeout=startup_timeout,
+    )
+
+    # Register cleanup
+    request.addfinalizer(_model_pool.shutdown)
+
+    return _model_pool
+
+
+@pytest.fixture
+def model_client(request: pytest.FixtureRequest, model_pool: "ModelPool"):
+    """Get OpenAI client for the model specified by @pytest.mark.model().
+
+    Usage:
+        @pytest.mark.model("llama-8b")
+        def test_chat(model_client):
+            response = model_client.chat.completions.create(...)
+    """
+    marker = request.node.get_closest_marker("model")
+    if marker is None:
+        pytest.fail(
+            "Test must be marked with @pytest.mark.model('model-id') to use model_client fixture"
+        )
+
+    model_id = marker.args[0]
+
+    try:
+        return model_pool.get_client(model_id)
+    except KeyError:
+        pytest.skip(f"Model {model_id} not available in model pool")
+
+
+@pytest.fixture
+def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> str:
+    """Get the base URL for the model specified by @pytest.mark.model().
+
+    Usage:
+        @pytest.mark.model("llama-8b")
+        def test_direct_http(model_base_url):
+            response = httpx.get(f"{model_base_url}/health")
+    """
+    marker = request.node.get_closest_marker("model")
+    if marker is None:
+        pytest.fail(
+            "Test must be marked with @pytest.mark.model('model-id') to use model_base_url fixture"
+        )
+
+    model_id = marker.args[0]
+
+    try:
+        return model_pool.get_base_url(model_id)
+    except KeyError:
+        pytest.skip(f"Model {model_id} not available in model pool")

sgl-model-gateway/e2e_test/infra/__init__.py

@@ -0,0 +1,31 @@
+"""Infrastructure for parallel GPU test execution."""
+
+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 .model_pool import ModelInstance, ModelPool
+from .model_specs import MODEL_SPECS
+
+__all__ = [
+    # 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",
+    # Model management
+    "ModelInstance",
+    "ModelPool",
+    "MODEL_SPECS",
+]

sgl-model-gateway/e2e_test/infra/gpu_allocator.py

@@ -0,0 +1,361 @@
+"""GPU detection and slot allocation for parallel test execution."""
+
+from __future__ import annotations
+
+import logging
+import os
+import socket
+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] = []
+
+    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")
+                    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]) -> list[GPUSlot]:
+        """Allocate GPU slots based on model memory requirements.
+
+        Uses a first-fit decreasing bin-packing algorithm:
+        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
+
+        Args:
+            model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys
+
+        Returns:
+            List of GPUSlots with assigned models
+        """
+        if not self.gpus:
+            logger.warning("No GPUs available for allocation")
+            return []
+
+        # Sort models by memory requirement (largest first for better packing)
+        sorted_models = sorted(
+            model_specs.items(),
+            key=lambda x: x[1].get("memory_gb", 0),
+            reverse=True,
+        )
+
+        # Track which GPUs are used
+        used_gpus: set[int] = set()
+        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]
+
+            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(),
+                        )
+                        slots.append(slot)
+                        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", memory_gb, model_id
+                    )
+            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(),
+                        )
+                        slots.append(slot)
+                        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",
+                        tp_size,
+                        memory_gb,
+                        model_id,
+                    )
+
+        self.slots = slots
+        return slots
+
+    def get_slot_for_model(self, model_id: str) -> GPUSlot | None:
+        """Get the slot assigned to a specific model."""
+        for slot in self.slots:
+            if slot.assigned_model == model_id:
+                return slot
+        return None
+
+    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"  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)

sgl-model-gateway/e2e_test/infra/model_pool.py

@@ -0,0 +1,280 @@
+"""Model pool for managing pre-loaded models across GPUs."""
+
+from __future__ import annotations
+
+import logging
+import os
+import signal
+import subprocess
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import httpx
+
+if TYPE_CHECKING:
+    import openai
+
+from .gpu_allocator import GPUAllocator, GPUSlot
+from .model_specs import MODEL_SPECS, get_model_spec
+
+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
+    model_path: str
+    base_url: str
+    process: subprocess.Popen
+    gpu_slot: GPUSlot
+    grpc_mode: bool = False
+
+    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."""
+        try:
+            resp = httpx.get(f"{self.base_url}/health", timeout=timeout)
+            return resp.status_code == 200
+        except (httpx.RequestError, httpx.TimeoutException):
+            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)
+
+        # 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)
+            self.process.kill()
+            self.process.wait()
+
+
+class ModelPool:
+    """Manages a pool of pre-loaded models across GPUs."""
+
+    def __init__(self, allocator: GPUAllocator | None = None):
+        """Initialize the model pool.
+
+        Args:
+            allocator: GPU allocator to use. If None, creates a new one.
+        """
+        self.allocator = allocator or GPUAllocator()
+        self.instances: dict[str, ModelInstance] = {}
+        self._startup_timeout = DEFAULT_STARTUP_TIMEOUT
+
+    def startup(
+        self,
+        model_ids: list[str] | None = None,
+        grpc_mode: bool = False,
+        startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
+    ) -> None:
+        """Spin up models in parallel on assigned GPU slots.
+
+        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.
+        """
+        self._startup_timeout = startup_timeout
+
+        # Determine which models to start
+        if model_ids is None:
+            model_ids = list(MODEL_SPECS.keys())
+
+        # Filter to models we have specs for
+        specs_to_start = {
+            mid: MODEL_SPECS[mid] for mid in model_ids if mid in MODEL_SPECS
+        }
+
+        if not specs_to_start:
+            logger.warning("No valid model specs to start")
+            return
+
+        # Allocate GPU slots
+        slots = self.allocator.allocate_slots(specs_to_start)
+
+        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)
+
+        # 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
+
+        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
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = slot.cuda_visible_devices()
+
+        # Build command
+        cmd = [
+            "python3",
+            "-m",
+            "sglang.launch_server",
+            "--model-path",
+            model_path,
+            "--port",
+            str(port),
+            "--tp-size",
+            str(tp_size),
+            "--log-level",
+            "warning",
+        ]
+
+        if grpc_mode:
+            cmd.append("--grpc-mode")
+
+        logger.info(
+            "Launching %s on GPUs %s port %d: %s",
+            model_id,
+            slot.gpu_ids,
+            port,
+            " ".join(cmd),
+        )
+
+        # Start the process
+        proc = subprocess.Popen(
+            cmd,
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            # Use process group for clean shutdown
+            start_new_session=True,
+        )
+
+        base_url = f"http://{DEFAULT_HOST}:{port}"
+        instance = ModelInstance(
+            model_id=model_id,
+            model_path=model_path,
+            base_url=base_url,
+            process=proc,
+            gpu_slot=slot,
+            grpc_mode=grpc_mode,
+        )
+        self.instances[model_id] = instance
+
+    def _wait_all_healthy(self) -> None:
+        """Wait for all model instances to become healthy."""
+        start_time = time.time()
+        pending = set(self.instances.keys())
+
+        while pending and (time.time() - start_time) < self._startup_timeout:
+            for model_id in list(pending):
+                instance = self.instances[model_id]
+
+                # Check if process died
+                if not instance.is_alive():
+                    logger.error(
+                        "Model %s (PID %d) died during startup",
+                        model_id,
+                        instance.process.pid,
+                    )
+                    # Read stderr for debugging
+                    if instance.process.stderr:
+                        stderr = instance.process.stderr.read()
+                        if stderr:
+                            logger.error("Stderr: %s", stderr.decode()[-2000:])
+                    pending.discard(model_id)
+                    continue
+
+                # Check health
+                if instance.health_check():
+                    logger.info(
+                        "Model %s is healthy at %s", model_id, instance.base_url
+                    )
+                    pending.discard(model_id)
+
+            if pending:
+                time.sleep(HEALTH_CHECK_INTERVAL)
+
+        if pending:
+            logger.error(
+                "Models failed to start within %ds: %s",
+                self._startup_timeout,
+                pending,
+            )
+            # Terminate failed instances
+            for model_id in pending:
+                self.instances[model_id].terminate()
+                del self.instances[model_id]
+
+    def get_client(self, model_id: str) -> "openai.OpenAI":
+        """Get OpenAI client for a specific model.
+
+        Args:
+            model_id: The model ID to get a client for.
+
+        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]
+        return openai.OpenAI(
+            base_url=f"{instance.base_url}/v1",
+            api_key="not-used",
+        )
+
+    def get_base_url(self, model_id: str) -> 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
+
+    def shutdown(self) -> None:
+        """Tear down all models."""
+        logger.info("Shutting down model pool (%d instances)", len(self.instances))
+        for instance in self.instances.values():
+            instance.terminate()
+        self.instances.clear()
+
+    def __enter__(self) -> "ModelPool":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.shutdown()

sgl-model-gateway/e2e_test/infra/model_specs.py

@@ -0,0 +1,109 @@
+"""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"],
+    },
+    # 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": 2,
+        "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")