refactor(e2e): unify RouterInstance into Gateway class, split conftest.py into modular fixtures (#16671)

sgl-model-gateway/e2e_test/fixtures/pool.py

@@ -0,0 +1,194 @@
+"""Model pool fixtures for E2E tests.
+
+This module provides session-scoped fixtures for managing SGLang worker processes.
+Workers are expensive to start (~30-60s each), so they're kept running across tests.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from infra import ModelPool
+
+from .hooks import get_pool_requirements
+
+logger = logging.getLogger(__name__)
+
+# Global model pool instance
+_model_pool: "ModelPool | None" = None
+
+
+@pytest.fixture(scope="session")
+def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
+    """Session-scoped fixture that manages SGLang worker processes.
+
+    Workers (sglang.launch_server) are expensive to start (~30-60s each due to
+    model loading). This fixture starts them ONCE per session and keeps them
+    running across all tests. The setup_backend fixture then launches cheap
+    routers (~1-2s) pointing to these workers.
+
+    Startup behavior:
+    - Scans test markers to determine required workers (model, mode, type, count)
+    - Launches workers in test collection order
+    - Waits for all workers to become healthy before returning
+
+    Test requirements are auto-detected from:
+    - @pytest.mark.parametrize("setup_backend", ["grpc", "http", "pd"])
+    - @pytest.mark.model("model-name")
+    - @pytest.mark.workers(count=N) for regular workers
+    - @pytest.mark.workers(prefill=N, decode=N) for PD workers
+
+    Environment variable overrides:
+    - E2E_MODELS: Comma-separated model IDs (e.g., "llama-8b,qwen-7b")
+    - E2E_BACKENDS: Comma-separated backends (e.g., "grpc,http")
+    - SKIP_MODEL_POOL: Set to "1" to skip worker startup
+    """
+    global _model_pool
+
+    from infra import (
+        DEFAULT_MODEL,
+        ENV_BACKENDS,
+        ENV_MODELS,
+        ENV_SKIP_MODEL_POOL,
+        ENV_STARTUP_TIMEOUT,
+        LOCAL_MODES,
+        MODEL_SPECS,
+        ConnectionMode,
+        GPUAllocator,
+        ModelPool,
+        WorkerIdentity,
+        WorkerType,
+    )
+
+    if _model_pool is not None:
+        return _model_pool
+
+    # Check if we should skip model startup
+    if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
+        logger.info("%s is set, skipping model pool startup", ENV_SKIP_MODEL_POOL)
+        _model_pool = ModelPool(GPUAllocator(gpus=[]))
+        return _model_pool
+
+    # Determine requirements from scanned tests or env vars
+    models_env = os.environ.get(ENV_MODELS, "")
+    backends_env = os.environ.get(ENV_BACKENDS, "")
+
+    if models_env or backends_env:
+        # Use env var overrides
+        models = (
+            {m.strip() for m in models_env.split(",") if m.strip()}
+            if models_env
+            else {DEFAULT_MODEL}
+        )
+
+        # Parse backend strings to ConnectionMode enums
+        backend_modes: set[ConnectionMode] = set()
+        if backends_env:
+            for b in backends_env.split(","):
+                b = b.strip()
+                if b:
+                    try:
+                        mode = ConnectionMode(b)
+                        if mode in LOCAL_MODES:
+                            backend_modes.add(mode)
+                    except ValueError:
+                        logger.warning("Unknown backend '%s', skipping", b)
+
+        # Default to HTTP if no valid backends
+        if not backend_modes:
+            backend_modes = {ConnectionMode.HTTP}
+
+        # Create WorkerIdentity objects (regular workers only from env vars)
+        requirements = [
+            WorkerIdentity(m, b, WorkerType.REGULAR, 0)
+            for m in models
+            for b in backend_modes
+        ]
+        logger.info("Using env var requirements: %s", [str(r) for r in requirements])
+    else:
+        # Use scanned requirements from test markers
+        requirements = get_pool_requirements()
+        logger.info("Using scanned requirements: %s", [str(r) for r in requirements])
+
+    # Filter to valid models
+    requirements = [r for r in requirements if r.model_id in MODEL_SPECS]
+
+    if not requirements:
+        logger.warning("No valid requirements, model pool will be empty")
+        _model_pool = ModelPool(GPUAllocator(gpus=[]))
+        return _model_pool
+
+    # Create and start the pool
+    allocator = GPUAllocator()
+    _model_pool = ModelPool(allocator)
+
+    startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300"))
+    _model_pool.startup(
+        requirements=requirements,
+        startup_timeout=startup_timeout,
+    )
+
+    # Log final GPU allocation summary
+    logger.info(_model_pool.allocator.summary())
+
+    # 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(...)
+    """
+    from infra import PARAM_MODEL
+
+    marker = request.node.get_closest_marker(PARAM_MODEL)
+    if marker is None:
+        pytest.fail(
+            f"Test must be marked with @pytest.mark.{PARAM_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")
+    """
+    from infra import PARAM_MODEL
+
+    marker = request.node.get_closest_marker(PARAM_MODEL)
+    if marker is None:
+        pytest.fail(
+            f"Test must be marked with @pytest.mark.{PARAM_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")