From 6037267f5bc14105965d00ecd712e8aa5e7aaff4 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Wed, 7 Jan 2026 13:25:41 -0800 Subject: [PATCH] [smg][ci] Add thread safety to ModelPool and GPUAllocator (#16674) --- .github/workflows/pr-test-rust.yml | 8 +- sgl-model-gateway/e2e_test/conftest.py | 42 +++- .../e2e_test/embeddings/test_correctness.py | 70 ++++--- .../e2e_test/fixtures/__init__.py | 4 + sgl-model-gateway/e2e_test/fixtures/hooks.py | 72 ++++++- sgl-model-gateway/e2e_test/fixtures/pool.py | 174 +++++++++------- .../e2e_test/fixtures/setup_backend.py | 86 ++++---- .../e2e_test/infra/gpu_allocator.py | 71 +++++-- .../e2e_test/infra/model_pool.py | 188 ++++++++++++++++-- sgl-model-gateway/e2e_test/pyproject.toml | 14 ++ 10 files changed, 534 insertions(+), 195 deletions(-) diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 5218123d2..6bbc75528 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -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() diff --git a/sgl-model-gateway/e2e_test/conftest.py b/sgl-model-gateway/e2e_test/conftest.py index b144ba627..da2c2b2ec 100644 --- a/sgl-model-gateway/e2e_test/conftest.py +++ b/sgl-model-gateway/e2e_test/conftest.py @@ -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", diff --git a/sgl-model-gateway/e2e_test/embeddings/test_correctness.py b/sgl-model-gateway/e2e_test/embeddings/test_correctness.py index 229863a50..00edb263b 100644 --- a/sgl-model-gateway/e2e_test/embeddings/test_correctness.py +++ b/sgl-model-gateway/e2e_test/embeddings/test_correctness.py @@ -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 diff --git a/sgl-model-gateway/e2e_test/fixtures/__init__.py b/sgl-model-gateway/e2e_test/fixtures/__init__.py index b391945ab..211c18e55 100644 --- a/sgl-model-gateway/e2e_test/fixtures/__init__.py +++ b/sgl-model-gateway/e2e_test/fixtures/__init__.py @@ -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", diff --git a/sgl-model-gateway/e2e_test/fixtures/hooks.py b/sgl-model-gateway/e2e_test/fixtures/hooks.py index 47cc36e42..e22180305 100644 --- a/sgl-model-gateway/e2e_test/fixtures/hooks.py +++ b/sgl-model-gateway/e2e_test/fixtures/hooks.py @@ -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}") diff --git a/sgl-model-gateway/e2e_test/fixtures/pool.py b/sgl-model-gateway/e2e_test/fixtures/pool.py index 98d19f231..2ee1beea9 100644 --- a/sgl-model-gateway/e2e_test/fixtures/pool.py +++ b/sgl-model-gateway/e2e_test/fixtures/pool.py @@ -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 diff --git a/sgl-model-gateway/e2e_test/fixtures/setup_backend.py b/sgl-model-gateway/e2e_test/fixtures/setup_backend.py index 078a3a7d9..53bdd4ed0 100644 --- a/sgl-model-gateway/e2e_test/fixtures/setup_backend.py +++ b/sgl-model-gateway/e2e_test/fixtures/setup_backend.py @@ -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], diff --git a/sgl-model-gateway/e2e_test/infra/gpu_allocator.py b/sgl-model-gateway/e2e_test/infra/gpu_allocator.py index 8fb5f2343..ca7b9bb96 100644 --- a/sgl-model-gateway/e2e_test/infra/gpu_allocator.py +++ b/sgl-model-gateway/e2e_test/infra/gpu_allocator.py @@ -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) diff --git a/sgl-model-gateway/e2e_test/infra/model_pool.py b/sgl-model-gateway/e2e_test/infra/model_pool.py index 1c51d2c7c..2b78c5ed9 100644 --- a/sgl-model-gateway/e2e_test/infra/model_pool.py +++ b/sgl-model-gateway/e2e_test/infra/model_pool.py @@ -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 diff --git a/sgl-model-gateway/e2e_test/pyproject.toml b/sgl-model-gateway/e2e_test/pyproject.toml index 00fb6a1d6..b39094284 100644 --- a/sgl-model-gateway/e2e_test/pyproject.toml +++ b/sgl-model-gateway/e2e_test/pyproject.toml @@ -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.