"""Pytest configuration for E2E tests. Markers ------- This module defines several pytest markers for configuring E2E tests: @pytest.mark.model(name) Specify which model to use for the test. Args: name: Model ID from MODEL_SPECS (e.g., "llama-8b", "qwen-7b") GPU Resource Management: When GPUs are limited (e.g., 4 GPUs, 6 models), the model pool uses MRU (Most Recently Used) eviction: 1. Models are pre-launched until GPUs are full 2. When a test needs a model that isn't running, MRU model is evicted (models just used are likely done, models not yet used are waiting) 3. The needed model is then launched on-demand Examples: @pytest.mark.model("llama-8b") @pytest.mark.model("qwen-72b") @pytest.mark.workers(count=1, prefill=None, decode=None) Configure worker topology for the test. Args: count: Number of regular workers (default: 1) prefill: Number of prefill workers for PD disaggregation decode: Number of decode workers for PD disaggregation Examples: @pytest.mark.workers(count=3) # 3 regular workers @pytest.mark.workers(prefill=2, decode=2) # PD mode @pytest.mark.gateway(policy="round_robin", timeout=None, extra_args=None) Configure the gateway/router. Args: policy: Routing policy ("round_robin", "random", etc.) timeout: Startup timeout in seconds extra_args: Additional CLI arguments for the router Examples: @pytest.mark.gateway(policy="random") @pytest.mark.gateway(extra_args=["--cache-routing"]) @pytest.mark.e2e Mark test as an end-to-end test requiring GPU workers. @pytest.mark.slow Mark test as slow-running. Fixtures -------- model_pool: Session-scoped fixture managing SGLang worker processes. setup_backend: Class-scoped fixture that launches gateway + provides client. Usage Examples -------------- Basic test with default model: @pytest.mark.e2e @pytest.mark.parametrize("setup_backend", ["http"], indirect=True) class TestBasic: def test_chat(self, setup_backend): backend, model, client, gateway = setup_backend response = client.chat.completions.create(...) Test with specific model and multiple backends: @pytest.mark.e2e @pytest.mark.model("qwen-7b") @pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True) class TestQwen: def test_generate(self, setup_backend): ... PD disaggregation mode: @pytest.mark.e2e @pytest.mark.workers(prefill=1, decode=1) @pytest.mark.parametrize("setup_backend", ["pd"], indirect=True) class TestPD: def test_pd_inference(self, setup_backend): ... """ from __future__ import annotations import logging import os import subprocess 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 # --------------------------------------------------------------------------- # Logging setup (clean output without pytest's "---- live log ----" dividers) # --------------------------------------------------------------------------- def _setup_logging() -> None: """Configure clean logging to stdout with timestamps.""" # Custom format: timestamp [logger] message fmt = "%(asctime)s.%(msecs)03d [%(name)s] %(message)s" datefmt = "%H:%M:%S" # Create handler for stdout handler = logging.StreamHandler(sys.stdout) handler.setFormatter(logging.Formatter(fmt, datefmt)) # Configure our e2e_test and infra modules for INFO level for logger_name in ("e2e_test", "infra"): log = logging.getLogger(logger_name) log.setLevel(logging.INFO) log.addHandler(handler) log.propagate = False # Don't double-log # Suppress noisy third-party loggers for logger_name in ("openai", "httpx", "httpcore", "numexpr"): logging.getLogger(logger_name).setLevel(logging.WARNING) _setup_logging() logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Test visibility hooks # --------------------------------------------------------------------------- def pytest_runtest_logstart(nodeid: str, location: tuple) -> None: """Print clear test header at start of each test.""" # Extract test name from nodeid (e.g., "test_mmlu.py::TestMMLU::test_mmlu_basic[grpc]") test_name = nodeid.split("::")[-1] if "::" in nodeid else nodeid print(f"\n{'='*60}") print(f"TEST: {test_name}") print(f"{'='*60}") # Path setup for imports _ROOT = Path(__file__).resolve().parents[1] # sgl-model-gateway/ _E2E_TEST = Path(__file__).resolve().parent # e2e_test/ _SRC = _ROOT / "bindings" / "python" # Add e2e_test to path so "from infra import ..." works if str(_E2E_TEST) not in sys.path: sys.path.insert(0, str(_E2E_TEST)) # Add bindings/python to path if the wheel is not installed (for local development) _wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None if not _wheel_installed and str(_SRC) not in sys.path: sys.path.insert(0, str(_SRC)) # Import constants after path setup from infra import ( DEFAULT_MODEL, DEFAULT_ROUTER_TIMEOUT, ENV_BACKENDS, ENV_MODEL, ENV_MODELS, ENV_SHOW_ROUTER_LOGS, ENV_SKIP_BACKEND_SETUP, ENV_SKIP_MODEL_POOL, ENV_STARTUP_TIMEOUT, LOCAL_MODES, PARAM_MODEL, PARAM_SETUP_BACKEND, ConnectionMode, ) # --------------------------------------------------------------------------- # Test collection: scan for required backends # --------------------------------------------------------------------------- # Global storage for scanned requirements _scanned_backends: set[str] = set() # {"grpc", "http", "openai", ...} _scanned_models: set[str] = set() # Models needed by tests _needs_default_model: bool = False # True if any e2e test lacks explicit model marker def pytest_collection_modifyitems( session: pytest.Session, config: pytest.Config, items: list[pytest.Item], ) -> None: """Scan collected tests to determine required backends and models. This runs after test collection but before tests execute. It extracts backend requirements from @pytest.mark.parametrize markers. """ global _scanned_backends, _scanned_models, _needs_default_model for item in items: # Track if this test has an explicit model marker has_model_marker = False # Scan parametrize markers for setup_backend for marker in item.iter_markers("parametrize"): if marker.args and len(marker.args) >= 2: param_name = marker.args[0] param_values = marker.args[1] if param_name == PARAM_SETUP_BACKEND: # Extract backend names from parametrize values if isinstance(param_values, (list, tuple)): _scanned_backends.update(param_values) elif param_name == PARAM_MODEL or PARAM_MODEL in param_name: # Extract model names from parametrize if isinstance(param_values, (list, tuple)): _scanned_models.update(param_values) has_model_marker = True # Check for @pytest.mark.model("name") markers model_marker = item.get_closest_marker(PARAM_MODEL) if model_marker and model_marker.args: model_name = model_marker.args[0] _scanned_models.add(model_name) has_model_marker = True # Check if this is an e2e test without an explicit model marker # Such tests need the DEFAULT_MODEL if not has_model_marker and item.get_closest_marker("e2e"): _needs_default_model = True logger.info( "Scanned test requirements - backends: %s, models: %s, needs default: %s", _scanned_backends or {"(none)"}, _scanned_models or {"(none)"}, _needs_default_model, ) def get_pool_requirements() -> list[tuple[str, ConnectionMode]]: """Build pool requirements from scanned test markers. Returns: List of (model_id, ConnectionMode) tuples to try to pre-launch. Models that don't fit will be launched on-demand. """ models = set(_scanned_models) # Add DEFAULT_MODEL if any e2e test lacks an explicit model marker, # or if no models were specified at all if _needs_default_model or not models: models.add(DEFAULT_MODEL) # Convert scanned string backends to ConnectionMode enums # Filter to local backends only (grpc, http) - cloud backends don't need workers local_modes: set[ConnectionMode] = set() for backend in _scanned_backends: try: mode = ConnectionMode(backend) if mode in LOCAL_MODES: local_modes.add(mode) except ValueError: # Not a ConnectionMode (e.g., "openai", "xai", "pd") - skip pass # Default to HTTP if no local backends specified if not local_modes: local_modes = {ConnectionMode.HTTP} # Build requirements: each model needs each mode requirements: list[tuple[str, ConnectionMode]] = [] for model in models: for mode in local_modes: requirements.append((model, mode)) return requirements # --------------------------------------------------------------------------- # 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 MODEL_SPECS", ) config.addinivalue_line( "markers", "backend(name): mark test to use a specific backend (grpc, http, openai, etc.)", ) config.addinivalue_line( "markers", "workers(count=1, prefill=None, decode=None): " "worker configuration - use count for regular workers, " "or prefill/decode for PD disaggregation mode", ) config.addinivalue_line( "markers", "gateway(policy='round_robin', timeout=None, extra_args=None): " "gateway/router configuration", ) 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 @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 (model, mode) combinations - Launches workers sequentially, but they boot up concurrently - Waits for all workers to become healthy before returning Test requirements are auto-detected from: - @pytest.mark.parametrize("setup_backend", ["grpc", "http"]) - @pytest.mark.model("model-name") 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 MODEL_SPECS, GPUAllocator, ModelPool 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} requirements = [(m, b) for m in models for b in backend_modes] logger.info("Using env var requirements: %s", requirements) else: # Use scanned requirements from test markers requirements = get_pool_requirements() logger.info("Using scanned requirements: %s", requirements) # Filter to valid models requirements = [(m, b) for m, b in requirements if m 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(...) """ 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") """ 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") # --------------------------------------------------------------------------- # Backend fixtures # --------------------------------------------------------------------------- def _get_marker_value( request: pytest.FixtureRequest, marker_name: str, arg_index: int = 0, default: any = None, ) -> any: """Get a value from a pytest marker. Args: request: The pytest fixture request. marker_name: Name of the marker to look for. arg_index: Index of positional argument to extract. default: Default value if marker not found. Returns: The marker argument value or default. """ marker = request.node.get_closest_marker(marker_name) if marker is None: return default if marker.args and len(marker.args) > arg_index: return marker.args[arg_index] return default def _get_marker_kwargs( request: pytest.FixtureRequest, marker_name: str, defaults: dict[str, any] | None = None, ) -> dict[str, any]: """Get keyword arguments from a pytest marker. Args: request: The pytest fixture request. marker_name: Name of the marker to look for. defaults: Default values if marker not found or missing kwargs. Returns: Dict of keyword arguments merged with defaults. """ result = dict(defaults) if defaults else {} marker = request.node.get_closest_marker(marker_name) if marker is not None: result.update(marker.kwargs) return result @pytest.fixture(scope="class") def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): """Class-scoped fixture that launches a router for each test class. Routers are cheap to start (~1-2s) compared to workers (~30-60s), so we launch a fresh router per test class for isolation while reusing the expensive workers from model_pool. Backend types: - "http", "grpc": Gets existing worker from model_pool, launches router - "pd": Launches prefill/decode workers via model_pool, launches PD router - "openai", "xai", etc.: Launches cloud router (no local workers) Configuration via markers: - @pytest.mark.model("model-id"): Override default model - @pytest.mark.workers(count=1): Number of regular workers behind router - @pytest.mark.workers(prefill=1, decode=1): PD worker configuration - @pytest.mark.gateway(policy="round_robin", timeout=60): Gateway configuration Returns: Tuple of (backend_name, model_path, openai_client, gateway) Usage: # Simple - uses defaults @pytest.mark.parametrize("setup_backend", ["http"], indirect=True) class TestBasic: ... # With model override @pytest.mark.model("qwen-7b") @pytest.mark.parametrize("setup_backend", ["http"], indirect=True) class TestWithModel: ... # Load balancing with multiple workers @pytest.mark.workers(count=3) @pytest.mark.gateway(policy="round_robin") @pytest.mark.parametrize("setup_backend", ["http"], indirect=True) class TestLoadBalancing: ... # PD with custom configuration @pytest.mark.workers(prefill=2, decode=2) @pytest.mark.gateway(policy="round_robin") @pytest.mark.parametrize("setup_backend", ["pd"], indirect=True) class TestPDScaling: ... """ import openai from infra import DEFAULT_ROUTER_TIMEOUT, Gateway, WorkerType backend_name = request.param # Skip if requested if os.environ.get(ENV_SKIP_BACKEND_SETUP, "").lower() in ("1", "true", "yes"): pytest.skip(f"{ENV_SKIP_BACKEND_SETUP} is set") # Get model from marker or env var or default model_id = _get_marker_value(request, "model") if model_id is None: model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL) # Get worker configuration from marker workers_config = _get_marker_kwargs( request, "workers", defaults={"count": 1, "prefill": None, "decode": None} ) # Get gateway configuration from marker gateway_config = _get_marker_kwargs( request, "gateway", defaults={ "policy": "round_robin", "timeout": DEFAULT_ROUTER_TIMEOUT, "extra_args": None, }, ) # PD disaggregation backend if backend_name == "pd": # Check PD requirements try: import sgl_kernel # noqa: F401 except ImportError: pytest.skip("sgl_kernel not available, required for PD disaggregation") try: import torch except ImportError: pytest.skip("torch not available") if not torch.cuda.is_available(): pytest.skip("CUDA not available") # Get PD configuration from workers marker num_prefill = workers_config.get("prefill") or 1 num_decode = workers_config.get("decode") or 1 # Check GPU requirements required_gpus = num_prefill + num_decode gpu_count = torch.cuda.device_count() if gpu_count < required_gpus: pytest.skip( f"PD tests require {required_gpus} GPUs " f"({num_prefill} prefill + {num_decode} decode), found {gpu_count}" ) # Try to use pre-launched PD workers, or launch new ones if needed existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL) existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE) if ( len(existing_prefills) >= num_prefill and len(existing_decodes) >= num_decode ): # Use pre-launched workers prefills = existing_prefills[:num_prefill] decodes = existing_decodes[:num_decode] logger.info( "Using pre-launched PD workers: %d prefill, %d decode", len(prefills), len(decodes), ) else: # Launch new PD workers (custom config or not pre-launched) prefills, decodes = model_pool.launch_pd_workers( model_id=model_id, num_prefill=num_prefill, num_decode=num_decode, startup_timeout=300, ) model_path = prefills[0].model_path if prefills else None # Launch PD gateway with configuration gateway = Gateway() gateway.start( prefill_workers=prefills, decode_workers=decodes, policy=gateway_config["policy"], timeout=gateway_config["timeout"], extra_args=gateway_config["extra_args"], ) client = openai.OpenAI( base_url=f"{gateway.base_url}/v1", api_key="not-used", ) logger.info( "Setup PD backend: model=%s, %d prefill + %d decode workers, " "gateway=%s, policy=%s", model_id, len(prefills), len(decodes), gateway.base_url, gateway_config["policy"], ) try: yield backend_name, model_path, client, gateway finally: logger.info("Tearing down PD gateway") gateway.shutdown() return # Check if this is a local backend (grpc, http) try: connection_mode = ConnectionMode(backend_name) is_local = connection_mode in LOCAL_MODES except ValueError: is_local = False connection_mode = None # Local backends: use worker from pool + launch gateway if is_local: # Get number of workers from marker num_workers = workers_config.get("count") or 1 try: instance = model_pool.get(model_id, connection_mode) except RuntimeError as e: pytest.fail(str(e)) # Build worker URLs list # For num_workers > 1, we need multiple workers from the pool # For now, we reuse the same worker URL (router will load balance) # TODO: Support launching multiple distinct workers for true LB testing worker_urls = [instance.worker_url] * num_workers # Launch gateway with configuration gateway = Gateway() gateway.start( worker_urls=worker_urls, model_path=instance.model_path, policy=gateway_config["policy"], timeout=gateway_config["timeout"], extra_args=gateway_config["extra_args"], ) client = openai.OpenAI( base_url=f"{gateway.base_url}/v1", api_key="not-used", ) logger.info( "Setup %s backend: model=%s, workers=%d, gateway=%s, policy=%s", backend_name, model_id, num_workers, gateway.base_url, gateway_config["policy"], ) try: yield backend_name, instance.model_path, client, gateway finally: logger.info("Tearing down gateway for %s backend", backend_name) gateway.shutdown() return # Cloud backends: launch cloud router from backends import CLOUD_BACKENDS, launch_cloud_backend if backend_name not in CLOUD_BACKENDS: pytest.fail(f"Unknown backend: {backend_name}") cfg = CLOUD_BACKENDS[backend_name] api_key_env = cfg.get("api_key_env") if api_key_env and not os.environ.get(api_key_env): pytest.skip(f"{api_key_env} not set, skipping {backend_name} tests") logger.info("Launching cloud backend: %s", backend_name) router = launch_cloud_backend(backend_name) api_key = os.environ.get(api_key_env) if api_key_env else "not-used" client = openai.OpenAI( base_url=f"{router.base_url}/v1", api_key=api_key, ) try: yield backend_name, cfg["model"], client finally: logger.info("Tearing down cloud backend: %s", backend_name) router.shutdown() @pytest.fixture def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"): """Function-scoped fixture for launching a fresh router per test. This launches a new Gateway for each test, pointing to workers from the pool. Use for tests that need isolated router state. Usage: @pytest.mark.parametrize("backend_router", ["grpc", "http"], indirect=True) def test_router_state(backend_router): gateway = backend_router # Test gateway-specific behavior """ from infra import Gateway backend_name = request.param model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL) # Convert string to ConnectionMode connection_mode = ConnectionMode(backend_name) try: instance = model_pool.get(model_id, connection_mode) except KeyError: pytest.skip(f"Model {model_id}:{backend_name} not available in pool") except RuntimeError as e: pytest.fail(str(e)) gateway = Gateway() gateway.start( worker_urls=[instance.worker_url], model_path=instance.model_path, ) try: yield gateway finally: gateway.shutdown()