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

[smg][ci] Add thread safety to ModelPool and GPUAllocator (#16674)

Browse Source
This commit is contained in:
Simo Lin
2026-01-07 13:25:41 -08:00
committed by GitHub
parent 0241e0460f
commit 6037267f5b
10 changed files with 534 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
.github/workflows/pr-test-rust.yml vendored
View File

@@ -172,6 +172,7 @@ jobs:
env_vars: ""
reruns: ""
upload_benchmarks: true
parallel_opts: "" # No parallel for benchmarks (performance measurement)
- name: response-api
timeout: 32
test_dirs: "e2e_test/e2e_response_api"
@@ -180,18 +181,21 @@ jobs:
reruns: "--reruns 3 --reruns-delay 2"
setup_oracle: true
setup_brave: true
parallel_opts: "" # Legacy tests, not yet migrated for parallel
- name: grpc
timeout: 32
test_dirs: "e2e_test/e2e_grpc"
extra_deps: ""
env_vars: "SHOW_ROUTER_LOGS=1"
reruns: "--reruns 3 --reruns-delay 2"
parallel_opts: "" # Legacy tests, not yet migrated for parallel
- name: e2e
timeout: 45
test_dirs: "e2e_test/router e2e_test/embeddings"
extra_deps: ""
extra_deps: "pytest-parallel py" # py is required for pytest-parallel with newer pytest
env_vars: "SHOW_WORKER_LOGS=0 SHOW_ROUTER_LOGS=1"
reruns: "--reruns 2 --reruns-delay 5"
parallel_opts: "--workers 1 --tests-per-worker 4" # Thread-based parallelism
runs-on: 4-gpu-a10
timeout-minutes: ${{ matrix.timeout }}
steps:
@@ -286,7 +290,7 @@ jobs:
bash scripts/killall_sglang.sh "nuk_gpus"
cd sgl-model-gateway
source "$HOME/.cargo/env"
${{ matrix.env_vars }} ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" pytest ${{ matrix.reruns }} ${{ matrix.test_dirs }} -s -vv -o log_cli=true --log-cli-level=INFO
${{ matrix.env_vars }} ROUTER_LOCAL_MODEL_PATH="/home/ubuntu/models" pytest ${{ matrix.reruns }} ${{ matrix.parallel_opts }} ${{ matrix.test_dirs }} -s -vv -o log_cli=true --log-cli-level=INFO
- name: Upload benchmark results
if: matrix.upload_benchmarks && success()

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

@@ -1,5 +1,17 @@
"""Pytest configuration for E2E tests.
Parallel Execution
------------------
Tests can run in parallel using pytest-parallel with shared worker processes.
Use --workers 1 --tests-per-worker N for N concurrent test threads:
pytest --workers 1 --tests-per-worker 4 e2e_test/router/
This leverages the thread-safe ModelPool and GPUAllocator classes to enable
true shared-worker parallelism where all threads share the same session-scoped
model_pool fixture. Tests marked with @pytest.mark.thread_unsafe will be
automatically skipped in parallel mode.
Markers
-------
This module defines several pytest markers for configuring E2E tests:
@@ -52,6 +64,18 @@ This module defines several pytest markers for configuring E2E tests:
@pytest.mark.slow
Mark test as slow-running.
@pytest.mark.thread_unsafe(reason=None)
Mark test as incompatible with parallel thread execution.
Tests with this marker are automatically skipped when running
with --tests-per-worker > 1.
Args:
reason: Optional explanation of why the test is thread-unsafe.
Examples:
@pytest.mark.thread_unsafe
@pytest.mark.thread_unsafe(reason="Modifies global state")
Fixtures
--------
model_pool: Session-scoped fixture managing SGLang worker processes.
@@ -119,8 +143,15 @@ if not _wheel_installed and str(_SRC) not in sys.path:
def _setup_logging() -> None:
"""Configure clean logging to stdout with timestamps."""
fmt = "%(asctime)s.%(msecs)03d [%(name)s] %(message)s"
"""Configure clean logging to stdout with timestamps and thread info.
In parallel mode (--tests-per-worker > 1), logs from different threads
would be interleaved. Including thread name helps identify which test
produced each log line.
"""
# Include thread name for parallel execution readability
# MainThread for sequential, Thread-N for parallel workers
fmt = "%(asctime)s.%(msecs)03d [%(threadName)s] [%(name)s] %(message)s"
datefmt = "%H:%M:%S"
handler = logging.StreamHandler(sys.stdout)
@@ -148,11 +179,14 @@ logger = logging.getLogger(__name__)
def pytest_runtest_logstart(nodeid: str, location: tuple) -> None:
"""Print clear test header at start of each test."""
import threading
from infra import LOG_SEPARATOR_WIDTH
test_name = nodeid.split("::")[-1] if "::" in nodeid else nodeid
thread_name = threading.current_thread().name
print(f"\n{'=' * LOG_SEPARATOR_WIDTH}")
print(f"TEST: {test_name}")
print(f"[{thread_name}] TEST: {test_name}")
print(f"{'=' * LOG_SEPARATOR_WIDTH}")
@@ -170,6 +204,7 @@ from fixtures import (
pytest_collection_finish,
pytest_collection_modifyitems,
pytest_configure,
pytest_runtest_setup,
setup_backend,
)
@@ -180,6 +215,7 @@ __all__ = [
"pytest_collection_modifyitems",
"pytest_collection_finish",
"pytest_configure",
"pytest_runtest_setup",
# Fixtures
"model_pool",
"model_client",

70
sgl-model-gateway/e2e_test/embeddings/test_correctness.py
View File

@@ -18,6 +18,7 @@ Requirements:
from __future__ import annotations
import logging
import threading
from typing import Any
import numpy as np
@@ -27,6 +28,10 @@ import torch.nn.functional as F
logger = logging.getLogger(__name__)
# Thread-safe storage for HF reference embeddings
_hf_embeddings_cache: dict[str, Any] | None = None
_hf_embeddings_lock = threading.Lock()
# Test data for semantic similarity checks
SEMANTIC_TEST_SETS: list[list[str]] = [
@@ -127,45 +132,54 @@ def get_input_texts(test_json: dict) -> list[str]:
return [doc["body"] for doc in test_json["sample_reference"]]
@pytest.fixture(scope="class")
@pytest.fixture(scope="session")
def hf_reference_embeddings(request):
"""Pre-compute HuggingFace reference embeddings on CPU.
This is done once per test class before launching workers to avoid
GPU memory conflicts in CI environments.
This is done once per session with thread-safe initialization to support
pytest-parallel execution. Uses CPU to avoid GPU memory conflicts.
"""
from infra.model_specs import MODEL_SPECS
global _hf_embeddings_cache
# Get model path from MODEL_SPECS for the embedding model
model_path = MODEL_SPECS.get("embedding", {}).get("model")
if model_path is None:
pytest.skip("Embedding model not found in MODEL_SPECS")
# Thread-safe initialization - only one thread computes embeddings
with _hf_embeddings_lock:
if _hf_embeddings_cache is not None:
return _hf_embeddings_cache
logger.info(
"Pre-computing HuggingFace reference embeddings (CPU) for %s", model_path
)
from infra.model_specs import MODEL_SPECS
# Flatten all test texts for semantic similarity
all_semantic_texts = []
for text_set in SEMANTIC_TEST_SETS:
all_semantic_texts.extend(text_set)
# Get model path from MODEL_SPECS for the embedding model
model_path = MODEL_SPECS.get("embedding", {}).get("model")
if model_path is None:
pytest.skip("Embedding model not found in MODEL_SPECS")
# Get relevance test texts
query = f"Instruct: Given a search query, retrieve relevant passages that answer the query\nQuery: {RELEVANCE_TEST_DATA['sample_query']}"
docs = get_input_texts(RELEVANCE_TEST_DATA)
logger.info(
"Pre-computing HuggingFace reference embeddings (CPU) for %s", model_path
)
# Compute all reference embeddings at once
hf_semantic = get_hf_st_embeddings(all_semantic_texts, model_path)
hf_query = get_hf_st_embeddings(query, model_path)
hf_docs = get_hf_st_embeddings(docs, model_path)
# Flatten all test texts for semantic similarity
all_semantic_texts = []
for text_set in SEMANTIC_TEST_SETS:
all_semantic_texts.extend(text_set)
logger.info("Reference embeddings computed on CPU")
# Get relevance test texts
query = f"Instruct: Given a search query, retrieve relevant passages that answer the query\nQuery: {RELEVANCE_TEST_DATA['sample_query']}"
docs = get_input_texts(RELEVANCE_TEST_DATA)
return {
"semantic": hf_semantic,
"query": hf_query,
"docs": hf_docs,
}
# Compute all reference embeddings at once
hf_semantic = get_hf_st_embeddings(all_semantic_texts, model_path)
hf_query = get_hf_st_embeddings(query, model_path)
hf_docs = get_hf_st_embeddings(docs, model_path)
logger.info("Reference embeddings computed on CPU")
_hf_embeddings_cache = {
"semantic": hf_semantic,
"query": hf_query,
"docs": hf_docs,
}
return _hf_embeddings_cache
@pytest.mark.e2e

4
sgl-model-gateway/e2e_test/fixtures/__init__.py
View File

@@ -14,9 +14,11 @@ Legacy modules (to be removed during e2e_response_api migration):
# Pytest hooks (imported by conftest.py via pytest_plugins)
from .hooks import (
get_pool_requirements,
is_parallel_execution,
pytest_collection_finish,
pytest_collection_modifyitems,
pytest_configure,
pytest_runtest_setup,
validate_gpu_requirements,
)
@@ -32,8 +34,10 @@ __all__ = [
"pytest_collection_modifyitems",
"pytest_collection_finish",
"pytest_configure",
"pytest_runtest_setup",
"get_pool_requirements",
"validate_gpu_requirements",
"is_parallel_execution",
# Pool fixtures
"model_pool",
"model_client",

72
sgl-model-gateway/e2e_test/fixtures/hooks.py
View File

@@ -269,21 +269,38 @@ def get_pool_requirements() -> list["WorkerIdentity"]:
# ---------------------------------------------------------------------------
def _count_gpus_without_cuda() -> int:
"""Count available GPUs without initializing CUDA.
Uses nvidia-smi to avoid CUDA initialization, which is critical for
pytest-parallel compatibility. CUDA cannot be re-initialized after a fork.
"""
import subprocess
try:
result = subprocess.run(
["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
capture_output=True,
text=True,
timeout=10,
)
if result.returncode == 0:
return len([line for line in result.stdout.strip().split("\n") if line])
except (subprocess.SubprocessError, FileNotFoundError, OSError):
pass
return 0
def validate_gpu_requirements() -> tuple[int, int]:
"""Check if there are enough GPUs for any single test.
Uses nvidia-smi instead of torch.cuda to avoid CUDA initialization,
which would break pytest-parallel (CUDA cannot be re-initialized after fork).
Returns:
Tuple of (max_required_gpus, available_gpus).
"""
available_gpus = 0
try:
import torch
if torch.cuda.is_available():
available_gpus = torch.cuda.device_count()
except ImportError:
pass
available_gpus = _count_gpus_without_cuda()
return _max_test_gpu_requirement, available_gpus
@@ -356,3 +373,40 @@ def pytest_configure(config: pytest.Config) -> None:
"markers",
"slow: mark test as slow-running",
)
config.addinivalue_line(
"markers",
"thread_unsafe: mark test as incompatible with parallel thread execution",
)
# ---------------------------------------------------------------------------
# Parallel execution support
# ---------------------------------------------------------------------------
def is_parallel_execution(config: pytest.Config) -> bool:
"""Check if tests are running in parallel mode (pytest-parallel).
Returns True if --tests-per-worker > 1, indicating concurrent thread execution.
"""
# pytest-parallel adds the 'tests_per_worker' option
tests_per_worker = getattr(config.option, "tests_per_worker", None)
if tests_per_worker is None:
return False
if tests_per_worker == "auto":
return True
try:
return int(tests_per_worker) > 1
except (ValueError, TypeError):
return False
def pytest_runtest_setup(item: pytest.Item) -> None:
"""Skip thread_unsafe tests when running in parallel mode."""
if is_parallel_execution(item.config):
marker = item.get_closest_marker("thread_unsafe")
if marker:
reason = marker.kwargs.get("reason", "Test is not thread-safe")
pytest.skip(f"Skipping in parallel mode: {reason}")

174
sgl-model-gateway/e2e_test/fixtures/pool.py
View File

@@ -6,8 +6,10 @@ Workers are expensive to start (~30-60s each), so they're kept running across te
from __future__ import annotations
import atexit
import logging
import os
import threading
from typing import TYPE_CHECKING
import pytest
@@ -19,8 +21,24 @@ from .hooks import get_pool_requirements
logger = logging.getLogger(__name__)
# Global model pool instance
# Global model pool instance with thread-safe initialization
_model_pool: "ModelPool | None" = None
_model_pool_lock = threading.Lock()
_shutdown_registered = False
def _shutdown_model_pool() -> None:
"""Shutdown the global model pool at process exit.
This is registered with atexit to ensure cleanup happens after all tests
complete, which is important for pytest-parallel where multiple threads
share the session-scoped fixture.
"""
global _model_pool
if _model_pool is not None:
logger.info("Shutting down model pool at process exit")
_model_pool.shutdown()
_model_pool = None
@pytest.fixture(scope="session")
@@ -65,82 +83,94 @@ def model_pool(request: pytest.FixtureRequest) -> "ModelPool":
WorkerType,
)
if _model_pool is not None:
return _model_pool
# Thread-safe initialization: use lock to ensure only one thread creates the pool
# This is critical for pytest-parallel which runs tests as concurrent threads
with _model_pool_lock:
if _model_pool is not None:
return _model_pool
# Check if we should skip model startup
if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
logger.info("%s is set, skipping model pool startup", ENV_SKIP_MODEL_POOL)
_model_pool = ModelPool(GPUAllocator(gpus=[]))
return _model_pool
# Check if we should skip model startup
if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
logger.info("%s is set, skipping model pool startup", ENV_SKIP_MODEL_POOL)
_model_pool = ModelPool(GPUAllocator(gpus=[]))
return _model_pool
# Determine requirements from scanned tests or env vars
models_env = os.environ.get(ENV_MODELS, "")
backends_env = os.environ.get(ENV_BACKENDS, "")
# Determine requirements from scanned tests or env vars
models_env = os.environ.get(ENV_MODELS, "")
backends_env = os.environ.get(ENV_BACKENDS, "")
if models_env or backends_env:
# Use env var overrides
models = (
{m.strip() for m in models_env.split(",") if m.strip()}
if models_env
else {DEFAULT_MODEL}
if models_env or backends_env:
# Use env var overrides
models = (
{m.strip() for m in models_env.split(",") if m.strip()}
if models_env
else {DEFAULT_MODEL}
)
# Parse backend strings to ConnectionMode enums
backend_modes: set[ConnectionMode] = set()
if backends_env:
for b in backends_env.split(","):
b = b.strip()
if b:
try:
mode = ConnectionMode(b)
if mode in LOCAL_MODES:
backend_modes.add(mode)
except ValueError:
logger.warning("Unknown backend '%s', skipping", b)
# Default to HTTP if no valid backends
if not backend_modes:
backend_modes = {ConnectionMode.HTTP}
# Create WorkerIdentity objects (regular workers only from env vars)
requirements = [
WorkerIdentity(m, b, WorkerType.REGULAR, 0)
for m in models
for b in backend_modes
]
logger.info(
"Using env var requirements: %s", [str(r) for r in requirements]
)
else:
# Use scanned requirements from test markers
requirements = get_pool_requirements()
logger.info(
"Using scanned requirements: %s", [str(r) for r in requirements]
)
# Filter to valid models
requirements = [r for r in requirements if r.model_id in MODEL_SPECS]
if not requirements:
logger.warning("No valid requirements, model pool will be empty")
_model_pool = ModelPool(GPUAllocator(gpus=[]))
return _model_pool
# Create and start the pool
allocator = GPUAllocator()
_model_pool = ModelPool(allocator)
startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300"))
_model_pool.startup(
requirements=requirements,
startup_timeout=startup_timeout,
)
# Parse backend strings to ConnectionMode enums
backend_modes: set[ConnectionMode] = set()
if backends_env:
for b in backends_env.split(","):
b = b.strip()
if b:
try:
mode = ConnectionMode(b)
if mode in LOCAL_MODES:
backend_modes.add(mode)
except ValueError:
logger.warning("Unknown backend '%s', skipping", b)
# Log final GPU allocation summary
logger.info(_model_pool.allocator.summary())
# Default to HTTP if no valid backends
if not backend_modes:
backend_modes = {ConnectionMode.HTTP}
# Register cleanup with atexit instead of request.addfinalizer
# This is critical for pytest-parallel where multiple threads share
# the session-scoped fixture - addfinalizer can fire too early
global _shutdown_registered
if not _shutdown_registered:
atexit.register(_shutdown_model_pool)
_shutdown_registered = True
# Create WorkerIdentity objects (regular workers only from env vars)
requirements = [
WorkerIdentity(m, b, WorkerType.REGULAR, 0)
for m in models
for b in backend_modes
]
logger.info("Using env var requirements: %s", [str(r) for r in requirements])
else:
# Use scanned requirements from test markers
requirements = get_pool_requirements()
logger.info("Using scanned requirements: %s", [str(r) for r in requirements])
# Filter to valid models
requirements = [r for r in requirements if r.model_id in MODEL_SPECS]
if not requirements:
logger.warning("No valid requirements, model pool will be empty")
_model_pool = ModelPool(GPUAllocator(gpus=[]))
return _model_pool
# Create and start the pool
allocator = GPUAllocator()
_model_pool = ModelPool(allocator)
startup_timeout = int(os.environ.get(ENV_STARTUP_TIMEOUT, "300"))
_model_pool.startup(
requirements=requirements,
startup_timeout=startup_timeout,
)
# Log final GPU allocation summary
logger.info(_model_pool.allocator.summary())
# Register cleanup
request.addfinalizer(_model_pool.shutdown)
return _model_pool
@pytest.fixture
def model_client(request: pytest.FixtureRequest, model_pool: "ModelPool"):
@@ -164,13 +194,11 @@ def model_client(request: pytest.FixtureRequest, model_pool: "ModelPool"):
model_id = marker.args[0]
try:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id)
except KeyError:
pytest.skip(f"Model {model_id} not available in model pool")
# Acquire reference to prevent eviction during test
instance.acquire()
client = openai.OpenAI(
base_url=f"{instance.base_url}/v1",
api_key="not-used",
@@ -203,13 +231,11 @@ def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> s
model_id = marker.args[0]
try:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id)
except KeyError:
pytest.skip(f"Model {model_id} not available in model pool")
# Acquire reference to prevent eviction during test
instance.acquire()
yield instance.base_url
# Release reference to allow eviction

86
sgl-model-gateway/e2e_test/fixtures/setup_backend.py
View File

@@ -130,34 +130,15 @@ def _setup_pd_backend(
import openai
from infra import ConnectionMode, Gateway, WorkerIdentity, WorkerType
# Check PD requirements
try:
import sgl_kernel # noqa: F401
except ImportError:
pytest.skip("sgl_kernel not available, required for PD disaggregation")
try:
import torch
except ImportError:
pytest.skip("torch not available")
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
logger.info("Setting up PD backend for model %s", model_id)
# Get PD configuration from workers marker
num_prefill = workers_config.get("prefill") or 1
num_decode = workers_config.get("decode") or 1
# Check GPU requirements
required_gpus = num_prefill + num_decode
gpu_count = torch.cuda.device_count()
if gpu_count < required_gpus:
pytest.skip(
f"PD tests require {required_gpus} GPUs "
f"({num_prefill} prefill + {num_decode} decode), found {gpu_count}"
)
logger.info("PD config: %d prefill, %d decode workers", num_prefill, num_decode)
# Try to use pre-launched PD workers, or launch additional ones if needed
# get_workers_by_type auto-acquires all returned workers
existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL)
existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE)
@@ -168,6 +149,11 @@ def _setup_pd_backend(
if missing_prefill == 0 and missing_decode == 0:
prefills = existing_prefills[:num_prefill]
decodes = existing_decodes[:num_decode]
# Release excess workers we won't use
for w in existing_prefills[num_prefill:]:
w.release()
for w in existing_decodes[num_decode:]:
w.release()
logger.info(
"Using pre-launched PD workers: %d prefill, %d decode",
len(prefills),
@@ -207,17 +193,36 @@ def _setup_pd_backend(
workers_to_launch, startup_timeout=300
)
if not new_instances:
# Release any existing workers we acquired
for w in existing_prefills + existing_decodes:
w.release()
pytest.fail(
f"Failed to launch PD workers: needed {len(workers_to_launch)} workers "
f"but could not allocate GPUs (all in use or timeout)"
)
# Acquire newly launched instances (launch_workers doesn't auto-acquire)
for inst in new_instances:
inst.acquire()
new_prefills = [w for w in new_instances if w.worker_type == WorkerType.PREFILL]
new_decodes = [w for w in new_instances if w.worker_type == WorkerType.DECODE]
prefills = existing_prefills + new_prefills
decodes = existing_decodes + new_decodes
# Acquire references to prevent eviction during test
all_workers = prefills + decodes
for worker in all_workers:
worker.acquire()
# All workers in prefills and decodes are now acquired
model_path = prefills[0].model_path if prefills else None
if not prefills or not decodes:
# This shouldn't happen but guard against it
for w in prefills + decodes:
w.release()
pytest.fail(
f"PD setup incomplete: have {len(prefills)} prefill, {len(decodes)} decode "
f"(need {num_prefill} prefill, {num_decode} decode)"
)
model_path = prefills[0].model_path
# Launch PD gateway
gateway = Gateway()
@@ -250,7 +255,7 @@ def _setup_pd_backend(
logger.info("Tearing down PD gateway")
gateway.shutdown()
# Release references to allow eviction
for worker in all_workers:
for worker in prefills + decodes:
worker.release()
@@ -272,11 +277,20 @@ def _setup_local_backend(
try:
if num_workers > 1:
existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR)
existing_for_mode = [w for w in existing if w.mode == connection_mode]
# get_workers_by_type auto-acquires all returned workers
all_existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR)
existing_for_mode = [w for w in all_existing if w.mode == connection_mode]
# Release workers we won't use (wrong mode)
for w in all_existing:
if w not in existing_for_mode:
w.release()
if len(existing_for_mode) >= num_workers:
instances = existing_for_mode[:num_workers]
# Release excess workers we won't use
for w in existing_for_mode[num_workers:]:
w.release()
else:
missing = num_workers - len(existing_for_mode)
workers_to_launch = [
@@ -291,6 +305,9 @@ def _setup_local_backend(
new_instances = model_pool.launch_workers(
workers_to_launch, startup_timeout=300
)
# Acquire newly launched instances
for inst in new_instances:
inst.acquire()
instances = existing_for_mode + new_instances
if not instances:
@@ -298,14 +315,11 @@ def _setup_local_backend(
worker_urls = [inst.worker_url for inst in instances]
model_path = instances[0].model_path
else:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id, connection_mode)
instances = [instance]
worker_urls = [instance.worker_url]
model_path = instance.model_path
# Acquire references to prevent eviction during test
for inst in instances:
inst.acquire()
except RuntimeError as e:
pytest.fail(str(e))
@@ -393,15 +407,13 @@ def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
connection_mode = ConnectionMode(backend_name)
try:
# get() auto-acquires the returned instance
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))
# Acquire reference to prevent eviction during test
instance.acquire()
gateway = Gateway()
gateway.start(
worker_urls=[instance.worker_url],

71
sgl-model-gateway/e2e_test/infra/gpu_allocator.py
View File

@@ -5,6 +5,7 @@ from __future__ import annotations
import logging
import os
import socket
import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass
@@ -200,6 +201,7 @@ class GPUAllocator:
self.gpus = gpus if gpus is not None else self._detect_gpus()
self.slots: list[GPUSlot] = []
self._used_gpus: set[int] = set() # Track GPUs used across all allocations
self._lock = threading.RLock() # Protects slots and _used_gpus
def _detect_gpus(self) -> list[GPUInfo]:
"""Auto-detect available GPUs via nvidia-ml-py (NVML)."""
@@ -261,6 +263,8 @@ class GPUAllocator:
Note: This method tracks used GPUs across multiple calls, so subsequent
allocations will use different GPUs than previous ones.
Thread-safe: Protected by internal lock.
Args:
model_specs: Dict of model_id -> spec dict with 'memory_gb' and 'tp' keys
preserve_order: If True, allocate in dict order (test order) instead
@@ -269,6 +273,13 @@ class GPUAllocator:
Returns:
List of GPUSlots with assigned models (only the newly allocated slots)
"""
with self._lock:
return self._allocate_slots_unlocked(model_specs, preserve_order)
def _allocate_slots_unlocked(
self, model_specs: dict[str, dict], preserve_order: bool = False
) -> list[GPUSlot]:
"""Internal allocation logic. Caller must hold _lock."""
if not self.gpus:
logger.warning("No GPUs available for allocation")
return []
@@ -362,23 +373,32 @@ class GPUAllocator:
return new_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
"""Get the slot assigned to a specific model.
Thread-safe: Protected by internal lock.
"""
with self._lock:
for slot in self.slots:
if slot.assigned_model == model_id:
return slot
return None
def release_gpus(self, gpu_ids: list[int]) -> None:
"""Release GPUs back to the available pool.
Thread-safe: Protected by internal lock.
Args:
gpu_ids: List of GPU IDs to release.
"""
for gpu_id in gpu_ids:
self._used_gpus.discard(gpu_id)
# Remove slots that used these GPUs
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)
with self._lock:
for gpu_id in gpu_ids:
self._used_gpus.discard(gpu_id)
# Remove slots that used these GPUs
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.
@@ -391,20 +411,27 @@ class GPUAllocator:
def available_gpus(self) -> list[int]:
"""Get list of available (unused) GPU IDs.
Thread-safe: Protected by internal lock.
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]
with self._lock:
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:"]
lines.append(f" Total GPUs: {len(self.gpus)}")
lines.append(f" Used GPUs: {sorted(self._used_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)
"""Return a summary of GPU allocations.
Thread-safe: Protected by internal lock.
"""
with self._lock:
lines = ["GPU Allocation Summary:"]
lines.append(f" Total GPUs: {len(self.gpus)}")
lines.append(f" Used GPUs: {sorted(self._used_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)

188
sgl-model-gateway/e2e_test/infra/model_pool.py
View File

@@ -293,6 +293,7 @@ class ModelPool:
self.allocator = allocator or GPUAllocator()
self.instances: dict[str, ModelInstance] = {} # key = "model_id:mode"
self._startup_timeout = DEFAULT_STARTUP_TIMEOUT
self._lock = threading.RLock() # Protects instances dict
def startup(
self,
@@ -309,11 +310,22 @@ class ModelPool:
Each WorkerIdentity uniquely identifies a worker by (model_id, mode,
worker_type, index).
Thread-safe: Protected by internal lock.
Args:
requirements: List of WorkerIdentity specifying what to start.
If None, starts default model in HTTP mode.
startup_timeout: Timeout in seconds for all models to become healthy.
"""
with self._lock:
self._startup_unlocked(requirements, startup_timeout)
def _startup_unlocked(
self,
requirements: list[WorkerIdentity] | None = None,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
) -> None:
"""Internal startup logic. Caller must hold _lock."""
self._startup_timeout = startup_timeout
if requirements is None:
@@ -615,22 +627,76 @@ class ModelPool:
model_id: str,
mode: ConnectionMode | str,
worker_type: WorkerType | str = WorkerType.REGULAR,
wait_for_gpus: bool = True,
gpu_wait_timeout: int = 300,
) -> ModelInstance:
"""Get a model instance by model_id, mode, and worker_type.
If the model is not running, it will be launched on-demand with MRU
eviction if GPU resources are constrained.
Thread-safe: Protected by internal lock. The returned instance has its
reference count incremented (via acquire()) to prevent eviction.
Caller MUST call release() on the instance when done.
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.
wait_for_gpus: If True, wait for GPUs to become available when all
are in use by other tests. Defaults to True.
gpu_wait_timeout: Max seconds to wait for GPUs (default 5 min).
Returns:
ModelInstance for the requested model/mode/worker_type.
ModelInstance for the requested model/mode/worker_type (already acquired).
Raises:
RuntimeError: If worker process died or failed health check.
RuntimeError: If worker process died, failed health check, or
timeout waiting for GPUs.
"""
deadline = time.time() + gpu_wait_timeout
poll_interval = 2.0 # seconds
while True:
with self._lock:
instance = self._get_unlocked(model_id, mode, worker_type)
if instance is not None:
# Acquire while holding lock to prevent race with eviction
instance.acquire()
return instance
# _get_unlocked returns None when GPUs unavailable after eviction
if not wait_for_gpus:
raise RuntimeError(
f"Cannot get {model_id}: GPUs unavailable and waiting disabled"
)
if time.time() >= deadline:
raise RuntimeError(
f"Timeout waiting for GPUs for {model_id} after {gpu_wait_timeout}s"
)
# Release lock while waiting so other tests can release workers
logger.info(
"All GPUs in use by other tests, waiting %.1fs for %s...",
poll_interval,
model_id,
)
time.sleep(poll_interval)
def _get_unlocked(
self,
model_id: str,
mode: ConnectionMode | str,
worker_type: WorkerType | str = WorkerType.REGULAR,
) -> ModelInstance | None:
"""Internal get logic. Caller must hold _lock.
Returns:
ModelInstance if successful, None if GPUs unavailable (signals retry).
Raises:
RuntimeError: If worker died or failed health check.
"""
# Accept both enum and string for convenience
if isinstance(mode, str):
@@ -649,7 +715,9 @@ class ModelPool:
"Model %s not running, launching on-demand with MRU eviction if needed",
key,
)
self._ensure_gpu_available(model_id)
if not self._ensure_gpu_available(model_id):
# GPUs not available after eviction - signal retry
return None
# Allocate GPU slot for this model
spec = get_model_spec(model_id)
@@ -752,14 +820,14 @@ class ModelPool:
if inst.gpu_slot:
freed_gpus += len(inst.gpu_slot.gpu_ids)
def _ensure_gpu_available(self, model_id: str) -> None:
def _ensure_gpu_available(self, model_id: str) -> bool:
"""Ensure GPU is available for a model, evicting if needed.
Args:
model_id: Model ID that needs GPU resources.
Raises:
RuntimeError: If not enough GPUs after eviction.
Returns:
True if GPUs are available, False if not (all in use by other tests).
"""
spec = get_model_spec(model_id)
required_gpus = spec.get("tp", 1)
@@ -774,10 +842,15 @@ class ModelPool:
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"
logger.info(
"Cannot launch %s: need %d GPUs, only %d available after eviction "
"(all workers in use by other tests)",
model_id,
required_gpus,
len(available),
)
return False
return True
def _evict_instance(self, key: str) -> None:
"""Evict a model instance and free its resources.
@@ -831,38 +904,96 @@ class ModelPool:
) -> list[ModelInstance]:
"""Get all workers of a specific type for a model.
Thread-safe: Protected by internal lock. All returned instances have their
reference count incremented (via acquire()) to prevent eviction.
Caller MUST call release() on each instance when done.
Args:
model_id: The model ID.
worker_type: The worker type to filter by.
Returns:
List of matching ModelInstance objects.
List of matching ModelInstance objects (already acquired).
"""
return [
inst
for inst in self.instances.values()
if inst.model_id == model_id and inst.worker_type == worker_type
]
with self._lock:
workers = [
inst
for inst in self.instances.values()
if inst.model_id == model_id and inst.worker_type == worker_type
]
# Acquire all while holding lock to prevent race with eviction
for worker in workers:
worker.acquire()
return workers
def launch_workers(
self,
workers: list[WorkerIdentity],
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
allow_eviction: bool = True,
wait_for_gpus: bool = True,
gpu_wait_timeout: int = 300,
) -> list[ModelInstance]:
"""Launch workers of any type.
This is the unified method for launching workers. It handles all worker
types (regular, prefill, decode) uniformly.
Thread-safe: Protected by internal lock.
Args:
workers: List of WorkerIdentity objects specifying workers to launch.
startup_timeout: Timeout for workers to become healthy.
allow_eviction: If True, evict MRU models to free GPUs.
wait_for_gpus: If True, wait for GPUs to become available when all
are in use by other tests (with eviction enabled).
gpu_wait_timeout: Max seconds to wait for GPUs (default 5 min).
Returns:
List of launched ModelInstance objects.
"""
deadline = time.time() + gpu_wait_timeout
poll_interval = 2.0 # seconds
while True:
with self._lock:
result = self._launch_workers_unlocked(
workers, startup_timeout, allow_eviction
)
if result is not None:
return result
# _launch_workers_unlocked returns None when GPUs unavailable
# after eviction attempt (all workers in use by other tests)
if not wait_for_gpus or not allow_eviction:
return []
if time.time() >= deadline:
logger.warning(
"Timeout waiting for GPUs after %ds, giving up",
gpu_wait_timeout,
)
return []
# Release lock while waiting so other tests can release workers
logger.info(
"All GPUs in use by other tests, waiting %.1fs for availability...",
poll_interval,
)
time.sleep(poll_interval)
def _launch_workers_unlocked(
self,
workers: list[WorkerIdentity],
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
allow_eviction: bool = True,
) -> list[ModelInstance] | None:
"""Internal launch logic. Caller must hold _lock.
Returns:
List of launched instances, empty list if no valid workers,
or None if GPUs unavailable (signals caller to wait and retry).
"""
if not workers:
return []
@@ -899,6 +1030,19 @@ class ModelPool:
len(available),
)
self._evict_for_gpus(total_gpus)
# Check again after eviction
available = self.allocator.available_gpus()
if len(available) < total_gpus:
# Still not enough - all workers are in use by other tests
# Return None to signal caller to wait and retry
logger.info(
"Still need %d GPUs, only %d available after eviction. "
"All workers in use by other tests.",
total_gpus,
len(available),
)
return None
else:
logger.warning(
"Need %d GPUs, only %d available. Skipping launch.",
@@ -976,11 +1120,15 @@ class ModelPool:
return self.get(model_id, mode).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()
"""Tear down all models.
Thread-safe: Protected by internal lock.
"""
with self._lock:
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

14
sgl-model-gateway/e2e_test/pyproject.toml
View File

@@ -9,7 +9,9 @@ dependencies = [
"grpcio-health-checking",
"httpx",
"openai",
"py", # Required for pytest-parallel with newer pytest versions
"pytest",
"pytest-parallel",
"pytest-rerunfailures",
]
@@ -23,8 +25,20 @@ testpaths = ["."]
markers = [
"e2e: mark test as end-to-end test requiring GPU workers",
"slow: mark test as slow-running",
"thread_unsafe: mark test as incompatible with parallel thread execution",
]
addopts = "-v -s"
# Explicitly disable live log to avoid "---- live log ----" dividers
# We configure logging manually in conftest.py
log_cli = false
# Parallel execution configuration:
# Use --workers 1 --tests-per-worker N to run N tests concurrently as threads
# within a single process. This enables true shared-worker parallelism where
# the session-scoped model_pool fixture is shared across all threads.
#
# Example usage:
# pytest --workers 1 --tests-per-worker 4 e2e_test/router/
#
# The thread-safe ModelPool and GPUAllocator classes enable safe concurrent
# access from multiple test threads.