minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
sgl-model-gateway/e2e_test/fixtures/__init__.py
+4
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",
sgl-model-gateway/e2e_test/fixtures/hooks.py
+63 -9
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}")
sgl-model-gateway/e2e_test/fixtures/pool.py
+100 -74
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
sgl-model-gateway/e2e_test/fixtures/setup_backend.py
+49 -37
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],