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

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