[CI] Deflake pooling shards with GPU teardown fixtures between tests (#57371)

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Kimi <noreply@moonshot.cn>

tests/models/language/pooling/conftest.py

@@ -2,12 +2,40 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Pytest configuration for vLLM language generation tests."""
 
+import gc
+
+import pytest
 import torch
 
+from tests.utils import wait_for_memory_to_settle
+from vllm.logger import init_logger
 from vllm.platforms import current_platform
 
+logger = init_logger(__name__)
+
 
 def pytest_sessionstart(session):
     """Configure ROCm-specific settings before test session starts."""
     if current_platform.is_rocm():
         torch.set_float32_matmul_precision("high")
+
+
+@pytest.fixture(autouse=True)
+def release_gpu_memory_between_tests():
+    """Reclaim cached GPU memory between tests in this shard.
+
+    These tests start many vLLM/HF engines sequentially in one pytest
+    process; without explicit reclamation between tests, VRAM fragmentation
+    accumulates and later tests fail at engine startup ("Engine core
+    initialization failed").
+    """
+    yield
+    gc.collect()
+    if torch.accelerator.is_available():
+        torch.accelerator.empty_cache()
+    try:
+        wait_for_memory_to_settle()
+    except ValueError as e:
+        # Longer-lived (class/module-scoped) engine fixtures may legitimately
+        # still hold VRAM; nothing further to reclaim in that case.
+        logger.info("Failed to clean GPU memory: %s", e)

tests/models/language/pooling/test_colbert.py

@@ -128,9 +128,12 @@ def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device
     """Download and return the ColBERT linear projection weight."""
     from safetensors.torch import load_file
 
-    from vllm.transformers_utils.repo_utils import hf_api
+    from vllm.transformers_utils.repo_utils import hf_api, with_retry
 
-    path = hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"])
+    path = with_retry(
+        lambda: hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]),
+        f"Error downloading {hf_spec['weights_file']} from {model_name}",
+    )
     weights = load_file(path)
     return weights[hf_spec["weights_key"]].to(device)
 

tests/models/language/pooling/test_nomic_max_model_len.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # ruff: noqa: SIM117
 
+from functools import partial
+
 import pytest
 
 from ...utils import EmbedModelInfo
@@ -26,6 +28,22 @@ original_max_position_embeddings = 2048
 max_model_len = int(original_max_position_embeddings * factor)
 
 
+@pytest.fixture(scope="module", autouse=True)
+def predownload_models():
+    """Warm the HF cache so engine startup never hits transient hub errors."""
+    from vllm.transformers_utils.repo_utils import hf_api, with_retry
+
+    for model_info in MODELS:
+        with_retry(
+            partial(
+                hf_api().snapshot_download,
+                model_info.name,
+                revision=model_info.revision,
+            ),
+            f"Error downloading {model_info.name}",
+        )
+
+
 @pytest.mark.parametrize("model_info", MODELS)
 def test_default(model_info, vllm_runner):
     with vllm_runner(

tests/models/language/pooling_mteb_test/conftest.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pytest configuration for vLLM pooling MTEB tests."""
+
+import gc
+
+import pytest
+import torch
+
+from tests.utils import wait_for_memory_to_settle
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def release_gpu_memory_between_tests():
+    """Reclaim cached GPU memory between tests in this shard.
+
+    These tests start many vLLM/HF engines sequentially in one pytest
+    process; without explicit reclamation between tests, VRAM fragmentation
+    accumulates and later tests fail at engine startup ("Engine core
+    initialization failed").
+    """
+    yield
+    gc.collect()
+    if torch.accelerator.is_available():
+        torch.accelerator.empty_cache()
+    try:
+        wait_for_memory_to_settle()
+    except ValueError as e:
+        # Longer-lived (class/module-scoped) engine fixtures may legitimately
+        # still hold VRAM; nothing further to reclaim in that case.
+        logger.info("Failed to clean GPU memory: %s", e)

tests/models/multimodal/pooling/conftest.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pytest configuration for vLLM multimodal pooling tests."""
+
+import gc
+
+import pytest
+import torch
+
+from tests.utils import wait_for_memory_to_settle
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def release_gpu_memory_between_tests():
+    """Reclaim cached GPU memory between tests in this shard.
+
+    These tests start many vLLM/HF engines sequentially in one pytest
+    process; without explicit reclamation between tests, VRAM fragmentation
+    accumulates and later tests fail at engine startup ("Engine core
+    initialization failed").
+    """
+    yield
+    gc.collect()
+    if torch.accelerator.is_available():
+        torch.accelerator.empty_cache()
+    try:
+        wait_for_memory_to_settle()
+    except ValueError as e:
+        # Longer-lived (class/module-scoped) engine fixtures may legitimately
+        # still hold VRAM; nothing further to reclaim in that case.
+        logger.info("Failed to clean GPU memory: %s", e)