[CI] Reclaim GPU memory between model initialization tests (#57362)

Signed-off-by: khluu <khluu000@gmail.com>

tests/models/test_initialization.py

@@ -1,12 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import gc
+from collections.abc import Generator
 from functools import partial
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm import LLM
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 from vllm.v1.core.kv_cache_utils import (
@@ -15,7 +20,12 @@ from vllm.v1.core.kv_cache_utils import (
 )
 from vllm.v1.engine.core import EngineCore as V1EngineCore
 
-from ..utils import create_new_process_for_each_test, requires_spawn_multiprocessing
+from ..utils import (
+    create_new_process_for_each_test,
+    requires_spawn_multiprocessing,
+    wait_for_gpu_memory_to_clear,
+    wait_for_memory_to_settle,
+)
 from .registry import (
     _TRANSFORMERS_BACKEND_MODELS,
     AUTO_EXAMPLE_MODELS,
@@ -24,6 +34,8 @@ from .registry import (
 )
 from .utils import dummy_hf_overrides
 
+logger = init_logger(__name__)
+
 # This minimal list of model architectures is smaller than the total list of
 # supported models. The intention is that in the "typical" regression testing
 # 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(
 )
 
 
+@pytest.fixture(autouse=True)
+def cleanup_gpu_memory_between_tests() -> Generator[None, None, None]:
+    """Reclaim GPU memory left over by each test's forked subprocess.
+
+    These tests are known to leave behind GPU memory (see the comment at
+    ``gpu_memory_utilization=0.80`` below). Each parametrized case builds an
+    LLM in a forked subprocess whose VRAM is only reclaimed asynchronously
+    when the subprocess exits, so the leftovers can OOM a later test in this
+    module (observed with the heavier multimodal archs, e.g.
+    InternVLChatModel). Mirrors ``clean_gpu_memory_between_tests`` in the
+    top-level conftest, but unconditionally for this module.
+    """
+    yield
+    gc.collect()
+    if not torch.cuda.is_available():
+        return
+    torch.accelerator.empty_cache()
+    if current_platform.is_rocm() or current_platform.is_xpu():
+        wait_for_memory_to_settle()
+        return
+    try:
+        wait_for_gpu_memory_to_clear(
+            devices=list(range(torch.accelerator.device_count())),
+            threshold_ratio=0.1,
+        )
+    except Exception as e:
+        # NVML is not usable on MIG slices (raises NVMLError_NoPermission);
+        # cleanup is best-effort, so never fail a test over it.
+        logger.info("Failed to clean GPU memory: %s", e)
+
+
 @create_new_process_for_each_test()
 def can_initialize(
     model_arch: str, monkeypatch: pytest.MonkeyPatch, EXAMPLE_MODELS: HfExampleModels