From 9d75dd454431ccbb3c249b172cf525a5ddc745ee Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 18 Sep 2026 20:30:02 -0700 Subject: [PATCH] [CI] Deflake pooling shards with GPU teardown fixtures between tests (#57371) Signed-off-by: khluu Co-authored-by: Kimi --- tests/models/language/pooling/conftest.py | 28 +++++++++++++++ tests/models/language/pooling/test_colbert.py | 7 ++-- .../pooling/test_nomic_max_model_len.py | 18 ++++++++++ .../language/pooling_mteb_test/conftest.py | 34 +++++++++++++++++++ tests/models/multimodal/pooling/conftest.py | 34 +++++++++++++++++++ 5 files changed, 119 insertions(+), 2 deletions(-) create mode 100644 tests/models/language/pooling_mteb_test/conftest.py create mode 100644 tests/models/multimodal/pooling/conftest.py diff --git a/tests/models/language/pooling/conftest.py b/tests/models/language/pooling/conftest.py index 0571f1cd6d..2da4c65133 100644 --- a/tests/models/language/pooling/conftest.py +++ b/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) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index 75323d54bf..8829558a06 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/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) diff --git a/tests/models/language/pooling/test_nomic_max_model_len.py b/tests/models/language/pooling/test_nomic_max_model_len.py index 6ea29f6d06..4ecbe77bc6 100644 --- a/tests/models/language/pooling/test_nomic_max_model_len.py +++ b/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( diff --git a/tests/models/language/pooling_mteb_test/conftest.py b/tests/models/language/pooling_mteb_test/conftest.py new file mode 100644 index 0000000000..5a31a5f415 --- /dev/null +++ b/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) diff --git a/tests/models/multimodal/pooling/conftest.py b/tests/models/multimodal/pooling/conftest.py new file mode 100644 index 0000000000..32f598b30a --- /dev/null +++ b/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)