diff --git a/sgl-model-gateway/e2e_test/conftest.py b/sgl-model-gateway/e2e_test/conftest.py index 5825125c7..11de3ec17 100644 --- a/sgl-model-gateway/e2e_test/conftest.py +++ b/sgl-model-gateway/e2e_test/conftest.py @@ -1,4 +1,103 @@ -"""Pytest configuration for E2E tests.""" +"""Pytest configuration for E2E tests. + +Markers +------- +This module defines several pytest markers for configuring E2E tests: + +@pytest.mark.model(name, scope="session") + Specify which model to use for the test. + + Args: + name: Model ID from MODEL_SPECS (e.g., "llama-8b", "qwen-7b") + scope: "session" (default) or "class" + - session: Pre-launched at test session start. Stays running. + - class: Launched on-demand when test class starts. + + GPU Resource Management: + When GPUs are limited (e.g., 4 GPUs, 6 models), the model pool uses + LRU (Least Recently Used) eviction: + 1. Session models are pre-launched until GPUs are full + 2. Overflow models are queued for on-demand launch + 3. When a queued model is needed, LRU model is evicted + 4. Evicted models go back to queue and can be re-launched later + + Examples: + @pytest.mark.model("llama-8b") # session scope, pre-launched + @pytest.mark.model("qwen-72b", scope="class") # on-demand only + +@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): + ... + +Large model loaded on-demand (class scope): + + @pytest.mark.e2e + @pytest.mark.model("llama-70b", scope="class") + @pytest.mark.parametrize("setup_backend", ["http"], indirect=True) + class TestLargeModel: + def test_inference(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 @@ -96,7 +195,8 @@ from infra import ( # Global storage for scanned requirements _scanned_backends: set[str] = set() # {"grpc", "http", "openai", ...} -_scanned_models: set[str] = set() # {"llama-8b", "qwen-7b", ...} +_session_models: set[str] = set() # Models to pre-launch at session start +_class_models: set[str] = set() # Models to launch on-demand per class def pytest_collection_modifyitems( @@ -108,8 +208,12 @@ def pytest_collection_modifyitems( This runs after test collection but before tests execute. It extracts backend requirements from @pytest.mark.parametrize markers. + + Models are categorized by scope: + - session: Pre-launched at session start (default) + - class: Launched on-demand when test class starts """ - global _scanned_backends, _scanned_models + global _scanned_backends, _session_models, _class_models for item in items: # Scan parametrize markers for setup_backend @@ -124,30 +228,45 @@ def pytest_collection_modifyitems( _scanned_backends.update(param_values) elif param_name == PARAM_MODEL or PARAM_MODEL in param_name: - # Extract model names + # Extract model names from parametrize - default to session scope if isinstance(param_values, (list, tuple)): - _scanned_models.update(param_values) + _session_models.update(param_values) - # Also check for @pytest.mark.model("name") markers + # Check for @pytest.mark.model("name", scope="...") markers model_marker = item.get_closest_marker(PARAM_MODEL) if model_marker and model_marker.args: - _scanned_models.add(model_marker.args[0]) + model_name = model_marker.args[0] + scope = model_marker.kwargs.get("scope", "session") + + if scope == "class": + _class_models.add(model_name) + else: + _session_models.add(model_name) + + # Remove class models from session models (class scope takes precedence if mixed) + # Actually, keep both - a model can be used by both session and class scoped tests + # The model_pool will handle this by keeping session models running logger.info( - "Scanned test requirements - backends: %s, models: %s", + "Scanned test requirements - backends: %s, session models: %s, class models: %s", _scanned_backends or {"(none)"}, - _scanned_models or {"(none)"}, + _session_models or {"(none)"}, + _class_models or {"(none)"}, ) def get_pool_requirements() -> list[tuple[str, ConnectionMode]]: """Build pool requirements from scanned test markers. + Only returns session-scoped models for pre-launching. + Class-scoped models are launched on-demand by model_pool.get(). + Returns: - List of (model_id, ConnectionMode) tuples needed by tests. + List of (model_id, ConnectionMode) tuples to pre-launch. """ + # Only pre-launch session-scoped models # Default model if none specified - models = _scanned_models or {DEFAULT_MODEL} + models = _session_models or {DEFAULT_MODEL} # Convert scanned string backends to ConnectionMode enums # Filter to local backends only (grpc, http) - cloud backends don't need workers @@ -174,6 +293,15 @@ def get_pool_requirements() -> list[tuple[str, ConnectionMode]]: return requirements +def get_class_scoped_models() -> set[str]: + """Get models that are class-scoped (launched on-demand). + + Returns: + Set of model IDs that should be launched on-demand. + """ + return _class_models.copy() + + # --------------------------------------------------------------------------- # Custom pytest markers # --------------------------------------------------------------------------- @@ -183,7 +311,8 @@ 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", + "model(name, scope='session'): mark test to use a specific model " + "(scope: 'session' for pre-launched, 'class' for on-demand)", ) config.addinivalue_line( "markers", @@ -191,11 +320,14 @@ def pytest_configure(config: pytest.Config) -> None: ) config.addinivalue_line( "markers", - "workers(n): number of workers to launch behind the router (default: 1)", + "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", - "pd(num_prefill=1, num_decode=1): PD disaggregation worker configuration", + "gateway(policy='round_robin', timeout=None, extra_args=None): " + "gateway/router configuration", ) config.addinivalue_line( "markers", @@ -299,6 +431,11 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool": allocator = GPUAllocator() _model_pool = ModelPool(allocator) + # Register class-scoped models for on-demand launching + class_models = get_class_scoped_models() + if class_models: + _model_pool.register_class_scoped_models(class_models) + startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300")) _model_pool.startup(requirements=requirements, startup_timeout=startup_timeout) @@ -306,7 +443,7 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool": if "pd" in _scanned_backends: logger.info("PD backend detected, pre-launching PD workers") # Use default model for PD workers - pd_model = next(iter(_scanned_models), DEFAULT_MODEL) + pd_model = next(iter(_session_models), DEFAULT_MODEL) if pd_model in MODEL_SPECS: try: _model_pool.launch_pd_workers( @@ -373,173 +510,6 @@ def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> s pytest.skip(f"Model {model_id} not available in model pool") -# --------------------------------------------------------------------------- -# Router launching helpers -# --------------------------------------------------------------------------- - - -def launch_local_router( - worker_urls: list[str], - model_path: str, - *, - policy: str = "round_robin", - router_args: list[str] | None = None, - timeout: float = DEFAULT_ROUTER_TIMEOUT, - show_output: bool | None = None, -) -> tuple[str, subprocess.Popen]: - """Launch a router pointing to pre-started workers. - - Args: - worker_urls: List of worker URLs (e.g., ["http://127.0.0.1:30000"]) - model_path: Model path for the router - policy: Routing policy - router_args: Additional router arguments - timeout: Startup timeout in seconds - show_output: Show subprocess output - - Returns: - Tuple of (base_url, router_process) - """ - from infra import get_open_port, wait_for_workers_ready - - if show_output is None: - show_output = os.environ.get(ENV_SHOW_ROUTER_LOGS, "0") == "1" - - router_port = get_open_port() - prometheus_port = get_open_port() - base_url = f"http://127.0.0.1:{router_port}" - - cmd = [ - "python3", - "-m", - "sglang_router.launch_router", - "--host", - "127.0.0.1", - "--port", - str(router_port), - "--prometheus-port", - str(prometheus_port), - "--policy", - policy, - "--model-path", - model_path, - "--log-level", - "warn", - "--worker-urls", - *worker_urls, - ] - - if router_args: - cmd.extend(router_args) - - logger.info("Starting router on port %d with workers: %s", router_port, worker_urls) - - router_proc = subprocess.Popen( - cmd, - stdout=None if show_output else subprocess.PIPE, - stderr=None if show_output else subprocess.PIPE, - start_new_session=True, - ) - - try: - wait_for_workers_ready(base_url, len(worker_urls), timeout=timeout) - except TimeoutError: - from infra import kill_process_tree - - kill_process_tree(router_proc.pid) - raise - - logger.info("Router ready at %s", base_url) - return base_url, router_proc - - -def launch_pd_router( - prefills: list, - decodes: list, - *, - policy: str = "round_robin", - router_args: list[str] | None = None, - timeout: float = DEFAULT_ROUTER_TIMEOUT, - show_output: bool | None = None, -) -> tuple[str, subprocess.Popen]: - """Launch a PD disaggregation router. - - Args: - prefills: List of prefill ModelInstance objects. - decodes: List of decode ModelInstance objects. - policy: Routing policy. - router_args: Additional router arguments. - timeout: Startup timeout in seconds. - show_output: Show subprocess output. - - Returns: - Tuple of (base_url, router_process) - """ - from infra import get_open_port, wait_for_health - - if show_output is None: - show_output = os.environ.get(ENV_SHOW_ROUTER_LOGS, "0") == "1" - - router_port = get_open_port() - prometheus_port = get_open_port() - base_url = f"http://127.0.0.1:{router_port}" - - cmd = [ - "python3", - "-m", - "sglang_router.launch_router", - "--host", - "127.0.0.1", - "--port", - str(router_port), - "--prometheus-port", - str(prometheus_port), - "--prometheus-host", - "127.0.0.1", - "--policy", - policy, - "--pd-disaggregation", - "--log-level", - "warn", - ] - - # Add prefill workers with bootstrap ports - for pf in prefills: - cmd += ["--prefill", pf.base_url, str(pf.bootstrap_port)] - - # Add decode workers - for dc in decodes: - cmd += ["--decode", dc.base_url] - - if router_args: - cmd.extend(router_args) - - logger.info( - "Starting PD router on port %d with %d prefill, %d decode workers", - router_port, - len(prefills), - len(decodes), - ) - - router_proc = subprocess.Popen( - cmd, - stdout=None if show_output else subprocess.PIPE, - stderr=None if show_output else subprocess.PIPE, - start_new_session=True, - ) - - try: - wait_for_health(base_url, timeout=timeout) - except TimeoutError: - from infra import kill_process_tree - - kill_process_tree(router_proc.pid) - raise - - logger.info("PD Router ready at %s", base_url) - return base_url, router_proc - - # --------------------------------------------------------------------------- # Backend fixtures # --------------------------------------------------------------------------- @@ -607,11 +577,12 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): Configuration via markers: - @pytest.mark.model("model-id"): Override default model - - @pytest.mark.workers(n): Number of workers behind router (default: 1) - - @pytest.mark.pd(num_prefill=1, num_decode=1): PD worker configuration + - @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) + Tuple of (backend_name, model_path, openai_client, gateway) Usage: # Simple - uses defaults @@ -626,19 +597,21 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): ... # Load balancing with multiple workers - @pytest.mark.workers(3) + @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.pd(num_prefill=2, num_decode=2) + @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 kill_process_tree + from infra import DEFAULT_ROUTER_TIMEOUT, Gateway, WorkerType backend_name = request.param @@ -651,6 +624,28 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): if model_id is None: model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL) + # Get model scope from marker (session or class) + model_marker = request.node.get_closest_marker("model") + model_scope = "session" + if model_marker: + model_scope = model_marker.kwargs.get("scope", "session") + + # 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 @@ -667,12 +662,9 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): if not torch.cuda.is_available(): pytest.skip("CUDA not available") - # Get PD configuration from marker - pd_config = _get_marker_kwargs( - request, "pd", defaults={"num_prefill": 1, "num_decode": 1} - ) - num_prefill = pd_config["num_prefill"] - num_decode = pd_config["num_decode"] + # 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 @@ -684,8 +676,6 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): ) # Try to use pre-launched PD workers, or launch new ones if needed - from infra import WorkerType - existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL) existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE) @@ -712,27 +702,36 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): model_path = prefills[0].model_path if prefills else None - # Launch PD router - base_url, router_proc = launch_pd_router(prefills, decodes) + # 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"{base_url}/v1", + base_url=f"{gateway.base_url}/v1", api_key="not-used", ) logger.info( - "Setup PD backend: model=%s, %d prefill + %d decode workers, router=%s", + "Setup PD backend: model=%s, %d prefill + %d decode workers, " + "gateway=%s, policy=%s", model_id, len(prefills), len(decodes), - base_url, + gateway.base_url, + gateway_config["policy"], ) try: - yield backend_name, model_path, client + yield backend_name, model_path, client, gateway finally: - logger.info("Tearing down PD router") - kill_process_tree(router_proc.pid) + logger.info("Tearing down PD gateway") + gateway.shutdown() return # Check if this is a local backend (grpc, http) @@ -743,15 +742,17 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): is_local = False connection_mode = None - # Local backends: use worker from pool + launch router + # Local backends: use worker from pool + launch gateway if is_local: # Get number of workers from marker - num_workers = _get_marker_value(request, "workers", default=1) + num_workers = workers_config.get("count") or 1 try: - instance = model_pool.get(model_id, connection_mode) + instance = model_pool.get(model_id, connection_mode, scope=model_scope) except KeyError: pytest.skip(f"Model {model_id}:{backend_name} not available in pool") + except RuntimeError as e: + pytest.fail(str(e)) # Build worker URLs list # For num_workers > 1, we need multiple workers from the pool @@ -759,30 +760,35 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): # TODO: Support launching multiple distinct workers for true LB testing worker_urls = [instance.worker_url] * num_workers - # Launch router pointing to the worker(s) - base_url, router_proc = launch_local_router( + # 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"{base_url}/v1", + base_url=f"{gateway.base_url}/v1", api_key="not-used", ) logger.info( - "Setup %s backend: model=%s, workers=%d, router=%s", + "Setup %s backend: model=%s, workers=%d, gateway=%s, policy=%s", backend_name, model_id, num_workers, - base_url, + gateway.base_url, + gateway_config["policy"], ) try: - yield backend_name, instance.model_path, client + yield backend_name, instance.model_path, client, gateway finally: - logger.info("Tearing down router for %s backend", backend_name) - kill_process_tree(router_proc.pid) + logger.info("Tearing down gateway for %s backend", backend_name) + gateway.shutdown() return # Cloud backends: launch cloud router @@ -817,16 +823,16 @@ def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"): def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"): """Function-scoped fixture for launching a fresh router per test. - This launches a new router for each test, pointing to workers from the pool. + 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): - base_url, router_proc = backend_router - # Test router-specific behavior + gateway = backend_router + # Test gateway-specific behavior """ - from infra import kill_process_tree + from infra import Gateway backend_name = request.param model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL) @@ -838,13 +844,16 @@ def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"): 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)) - base_url, router_proc = launch_local_router( + gateway = Gateway() + gateway.start( worker_urls=[instance.worker_url], model_path=instance.model_path, ) try: - yield base_url, router_proc + yield gateway finally: - kill_process_tree(router_proc.pid) + gateway.shutdown() diff --git a/sgl-model-gateway/e2e_test/infra/__init__.py b/sgl-model-gateway/e2e_test/infra/__init__.py index 97f3ccf98..6b92a9885 100644 --- a/sgl-model-gateway/e2e_test/infra/__init__.py +++ b/sgl-model-gateway/e2e_test/infra/__init__.py @@ -24,6 +24,7 @@ from .constants import ( # Enums; Convenience sets; Fixture parameters; Default Runtime, WorkerType, ) +from .gateway import Gateway, WorkerInfo from .gpu_allocator import ( GPUAllocator, GPUInfo, @@ -107,6 +108,9 @@ __all__ = [ "ModelInstance", "ModelPool", "MODEL_SPECS", + # Gateway + "Gateway", + "WorkerInfo", # Default model paths "DEFAULT_MODEL_PATH", "DEFAULT_SMALL_MODEL_PATH", diff --git a/sgl-model-gateway/e2e_test/infra/gateway.py b/sgl-model-gateway/e2e_test/infra/gateway.py new file mode 100644 index 000000000..894b204ca --- /dev/null +++ b/sgl-model-gateway/e2e_test/infra/gateway.py @@ -0,0 +1,503 @@ +"""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 + + Three 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 + + 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() + """ + + 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._started: bool = False + + @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, + # 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 three 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() + + 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). + 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 + + # Validate mode exclusivity + modes_specified = sum([is_pd_mode, is_regular_mode, is_igw_mode]) + if modes_specified > 1: + raise ValueError( + "Cannot specify multiple modes. Choose one of: " + "worker_urls (regular), prefill/decode_workers (PD), or igw_mode" + ) + + if modes_specified == 0: + raise ValueError( + "Must specify one mode: worker_urls (regular), " + "prefill/decode_workers (PD), or igw_mode=True" + ) + + 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)", + ) + 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) + + self.process = subprocess.Popen( + cmd, + 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 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() + workers = [] + for w in data.get("workers", []): + # Map API fields to WorkerInfo + status = "healthy" if w.get("is_healthy", False) else "unhealthy" + workers.append( + 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"), + }, + ) + ) + return 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: + w = resp.json() + 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"), + }, + ) + 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() diff --git a/sgl-model-gateway/e2e_test/infra/gpu_allocator.py b/sgl-model-gateway/e2e_test/infra/gpu_allocator.py index 517a69a80..9e5d63371 100644 --- a/sgl-model-gateway/e2e_test/infra/gpu_allocator.py +++ b/sgl-model-gateway/e2e_test/infra/gpu_allocator.py @@ -368,6 +368,22 @@ class GPUAllocator: 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. + + Returns: + List of GPU IDs that are not currently allocated. + """ + 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.""" lines = ["GPU Allocation Summary:"] diff --git a/sgl-model-gateway/e2e_test/infra/model_pool.py b/sgl-model-gateway/e2e_test/infra/model_pool.py index b41c2e527..98e4b5e9b 100644 --- a/sgl-model-gateway/e2e_test/infra/model_pool.py +++ b/sgl-model-gateway/e2e_test/infra/model_pool.py @@ -44,6 +44,9 @@ class ModelInstance: gpu_slot: GPUSlot | None worker_type: WorkerType = WorkerType.REGULAR bootstrap_port: int | None = None # For prefill workers in PD mode + scope: str = "session" # "session" or "class" + last_used: float = 0.0 # Timestamp for LRU eviction + _healthy: bool = False # Track if initial health check passed @property def key(self) -> str: @@ -165,10 +168,14 @@ class ModelPool: keeps them running and allows reuse across multiple tests. Routers can then be launched cheaply (~1-2s) pointing to these workers. + Model scopes: + - session: Pre-launched at session start, never evicted + - class: Launched on-demand, can be evicted when GPUs are needed + Startup behavior: - - Workers are launched sequentially (one subprocess.Popen at a time) - - But they boot up concurrently (overlapping model loading) - - _wait_all_healthy() blocks until all workers respond to health checks + - Session-scoped workers are launched at startup + - Class-scoped workers are launched on-demand via get() + - When GPUs are full, class-scoped workers are evicted (LRU) Instance keys: - Regular workers: "model_id:mode" (e.g., "llama-8b:http") @@ -176,16 +183,18 @@ class ModelPool: Limitations: - Currently one worker instance per (model_id, mode) combination - - @pytest.mark.workers(n) duplicates URLs to router, not distinct workers + - @pytest.mark.workers(count=n) duplicates URLs to router, not distinct workers - For true multi-worker LB testing, extend to support multiple instances Usage: pool = ModelPool() pool.startup(requirements=[("llama-8b", ConnectionMode.HTTP)]) + # Session-scoped (pre-launched) instance = pool.get("llama-8b", "http") - # instance.base_url -> "http://127.0.0.1:30000" - # instance.worker_url -> URL for router to connect to + + # Class-scoped (on-demand) + instance = pool.get("qwen-7b", "http", scope="class") """ def __init__(self, allocator: GPUAllocator | None = None): @@ -197,6 +206,22 @@ class ModelPool: self.allocator = allocator or GPUAllocator() self.instances: dict[str, ModelInstance] = {} # key = "model_id:mode" self._startup_timeout = DEFAULT_STARTUP_TIMEOUT + self._class_scoped_models: set[str] = ( + set() + ) # Models that can be launched on-demand + self._queued_models: set[str] = ( + set() + ) # Session models that couldn't be pre-launched + + def register_class_scoped_models(self, models: set[str]) -> None: + """Register models that may be launched on-demand. + + Args: + models: Set of model IDs that are class-scoped. + """ + self._class_scoped_models = models + if models: + logger.info("Registered class-scoped models: %s", models) def startup( self, @@ -253,11 +278,15 @@ class ModelPool: # Allocate GPU slots slots = self.allocator.allocate_slots(allocation_specs) + # Track which models got slots + launched_keys = set() + if not slots: logger.warning("No GPU slots allocated, launching without GPU assignment") # Fallback: launch without specific GPU assignment for model_id, mode in valid_requirements: self._launch_model(model_id, mode, gpu_slot=None) + launched_keys.add(f"{model_id}:{mode.value}") else: # Launch on allocated slots for slot in slots: @@ -266,8 +295,20 @@ class ModelPool: model_id, mode_str = slot.assigned_model.rsplit(":", 1) mode = ConnectionMode(mode_str) self._launch_model(model_id, mode, gpu_slot=slot) + launched_keys.add(slot.assigned_model) - # Wait for all to be healthy + # Track queued models (requested but couldn't be launched due to GPU constraints) + all_keys = set(allocation_specs.keys()) + queued_keys = all_keys - launched_keys + if queued_keys: + self._queued_models.update(queued_keys) + logger.info( + "Queued %d models for on-demand launch (GPU constraints): %s", + len(queued_keys), + queued_keys, + ) + + # Wait for all launched models to be healthy self._wait_all_healthy() def _launch_model( @@ -278,6 +319,7 @@ class ModelPool: worker_type: WorkerType = WorkerType.REGULAR, bootstrap_port: int | None = None, ib_device: str | None = None, + scope: str = "session", ) -> ModelInstance: """Launch a model instance. @@ -288,6 +330,7 @@ class ModelPool: worker_type: Worker type (REGULAR, PREFILL, or DECODE). bootstrap_port: Bootstrap port for prefill workers in PD mode. ib_device: InfiniBand device for PD disaggregation. + scope: Model scope ("session" or "class"). Returns: The launched ModelInstance. @@ -367,16 +410,27 @@ class ModelPool: gpu_slot=gpu_slot, worker_type=worker_type, bootstrap_port=bootstrap_port, + scope=scope, + last_used=time.time(), ) self.instances[key] = instance return instance def _wait_all_healthy(self) -> None: - """Wait for all model instances to become healthy.""" + """Wait for all model instances to become healthy. + + Only checks workers that haven't been marked healthy yet, + avoiding redundant health checks on already-verified workers. + """ start_time = time.time() - pending = set(self.instances.keys()) + # Only wait for workers that haven't been verified healthy yet + pending = {key for key, inst in self.instances.items() if not inst._healthy} check_count = 0 + if not pending: + logger.info("All workers already healthy, skipping health check") + return + logger.info( "Waiting for %d workers to become healthy (timeout: %ds)...", len(pending), @@ -415,6 +469,7 @@ class ModelPool: instance.base_url, check_count, ) + instance._healthy = True pending.discard(key) if pending: @@ -454,19 +509,26 @@ class ModelPool: model_id: str, mode: ConnectionMode | str, worker_type: WorkerType | str = WorkerType.REGULAR, + scope: str = "session", ) -> ModelInstance: """Get a model instance by model_id, mode, and worker_type. + For session-scoped models, raises KeyError if not pre-launched. + For class-scoped models, launches on-demand if not running. + Args: model_id: The model ID (e.g., "llama-8b") mode: The mode (ConnectionMode.HTTP or ConnectionMode.GRPC, or string) worker_type: The worker type (REGULAR, PREFILL, DECODE). Defaults to REGULAR. + scope: Model scope ("session" or "class"). Class-scoped models are + launched on-demand if not running. Returns: ModelInstance for the requested model/mode/worker_type. Raises: - KeyError: If model/mode/worker_type combination is not running. + KeyError: If session-scoped model is not running. + RuntimeError: If worker process died or failed health check. """ # Accept both enum and string for convenience if isinstance(mode, str): @@ -479,13 +541,36 @@ class ModelPool: else: key = f"{model_id}:{mode.value}:{worker_type.value}" + # Check if instance exists if key not in self.instances: - raise KeyError( - f"{key} not running. Available: {list(self.instances.keys())}" - ) + # Check if this model can be launched on-demand + is_class_scoped = scope == "class" or model_id in self._class_scoped_models + is_queued = key in self._queued_models + + if is_class_scoped or is_queued: + launch_scope = "class" if is_class_scoped else "session" + logger.info( + "Launching %s model %s on-demand (queued=%s)", + launch_scope, + key, + is_queued, + ) + self._ensure_gpu_available(model_id) + self._launch_model(model_id, mode, scope=launch_scope) + self._wait_for_instance(key) + + # Remove from queued if it was there + self._queued_models.discard(key) + else: + raise KeyError( + f"{key} not running. Available: {list(self.instances.keys())}" + ) instance = self.instances[key] + # Update last_used timestamp + instance.last_used = time.time() + # Verify worker is still alive and healthy if not instance.is_alive(): raise RuntimeError(f"Worker {key} process died (was healthy at startup)") @@ -499,6 +584,104 @@ class ModelPool: logger.info("Worker %s passed deep health check", key) return instance + def _ensure_gpu_available(self, model_id: str) -> None: + """Ensure GPU is available, evicting models if needed (LRU). + + All models can be evicted when GPU resources are needed. + Uses LRU (least recently used) eviction strategy. + + Args: + model_id: Model ID that needs GPU resources. + """ + spec = get_model_spec(model_id) + required_gpus = spec.get("tp", 1) + + # Check if we have enough free GPUs + available = self.allocator.available_gpus() + if len(available) >= required_gpus: + return # Enough GPUs available + + # Need to evict models to free up GPUs + # Sort by last_used (LRU eviction) - evict least recently used first + evictable = [ + inst + for inst in self.instances.values() + if inst.worker_type == WorkerType.REGULAR + ] + evictable.sort(key=lambda x: x.last_used) + + freed_gpus = 0 + for inst in evictable: + if freed_gpus >= required_gpus: + break + + logger.info( + "Evicting model %s (LRU) to free GPUs for %s", inst.key, model_id + ) + self._evict_instance(inst.key) + if inst.gpu_slot: + freed_gpus += len(inst.gpu_slot.gpu_ids) + + # Recheck available GPUs + available = self.allocator.available_gpus() + if len(available) < required_gpus: + raise RuntimeError( + f"Cannot launch {model_id}: need {required_gpus} GPUs, " + f"only {len(available)} available after eviction" + ) + + def _evict_instance(self, key: str) -> None: + """Evict a model instance and free its resources. + + Evicted models are added back to the queue for potential re-launch. + + Args: + key: Instance key to evict. + """ + if key not in self.instances: + return + + instance = self.instances[key] + instance.terminate() + + # Release GPU slot back to allocator + if instance.gpu_slot: + self.allocator.release_slot(instance.gpu_slot) + + # Add to queued so it can be re-launched on-demand + self._queued_models.add(key) + + del self.instances[key] + logger.info("Evicted instance %s (added to queue for re-launch)", key) + + def _wait_for_instance(self, key: str, timeout: float | None = None) -> None: + """Wait for a specific instance to become healthy. + + Args: + key: Instance key to wait for. + timeout: Timeout in seconds. Defaults to _startup_timeout. + """ + if timeout is None: + timeout = self._startup_timeout + + start_time = time.time() + instance = self.instances.get(key) + if not instance: + raise KeyError(f"Instance {key} not found") + + while (time.time() - start_time) < timeout: + if not instance.is_alive(): + raise RuntimeError(f"Worker {key} died during startup") + + if instance.health_check(): + logger.info("Instance %s is healthy", key) + instance._healthy = True + return + + time.sleep(HEALTH_CHECK_INTERVAL) + + raise TimeoutError(f"Instance {key} did not become healthy within {timeout}s") + def get_workers_by_type( self, model_id: str, worker_type: WorkerType ) -> list[ModelInstance]: diff --git a/sgl-model-gateway/e2e_test/infra/simple_eval_common.py b/sgl-model-gateway/e2e_test/infra/simple_eval_common.py index 17560377b..86855dc9c 100644 --- a/sgl-model-gateway/e2e_test/infra/simple_eval_common.py +++ b/sgl-model-gateway/e2e_test/infra/simple_eval_common.py @@ -141,7 +141,7 @@ class ChatCompletionSampler(SamplerBase): self._pack_message("system", self.system_message) ] + message_list trial = 0 - while trial < 6: # 126 seconds in total + while trial < 6: # Max 63 seconds backoff (1+2+4+8+16+32) try: response = self.client.chat.completions.create( model=self.model, @@ -157,15 +157,20 @@ class ChatCompletionSampler(SamplerBase): return "" except Exception as e: exception_backoff = 2**trial # exponential back off - logger.debug( - "Rate limit, retry %d after %ds: %s", - trial, + # 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, + 6, exception_backoff, e, ) time.sleep(exception_backoff) trial += 1 - logger.warning("All retry attempts exhausted, returning empty response") + logger.warning( + "All retry attempts exhausted after 6 retries, returning empty response" + ) return "" diff --git a/sgl-model-gateway/e2e_test/router/test_mmlu.py b/sgl-model-gateway/e2e_test/router/test_mmlu.py index 13d491e44..a9a6ab11e 100644 --- a/sgl-model-gateway/e2e_test/router/test_mmlu.py +++ b/sgl-model-gateway/e2e_test/router/test_mmlu.py @@ -35,7 +35,7 @@ class TestMMLU: Note: setup_backend fixture already waits for workers to be ready. """ - backend, model, client = setup_backend + backend, model, client, *_ = setup_backend base_url = str(client.base_url).rstrip("/v1") args = SimpleNamespace( @@ -59,7 +59,7 @@ class TestMMLU: Runs MMLU with 128 examples for more statistically significant results. """ - backend, model, client = setup_backend + backend, model, client, *_ = setup_backend base_url = str(client.base_url).rstrip("/v1") args = SimpleNamespace( diff --git a/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py b/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py index 880b73756..e6b740bbd 100644 --- a/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py +++ b/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py @@ -10,7 +10,8 @@ Requirements: Configuration via markers: @pytest.mark.model("model-id") # Override default model - @pytest.mark.pd(num_prefill=2, num_decode=2) # Custom worker counts + @pytest.mark.workers(prefill=2, decode=2) # Custom worker counts + @pytest.mark.gateway(policy="round_robin") # Gateway configuration Usage: # Basic (1 prefill + 1 decode) @@ -42,7 +43,7 @@ class TestPDMMLU: Runs MMLU with 1 prefill + 1 decode worker and validates accuracy meets threshold (>= 0.65). """ - backend, model, client = setup_backend + backend, model, client, *_ = setup_backend base_url = str(client.base_url).rstrip("/v1") args = SimpleNamespace( diff --git a/sgl-model-gateway/e2e_test/router/test_worker_api.py b/sgl-model-gateway/e2e_test/router/test_worker_api.py new file mode 100644 index 000000000..ff356b38b --- /dev/null +++ b/sgl-model-gateway/e2e_test/router/test_worker_api.py @@ -0,0 +1,156 @@ +"""Tests for gateway worker management APIs. + +Tests the gateway's worker management endpoints: +- GET /workers - List all workers +- POST /add_worker - Add a worker dynamically +- POST /remove_worker - Remove a worker dynamically +- GET /v1/models - List available models + +Usage: + pytest e2e_test/router/test_worker_api.py -v +""" + +from __future__ import annotations + +import logging + +import pytest +from infra import ConnectionMode, Gateway, ModelPool + +logger = logging.getLogger(__name__) + + +@pytest.mark.e2e +@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True) +class TestWorkerAPI: + """Tests for worker management APIs using setup_backend fixture.""" + + def test_list_workers(self, setup_backend): + """Test listing workers via /workers endpoint.""" + backend, model, client, gateway = setup_backend + + workers = gateway.list_workers() + assert len(workers) >= 1, "Expected at least one worker" + logger.info("Found %d workers", len(workers)) + + for worker in workers: + logger.info( + "Worker: id=%s, url=%s, status=%s", + worker.id, + worker.url, + worker.status, + ) + assert worker.url, "Worker should have a URL" + + def test_list_models(self, setup_backend): + """Test listing models via /v1/models endpoint.""" + backend, model, client, gateway = setup_backend + + models = gateway.list_models() + assert len(models) >= 1, "Expected at least one model" + logger.info("Found %d models", len(models)) + + for m in models: + logger.info("Model: %s", m.get("id", "unknown")) + assert "id" in m, "Model should have an id" + + def test_health_endpoint(self, setup_backend): + """Test health check endpoint.""" + backend, model, client, gateway = setup_backend + + assert gateway.health(), "Gateway should be healthy" + logger.info("Gateway health check passed") + + +@pytest.mark.e2e +class TestIGWMode: + """Tests for IGW mode - start gateway empty, add workers via API.""" + + def test_igw_start_empty(self, model_pool: ModelPool): + """Test starting gateway in IGW mode with no workers.""" + gateway = Gateway() + gateway.start(igw_mode=True) + + try: + assert gateway.health(), "Gateway should be healthy" + assert gateway.igw_mode, "Gateway should be in IGW mode" + + workers = gateway.list_workers() + logger.info("IGW gateway started with %d workers", len(workers)) + finally: + gateway.shutdown() + + def test_igw_add_worker(self, model_pool: ModelPool): + """Test adding a worker to IGW gateway.""" + http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP) + + gateway = Gateway() + gateway.start(igw_mode=True) + + try: + # Add worker + success, result = gateway.add_worker(http_instance.worker_url) + assert success, f"Failed to add worker: {result}" + logger.info("Added worker: %s", result) + + # Verify worker was added + workers = gateway.list_workers() + assert len(workers) >= 1, "Expected at least one worker" + logger.info("Worker count: %d", len(workers)) + + # Verify models are available + models = gateway.list_models() + logger.info("Models available: %d", len(models)) + finally: + gateway.shutdown() + + def test_igw_add_and_remove_worker(self, model_pool: ModelPool): + """Test adding and removing workers dynamically.""" + http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP) + + gateway = Gateway() + gateway.start(igw_mode=True) + + try: + # Add worker + success, _ = gateway.add_worker(http_instance.worker_url) + assert success, "Failed to add worker" + + initial_count = len(gateway.list_workers()) + logger.info("Worker count after add: %d", initial_count) + + # Remove worker + success, msg = gateway.remove_worker(http_instance.worker_url) + if success: + logger.info("Removed worker: %s", msg) + final_count = len(gateway.list_workers()) + logger.info("Worker count after remove: %d", final_count) + else: + logger.warning("Remove worker not supported: %s", msg) + finally: + gateway.shutdown() + + def test_igw_multiple_workers(self, model_pool: ModelPool): + """Test adding multiple workers to IGW gateway.""" + http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP) + grpc_instance = model_pool.get("llama-8b", ConnectionMode.GRPC) + + gateway = Gateway() + gateway.start(igw_mode=True) + + try: + # Add both workers + success1, _ = gateway.add_worker(http_instance.worker_url) + success2, _ = gateway.add_worker(grpc_instance.worker_url) + + if not success1 or not success2: + pytest.skip("Dynamic worker management not fully supported") + + workers = gateway.list_workers() + logger.info("Worker count: %d", len(workers)) + assert len(workers) >= 2, "Expected at least 2 workers" + + for w in workers: + logger.info("Worker: id=%s, url=%s", w.id, w.url) + finally: + gateway.shutdown() diff --git a/sgl-model-gateway/e2e_test/utils.py b/sgl-model-gateway/e2e_test/utils.py index 7b1bbc23e..6fec3b424 100644 --- a/sgl-model-gateway/e2e_test/utils.py +++ b/sgl-model-gateway/e2e_test/utils.py @@ -21,7 +21,7 @@ from typing import TYPE_CHECKING logger = logging.getLogger(__name__) # Re-export commonly used items from submodules -from backends import kill_process_tree # noqa: F401 +from infra import kill_process_tree # noqa: F401 from infra.model_specs import ( # noqa: F401; Default model paths DEFAULT_EMBEDDING_MODEL_PATH, DEFAULT_ENABLE_THINKING_MODEL_PATH,