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

[smg][ci] fix model pool GPU cleanup and add startup reliability improvements (#16745)

Browse Source
This commit is contained in:
Simo Lin
2026-01-08 08:55:18 -08:00
committed by GitHub
parent aecd5f5f3e
commit 8a45a9c6a9
3 changed files with 51 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
sgl-model-gateway/e2e_test/infra/constants.py
View File

@@ -57,7 +57,11 @@ DEFAULT_HOST = "127.0.0.1"
# Timeouts (seconds)
DEFAULT_STARTUP_TIMEOUT = 300
DEFAULT_ROUTER_TIMEOUT = 60
HEALTH_CHECK_INTERVAL = 5
HEALTH_CHECK_INTERVAL = 2 # Check every 2s (was 5s)
# Model loading configuration
INITIAL_GRACE_PERIOD = 30 # Wait before first health check (model loading time)
LAUNCH_STAGGER_DELAY = 5 # Delay between launching multiple workers
# Retry configuration
MAX_RETRY_ATTEMPTS = (

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

@@ -21,6 +21,8 @@ from .constants import (
DEFAULT_STARTUP_TIMEOUT,
ENV_SHOW_WORKER_LOGS,
HEALTH_CHECK_INTERVAL,
INITIAL_GRACE_PERIOD,
LAUNCH_STAGGER_DELAY,
LOCAL_MODES,
ConnectionMode,
WorkerType,
@@ -361,6 +363,7 @@ class ModelPool:
logger.info("Detected InfiniBand device: %s", ib_device)
deferred: list[str] = []
launched_count = 0
# Process requirements in order - all workers treated uniformly
for identity in valid_requirements:
@@ -395,6 +398,14 @@ class ModelPool:
# Each prefill worker needs its own bootstrap port for PD communication
bootstrap_port = get_open_port() if identity.is_prefill else None
# Stagger launches to avoid resource contention during model loading
if launched_count > 0 and LAUNCH_STAGGER_DELAY > 0:
logger.info(
"Staggering launch by %ds to reduce resource contention",
LAUNCH_STAGGER_DELAY,
)
time.sleep(LAUNCH_STAGGER_DELAY)
# Launch the worker
self._launch_model(
model_id=identity.model_id,
@@ -407,6 +418,7 @@ class ModelPool:
),
instance_key=identity.key,
)
launched_count += 1
# Log deferred workers
if deferred:
@@ -559,6 +571,14 @@ class ModelPool:
self._startup_timeout,
)
# Initial grace period to allow models to load before health checks
if INITIAL_GRACE_PERIOD > 0:
logger.info(
"Waiting %ds for initial model loading before health checks...",
INITIAL_GRACE_PERIOD,
)
time.sleep(INITIAL_GRACE_PERIOD)
while pending and (time.time() - start_time) < self._startup_timeout:
check_count += 1
elapsed = time.time() - start_time
@@ -579,6 +599,8 @@ class ModelPool:
stderr = instance.process.stderr.read()
if stderr:
logger.error("Stderr: %s", stderr.decode()[-2000:])
# Evict dead instance and release GPUs
self._evict_instance(key)
pending.discard(key)
continue
@@ -614,10 +636,31 @@ class ModelPool:
self._startup_timeout,
pending,
)
# Terminate failed instances
# Log stderr from failed workers for debugging
for key in pending:
self.instances[key].terminate()
del self.instances[key]
instance = self.instances.get(key)
if instance and instance.process.stderr:
try:
import select
# Use select for non-blocking read with short timeout
# to avoid hanging if worker is unresponsive
ready, _, _ = select.select(
[instance.process.stderr], [], [], 0.1
)
if ready:
stderr = instance.process.stderr.read()
if stderr:
logger.error(
"[%s] Last stderr output:\n%s",
key,
stderr.decode(errors="replace")[-3000:],
)
except Exception as e:
logger.error("[%s] Could not read stderr: %s", key, e)
# Terminate failed instances and release their GPUs
for key in pending:
self._evict_instance(key)
else:
elapsed = time.time() - start_time
logger.info(

183
sgl-model-gateway/e2e_test/utils.py
View File

@@ -1,183 +0,0 @@
"""Consolidated utilities for E2E tests.
This module provides common utilities used across E2E tests:
- Tokenizer loading (get_tokenizer)
- Test base classes (CustomTestCase for unittest compatibility)
- Model path resolution
- Process management utilities
Import examples:
from utils import get_tokenizer, CustomTestCase
from utils import DEFAULT_MODEL_PATH, DEFAULT_TIMEOUT
"""
from __future__ import annotations
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING
logger = logging.getLogger(__name__)
# Re-export commonly used items from submodules
from infra import kill_process_tree # noqa: F401
from infra.model_specs import ( # noqa: F401; Default model paths
DEFAULT_EMBEDDING_MODEL_PATH,
DEFAULT_ENABLE_THINKING_MODEL_PATH,
DEFAULT_GPT_OSS_MODEL_PATH,
DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_MODEL_PATH,
DEFAULT_QWEN_FUNCTION_CALLING_MODEL_PATH,
DEFAULT_REASONING_MODEL_PATH,
DEFAULT_SMALL_MODEL_PATH,
MODEL_SPECS,
ROUTER_LOCAL_MODEL_PATH,
_resolve_model_path,
)
# =============================================================================
# Tokenizer Utilities
# =============================================================================
# Lazy import transformers to avoid import errors in environments without it
_transformers_available = None
_AutoTokenizer = None
_PreTrainedTokenizer = None
_PreTrainedTokenizerBase = None
_PreTrainedTokenizerFast = None
def _ensure_transformers():
"""Lazy load transformers module."""
global _transformers_available, _AutoTokenizer
global _PreTrainedTokenizer, _PreTrainedTokenizerBase, _PreTrainedTokenizerFast
if _transformers_available is not None:
return _transformers_available
try:
from transformers import (
AutoTokenizer,
PreTrainedTokenizer,
PreTrainedTokenizerBase,
PreTrainedTokenizerFast,
)
_AutoTokenizer = AutoTokenizer
_PreTrainedTokenizer = PreTrainedTokenizer
_PreTrainedTokenizerBase = PreTrainedTokenizerBase
_PreTrainedTokenizerFast = PreTrainedTokenizerFast
_transformers_available = True
except ImportError:
_transformers_available = False
return _transformers_available
def check_gguf_file(model_path: str) -> bool:
"""Check if the model path points to a GGUF file."""
if not isinstance(model_path, str):
return False
return model_path.endswith(".gguf")
def is_remote_url(path: str) -> bool:
"""Check if the path is a remote URL."""
if not isinstance(path, str):
return False
return path.startswith("http://") or path.startswith("https://")
def get_tokenizer(
tokenizer_name: str,
*args,
tokenizer_mode: str = "auto",
trust_remote_code: bool = False,
tokenizer_revision: str | None = None,
**kwargs,
):
"""Gets a tokenizer for the given model name via Huggingface.
Args:
tokenizer_name: Name or path of the tokenizer
tokenizer_mode: Mode for tokenizer loading ("auto", "slow")
trust_remote_code: Whether to trust remote code
tokenizer_revision: Specific revision to use
**kwargs: Additional arguments passed to AutoTokenizer.from_pretrained
Returns:
Loaded tokenizer instance
Raises:
ImportError: If transformers is not installed
RuntimeError: If tokenizer loading fails
"""
if not _ensure_transformers():
raise ImportError(
"transformers is required for tokenizer utilities. "
"Install with: pip install transformers"
)
if tokenizer_mode == "slow":
if kwargs.get("use_fast", False):
raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
kwargs["use_fast"] = False
# Handle special model name mapping
if tokenizer_name == "mistralai/Devstral-Small-2505":
tokenizer_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
is_gguf = check_gguf_file(tokenizer_name)
if is_gguf:
kwargs["gguf_file"] = tokenizer_name
tokenizer_name = str(Path(tokenizer_name).parent)
try:
tokenizer = _AutoTokenizer.from_pretrained(
tokenizer_name,
*args,
trust_remote_code=trust_remote_code,
tokenizer_revision=tokenizer_revision,
**kwargs,
)
except TypeError as e:
err_msg = (
"Failed to load the tokenizer. If you are running a model with "
"a custom tokenizer, please set the --trust-remote-code flag."
)
raise RuntimeError(err_msg) from e
if not isinstance(tokenizer, _PreTrainedTokenizerFast):
logger.warning(
"Using a slow tokenizer. This might cause a performance "
"degradation. Consider using a fast tokenizer instead."
)
return tokenizer
def get_tokenizer_from_processor(processor):
"""Extract tokenizer from a processor object."""
if not _ensure_transformers():
raise ImportError("transformers is required for tokenizer utilities.")
if isinstance(processor, _PreTrainedTokenizerBase):
return processor
return processor.tokenizer
# =============================================================================
# Environment Utilities
# =============================================================================
def is_ci_environment() -> bool:
"""Check if running in CI environment."""
ci_vars = ["CI", "GITHUB_ACTIONS", "JENKINS_URL", "GITLAB_CI", "CIRCLECI"]
return any(os.environ.get(var) for var in ci_vars)
def get_test_timeout() -> int:
"""Get test timeout from environment or default (600 seconds)."""
return int(os.environ.get("E2E_TEST_TIMEOUT", "600"))