@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6f5478298d
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Use `aiohttp` connection pool for benchmarking (#21981)

Signed-off-by: Seiji Eicher <seiji@anyscale.com>

Seiji Eicher committed Aug 3, 2025 at 19:23 UTC 6f5478298ddd8e6aa330f171c70811f667b8699b
3 files changed +276 -247
vllm/benchmarks/lib/endpoint_request_func.py
+238 -241
@@ -50,6 +50,7 @@ class RequestFuncOutput:
50
51 async def async_request_openai_completions(
52 request_func_input: RequestFuncInput,
53 + session: aiohttp.ClientSession,
54 pbar: Optional[tqdm] = None,
55 ) -> RequestFuncOutput:
56 """The async request function for the OpenAI Completions API.
@@ -66,96 +67,94 @@ async def async_request_openai_completions(
67 ("completions", "profile")
68 ), "OpenAI Completions API URL must end with 'completions' or 'profile'."
69
69 - async with aiohttp.ClientSession(trust_env=True,
70 - timeout=AIOHTTP_TIMEOUT) as session:
71 - payload = {
72 - "model": request_func_input.model_name \
73 - if request_func_input.model_name else request_func_input.model,
74 - "prompt": request_func_input.prompt,
75 - "temperature": 0.0,
76 - "repetition_penalty": 1.0,
77 - "max_tokens": request_func_input.output_len,
78 - "logprobs": request_func_input.logprobs,
79 - "stream": True,
80 - "stream_options": {
81 - "include_usage": True,
82 - },
83 - }
84 - if request_func_input.ignore_eos:
85 - payload["ignore_eos"] = request_func_input.ignore_eos
86 - if request_func_input.extra_body:
87 - payload.update(request_func_input.extra_body)
88 - headers = {
89 - "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
90 - }
91 -
92 - output = RequestFuncOutput()
93 - output.prompt_len = request_func_input.prompt_len
94 -
95 - generated_text = ""
96 - st = time.perf_counter()
97 - most_recent_timestamp = st
98 - try:
99 - async with session.post(url=api_url, json=payload,
100 - headers=headers) as response:
101 - if response.status == 200:
102 - first_chunk_received = False
103 - async for chunk_bytes in response.content:
104 - chunk_bytes = chunk_bytes.strip()
105 - if not chunk_bytes:
106 - continue
107 - chunk_bytes = chunk_bytes.decode("utf-8")
108 - # NOTE: SSE comments (often used as pings) start with
109 - # a colon. These are not JSON data payload and should
110 - # be skipped.
111 - if chunk_bytes.startswith(":"):
112 - continue
113 -
114 - chunk = chunk_bytes.removeprefix("data: ")
115 -
116 - if chunk != "[DONE]":
117 - data = json.loads(chunk)
118 -
119 - # NOTE: Some completion API might have a last
120 - # usage summary response without a token so we
121 - # want to check a token was generated
122 - if choices := data.get("choices"):
123 - # Note that text could be empty here
124 - # e.g. for special tokens
125 - text = choices[0].get("text")
126 - timestamp = time.perf_counter()
127 - # First token
128 - if not first_chunk_received:
129 - first_chunk_received = True
130 - ttft = time.perf_counter() - st
131 - output.ttft = ttft
70 + payload = {
71 + "model": request_func_input.model_name \
72 + if request_func_input.model_name else request_func_input.model,
73 + "prompt": request_func_input.prompt,
74 + "temperature": 0.0,
75 + "repetition_penalty": 1.0,
76 + "max_tokens": request_func_input.output_len,
77 + "logprobs": request_func_input.logprobs,
78 + "stream": True,
79 + "stream_options": {
80 + "include_usage": True,
81 + },
82 + }
83 + if request_func_input.ignore_eos:
84 + payload["ignore_eos"] = request_func_input.ignore_eos
85 + if request_func_input.extra_body:
86 + payload.update(request_func_input.extra_body)
87 + headers = {
88 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
89 + }
90 +
91 + output = RequestFuncOutput()
92 + output.prompt_len = request_func_input.prompt_len
93 +
94 + generated_text = ""
95 + st = time.perf_counter()
96 + most_recent_timestamp = st
97 + try:
98 + async with session.post(url=api_url, json=payload,
99 + headers=headers) as response:
100 + if response.status == 200:
101 + first_chunk_received = False
102 + async for chunk_bytes in response.content:
103 + chunk_bytes = chunk_bytes.strip()
104 + if not chunk_bytes:
105 + continue
106 + chunk_bytes = chunk_bytes.decode("utf-8")
107 + # NOTE: SSE comments (often used as pings) start with
108 + # a colon. These are not JSON data payload and should
109 + # be skipped.
110 + if chunk_bytes.startswith(":"):
111 + continue
112 +
113 + chunk = chunk_bytes.removeprefix("data: ")
114 +
115 + if chunk != "[DONE]":
116 + data = json.loads(chunk)
117 +
118 + # NOTE: Some completion API might have a last
119 + # usage summary response without a token so we
120 + # want to check a token was generated
121 + if choices := data.get("choices"):
122 + # Note that text could be empty here
123 + # e.g. for special tokens
124 + text = choices[0].get("text")
125 + timestamp = time.perf_counter()
126 + # First token
127 + if not first_chunk_received:
128 + first_chunk_received = True
129 + ttft = time.perf_counter() - st
130 + output.ttft = ttft
131
133 - # Decoding phase
134 - else:
135 - output.itl.append(timestamp -
136 - most_recent_timestamp)
132 + # Decoding phase
133 + else:
134 + output.itl.append(timestamp -
135 + most_recent_timestamp)
136
138 - most_recent_timestamp = timestamp
139 - generated_text += text or ""
140 - elif usage := data.get("usage"):
141 - output.output_tokens = usage.get(
142 - "completion_tokens")
143 - if first_chunk_received:
144 - output.success = True
145 - else:
146 - output.success = False
147 - output.error = (
148 - "Never received a valid chunk to calculate TTFT."
149 - "This response will be marked as failed!")
150 - output.generated_text = generated_text
151 - output.latency = most_recent_timestamp - st
137 + most_recent_timestamp = timestamp
138 + generated_text += text or ""
139 + elif usage := data.get("usage"):
140 + output.output_tokens = usage.get(
141 + "completion_tokens")
142 + if first_chunk_received:
143 + output.success = True
144 else:
153 - output.error = response.reason or ""
145 output.success = False
155 - except Exception:
156 - output.success = False
157 - exc_info = sys.exc_info()
158 - output.error = "".join(traceback.format_exception(*exc_info))
146 + output.error = (
147 + "Never received a valid chunk to calculate TTFT."
148 + "This response will be marked as failed!")
149 + output.generated_text = generated_text
150 + output.latency = most_recent_timestamp - st
151 + else:
152 + output.error = response.reason or ""
153 + output.success = False
154 + except Exception:
155 + output.success = False
156 + exc_info = sys.exc_info()
157 + output.error = "".join(traceback.format_exception(*exc_info))
158
159 if pbar:
160 pbar.update(1)
@@ -164,45 +163,158 @@ async def async_request_openai_completions(
163
164 async def async_request_openai_chat_completions(
165 request_func_input: RequestFuncInput,
166 + session: aiohttp.ClientSession,
167 pbar: Optional[tqdm] = None,
168 ) -> RequestFuncOutput:
169 api_url = request_func_input.api_url
170 assert api_url.endswith(("chat/completions", "profile")), (
171 "OpenAI Chat Completions API URL must end with 'chat/completions'.")
172
173 - async with aiohttp.ClientSession(trust_env=True,
174 - timeout=AIOHTTP_TIMEOUT) as session:
175 - content = [{"type": "text", "text": request_func_input.prompt}]
176 - if request_func_input.multi_modal_content:
177 - content.append(request_func_input.multi_modal_content)
178 - payload = {
179 - "model":
180 - request_func_input.model_name
181 - if request_func_input.model_name else request_func_input.model,
182 - "messages": [
183 - {
184 - "role": "user",
185 - "content": content
186 - },
187 - ],
188 - "temperature":
189 - 0.0,
190 - "max_completion_tokens":
191 - request_func_input.output_len,
192 - "stream":
193 - True,
194 - "stream_options": {
195 - "include_usage": True,
173 + content = [{"type": "text", "text": request_func_input.prompt}]
174 + if request_func_input.multi_modal_content:
175 + content.append(request_func_input.multi_modal_content)
176 + payload = {
177 + "model":
178 + request_func_input.model_name
179 + if request_func_input.model_name else request_func_input.model,
180 + "messages": [
181 + {
182 + "role": "user",
183 + "content": content
184 },
197 - }
198 - if request_func_input.ignore_eos:
199 - payload["ignore_eos"] = request_func_input.ignore_eos
200 - if request_func_input.extra_body:
201 - payload.update(request_func_input.extra_body)
202 - headers = {
203 - "Content-Type": "application/json",
204 - "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
205 - }
185 + ],
186 + "temperature":
187 + 0.0,
188 + "max_completion_tokens":
189 + request_func_input.output_len,
190 + "stream":
191 + True,
192 + "stream_options": {
193 + "include_usage": True,
194 + },
195 + }
196 + if request_func_input.ignore_eos:
197 + payload["ignore_eos"] = request_func_input.ignore_eos
198 + if request_func_input.extra_body:
199 + payload.update(request_func_input.extra_body)
200 + headers = {
201 + "Content-Type": "application/json",
202 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
203 + }
204 +
205 + output = RequestFuncOutput()
206 + output.prompt_len = request_func_input.prompt_len
207 +
208 + generated_text = ""
209 + ttft = 0.0
210 + st = time.perf_counter()
211 + most_recent_timestamp = st
212 + try:
213 + async with session.post(url=api_url, json=payload,
214 + headers=headers) as response:
215 + if response.status == 200:
216 + async for chunk_bytes in response.content:
217 + chunk_bytes = chunk_bytes.strip()
218 + if not chunk_bytes:
219 + continue
220 + chunk_bytes = chunk_bytes.decode("utf-8")
221 + # NOTE: SSE comments (often used as pings) start with
222 + # a colon. These are not JSON data payload and should
223 + # be skipped.
224 + if chunk_bytes.startswith(":"):
225 + continue
226 +
227 + chunk = chunk_bytes.removeprefix("data: ")
228 +
229 + if chunk != "[DONE]":
230 + timestamp = time.perf_counter()
231 + data = json.loads(chunk)
232 +
233 + if choices := data.get("choices"):
234 + content = choices[0]["delta"].get("content")
235 + # First token
236 + if ttft == 0.0:
237 + ttft = timestamp - st
238 + output.ttft = ttft
239 +
240 + # Decoding phase
241 + else:
242 + output.itl.append(timestamp -
243 + most_recent_timestamp)
244 +
245 + generated_text += content or ""
246 + elif usage := data.get("usage"):
247 + output.output_tokens = usage.get(
248 + "completion_tokens")
249 +
250 + most_recent_timestamp = timestamp
251 +
252 + output.generated_text = generated_text
253 + output.success = True
254 + output.latency = most_recent_timestamp - st
255 + else:
256 + output.error = response.reason or ""
257 + output.success = False
258 + except Exception:
259 + output.success = False
260 + exc_info = sys.exc_info()
261 + output.error = "".join(traceback.format_exception(*exc_info))
262 +
263 + if pbar:
264 + pbar.update(1)
265 + return output
266 +
267 +
268 +async def async_request_openai_audio(
269 + request_func_input: RequestFuncInput,
270 + session: aiohttp.ClientSession,
271 + pbar: Optional[tqdm] = None,
272 +) -> RequestFuncOutput:
273 + # Lazy import without PlaceholderModule to avoid vllm dep.
274 + import soundfile
275 +
276 + api_url = request_func_input.api_url
277 + assert api_url.endswith(("transcriptions", "translations")), (
278 + "OpenAI Chat Completions API URL must end with 'transcriptions' ")
279 + "or `translations`."
280 +
281 + content = [{"type": "text", "text": request_func_input.prompt}]
282 + payload = {
283 + "model":
284 + request_func_input.model_name
285 + if request_func_input.model_name else request_func_input.model,
286 + "temperature":
287 + 0.0,
288 + "max_completion_tokens":
289 + request_func_input.output_len,
290 + "stream":
291 + True,
292 + "language":
293 + "en",
294 + # Flattened due to multipart/form-data
295 + "stream_include_usage":
296 + True,
297 + "stream_continuous_usage_stats":
298 + True,
299 + }
300 + if request_func_input.extra_body:
301 + payload.update(request_func_input.extra_body)
302 + headers = {
303 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
304 + }
305 +
306 + # Send audio file
307 + def to_bytes(y, sr):
308 + buffer = io.BytesIO()
309 + soundfile.write(buffer, y, sr, format="WAV")
310 + buffer.seek(0)
311 + return buffer
312 +
313 + with to_bytes(*request_func_input.multi_modal_content["audio"]) as f:
314 + form = aiohttp.FormData()
315 + form.add_field("file", f, content_type="audio/wav")
316 + for key, value in payload.items():
317 + form.add_field(key, str(value))
318
319 output = RequestFuncOutput()
320 output.prompt_len = request_func_input.prompt_len
@@ -212,28 +324,24 @@ async def async_request_openai_chat_completions(
324 st = time.perf_counter()
325 most_recent_timestamp = st
326 try:
215 - async with session.post(url=api_url, json=payload,
327 + async with session.post(url=api_url,
328 + data=form,
329 headers=headers) as response:
330 if response.status == 200:
331 async for chunk_bytes in response.content:
332 chunk_bytes = chunk_bytes.strip()
333 if not chunk_bytes:
334 continue
222 - chunk_bytes = chunk_bytes.decode("utf-8")
223 - # NOTE: SSE comments (often used as pings) start with
224 - # a colon. These are not JSON data payload and should
225 - # be skipped.
226 - if chunk_bytes.startswith(":"):
227 - continue
228 -
229 - chunk = chunk_bytes.removeprefix("data: ")
335
336 + chunk = chunk_bytes.decode("utf-8").removeprefix(
337 + "data: ")
338 if chunk != "[DONE]":
339 timestamp = time.perf_counter()
340 data = json.loads(chunk)
341
342 if choices := data.get("choices"):
236 - content = choices[0]["delta"].get("content")
343 + content = choices[0]["delta"].get(
344 + "content")
345 # First token
346 if ttft == 0.0:
347 ttft = timestamp - st
@@ -241,8 +349,8 @@ async def async_request_openai_chat_completions(
349
350 # Decoding phase
351 else:
244 - output.itl.append(timestamp -
245 - most_recent_timestamp)
352 + output.itl.append(
353 + timestamp - most_recent_timestamp)
354
355 generated_text += content or ""
356 elif usage := data.get("usage"):
@@ -267,117 +375,6 @@ async def async_request_openai_chat_completions(
375 return output
376
377
270 -async def async_request_openai_audio(
271 - request_func_input: RequestFuncInput,
272 - pbar: Optional[tqdm] = None,
273 -) -> RequestFuncOutput:
274 - # Lazy import without PlaceholderModule to avoid vllm dep.
275 - import soundfile
276 -
277 - api_url = request_func_input.api_url
278 - assert api_url.endswith(("transcriptions", "translations")), (
279 - "OpenAI Chat Completions API URL must end with 'transcriptions' ")
280 - "or `translations`."
281 -
282 - async with aiohttp.ClientSession(trust_env=True,
283 - timeout=AIOHTTP_TIMEOUT) as session:
284 - content = [{"type": "text", "text": request_func_input.prompt}]
285 - payload = {
286 - "model":
287 - request_func_input.model_name
288 - if request_func_input.model_name else request_func_input.model,
289 - "temperature":
290 - 0.0,
291 - "max_completion_tokens":
292 - request_func_input.output_len,
293 - "stream":
294 - True,
295 - "language":
296 - "en",
297 - # Flattened due to multipart/form-data
298 - "stream_include_usage":
299 - True,
300 - "stream_continuous_usage_stats":
301 - True,
302 - }
303 - if request_func_input.extra_body:
304 - payload.update(request_func_input.extra_body)
305 - headers = {
306 - "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
307 - }
308 -
309 - # Send audio file
310 - def to_bytes(y, sr):
311 - buffer = io.BytesIO()
312 - soundfile.write(buffer, y, sr, format="WAV")
313 - buffer.seek(0)
314 - return buffer
315 -
316 - with to_bytes(*request_func_input.multi_modal_content["audio"]) as f:
317 - form = aiohttp.FormData()
318 - form.add_field("file", f, content_type="audio/wav")
319 - for key, value in payload.items():
320 - form.add_field(key, str(value))
321 -
322 - output = RequestFuncOutput()
323 - output.prompt_len = request_func_input.prompt_len
324 -
325 - generated_text = ""
326 - ttft = 0.0
327 - st = time.perf_counter()
328 - most_recent_timestamp = st
329 - try:
330 - async with session.post(url=api_url,
331 - data=form,
332 - headers=headers) as response:
333 - if response.status == 200:
334 - async for chunk_bytes in response.content:
335 - chunk_bytes = chunk_bytes.strip()
336 - if not chunk_bytes:
337 - continue
338 -
339 - chunk = chunk_bytes.decode("utf-8").removeprefix(
340 - "data: ")
341 - if chunk != "[DONE]":
342 - timestamp = time.perf_counter()
343 - data = json.loads(chunk)
344 -
345 - if choices := data.get("choices"):
346 - content = choices[0]["delta"].get(
347 - "content")
348 - # First token
349 - if ttft == 0.0:
350 - ttft = timestamp - st
351 - output.ttft = ttft
352 -
353 - # Decoding phase
354 - else:
355 - output.itl.append(
356 - timestamp - most_recent_timestamp)
357 -
358 - generated_text += content or ""
359 - elif usage := data.get("usage"):
360 - output.output_tokens = usage.get(
361 - "completion_tokens")
362 -
363 - most_recent_timestamp = timestamp
364 -
365 - output.generated_text = generated_text
366 - output.success = True
367 - output.latency = most_recent_timestamp - st
368 - else:
369 - output.error = response.reason or ""
370 - output.success = False
371 - except Exception:
372 - output.success = False
373 - exc_info = sys.exc_info()
374 - output.error = "".join(traceback.format_exception(*exc_info))
375 -
376 - if pbar:
377 - pbar.update(1)
378 - return output
379 -
380 -
378 # TODO: Add more request functions for different API protocols.
379 ASYNC_REQUEST_FUNCS = {
380 "vllm": async_request_openai_completions,
vllm/benchmarks/lib/ready_checker.py
+3 -1
@@ -14,6 +14,7 @@ from .endpoint_request_func import RequestFuncInput, RequestFuncOutput
14 async def wait_for_endpoint(
15 request_func,
16 test_input: RequestFuncInput,
17 + session: aiohttp.ClientSession,
18 timeout_seconds: int = 600,
19 retry_interval: int = 5,
20 ) -> RequestFuncOutput:
@@ -55,7 +56,8 @@ async def wait_for_endpoint(
56
57 # ping the endpoint using request_func
58 try:
58 - output = await request_func(request_func_input=test_input)
59 + output = await request_func(
60 + request_func_input=test_input, session=session)
61 if output.success:
62 pbar.close()
63 return output
vllm/benchmarks/serve.py
+35 -5
@@ -28,6 +28,7 @@ from dataclasses import dataclass
28 from datetime import datetime
29 from typing import Any, Literal, Optional
30
31 +import aiohttp
32 import numpy as np
33 from tqdm.asyncio import tqdm
34 from transformers import PreTrainedTokenizerBase
@@ -338,6 +339,24 @@ async def benchmark(
339 else:
340 raise ValueError(f"Unknown endpoint_type: {endpoint_type}")
341
342 + # Reuses connections across requests to reduce TLS handshake overhead.
343 + connector = aiohttp.TCPConnector(
344 + limit=max_concurrency or 0,
345 + limit_per_host=max_concurrency or 0,
346 + ttl_dns_cache=300,
347 + use_dns_cache=True,
348 + keepalive_timeout=60,
349 + enable_cleanup_closed=True,
350 + force_close=False,
351 + ssl=("https://" in api_url),
352 + )
353 +
354 + session = aiohttp.ClientSession(
355 + connector=connector,
356 + trust_env=True,
357 + timeout=aiohttp.ClientTimeout(total=6 * 60 * 60),
358 + )
359 +
360 print("Starting initial single prompt test run...")
361 test_prompt, test_prompt_len, test_output_len, test_mm_content = (
362 input_requests[0].prompt,
@@ -361,7 +380,11 @@ async def benchmark(
380 )
381
382 test_output = await wait_for_endpoint(
364 - request_func, test_input, timeout_seconds=ready_check_timeout_sec)
383 + request_func,
384 + test_input,
385 + session,
386 + timeout_seconds=ready_check_timeout_sec,
387 + )
388 if not test_output.success:
389 raise ValueError(
390 "Initial test run failed - Please make sure benchmark arguments "
@@ -386,7 +409,8 @@ async def benchmark(
409 multi_modal_content=test_mm_content,
410 ignore_eos=ignore_eos,
411 extra_body=extra_body)
389 - profile_output = await request_func(request_func_input=profile_input)
412 + profile_output = await request_func(
413 + request_func_input=profile_input, session=session)
414 if profile_output.success:
415 print("Profiler started")
416
@@ -412,12 +436,14 @@ async def benchmark(
436 semaphore = (asyncio.Semaphore(max_concurrency)
437 if max_concurrency else None)
438
415 - async def limited_request_func(request_func_input, pbar):
439 + async def limited_request_func(request_func_input, session, pbar):
440 if semaphore is None:
441 return await request_func(request_func_input=request_func_input,
442 + session=session,
443 pbar=pbar)
444 async with semaphore:
420 - return await request_func(request_func_input=request_func_input,
445 + return await request_func(request_func_input=request_func_input,
446 + session=session,
447 pbar=pbar)
448
449 benchmark_start_time = time.perf_counter()
@@ -469,6 +495,7 @@ async def benchmark(
495 tasks.append(
496 asyncio.create_task(
497 limited_request_func(request_func_input=request_func_input,
498 + session=session,
499 pbar=pbar)))
500 outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks)
501
@@ -580,9 +607,12 @@ async def benchmark(
607 output_len=test_output_len,
608 logprobs=logprobs,
609 )
583 - profile_output = await request_func(request_func_input=profile_input)
610 + profile_output = await request_func(
611 + request_func_input=profile_input, session=session)
612 if profile_output.success:
613 print("Profiler stopped")
614 +
615 + await session.close()
616 return result
617
618