refactor(benchmarks): add type annotations to wait_for_endpoint parameters (#25218)

Signed-off-by: samzong <samzong.lu@gmail.com>
This commit is contained in:
samzong
2025-09-20 00:36:52 +08:00
committed by GitHub
parent aed16879a9
commit ce75e15373
2 changed files with 16 additions and 4 deletions

View File

@@ -8,8 +8,9 @@ import os
import sys
import time
import traceback
from collections.abc import Awaitable
from dataclasses import dataclass, field
from typing import Optional, Union
from typing import Optional, Protocol, Union
import aiohttp
from tqdm.asyncio import tqdm
@@ -92,6 +93,16 @@ class RequestFuncOutput:
start_time: float = 0.0
class RequestFunc(Protocol):
def __call__(
self,
request_func_input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: Optional[tqdm] = None,
) -> Awaitable[RequestFuncOutput]:
...
async def async_request_openai_completions(
request_func_input: RequestFuncInput,
session: aiohttp.ClientSession,
@@ -507,7 +518,7 @@ async def async_request_openai_embeddings(
# TODO: Add more request functions for different API protocols.
ASYNC_REQUEST_FUNCS = {
ASYNC_REQUEST_FUNCS: dict[str, RequestFunc] = {
"vllm": async_request_openai_completions,
"openai": async_request_openai_completions,
"openai-chat": async_request_openai_chat_completions,