Add reference counting to ModelInstance for parallel test safety (#16672)

sgl-model-gateway/e2e_test/fixtures/pool.py

@@ -151,6 +151,7 @@ def model_client(request: pytest.FixtureRequest, model_pool: "ModelPool"):
         def test_chat(model_client):
             response = model_client.chat.completions.create(...)
     """
+    import openai
     from infra import PARAM_MODEL
 
     marker = request.node.get_closest_marker(PARAM_MODEL)
@@ -163,10 +164,23 @@ def model_client(request: pytest.FixtureRequest, model_pool: "ModelPool"):
     model_id = marker.args[0]
 
     try:
-        return model_pool.get_client(model_id)
+        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",
+    )
+
+    yield client
+
+    # Release reference to allow eviction
+    instance.release()
+
 
 @pytest.fixture
 def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> str:
@@ -189,6 +203,14 @@ def model_base_url(request: pytest.FixtureRequest, model_pool: "ModelPool") -> s
     model_id = marker.args[0]
 
     try:
-        return model_pool.get_base_url(model_id)
+        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
+    instance.release()

sgl-model-gateway/e2e_test/fixtures/setup_backend.py

@@ -212,6 +212,11 @@ def _setup_pd_backend(
         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()
+
     model_path = prefills[0].model_path if prefills else None
 
     # Launch PD gateway
@@ -244,6 +249,9 @@ def _setup_pd_backend(
     finally:
         logger.info("Tearing down PD gateway")
         gateway.shutdown()
+        # Release references to allow eviction
+        for worker in all_workers:
+            worker.release()
 
 
 def _setup_local_backend(
@@ -260,6 +268,7 @@ def _setup_local_backend(
     from infra import Gateway, WorkerIdentity, WorkerType
 
     num_workers = workers_config.get("count") or 1
+    instances: list = []  # Track instances for reference counting
 
     try:
         if num_workers > 1:
@@ -290,8 +299,13 @@ def _setup_local_backend(
             model_path = instances[0].model_path
         else:
             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))
 
@@ -324,6 +338,9 @@ def _setup_local_backend(
     finally:
         logger.info("Tearing down gateway for %s backend", backend_name)
         gateway.shutdown()
+        # Release references to allow eviction
+        for inst in instances:
+            inst.release()
 
 
 def _setup_cloud_backend(backend_name: str):
@@ -382,6 +399,9 @@ def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
     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],
@@ -392,3 +412,5 @@ def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
         yield gateway
     finally:
         gateway.shutdown()
+        # Release reference to allow eviction
+        instance.release()

sgl-model-gateway/e2e_test/infra/model_pool.py

@@ -5,8 +5,9 @@ from __future__ import annotations
 import logging
 import os
 import subprocess
+import threading
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 import httpx
@@ -102,6 +103,10 @@ class ModelInstance:
     last_used: float = 0.0  # Timestamp for MRU eviction
     _healthy: bool = False  # Track if initial health check passed
 
+    # Reference counting for safe parallel test execution
+    _ref_count: int = 0
+    _ref_lock: threading.Lock = field(default_factory=threading.Lock)
+
     @property
     def identity(self) -> WorkerIdentity:
         """Get the identity (model_id, mode, worker_type) of this instance."""
@@ -111,6 +116,44 @@ class ModelInstance:
             worker_type=self.worker_type,
         )
 
+    @property
+    def is_in_use(self) -> bool:
+        """Check if this instance has active references (tests using it)."""
+        with self._ref_lock:
+            return self._ref_count > 0
+
+    def acquire(self) -> None:
+        """Acquire a reference to this instance.
+
+        Call this before using the instance in a test to prevent eviction.
+        Must be paired with a release() call when done.
+        Also updates last_used timestamp atomically with ref count.
+        """
+        with self._ref_lock:
+            self._ref_count += 1
+            self.last_used = time.time()
+            logger.debug(
+                "Acquired reference to %s (ref_count=%d)", self.key, self._ref_count
+            )
+
+    def release(self) -> None:
+        """Release a reference to this instance.
+
+        Call this when done using the instance in a test.
+        """
+        with self._ref_lock:
+            if self._ref_count > 0:
+                self._ref_count -= 1
+                logger.debug(
+                    "Released reference to %s (ref_count=%d)",
+                    self.key,
+                    self._ref_count,
+                )
+            else:
+                logger.warning(
+                    "Attempted to release reference to %s with ref_count=0", self.key
+                )
+
     @property
     def worker_url(self) -> str:
         """URL to use when connecting router to this worker."""
@@ -629,8 +672,8 @@ class ModelPool:
 
         instance = self.instances[key]
 
-        # Update last_used timestamp
-        instance.last_used = time.time()
+        # Note: last_used is updated in acquire() which should be called by fixtures
+        # to prevent eviction during test execution
 
         # Verify worker is still alive and healthy
         if not instance.is_alive():
@@ -671,8 +714,15 @@ class ModelPool:
 
         # Sort by last_used descending (MRU eviction) - evict most recently used first
         # Store (dict_key, instance) tuples to preserve the actual key for eviction
+        # Note: Make a copy of items to avoid RuntimeError if dict is modified during iteration
         evictable: list[tuple[str, ModelInstance]] = []
-        for dict_key, inst in self.instances.items():
+        for dict_key, inst in list(self.instances.items()):
+            # Skip instances with active references (tests using them)
+            if inst.is_in_use:
+                logger.debug(
+                    "Skipping eviction of %s - has active references", dict_key
+                )
+                continue
             if exclude_worker_types is not None:
                 # Precise matching with worker types
                 # Must match model_id AND worker_type, mode is optional