chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Simo Lin
2026-01-04 23:30:32 -08:00
committed by GitHub
parent 399ca037b1
commit 7f6a678f8f
5 changed files with 957 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

176
sgl-model-gateway/e2e_test/conftest.py
View File

@@ -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")

31
sgl-model-gateway/e2e_test/infra/__init__.py Normal file
View File

@@ -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",
]

361
sgl-model-gateway/e2e_test/infra/gpu_allocator.py Normal file
View File

@@ -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)

280
sgl-model-gateway/e2e_test/infra/model_pool.py Normal file
View File

@@ -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()

109
sgl-model-gateway/e2e_test/infra/model_specs.py Normal file
View File

@@ -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")