[CI] Reclaim GPU memory between model initialization tests (#57362)
Signed-off-by: khluu <khluu000@gmail.com>
Kevin H. Luu committed
Sep 18, 2026 at 17:04 UTC
729ebac4983e1510e035bc579142b0fc210a49a3
1 file changed
+44
-1
tests/models/test_initialization.py
+44
-1
@@ -1,12 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import gc
5
+from collections.abc import Generator
6
from functools import partial
7
from unittest.mock import patch
8
9
import pytest
10
+import torch
11
12
from vllm import LLM
13
+from vllm.logger import init_logger
14
+from vllm.platforms import current_platform
15
from vllm.utils.mem_constants import GiB_bytes
16
from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
17
from vllm.v1.core.kv_cache_utils import (
@@ -15,7 +20,12 @@ from vllm.v1.core.kv_cache_utils import (
20
)
21
from vllm.v1.engine.core import EngineCore as V1EngineCore
22
18
-from ..utils import create_new_process_for_each_test, requires_spawn_multiprocessing
23
+from ..utils import (
24
+ create_new_process_for_each_test,
25
+ requires_spawn_multiprocessing,
26
+ wait_for_gpu_memory_to_clear,
27
+ wait_for_memory_to_settle,
28
+)
29
from .registry import (
30
_TRANSFORMERS_BACKEND_MODELS,
31
AUTO_EXAMPLE_MODELS,
@@ -24,6 +34,8 @@ from .registry import (
34
)
35
from .utils import dummy_hf_overrides
36
37
+logger = init_logger(__name__)
38
+
39
# This minimal list of model architectures is smaller than the total list of
40
# supported models. The intention is that in the "typical" regression testing
41
# scenario, we only test initializing these models. This subset was chosen
@@ -53,6 +65,37 @@ OTHER_MODEL_ARCH_LIST = set(HF_EXAMPLE_MODELS.get_supported_archs()) - set(
65
)
66
67
68
+@pytest.fixture(autouse=True)
69
+def cleanup_gpu_memory_between_tests() -> Generator[None, None, None]:
70
+ """Reclaim GPU memory left over by each test's forked subprocess.
71
+
72
+ These tests are known to leave behind GPU memory (see the comment at
73
+ ``gpu_memory_utilization=0.80`` below). Each parametrized case builds an
74
+ LLM in a forked subprocess whose VRAM is only reclaimed asynchronously
75
+ when the subprocess exits, so the leftovers can OOM a later test in this
76
+ module (observed with the heavier multimodal archs, e.g.
77
+ InternVLChatModel). Mirrors ``clean_gpu_memory_between_tests`` in the
78
+ top-level conftest, but unconditionally for this module.
79
+ """
80
+ yield
81
+ gc.collect()
82
+ if not torch.cuda.is_available():
83
+ return
84
+ torch.accelerator.empty_cache()
85
+ if current_platform.is_rocm() or current_platform.is_xpu():
86
+ wait_for_memory_to_settle()
87
+ return
88
+ try:
89
+ wait_for_gpu_memory_to_clear(
90
+ devices=list(range(torch.accelerator.device_count())),
91
+ threshold_ratio=0.1,
92
+ )
93
+ except Exception as e:
94
+ # NVML is not usable on MIG slices (raises NVMLError_NoPermission);
95
+ # cleanup is best-effort, so never fail a test over it.
96
+ logger.info("Failed to clean GPU memory: %s", e)
97
+
98
+
99
@create_new_process_for_each_test()
100
def can_initialize(
101
model_arch: str, monkeypatch: pytest.MonkeyPatch, EXAMPLE_MODELS: HfExampleModels