From 8a45a9c6a92203247c3202c08dd0371ec5b047b4 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Thu, 8 Jan 2026 08:55:18 -0800 Subject: [PATCH] [smg][ci] fix model pool GPU cleanup and add startup reliability improvements (#16745) --- sgl-model-gateway/e2e_test/infra/constants.py | 6 +- .../e2e_test/infra/model_pool.py | 49 ++++- sgl-model-gateway/e2e_test/utils.py | 183 ------------------ 3 files changed, 51 insertions(+), 187 deletions(-) delete mode 100644 sgl-model-gateway/e2e_test/utils.py diff --git a/sgl-model-gateway/e2e_test/infra/constants.py b/sgl-model-gateway/e2e_test/infra/constants.py index e12f8bb34..f7f2c6fc3 100644 --- a/sgl-model-gateway/e2e_test/infra/constants.py +++ b/sgl-model-gateway/e2e_test/infra/constants.py @@ -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 = ( diff --git a/sgl-model-gateway/e2e_test/infra/model_pool.py b/sgl-model-gateway/e2e_test/infra/model_pool.py index 92bcb25ed..9d1db1f02 100644 --- a/sgl-model-gateway/e2e_test/infra/model_pool.py +++ b/sgl-model-gateway/e2e_test/infra/model_pool.py @@ -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( diff --git a/sgl-model-gateway/e2e_test/utils.py b/sgl-model-gateway/e2e_test/utils.py deleted file mode 100644 index b7b18f23e..000000000 --- a/sgl-model-gateway/e2e_test/utils.py +++ /dev/null @@ -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"))