@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 183 lines 5.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import gc
5 import random
6 import string
7 import sys
8 import weakref
9
10 import pytest
11 import torch
12
13 from tests.models.registry import HF_EXAMPLE_MODELS
14 from vllm import LLM, SamplingParams
15 from vllm.distributed import cleanup_dist_env_and_memory
16 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
17 from vllm.platforms import current_platform
18 from vllm.utils.mem_constants import KiB_bytes, MiB_bytes
19 from vllm.utils.mem_utils import format_mib
20
21 MODEL_NAME = "Qwen/Qwen3-VL-4B-Instruct"
22 RANDOM_PREFIX_LEN = 100
23 TEST_IMAGE_NAMES = [
24 "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg",
25 "Grayscale_8bits_palette_sample_image.png",
26 ]
27 MAX_MODEL_LEN = 8192
28 REQUESTS_PER_ROUND = 4
29 WARMUP_ROUNDS = 2
30 MEASURED_ROUNDS = 16
31 GPU_GROWTH_THRESHOLD_MIB = 0
32 CPU_PEAK_GROWTH_THRESHOLD_MIB = 0
33
34 SAMPLING_PARAMS = SamplingParams(
35 temperature=0.0,
36 max_tokens=16,
37 )
38
39
40 def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]:
41 # Avoid obscuring memory leaks because of prefix caching
42 random_text = "".join(random.choices(string.ascii_uppercase, k=RANDOM_PREFIX_LEN))
43
44 return [
45 {
46 "role": "user",
47 "content": [
48 {
49 "type": "text",
50 "text": f"Ignore this random string: {random_text}",
51 },
52 {"type": "image_url", "image_url": {"url": image_url}},
53 {
54 "type": "text",
55 "text": "Describe this image in one short sentence.",
56 },
57 ],
58 }
59 ]
60
61
62 def _build_request_batch(
63 image_urls: list[str],
64 ) -> list[list[ChatCompletionMessageParam]]:
65 return [
66 _make_messages(image_urls[i % len(image_urls)])
67 for i in range(REQUESTS_PER_ROUND)
68 ]
69
70
71 def _ru_maxrss_bytes() -> int | None:
72 try:
73 import resource
74 except ImportError:
75 return None
76
77 rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
78 if rss <= 0:
79 return 0
80
81 # Linux reports kilobytes, macOS reports bytes.
82 return rss if sys.platform == "darwin" else rss * KiB_bytes
83
84
85 def _gpu_used_bytes() -> int:
86 torch.accelerator.synchronize()
87 free_bytes, total_bytes = torch.accelerator.get_memory_info()
88 return int(total_bytes - free_bytes)
89
90
91 def _format_mib(num_bytes: int | None) -> str:
92 if num_bytes is None:
93 return "n/a"
94
95 return f"{format_mib(num_bytes)} MiB"
96
97
98 @pytest.fixture(scope="function")
99 def llm(monkeypatch):
100 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
101
102 # pytest caches the fixture so we use weakref.proxy to
103 # enable garbage collection
104 llm_kwargs = dict(
105 model=MODEL_NAME,
106 enforce_eager=True,
107 max_model_len=MAX_MODEL_LEN,
108 max_num_seqs=REQUESTS_PER_ROUND,
109 limit_mm_per_prompt={"image": 1},
110 seed=0,
111 disable_log_stats=True,
112 gpu_memory_utilization=0.8,
113 )
114 if current_platform.is_rocm():
115 llm_kwargs["attention_backend"] = "TRITON_ATTN"
116
117 llm = LLM(**llm_kwargs)
118
119 yield weakref.proxy(llm)
120
121 del llm
122
123 cleanup_dist_env_and_memory()
124
125
126 @pytest.mark.core_model
127 @pytest.mark.parametrize("image_urls", [TEST_IMAGE_NAMES], indirect=True)
128 def test_no_memory_leak(llm, image_urls: list[str]) -> None:
129 model_info = HF_EXAMPLE_MODELS.find_hf_info(MODEL_NAME)
130 model_info.check_available_online(on_fail="skip")
131 model_info.check_transformers_version(on_fail="skip")
132
133 request_batch = _build_request_batch(image_urls)
134
135 # Establish a warmup baseline after model load and the first multimodal
136 # requests complete. Later rounds should remain near this steady state.
137 for _ in range(WARMUP_ROUNDS):
138 outputs = llm.chat(request_batch, sampling_params=SAMPLING_PARAMS)
139 assert len(outputs) == len(request_batch)
140 assert llm.llm_engine.get_num_unfinished_requests() == 0
141 del outputs
142
143 gc.collect()
144 warmup_gpu = _gpu_used_bytes()
145 warmup_cpu_peak = _ru_maxrss_bytes()
146
147 post_warmup_gpu_samples: list[int] = []
148 post_warmup_cpu_peak_samples: list[int] = []
149
150 for _ in range(MEASURED_ROUNDS):
151 outputs = llm.chat(request_batch, sampling_params=SAMPLING_PARAMS)
152 assert len(outputs) == len(request_batch)
153 assert llm.llm_engine.get_num_unfinished_requests() == 0
154 del outputs
155
156 gc.collect()
157 post_warmup_gpu_samples.append(_gpu_used_bytes())
158 cpu_peak = _ru_maxrss_bytes()
159 if cpu_peak is not None:
160 post_warmup_cpu_peak_samples.append(cpu_peak)
161
162 gpu_growth = max(post_warmup_gpu_samples) - warmup_gpu
163 gpu_threshold = GPU_GROWTH_THRESHOLD_MIB * MiB_bytes
164
165 assert gpu_growth <= gpu_threshold, (
166 "Qwen3-VL GPU memory kept growing after warmup. "
167 f"warmup_baseline={_format_mib(warmup_gpu)}, "
168 f"post_warmup_samples={[_format_mib(x) for x in post_warmup_gpu_samples]}, "
169 f"gpu_growth={_format_mib(gpu_growth)}, "
170 f"gpu_threshold={GPU_GROWTH_THRESHOLD_MIB} MiB"
171 )
172
173 if warmup_cpu_peak is not None and post_warmup_cpu_peak_samples:
174 cpu_peak_growth = max(post_warmup_cpu_peak_samples) - warmup_cpu_peak
175 cpu_threshold = CPU_PEAK_GROWTH_THRESHOLD_MIB * MiB_bytes
176
177 assert cpu_peak_growth <= cpu_threshold, (
178 "Qwen3-VL CPU peak RSS kept growing after warmup. "
179 f"warmup_ru_maxrss={_format_mib(warmup_cpu_peak)}, "
180 f"post_warmup_ru_maxrss={[_format_mib(x) for x in post_warmup_cpu_peak_samples]}, " # noqa: E501
181 f"cpu_peak_growth={_format_mib(cpu_peak_growth)}, "
182 f"cpu_peak_threshold={CPU_PEAK_GROWTH_THRESHOLD_MIB} MiB"
183 )