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
refactor(benchmarks): add type annotations to wait_for_endpoint parameters (#25218)
Signed-off-by: samzong <samzong.lu@gmail.com>
samzong committed
Sep 20, 2025 at 00:36 UTC
ce75e153735363eca01ce67b9f69e7b9ea440c63
2 files changed
+16
-4
vllm/benchmarks/lib/endpoint_request_func.py
+13
-2
@@ -8,8 +8,9 @@ import os
8
import sys
9
import time
10
import traceback
11
+from collections.abc import Awaitable
12
from dataclasses import dataclass, field
12
-from typing import Optional, Union
13
+from typing import Optional, Protocol, Union
14
15
import aiohttp
16
from tqdm.asyncio import tqdm
@@ -92,6 +93,16 @@ class RequestFuncOutput:
93
start_time: float = 0.0
94
95
96
+class RequestFunc(Protocol):
97
+ def __call__(
98
+ self,
99
+ request_func_input: RequestFuncInput,
100
+ session: aiohttp.ClientSession,
101
+ pbar: Optional[tqdm] = None,
102
+ ) -> Awaitable[RequestFuncOutput]:
103
+ ...
104
+
105
+
106
async def async_request_openai_completions(
107
request_func_input: RequestFuncInput,
108
session: aiohttp.ClientSession,
@@ -507,7 +518,7 @@ async def async_request_openai_embeddings(
518
519
520
# TODO: Add more request functions for different API protocols.
510
-ASYNC_REQUEST_FUNCS = {
521
+ASYNC_REQUEST_FUNCS: dict[str, RequestFunc] = {
522
"vllm": async_request_openai_completions,
523
"openai": async_request_openai_completions,
524
"openai-chat": async_request_openai_chat_completions,
vllm/benchmarks/lib/ready_checker.py
+3
-2
@@ -8,11 +8,12 @@ import time
8
import aiohttp
9
from tqdm.asyncio import tqdm
10
11
-from .endpoint_request_func import RequestFuncInput, RequestFuncOutput
11
+from .endpoint_request_func import (RequestFunc, RequestFuncInput,
12
+ RequestFuncOutput)
13
14
15
async def wait_for_endpoint(
15
- request_func,
16
+ request_func: RequestFunc,
17
test_input: RequestFuncInput,
18
session: aiohttp.ClientSession,
19
timeout_seconds: int = 600,