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

.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()

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."""
+    """Configure clean logging to stdout with timestamps and thread info.
-    fmt = "%(asctime)s.%(msecs)03d [%(name)s] %(message)s"
+
+    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",

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
+    This is done once per session with thread-safe initialization to support
-    GPU memory conflicts in CI environments.
+    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
+    # Thread-safe initialization - only one thread computes embeddings
-    model_path = MODEL_SPECS.get("embedding", {}).get("model")
+    with _hf_embeddings_lock:
-    if model_path is None:
+        if _hf_embeddings_cache is not None:
-        pytest.skip("Embedding model not found in MODEL_SPECS")
+            return _hf_embeddings_cache
 
-    logger.info(
+        from infra.model_specs import MODEL_SPECS
-        "Pre-computing HuggingFace reference embeddings (CPU) for %s", model_path
-    )
 
-    # Flatten all test texts for semantic similarity
+        # Get model path from MODEL_SPECS for the embedding model
-    all_semantic_texts = []
+        model_path = MODEL_SPECS.get("embedding", {}).get("model")
-    for text_set in SEMANTIC_TEST_SETS:
+        if model_path is None:
-        all_semantic_texts.extend(text_set)
+            pytest.skip("Embedding model not found in MODEL_SPECS")
 
-    # Get relevance test texts
+        logger.info(
-    query = f"Instruct: Given a search query, retrieve relevant passages that answer the query\nQuery: {RELEVANCE_TEST_DATA['sample_query']}"
+            "Pre-computing HuggingFace reference embeddings (CPU) for %s", model_path
-    docs = get_input_texts(RELEVANCE_TEST_DATA)
+        )
 
-    # Compute all reference embeddings at once
+        # Flatten all test texts for semantic similarity
-    hf_semantic = get_hf_st_embeddings(all_semantic_texts, model_path)
+        all_semantic_texts = []
-    hf_query = get_hf_st_embeddings(query, model_path)
+        for text_set in SEMANTIC_TEST_SETS:
-    hf_docs = get_hf_st_embeddings(docs, model_path)
+            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 {
+        # Compute all reference embeddings at once
-        "semantic": hf_semantic,
+        hf_semantic = get_hf_st_embeddings(all_semantic_texts, model_path)
-        "query": hf_query,
+        hf_query = get_hf_st_embeddings(query, model_path)
-        "docs": hf_docs,
+        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

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",

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
+    available_gpus = _count_gpus_without_cuda()
-    try:
-        import torch
-
-        if torch.cuda.is_available():
-            available_gpus = torch.cuda.device_count()
-    except ImportError:
-        pass
-
     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

@@ -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:
+    # Thread-safe initialization: use lock to ensure only one thread creates the pool
-        return _model_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
+        # Check if we should skip model startup
-    if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
+        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)
+            logger.info("%s is set, skipping model pool startup", ENV_SKIP_MODEL_POOL)
-        _model_pool = ModelPool(GPUAllocator(gpus=[]))
+            _model_pool = ModelPool(GPUAllocator(gpus=[]))
-        return _model_pool
+            return _model_pool
 
-    # Determine requirements from scanned tests or env vars
+        # Determine requirements from scanned tests or env vars
-    models_env = os.environ.get(ENV_MODELS, "")
+        models_env = os.environ.get(ENV_MODELS, "")
-    backends_env = os.environ.get(ENV_BACKENDS, "")
+        backends_env = os.environ.get(ENV_BACKENDS, "")
 
-    if models_env or backends_env:
+        if models_env or backends_env:
-        # Use env var overrides
+            # Use env var overrides
-        models = (
+            models = (
-            {m.strip() for m in models_env.split(",") if m.strip()}
+                {m.strip() for m in models_env.split(",") if m.strip()}
-            if models_env
+                if models_env
-            else {DEFAULT_MODEL}
+                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
+        # Log final GPU allocation summary
-        backend_modes: set[ConnectionMode] = set()
+        logger.info(_model_pool.allocator.summary())
-        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
+        # Register cleanup with atexit instead of request.addfinalizer
-        if not backend_modes:
+        # This is critical for pytest-parallel where multiple threads share
-            backend_modes = {ConnectionMode.HTTP}
+        # 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

@@ -130,34 +130,15 @@ def _setup_pd_backend(
     import openai
     from infra import ConnectionMode, Gateway, WorkerIdentity, WorkerType
 
-    # Check PD requirements
+    logger.info("Setting up PD backend for model %s", model_id)
-    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")
 
     # Get PD configuration from workers marker
     num_prefill = workers_config.get("prefill") or 1
     num_decode = workers_config.get("decode") or 1
-
+    logger.info("PD config: %d prefill, %d decode workers", num_prefill, num_decode)
-    # 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}"
-        )
 
     # 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 in prefills and decodes are now acquired
-    all_workers = prefills + decodes
-    for worker in all_workers:
-        worker.acquire()
 
-    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)
+            # get_workers_by_type auto-acquires all returned workers
-            existing_for_mode = [w for w in existing if w.mode == connection_mode]
+            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],

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."""
+        """Get the slot assigned to a specific model.
-        for slot in self.slots:
+
-            if slot.assigned_model == model_id:
+        Thread-safe: Protected by internal lock.
-                return slot
+        """
-        return None
+        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:
+        with self._lock:
-            self._used_gpus.discard(gpu_id)
+            for gpu_id in gpu_ids:
-        # Remove slots that used these GPUs
+                self._used_gpus.discard(gpu_id)
-        self.slots = [s for s in self.slots if not any(g in gpu_ids for g in s.gpu_ids)]
+            # Remove slots that used these GPUs
-        logger.info("Released GPUs %s, now used: %s", gpu_ids, self._used_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."""
+        """Return a summary of GPU allocations.
-        lines = ["GPU Allocation Summary:"]
+
-        lines.append(f"  Total GPUs: {len(self.gpus)}")
+        Thread-safe: Protected by internal lock.
-        lines.append(f"  Used GPUs: {sorted(self._used_gpus)}")
+        """
-        lines.append(f"  Allocated Slots: {len(self.slots)}")
+        with self._lock:
-        for slot in self.slots:
+            lines = ["GPU Allocation Summary:"]
-            lines.append(
+            lines.append(f"  Total GPUs: {len(self.gpus)}")
-                f"    - {slot.assigned_model}: GPUs {slot.gpu_ids} "
+            lines.append(f"  Used GPUs: {sorted(self._used_gpus)}")
-                f"({slot.total_memory_gb:.1f}GB) port={slot.port}"
+            lines.append(f"  Allocated Slots: {len(self.slots)}")
-            )
+            for slot in self.slots:
-        return "\n".join(lines)
+                lines.append(
+                    f"    - {slot.assigned_model}: GPUs {slot.gpu_ids} "
+                    f"({slot.total_memory_gb:.1f}GB) port={slot.port}"
+                )
+            return "\n".join(lines)

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:
+        Returns:
-            RuntimeError: If not enough GPUs after eviction.
+            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(
+            logger.info(
-                f"Cannot launch {model_id}: need {required_gpus} GPUs, "
+                "Cannot launch %s: need %d GPUs, only %d available after eviction "
-                f"only {len(available)} 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 [
+        with self._lock:
-            inst
+            workers = [
-            for inst in self.instances.values()
+                inst
-            if inst.model_id == model_id and inst.worker_type == worker_type
+                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."""
+        """Tear down all models.
-        logger.info("Shutting down model pool (%d instances)", len(self.instances))
+
-        for instance in self.instances.values():
+        Thread-safe: Protected by internal lock.
-            instance.terminate()
+        """
-        self.instances.clear()
+        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

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.