[ROCm][CI] Explicitly tear down multimodal offline LLMs (#46999)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 29, 2026 at 02:59 UTC
db28ae2d078da82d01f8e7fad05fb27a52ce37ef
6 files changed
+91
-29
tests/conftest.py
+7
-1
@@ -1587,7 +1587,13 @@ class AssetHandler(http.server.BaseHTTPRequestHandler):
1587
self.send_header("Content-Type", ctype)
1588
self.send_header("Content-Length", str(len(data)))
1589
self.end_headers()
1590
- self.wfile.write(data)
1590
+ try:
1591
+ self.wfile.write(data)
1592
+ except (BrokenPipeError, ConnectionResetError) as e:
1593
+ logger.debug(
1594
+ "Client disconnected while serving test asset %s: %r", filename, e
1595
+ )
1596
+ self.close_connection = True
1597
1598
1599
def _find_free_port() -> int:
tests/entrypoints/multimodal/conftest.py
+72
@@ -1,5 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Callable, Iterator
4
+from contextlib import contextmanager
5
+from typing import Any
6
+
7
+import pytest
8
9
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
10
TEST_IMAGE_ASSETS = [
@@ -8,3 +13,70 @@ TEST_IMAGE_ASSETS = [
13
"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",
14
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
15
]
16
+
17
+
18
+def _shutdown_llm(llm: Any, gpu_memory_utilization: float) -> None:
19
+ from vllm.distributed import cleanup_dist_env_and_memory
20
+ from vllm.platforms import current_platform
21
+
22
+ try:
23
+ shutdown_timeout = 60.0 if current_platform.is_rocm() else None
24
+ llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout)
25
+ except Exception:
26
+ pass
27
+
28
+ del llm
29
+
30
+ try:
31
+ import torch
32
+
33
+ torch._dynamo.reset()
34
+ except Exception:
35
+ pass
36
+
37
+ cleanup_dist_env_and_memory()
38
+
39
+ if current_platform.is_rocm():
40
+ from tests.utils import wait_for_rocm_memory_to_settle
41
+
42
+ wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
43
+
44
+
45
+@contextmanager
46
+def managed_llm(*args: Any, **kwargs: Any) -> Iterator[Any]:
47
+ from vllm import LLM
48
+
49
+ llm = LLM(*args, **kwargs)
50
+ gpu_memory_utilization = (
51
+ llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
52
+ )
53
+ try:
54
+ yield llm
55
+ finally:
56
+ _shutdown_llm(llm, gpu_memory_utilization)
57
+
58
+
59
+def _make_managed_llm_factory() -> Iterator[Callable[..., Any]]:
60
+ from vllm import LLM
61
+
62
+ llms: list[tuple[Any, float]] = []
63
+
64
+ def make_llm(*args: Any, **kwargs: Any) -> Any:
65
+ llm = LLM(*args, **kwargs)
66
+ gpu_memory_utilization = (
67
+ llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
68
+ )
69
+ llms.append((llm, gpu_memory_utilization))
70
+ return llm
71
+
72
+ try:
73
+ yield make_llm
74
+ finally:
75
+ while llms:
76
+ llm, gpu_memory_utilization = llms.pop()
77
+ _shutdown_llm(llm, gpu_memory_utilization)
78
+
79
+
80
+@pytest.fixture
81
+def multimodal_llm_factory() -> Iterator[Callable[..., Any]]:
82
+ yield from _make_managed_llm_factory()
tests/entrypoints/multimodal/llm/test_chat.py
+2
-14
@@ -1,19 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import weakref
4
-
3
import pytest
4
5
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
8
-from vllm import LLM
9
-from vllm.distributed import cleanup_dist_env_and_memory
6
7
8
@pytest.fixture(scope="function")
13
-def vision_llm():
14
- # pytest caches the fixture so we use weakref.proxy to
15
- # enable garbage collection
16
- llm = LLM(
9
+def vision_llm(multimodal_llm_factory):
10
+ return multimodal_llm_factory(
11
model="microsoft/Phi-3.5-vision-instruct",
12
max_model_len=4096,
13
max_num_seqs=5,
@@ -23,12 +17,6 @@ def vision_llm():
17
seed=0,
18
)
19
26
- yield weakref.proxy(llm)
27
-
28
- del llm
29
-
30
- cleanup_dist_env_and_memory()
31
-
20
21
@pytest.mark.parametrize(
22
"image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True
tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py
+4
-2
@@ -69,6 +69,7 @@ def test_inject_into_mm_cache(
69
image_urls,
70
mm_processor_cache_type,
71
caplog_vllm,
72
+ multimodal_llm_factory,
73
):
74
"""Test that inject_into_mm_cache() injects pre-processed mm_kwargs into
75
the processor cache and MM cache hit metrics are updated correctly.
@@ -78,7 +79,7 @@ def test_inject_into_mm_cache(
79
2. Extract cached kwargs, call inject_into_mm_cache with a new hash,
80
then generate with a pre-rendered input -> verifies injection works
81
"""
81
- llm = LLM(
82
+ llm = multimodal_llm_factory(
83
model="llava-hf/llava-1.5-7b-hf",
84
max_model_len=4096,
85
max_num_seqs=5,
@@ -145,11 +146,12 @@ def test_inject_into_mm_cache(
146
def test_inject_into_mm_cache_without_cache(
147
num_gpus_available,
148
image_urls,
149
+ multimodal_llm_factory,
150
):
151
"""Test that inject_into_mm_cache works gracefully when processor cache
152
is disabled (mm_processor_cache_gb=0). Should not crash.
153
"""
152
- llm = LLM(
154
+ llm = multimodal_llm_factory(
155
model="llava-hf/llava-1.5-7b-hf",
156
max_model_len=4096,
157
max_num_seqs=5,
tests/entrypoints/multimodal/llm/test_mm_cache_stats.py
+2
-1
@@ -61,8 +61,9 @@ def test_mm_cache_stats(
61
image_urls,
62
mm_processor_cache_type,
63
caplog_vllm,
64
+ multimodal_llm_factory,
65
):
65
- llm = LLM(
66
+ llm = multimodal_llm_factory(
67
model="llava-hf/llava-1.5-7b-hf",
68
max_model_len=4096,
69
max_num_seqs=5,
tests/entrypoints/multimodal/llm/test_mm_embeds_only.py
+4
-11
@@ -1,13 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import weakref
5
-
4
import pytest
5
6
+from tests.entrypoints.multimodal.conftest import managed_llm
7
from vllm import LLM, SamplingParams
8
from vllm.assets.image import ImageAsset
10
-from vllm.distributed import cleanup_dist_env_and_memory
9
10
MODEL = "llava-hf/llava-1.5-7b-hf"
11
PROMPT = "USER: <image>\nDescribe this image briefly.\nASSISTANT:"
@@ -17,20 +15,15 @@ TEXT_ONLY_PROMPT = "USER: What is 2 + 2?\nASSISTANT:"
15
@pytest.fixture(scope="module")
16
def llm():
17
"""LLM with enable_mm_embeds=True and all modality limits zeroed out."""
20
- llm = LLM(
18
+ with managed_llm(
19
model=MODEL,
20
max_model_len=2048,
21
enforce_eager=True,
22
gpu_memory_utilization=0.8,
23
enable_mm_embeds=True,
24
limit_mm_per_prompt={"image": 0},
27
- )
28
-
29
- yield weakref.proxy(llm)
30
-
31
- del llm
32
-
33
- cleanup_dist_env_and_memory()
25
+ ) as llm:
26
+ yield llm
27
28
29
@pytest.mark.skip_global_cleanup