[ROCm][CI] Explicitly tear down multimodal offline LLMs (#46999)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

tests/conftest.py

@@ -1587,7 +1587,13 @@ class AssetHandler(http.server.BaseHTTPRequestHandler):
         self.send_header("Content-Type", ctype)
         self.send_header("Content-Length", str(len(data)))
         self.end_headers()
-        self.wfile.write(data)
+        try:
+            self.wfile.write(data)
+        except (BrokenPipeError, ConnectionResetError) as e:
+            logger.debug(
+                "Client disconnected while serving test asset %s: %r", filename, e
+            )
+            self.close_connection = True
 
 
 def _find_free_port() -> int:

tests/entrypoints/multimodal/conftest.py

@@ -1,5 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from typing import Any
+
+import pytest
 
 # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
 TEST_IMAGE_ASSETS = [
@@ -8,3 +13,70 @@ TEST_IMAGE_ASSETS = [
     "1280px-Venn_diagram_rgb.svg.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
     "RGBA_comp.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
 ]
+
+
+def _shutdown_llm(llm: Any, gpu_memory_utilization: float) -> None:
+    from vllm.distributed import cleanup_dist_env_and_memory
+    from vllm.platforms import current_platform
+
+    try:
+        shutdown_timeout = 60.0 if current_platform.is_rocm() else None
+        llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout)
+    except Exception:
+        pass
+
+    del llm
+
+    try:
+        import torch
+
+        torch._dynamo.reset()
+    except Exception:
+        pass
+
+    cleanup_dist_env_and_memory()
+
+    if current_platform.is_rocm():
+        from tests.utils import wait_for_rocm_memory_to_settle
+
+        wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
+
+
+@contextmanager
+def managed_llm(*args: Any, **kwargs: Any) -> Iterator[Any]:
+    from vllm import LLM
+
+    llm = LLM(*args, **kwargs)
+    gpu_memory_utilization = (
+        llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
+    )
+    try:
+        yield llm
+    finally:
+        _shutdown_llm(llm, gpu_memory_utilization)
+
+
+def _make_managed_llm_factory() -> Iterator[Callable[..., Any]]:
+    from vllm import LLM
+
+    llms: list[tuple[Any, float]] = []
+
+    def make_llm(*args: Any, **kwargs: Any) -> Any:
+        llm = LLM(*args, **kwargs)
+        gpu_memory_utilization = (
+            llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
+        )
+        llms.append((llm, gpu_memory_utilization))
+        return llm
+
+    try:
+        yield make_llm
+    finally:
+        while llms:
+            llm, gpu_memory_utilization = llms.pop()
+            _shutdown_llm(llm, gpu_memory_utilization)
+
+
+@pytest.fixture
+def multimodal_llm_factory() -> Iterator[Callable[..., Any]]:
+    yield from _make_managed_llm_factory()

tests/entrypoints/multimodal/llm/test_chat.py

@@ -1,19 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import weakref
-
 import pytest
 
 from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
-from vllm import LLM
-from vllm.distributed import cleanup_dist_env_and_memory
 
 
 @pytest.fixture(scope="function")
-def vision_llm():
-    # pytest caches the fixture so we use weakref.proxy to
-    # enable garbage collection
-    llm = LLM(
+def vision_llm(multimodal_llm_factory):
+    return multimodal_llm_factory(
         model="microsoft/Phi-3.5-vision-instruct",
         max_model_len=4096,
         max_num_seqs=5,
@@ -23,12 +17,6 @@ def vision_llm():
         seed=0,
     )
 
-    yield weakref.proxy(llm)
-
-    del llm
-
-    cleanup_dist_env_and_memory()
-
 
 @pytest.mark.parametrize(
     "image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True

tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py

@@ -69,6 +69,7 @@ def test_inject_into_mm_cache(
     image_urls,
     mm_processor_cache_type,
     caplog_vllm,
+    multimodal_llm_factory,
 ):
     """Test that inject_into_mm_cache() injects pre-processed mm_kwargs into
     the processor cache and MM cache hit metrics are updated correctly.
@@ -78,7 +79,7 @@ def test_inject_into_mm_cache(
     2. Extract cached kwargs, call inject_into_mm_cache with a new hash,
        then generate with a pre-rendered input -> verifies injection works
     """
-    llm = LLM(
+    llm = multimodal_llm_factory(
         model="llava-hf/llava-1.5-7b-hf",
         max_model_len=4096,
         max_num_seqs=5,
@@ -145,11 +146,12 @@ def test_inject_into_mm_cache(
 def test_inject_into_mm_cache_without_cache(
     num_gpus_available,
     image_urls,
+    multimodal_llm_factory,
 ):
     """Test that inject_into_mm_cache works gracefully when processor cache
     is disabled (mm_processor_cache_gb=0). Should not crash.
     """
-    llm = LLM(
+    llm = multimodal_llm_factory(
         model="llava-hf/llava-1.5-7b-hf",
         max_model_len=4096,
         max_num_seqs=5,

tests/entrypoints/multimodal/llm/test_mm_cache_stats.py

@@ -61,8 +61,9 @@ def test_mm_cache_stats(
     image_urls,
     mm_processor_cache_type,
     caplog_vllm,
+    multimodal_llm_factory,
 ):
-    llm = LLM(
+    llm = multimodal_llm_factory(
         model="llava-hf/llava-1.5-7b-hf",
         max_model_len=4096,
         max_num_seqs=5,

tests/entrypoints/multimodal/llm/test_mm_embeds_only.py

@@ -1,13 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import weakref
-
 import pytest
 
+from tests.entrypoints.multimodal.conftest import managed_llm
 from vllm import LLM, SamplingParams
 from vllm.assets.image import ImageAsset
-from vllm.distributed import cleanup_dist_env_and_memory
 
 MODEL = "llava-hf/llava-1.5-7b-hf"
 PROMPT = "USER: <image>\nDescribe this image briefly.\nASSISTANT:"
@@ -17,20 +15,15 @@ TEXT_ONLY_PROMPT = "USER: What is 2 + 2?\nASSISTANT:"
 @pytest.fixture(scope="module")
 def llm():
     """LLM with enable_mm_embeds=True and all modality limits zeroed out."""
-    llm = LLM(
+    with managed_llm(
         model=MODEL,
         max_model_len=2048,
         enforce_eager=True,
         gpu_memory_utilization=0.8,
         enable_mm_embeds=True,
         limit_mm_per_prompt={"image": 0},
-    )
-
-    yield weakref.proxy(llm)
-
-    del llm
-
-    cleanup_dist_env_and_memory()
+    ) as llm:
+        yield llm
 
 
 @pytest.mark.skip_global_cleanup