diff --git a/sgl-model-gateway/e2e_test/conftest.py b/sgl-model-gateway/e2e_test/conftest.py index fe4fc6994..b41906fe0 100644 --- a/sgl-model-gateway/e2e_test/conftest.py +++ b/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") diff --git a/sgl-model-gateway/e2e_test/infra/__init__.py b/sgl-model-gateway/e2e_test/infra/__init__.py new file mode 100644 index 000000000..0c06d7c01 --- /dev/null +++ b/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", +] diff --git a/sgl-model-gateway/e2e_test/infra/gpu_allocator.py b/sgl-model-gateway/e2e_test/infra/gpu_allocator.py new file mode 100644 index 000000000..40944d3ec --- /dev/null +++ b/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) diff --git a/sgl-model-gateway/e2e_test/infra/model_pool.py b/sgl-model-gateway/e2e_test/infra/model_pool.py new file mode 100644 index 000000000..07c61f670 --- /dev/null +++ b/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() diff --git a/sgl-model-gateway/e2e_test/infra/model_specs.py b/sgl-model-gateway/e2e_test/infra/model_specs.py new file mode 100644 index 000000000..519a3a7f2 --- /dev/null +++ b/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")