[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>
Kevin H. Luu committed
Sep 18, 2026 at 20:30 UTC
9d75dd454431ccbb3c249b172cf525a5ddc745ee
5 files changed
+119
-2
tests/models/language/pooling/conftest.py
+28
@@ -2,12 +2,40 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Pytest configuration for vLLM language generation tests."""
4
5
+import gc
6
+
7
+import pytest
8
import torch
9
10
+from tests.utils import wait_for_memory_to_settle
11
+from vllm.logger import init_logger
12
from vllm.platforms import current_platform
13
14
+logger = init_logger(__name__)
15
+
16
17
def pytest_sessionstart(session):
18
"""Configure ROCm-specific settings before test session starts."""
19
if current_platform.is_rocm():
20
torch.set_float32_matmul_precision("high")
21
+
22
+
23
+@pytest.fixture(autouse=True)
24
+def release_gpu_memory_between_tests():
25
+ """Reclaim cached GPU memory between tests in this shard.
26
+
27
+ These tests start many vLLM/HF engines sequentially in one pytest
28
+ process; without explicit reclamation between tests, VRAM fragmentation
29
+ accumulates and later tests fail at engine startup ("Engine core
30
+ initialization failed").
31
+ """
32
+ yield
33
+ gc.collect()
34
+ if torch.accelerator.is_available():
35
+ torch.accelerator.empty_cache()
36
+ try:
37
+ wait_for_memory_to_settle()
38
+ except ValueError as e:
39
+ # Longer-lived (class/module-scoped) engine fixtures may legitimately
40
+ # still hold VRAM; nothing further to reclaim in that case.
41
+ logger.info("Failed to clean GPU memory: %s", e)
tests/models/language/pooling/test_colbert.py
+5
-2
@@ -128,9 +128,12 @@ def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device
128
"""Download and return the ColBERT linear projection weight."""
129
from safetensors.torch import load_file
130
131
- from vllm.transformers_utils.repo_utils import hf_api
131
+ from vllm.transformers_utils.repo_utils import hf_api, with_retry
132
133
- path = hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"])
133
+ path = with_retry(
134
+ lambda: hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]),
135
+ f"Error downloading {hf_spec['weights_file']} from {model_name}",
136
+ )
137
weights = load_file(path)
138
return weights[hf_spec["weights_key"]].to(device)
139
tests/models/language/pooling/test_nomic_max_model_len.py
+18
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# ruff: noqa: SIM117
4
5
+from functools import partial
6
+
7
import pytest
8
9
from ...utils import EmbedModelInfo
@@ -26,6 +28,22 @@ original_max_position_embeddings = 2048
28
max_model_len = int(original_max_position_embeddings * factor)
29
30
31
+@pytest.fixture(scope="module", autouse=True)
32
+def predownload_models():
33
+ """Warm the HF cache so engine startup never hits transient hub errors."""
34
+ from vllm.transformers_utils.repo_utils import hf_api, with_retry
35
+
36
+ for model_info in MODELS:
37
+ with_retry(
38
+ partial(
39
+ hf_api().snapshot_download,
40
+ model_info.name,
41
+ revision=model_info.revision,
42
+ ),
43
+ f"Error downloading {model_info.name}",
44
+ )
45
+
46
+
47
@pytest.mark.parametrize("model_info", MODELS)
48
def test_default(model_info, vllm_runner):
49
with vllm_runner(
tests/models/language/pooling_mteb_test/conftest.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Pytest configuration for vLLM pooling MTEB tests."""
4
+
5
+import gc
6
+
7
+import pytest
8
+import torch
9
+
10
+from tests.utils import wait_for_memory_to_settle
11
+from vllm.logger import init_logger
12
+
13
+logger = init_logger(__name__)
14
+
15
+
16
+@pytest.fixture(autouse=True)
17
+def release_gpu_memory_between_tests():
18
+ """Reclaim cached GPU memory between tests in this shard.
19
+
20
+ These tests start many vLLM/HF engines sequentially in one pytest
21
+ process; without explicit reclamation between tests, VRAM fragmentation
22
+ accumulates and later tests fail at engine startup ("Engine core
23
+ initialization failed").
24
+ """
25
+ yield
26
+ gc.collect()
27
+ if torch.accelerator.is_available():
28
+ torch.accelerator.empty_cache()
29
+ try:
30
+ wait_for_memory_to_settle()
31
+ except ValueError as e:
32
+ # Longer-lived (class/module-scoped) engine fixtures may legitimately
33
+ # still hold VRAM; nothing further to reclaim in that case.
34
+ logger.info("Failed to clean GPU memory: %s", e)
tests/models/multimodal/pooling/conftest.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Pytest configuration for vLLM multimodal pooling tests."""
4
+
5
+import gc
6
+
7
+import pytest
8
+import torch
9
+
10
+from tests.utils import wait_for_memory_to_settle
11
+from vllm.logger import init_logger
12
+
13
+logger = init_logger(__name__)
14
+
15
+
16
+@pytest.fixture(autouse=True)
17
+def release_gpu_memory_between_tests():
18
+ """Reclaim cached GPU memory between tests in this shard.
19
+
20
+ These tests start many vLLM/HF engines sequentially in one pytest
21
+ process; without explicit reclamation between tests, VRAM fragmentation
22
+ accumulates and later tests fail at engine startup ("Engine core
23
+ initialization failed").
24
+ """
25
+ yield
26
+ gc.collect()
27
+ if torch.accelerator.is_available():
28
+ torch.accelerator.empty_cache()
29
+ try:
30
+ wait_for_memory_to_settle()
31
+ except ValueError as e:
32
+ # Longer-lived (class/module-scoped) engine fixtures may legitimately
33
+ # still hold VRAM; nothing further to reclaim in that case.
34
+ logger.info("Failed to clean GPU memory: %s", e)