1153 lines
40 KiB
Python
1153 lines
40 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""The request function for API endpoints."""
|
|
|
|
import codecs
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from collections.abc import Awaitable
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Literal, Protocol
|
|
|
|
import aiohttp
|
|
import regex as re
|
|
from tqdm.asyncio import tqdm
|
|
|
|
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
|
|
|
|
|
|
class StreamedResponseHandler:
|
|
"""Handles streaming HTTP responses by accumulating chunks until complete
|
|
messages are available."""
|
|
|
|
def __init__(self):
|
|
self.buffer = ""
|
|
self._decoder = codecs.getincrementaldecoder("utf-8")()
|
|
|
|
def add_chunk(self, chunk_bytes: bytes) -> list[str]:
|
|
"""Add a chunk of bytes to the buffer and return any complete
|
|
messages."""
|
|
chunk_str = self._decoder.decode(chunk_bytes)
|
|
self.buffer += chunk_str
|
|
|
|
messages = []
|
|
|
|
# Split by double newlines (SSE message separator)
|
|
while "\n\n" in self.buffer:
|
|
message, self.buffer = self.buffer.split("\n\n", 1)
|
|
message = message.strip()
|
|
if message:
|
|
messages.append(message)
|
|
|
|
# if self.buffer is not empty, check if it is a complete message
|
|
# by removing data: prefix and check if it is a valid JSON
|
|
if self.buffer.startswith("data: "):
|
|
message_content = self.buffer.removeprefix("data: ").strip()
|
|
if message_content == "[DONE]":
|
|
messages.append(self.buffer.strip())
|
|
self.buffer = ""
|
|
elif message_content:
|
|
try:
|
|
json.loads(message_content)
|
|
messages.append(self.buffer.strip())
|
|
self.buffer = ""
|
|
except json.JSONDecodeError:
|
|
# Incomplete JSON, wait for more chunks.
|
|
pass
|
|
|
|
return messages
|
|
|
|
|
|
@dataclass
|
|
class RequestFuncInput:
|
|
"""The input for the request function."""
|
|
|
|
prompt: str | list[str] | list[int] | list[list[int]] | list[dict[str, Any]]
|
|
api_url: str
|
|
prompt_len: int
|
|
output_len: int
|
|
model: str
|
|
model_name: str | None = None
|
|
logprobs: int | None = None
|
|
extra_headers: dict | None = None
|
|
extra_body: dict | None = None
|
|
multi_modal_content: dict[str, Any] | list[dict[str, Any]] | None = None
|
|
ignore_eos: bool = False
|
|
language: str | None = None
|
|
request_id: str | None = None
|
|
# Pre-built chat messages. When set, `async_request_openai_chat_completions`
|
|
# uses this list directly and skips building messages from `prompt` and
|
|
# `multi_modal_content`.
|
|
chat_messages: list[dict[str, Any]] | None = None
|
|
|
|
|
|
@dataclass
|
|
class RequestFuncOutput:
|
|
"""The output of the request function including metrics."""
|
|
|
|
generated_text: str = ""
|
|
success: bool = False
|
|
latency: float = 0.0
|
|
output_tokens: int = 0
|
|
ttft: float = 0.0 # Time to first token
|
|
itl: list[float] = field(default_factory=list) # list of inter-token latencies
|
|
tpot: float = 0.0 # avg next-token latencies
|
|
prompt_len: int = 0
|
|
error: str = ""
|
|
start_time: float = 0.0
|
|
# Time spent awaiting the benchmark client concurrency semaphore.
|
|
client_queue_time: float = 0.0
|
|
input_audio_duration: float = 0.0 # in seconds
|
|
num_input_sequences: int = 1
|
|
|
|
|
|
class RequestFunc(Protocol):
|
|
def __call__(
|
|
self,
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> Awaitable[RequestFuncOutput]: ...
|
|
|
|
|
|
def _validate_api_url(
|
|
api_url: str,
|
|
api_name: str,
|
|
expected_suffixes: str | set[str],
|
|
) -> None:
|
|
if isinstance(expected_suffixes, str):
|
|
expected_suffixes = {expected_suffixes}
|
|
|
|
if not api_url.endswith(tuple(expected_suffixes)):
|
|
raise ValueError(f"{api_name} URL must end with one of: {expected_suffixes}.")
|
|
|
|
|
|
def _update_payload_common(
|
|
payload: dict[str, Any],
|
|
request_func_input: RequestFuncInput,
|
|
) -> None:
|
|
if request_func_input.ignore_eos:
|
|
payload["ignore_eos"] = request_func_input.ignore_eos
|
|
if request_func_input.extra_body:
|
|
payload.update(request_func_input.extra_body)
|
|
|
|
|
|
def _update_headers_common(
|
|
headers: dict[str, Any],
|
|
request_func_input: RequestFuncInput,
|
|
) -> None:
|
|
if request_func_input.extra_headers:
|
|
headers |= request_func_input.extra_headers
|
|
if request_func_input.request_id:
|
|
headers["x-request-id"] = request_func_input.request_id
|
|
|
|
|
|
def _get_headers(content_type: str | None = None) -> dict[str, str]:
|
|
headers = {}
|
|
if content_type:
|
|
headers["Content-Type"] = content_type
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
return headers
|
|
|
|
|
|
async def async_request_profile(
|
|
api_url: str,
|
|
session: aiohttp.ClientSession,
|
|
extra_headers: dict | None = None,
|
|
) -> RequestFuncOutput:
|
|
_validate_api_url(
|
|
api_url,
|
|
"Profiling API",
|
|
{"start_profile", "stop_profile"},
|
|
)
|
|
|
|
headers = _get_headers()
|
|
if extra_headers:
|
|
headers |= extra_headers
|
|
|
|
output = RequestFuncOutput()
|
|
try:
|
|
async with session.post(url=api_url, headers=headers) as response:
|
|
output.success = response.status == 200
|
|
if not output.success:
|
|
output.error = response.reason or ""
|
|
except Exception:
|
|
exc_info = sys.exc_info()
|
|
output.error = "".join(traceback.format_exception(*exc_info))
|
|
|
|
return output
|
|
|
|
|
|
async def async_request_openai_completions(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
"""The async request function for the OpenAI Completions API.
|
|
|
|
Args:
|
|
request_func_input: The input for the request function.
|
|
session: The aiohttp session used to issue the request.
|
|
pbar: The progress bar to display the progress.
|
|
|
|
Returns:
|
|
The output of the request function.
|
|
|
|
"""
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "OpenAI Completions API", "completions")
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"prompt": request_func_input.prompt,
|
|
"repetition_penalty": 1.0,
|
|
"max_tokens": request_func_input.output_len,
|
|
"logprobs": request_func_input.logprobs,
|
|
"stream": True,
|
|
"stream_options": {
|
|
"include_usage": True,
|
|
},
|
|
}
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers()
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
output = RequestFuncOutput()
|
|
output.prompt_len = request_func_input.prompt_len
|
|
|
|
generated_text = ""
|
|
st = time.perf_counter()
|
|
output.start_time = st
|
|
most_recent_timestamp = st
|
|
try:
|
|
async with session.post(url=api_url, json=payload, headers=headers) as response:
|
|
if response.status == 200:
|
|
first_chunk_received = False
|
|
handler = StreamedResponseHandler()
|
|
|
|
async for chunk_bytes in response.content.iter_any():
|
|
if not chunk_bytes:
|
|
continue
|
|
|
|
messages = handler.add_chunk(chunk_bytes)
|
|
for message in messages:
|
|
# NOTE: SSE comments (often used as pings) start with
|
|
# a colon. These are not JSON data payload and should
|
|
# be skipped.
|
|
if message.startswith(":"):
|
|
continue
|
|
|
|
chunk = message.removeprefix("data: ")
|
|
|
|
if chunk != "[DONE]":
|
|
data = json.loads(chunk)
|
|
|
|
# NOTE: Some completion API might have a last
|
|
# usage summary response without a token so we
|
|
# want to check a token was generated
|
|
if choices := data.get("choices"):
|
|
# Note that text could be empty here
|
|
# e.g. for special tokens
|
|
text = choices[0].get("text")
|
|
timestamp = time.perf_counter()
|
|
# First token
|
|
if not first_chunk_received:
|
|
first_chunk_received = True
|
|
output.ttft = timestamp - st
|
|
|
|
# Decoding phase
|
|
else:
|
|
output.itl.append(timestamp - most_recent_timestamp)
|
|
|
|
most_recent_timestamp = timestamp
|
|
generated_text += text or ""
|
|
elif usage := data.get("usage"):
|
|
output.output_tokens = usage.get("completion_tokens")
|
|
if (pt := usage.get("prompt_tokens")) is not None:
|
|
output.prompt_len = pt
|
|
if first_chunk_received:
|
|
output.success = True
|
|
else:
|
|
output.success = False
|
|
output.error = (
|
|
"Never received a valid chunk to calculate TTFT."
|
|
"This response will be marked as failed!"
|
|
)
|
|
output.generated_text = generated_text
|
|
output.latency = most_recent_timestamp - st
|
|
else:
|
|
output.error = response.reason or ""
|
|
output.success = False
|
|
except Exception:
|
|
output.success = False
|
|
exc_info = sys.exc_info()
|
|
output.error = "".join(traceback.format_exception(*exc_info))
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
def _get_chat_content(
|
|
request_func_input: RequestFuncInput,
|
|
mm_position: Literal["first", "last"] = "last",
|
|
) -> list[dict[str, Any]]:
|
|
mm_contents: list[dict[str, Any]] = []
|
|
if request_func_input.multi_modal_content:
|
|
mm_content = request_func_input.multi_modal_content
|
|
if isinstance(mm_content, list):
|
|
mm_contents.extend(mm_content)
|
|
elif isinstance(mm_content, dict):
|
|
mm_contents.append(mm_content)
|
|
else:
|
|
raise TypeError(
|
|
"multi_modal_content must be a dict or list[dict] for openai-chat"
|
|
)
|
|
|
|
prompt = request_func_input.prompt
|
|
if (
|
|
isinstance(prompt, list)
|
|
and prompt
|
|
and all(
|
|
isinstance(item, dict) and isinstance(item.get("type"), str)
|
|
for item in prompt
|
|
)
|
|
):
|
|
prompt_dicts: list[dict[str, Any]] = prompt # type: ignore[assignment]
|
|
if mm_position == "first":
|
|
return mm_contents + prompt_dicts
|
|
|
|
return prompt_dicts + mm_contents
|
|
|
|
text_contents = [{"type": "text", "text": prompt}]
|
|
|
|
if mm_position == "first":
|
|
return mm_contents + text_contents
|
|
|
|
return text_contents + mm_contents
|
|
|
|
|
|
def _is_chat_messages(prompt: Any) -> bool:
|
|
if not isinstance(prompt, list):
|
|
return False
|
|
if not prompt:
|
|
return False
|
|
return all(
|
|
isinstance(item, dict)
|
|
and isinstance(item.get("role"), str)
|
|
and isinstance(item.get("content"), (str, list))
|
|
for item in prompt
|
|
)
|
|
|
|
|
|
def _get_chat_messages(
|
|
request_func_input: RequestFuncInput,
|
|
mm_position: Literal["first", "last"] = "last",
|
|
) -> list[dict[str, Any]]:
|
|
prompt = request_func_input.prompt
|
|
if _is_chat_messages(prompt):
|
|
return prompt # type: ignore[return-value]
|
|
|
|
return [
|
|
{
|
|
"role": "user",
|
|
"content": _get_chat_content(
|
|
request_func_input,
|
|
mm_position=mm_position,
|
|
),
|
|
}
|
|
]
|
|
|
|
|
|
async def async_request_openai_chat_completions(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
mm_position: Literal["first", "last"] = "last",
|
|
) -> RequestFuncOutput:
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions")
|
|
|
|
if request_func_input.chat_messages is not None:
|
|
messages = request_func_input.chat_messages
|
|
else:
|
|
messages = _get_chat_messages(request_func_input, mm_position=mm_position)
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"messages": messages,
|
|
"max_completion_tokens": request_func_input.output_len,
|
|
"stream": True,
|
|
"stream_options": {
|
|
"include_usage": True,
|
|
},
|
|
}
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers("application/json")
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
output = RequestFuncOutput()
|
|
output.prompt_len = request_func_input.prompt_len
|
|
|
|
generated_text = ""
|
|
st = time.perf_counter()
|
|
output.start_time = st
|
|
most_recent_timestamp = st
|
|
try:
|
|
async with session.post(url=api_url, json=payload, headers=headers) as response:
|
|
if response.status == 200:
|
|
first_chunk_received = False
|
|
handler = StreamedResponseHandler()
|
|
async for chunk_bytes in response.content.iter_any():
|
|
if not chunk_bytes:
|
|
continue
|
|
|
|
message_strings = handler.add_chunk(chunk_bytes)
|
|
for message in message_strings:
|
|
# NOTE: SSE comments (often used as pings) start with
|
|
# a colon. These are not JSON data payload and should
|
|
# be skipped.
|
|
if message.startswith(":"):
|
|
continue
|
|
|
|
chunk = message.removeprefix("data: ")
|
|
|
|
if chunk != "[DONE]":
|
|
timestamp = time.perf_counter()
|
|
data = json.loads(chunk)
|
|
|
|
if choices := data.get("choices"):
|
|
content = choices[0]["delta"].get("content")
|
|
# First token
|
|
if not first_chunk_received:
|
|
first_chunk_received = True
|
|
output.ttft = timestamp - st
|
|
|
|
# Decoding phase
|
|
else:
|
|
output.itl.append(timestamp - most_recent_timestamp)
|
|
|
|
generated_text += content or ""
|
|
# Only token chunks advance the request end;
|
|
# the trailing usage chunk carries no token.
|
|
most_recent_timestamp = timestamp
|
|
elif usage := data.get("usage"):
|
|
output.output_tokens = usage.get("completion_tokens")
|
|
if (pt := usage.get("prompt_tokens")) is not None:
|
|
output.prompt_len = pt
|
|
|
|
output.generated_text = generated_text
|
|
if first_chunk_received:
|
|
output.success = True
|
|
else:
|
|
output.success = False
|
|
output.error = (
|
|
"Never received a valid chunk to calculate TTFT."
|
|
"This response will be marked as failed!"
|
|
)
|
|
output.latency = most_recent_timestamp - st
|
|
else:
|
|
output.error = response.reason or ""
|
|
output.success = False
|
|
except Exception:
|
|
output.success = False
|
|
exc_info = sys.exc_info()
|
|
output.error = "".join(traceback.format_exception(*exc_info))
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
def _extract_sse_data(message: str) -> str | None:
|
|
"""Extract the ``data`` payload of a single SSE message.
|
|
|
|
The Responses API stream labels every message with an ``event:`` line
|
|
before its ``data:`` line, so the payload cannot be recovered by
|
|
stripping a leading ``data: `` the way the Completions and Chat
|
|
Completions streams allow.
|
|
|
|
Args:
|
|
message: One complete SSE message, without its trailing blank line.
|
|
|
|
Returns:
|
|
The payload, with multi-line payloads joined by newlines as required
|
|
by the SSE spec, or None if the message carries no ``data:`` line.
|
|
|
|
"""
|
|
# NOTE: split on "\n" rather than using splitlines(), which also breaks on
|
|
# U+2028, U+2029 and U+0085. Those are legal inside a JSON string and the
|
|
# server emits them unescaped, so splitlines() would truncate the payload.
|
|
data_lines = [
|
|
line.removeprefix("data:").removeprefix(" ")
|
|
for line in message.split("\n")
|
|
if line.startswith("data:")
|
|
]
|
|
if not data_lines:
|
|
return None
|
|
|
|
return "\n".join(data_lines)
|
|
|
|
|
|
# Responses events that carry generated tokens. Reasoning deltas are included
|
|
# because the server is already decoding when it emits them, so excluding them
|
|
# would report the whole reasoning phase as time to first token.
|
|
_RESPONSES_TOKEN_DELTA_EVENTS = frozenset(
|
|
{
|
|
"response.output_text.delta",
|
|
"response.reasoning_text.delta",
|
|
}
|
|
)
|
|
|
|
# Responses events that end a request. Truncation at the output length limit
|
|
# is a normal benchmark outcome: vLLM reports it as `response.completed` with
|
|
# `status: "incomplete"`, while the spec also defines a dedicated
|
|
# `response.incomplete` terminal event.
|
|
_RESPONSES_TERMINAL_EVENTS = frozenset(
|
|
{
|
|
"response.completed",
|
|
"response.incomplete",
|
|
}
|
|
)
|
|
|
|
# The spec's failure terminals. vLLM does not emit them today, but
|
|
# `vllm bench serve` can be pointed at other Responses servers.
|
|
_RESPONSES_ERROR_EVENTS = frozenset({"error", "response.failed"})
|
|
|
|
|
|
async def async_request_openai_responses(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "OpenAI Responses API", "responses")
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"input": request_func_input.prompt,
|
|
"max_output_tokens": request_func_input.output_len,
|
|
"stream": True,
|
|
# `store` defaults to true. A server started with
|
|
# VLLM_ENABLE_RESPONSES_API_STORE=1 keeps every stored response for
|
|
# the lifetime of the process, which would grow its heap over a long
|
|
# benchmark run.
|
|
"store": False,
|
|
}
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers("application/json")
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
output = RequestFuncOutput()
|
|
output.prompt_len = request_func_input.prompt_len
|
|
|
|
# Only output text is accumulated. Reasoning text is timed but not
|
|
# returned, matching how `openai-chat` handles `DeltaMessage.reasoning`.
|
|
generated_text = ""
|
|
stream_error: str | None = None
|
|
saw_terminal_event = False
|
|
first_chunk_received = False
|
|
st = time.perf_counter()
|
|
output.start_time = st
|
|
most_recent_timestamp = st
|
|
try:
|
|
async with session.post(url=api_url, json=payload, headers=headers) as response:
|
|
if response.status == 200:
|
|
handler = StreamedResponseHandler()
|
|
|
|
async for chunk_bytes in response.content.iter_any():
|
|
for message in handler.add_chunk(chunk_bytes):
|
|
# NOTE: `_extract_sse_data` returns None for SSE
|
|
# comments, which servers may send as keep-alives.
|
|
chunk = _extract_sse_data(message)
|
|
if not chunk or chunk == "[DONE]":
|
|
continue
|
|
|
|
timestamp = time.perf_counter()
|
|
data = json.loads(chunk)
|
|
event_type = data.get("type")
|
|
|
|
if event_type in _RESPONSES_TOKEN_DELTA_EVENTS:
|
|
# First token
|
|
if not first_chunk_received:
|
|
first_chunk_received = True
|
|
output.ttft = timestamp - st
|
|
|
|
# Decoding phase
|
|
else:
|
|
output.itl.append(timestamp - most_recent_timestamp)
|
|
|
|
if event_type == "response.output_text.delta":
|
|
generated_text += data.get("delta") or ""
|
|
|
|
# Only token events advance the end of the
|
|
# request, so `latency - ttft` stays equal to
|
|
# `sum(itl)` and TPOT means what serve.py says.
|
|
most_recent_timestamp = timestamp
|
|
elif event_type in _RESPONSES_TERMINAL_EVENTS:
|
|
saw_terminal_event = True
|
|
usage = (data.get("response") or {}).get("usage") or {}
|
|
if (n_output := usage.get("output_tokens")) is not None:
|
|
output.output_tokens = n_output
|
|
if (n_input := usage.get("input_tokens")) is not None:
|
|
output.prompt_len = n_input
|
|
elif event_type in _RESPONSES_ERROR_EVENTS or data.get("error"):
|
|
stream_error = json.dumps(data.get("error") or data)
|
|
|
|
output.generated_text = generated_text
|
|
output.latency = most_recent_timestamp - st
|
|
if stream_error is not None:
|
|
output.success = False
|
|
output.error = stream_error
|
|
elif not first_chunk_received:
|
|
output.success = False
|
|
output.error = (
|
|
"Never received a token delta to calculate TTFT."
|
|
"This response will be marked as failed!"
|
|
)
|
|
elif not saw_terminal_event:
|
|
# Without the terminal event the response is truncated and
|
|
# its usage block is missing.
|
|
output.success = False
|
|
output.error = (
|
|
"Stream ended without a terminal event."
|
|
"This response will be marked as failed!"
|
|
)
|
|
else:
|
|
output.success = True
|
|
else:
|
|
output.error = response.reason or ""
|
|
output.success = False
|
|
except Exception:
|
|
output.success = False
|
|
exc_info = sys.exc_info()
|
|
output.error = "".join(traceback.format_exception(*exc_info))
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
async def async_request_openai_audio(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
# Lazy import without PlaceholderModule to avoid vllm dep.
|
|
import soundfile
|
|
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "OpenAI Audio API", {"transcriptions", "translations"})
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"max_completion_tokens": request_func_input.output_len,
|
|
"stream": True,
|
|
"language": "en",
|
|
# Flattened due to multipart/form-data
|
|
"stream_include_usage": True,
|
|
"stream_continuous_usage_stats": True,
|
|
}
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers()
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
# Send audio file
|
|
def to_bytes(y, sr):
|
|
buffer = io.BytesIO()
|
|
soundfile.write(buffer, y, sr, format="WAV")
|
|
buffer.seek(0)
|
|
return buffer
|
|
|
|
async def send_audio_file(
|
|
audio_file: io.BytesIO | Any,
|
|
*,
|
|
input_audio_duration: float,
|
|
filename: str | None = None,
|
|
content_type: str | None = None,
|
|
) -> RequestFuncOutput:
|
|
form = aiohttp.FormData()
|
|
add_field_kwargs: dict[str, str] = {}
|
|
if filename is not None:
|
|
add_field_kwargs["filename"] = filename
|
|
if content_type is not None:
|
|
add_field_kwargs["content_type"] = content_type
|
|
form.add_field("file", audio_file, **add_field_kwargs)
|
|
for key, value in payload.items():
|
|
form.add_field(key, str(value))
|
|
|
|
output = RequestFuncOutput()
|
|
output.prompt_len = request_func_input.prompt_len
|
|
output.input_audio_duration = input_audio_duration
|
|
|
|
generated_text = ""
|
|
st = time.perf_counter()
|
|
output.start_time = st
|
|
most_recent_timestamp = st
|
|
try:
|
|
async with session.post(
|
|
url=api_url, data=form, headers=headers
|
|
) as response:
|
|
if response.status == 200:
|
|
first_chunk_received = False
|
|
handler = StreamedResponseHandler()
|
|
|
|
async for chunk_bytes in response.content.iter_any():
|
|
if not chunk_bytes:
|
|
continue
|
|
|
|
messages = handler.add_chunk(chunk_bytes)
|
|
for message in messages:
|
|
if type(message) is bytes:
|
|
message = message.decode("utf-8")
|
|
chunk = message.removeprefix("data: ")
|
|
if chunk != "[DONE]":
|
|
timestamp = time.perf_counter()
|
|
data = json.loads(chunk)
|
|
|
|
if choices := data.get("choices"):
|
|
content = choices[0]["delta"].get("content")
|
|
# First token
|
|
if not first_chunk_received:
|
|
first_chunk_received = True
|
|
output.ttft = timestamp - st
|
|
|
|
# Decoding phase
|
|
else:
|
|
output.itl.append(
|
|
timestamp - most_recent_timestamp
|
|
)
|
|
|
|
generated_text += content or ""
|
|
# Only token chunks advance the request
|
|
# end; the trailing usage chunk carries no
|
|
# token.
|
|
most_recent_timestamp = timestamp
|
|
elif usage := data.get("usage"):
|
|
output.output_tokens = usage.get(
|
|
"completion_tokens"
|
|
)
|
|
|
|
output.generated_text = generated_text
|
|
if first_chunk_received:
|
|
output.success = True
|
|
else:
|
|
output.success = False
|
|
output.error = (
|
|
"Never received a valid chunk to calculate TTFT."
|
|
"This response will be marked as failed!"
|
|
)
|
|
output.latency = most_recent_timestamp - st
|
|
else:
|
|
output.error = response.reason or ""
|
|
output.success = False
|
|
except Exception:
|
|
output.success = False
|
|
exc_info = sys.exc_info()
|
|
output.error = "".join(traceback.format_exception(*exc_info))
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
mm_audio = request_func_input.multi_modal_content
|
|
if not isinstance(mm_audio, dict):
|
|
raise TypeError(
|
|
"multi_modal_content must be a dict containing 'audio' or 'audio_path'"
|
|
)
|
|
if "audio" in mm_audio:
|
|
with to_bytes(*mm_audio["audio"]) as f:
|
|
input_audio_duration = soundfile.info(f).duration
|
|
f.seek(0)
|
|
return await send_audio_file(
|
|
f,
|
|
input_audio_duration=input_audio_duration,
|
|
filename="audio.wav",
|
|
content_type="audio/wav",
|
|
)
|
|
if "audio_path" in mm_audio:
|
|
audio_path = mm_audio["audio_path"]
|
|
with open(audio_path, "rb") as f:
|
|
return await send_audio_file(
|
|
f,
|
|
input_audio_duration=soundfile.info(audio_path).duration,
|
|
filename=os.path.basename(audio_path),
|
|
)
|
|
raise TypeError(
|
|
"multi_modal_content must be a dict containing 'audio' or 'audio_path'"
|
|
)
|
|
|
|
|
|
async def _run_pooling_request(
|
|
session: aiohttp.ClientSession,
|
|
api_url: str,
|
|
payload: dict[str, Any],
|
|
headers: dict[str, Any],
|
|
pbar: tqdm | None = None,
|
|
num_input_sequences: int = 1,
|
|
prompt_len: int = 0,
|
|
) -> RequestFuncOutput:
|
|
output = RequestFuncOutput(num_input_sequences=num_input_sequences)
|
|
st = time.perf_counter()
|
|
output.start_time = st
|
|
try:
|
|
async with session.post(url=api_url, headers=headers, json=payload) as response:
|
|
if response.status == 200:
|
|
encoding_format = payload.get("encoding_format", "float")
|
|
if encoding_format in ("bytes", "bytes_only"):
|
|
async for _ in response.content.iter_any():
|
|
pass
|
|
else:
|
|
await response.read()
|
|
output.ttft = output.latency = time.perf_counter() - st
|
|
|
|
if encoding_format == "bytes":
|
|
metadata = json.loads(response.headers["metadata"])
|
|
usage = metadata.get("usage", {})
|
|
elif encoding_format == "bytes_only":
|
|
usage = {"prompt_tokens": prompt_len}
|
|
else:
|
|
data = await response.json()
|
|
usage = data.get("usage", {})
|
|
|
|
output.success = True
|
|
output.generated_text = ""
|
|
output.prompt_len = usage.get("prompt_tokens", 0)
|
|
else:
|
|
output.success = False
|
|
output.error = response.reason or ""
|
|
except Exception as e:
|
|
output.success = False
|
|
output.error = str(e)
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
return output
|
|
|
|
|
|
def _get_num_input_sequences(prompt: Any) -> int:
|
|
if prompt and isinstance(prompt, list) and isinstance(prompt[0], (str, list)):
|
|
return len(prompt)
|
|
return 1
|
|
|
|
|
|
async def async_request_openai_embeddings(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "OpenAI Embeddings API", "embeddings")
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"input": request_func_input.prompt,
|
|
# Many embedding models have short context length,
|
|
# this is to avoid dropping some of the requests.
|
|
"truncate_prompt_tokens": -1,
|
|
}
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers("application/json")
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
return await _run_pooling_request(
|
|
session,
|
|
api_url,
|
|
payload=payload,
|
|
headers=headers,
|
|
pbar=pbar,
|
|
num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
|
|
prompt_len=request_func_input.prompt_len,
|
|
)
|
|
|
|
|
|
async def async_request_vllm_rerank(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "vLLM score API", "rerank")
|
|
|
|
assert (
|
|
isinstance(request_func_input.prompt, list)
|
|
and len(request_func_input.prompt) > 1
|
|
)
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"query": request_func_input.prompt[0],
|
|
"documents": request_func_input.prompt[1:],
|
|
# Many reranker models have short context length,
|
|
# this is to avoid dropping some of the requests.
|
|
"truncate_prompt_tokens": -1,
|
|
}
|
|
|
|
headers = _get_headers("application/json")
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
return await _run_pooling_request(
|
|
session,
|
|
api_url,
|
|
payload=payload,
|
|
headers=headers,
|
|
pbar=pbar,
|
|
num_input_sequences=len(request_func_input.prompt) - 1,
|
|
prompt_len=request_func_input.prompt_len,
|
|
)
|
|
|
|
|
|
async def async_request_openai_embeddings_chat(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
mm_position: Literal["first", "last"] = "last",
|
|
) -> RequestFuncOutput:
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "OpenAI Embeddings API", "embeddings")
|
|
|
|
messages = _get_chat_messages(request_func_input, mm_position=mm_position)
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"messages": messages,
|
|
# Many embedding models have short context length,
|
|
# this is to avoid dropping some of the requests.
|
|
"truncate_prompt_tokens": -1,
|
|
}
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers("application/json")
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
return await _run_pooling_request(
|
|
session,
|
|
api_url,
|
|
payload=payload,
|
|
headers=headers,
|
|
pbar=pbar,
|
|
prompt_len=request_func_input.prompt_len,
|
|
)
|
|
|
|
|
|
def _try_extract_request_idx(request_func_input: RequestFuncInput):
|
|
if request_func_input.request_id:
|
|
match = re.search(r"(\d+)$", request_func_input.request_id)
|
|
if match:
|
|
try:
|
|
return int(match.group(1))
|
|
except ValueError:
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
def _preprocess_clip(request_func_input: RequestFuncInput):
|
|
if request_func_input.multi_modal_content:
|
|
# Image input
|
|
request_func_input.prompt = ""
|
|
|
|
|
|
def _preprocess_vlm2vec(request_func_input: RequestFuncInput):
|
|
if request_func_input.multi_modal_content:
|
|
request_idx = _try_extract_request_idx(request_func_input)
|
|
|
|
# Adjust the ratio manually if needed.
|
|
use_image_only_prompt = request_idx is None or request_idx % 2 == 0
|
|
|
|
if use_image_only_prompt:
|
|
# Image input
|
|
request_func_input.prompt = "Represent the given image."
|
|
else:
|
|
# Text+Image input
|
|
request_func_input.prompt = (
|
|
f"Represent the given image with the following question: "
|
|
f"{request_func_input.prompt}"
|
|
)
|
|
|
|
|
|
async def async_request_openai_embeddings_clip(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
_preprocess_clip(request_func_input)
|
|
|
|
return await async_request_openai_embeddings_chat(
|
|
request_func_input,
|
|
session,
|
|
pbar=pbar,
|
|
)
|
|
|
|
|
|
async def async_request_openai_embeddings_vlm2vec(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
_preprocess_vlm2vec(request_func_input)
|
|
|
|
return await async_request_openai_embeddings_chat(
|
|
request_func_input,
|
|
session,
|
|
pbar=pbar,
|
|
mm_position="first",
|
|
)
|
|
|
|
|
|
async def async_request_infinity_embeddings(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "Infinity Embeddings API", "embeddings")
|
|
|
|
payload: dict[str, Any] = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
}
|
|
|
|
if request_func_input.prompt:
|
|
payload["input"] = request_func_input.prompt
|
|
else:
|
|
mm_content = request_func_input.multi_modal_content
|
|
assert isinstance(mm_content, dict)
|
|
|
|
mm_type = mm_content["type"]
|
|
payload["input"] = mm_content[mm_type]["url"]
|
|
payload["modality"] = mm_type.split("_", 1)[0]
|
|
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers("application/json")
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
return await _run_pooling_request(
|
|
session,
|
|
api_url,
|
|
payload=payload,
|
|
headers=headers,
|
|
pbar=pbar,
|
|
num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
|
|
prompt_len=request_func_input.prompt_len,
|
|
)
|
|
|
|
|
|
async def async_request_infinity_embeddings_clip(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
_preprocess_clip(request_func_input)
|
|
|
|
return await async_request_infinity_embeddings(
|
|
request_func_input,
|
|
session,
|
|
pbar=pbar,
|
|
)
|
|
|
|
|
|
async def async_request_vllm_pooling(
|
|
request_func_input: RequestFuncInput,
|
|
session: aiohttp.ClientSession,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestFuncOutput:
|
|
api_url = request_func_input.api_url
|
|
_validate_api_url(api_url, "vLLM Pooling API", "pooling")
|
|
|
|
payload = {
|
|
"model": request_func_input.model_name
|
|
if request_func_input.model_name
|
|
else request_func_input.model,
|
|
"truncate_prompt_tokens": -1,
|
|
}
|
|
|
|
if isinstance(request_func_input.prompt, dict):
|
|
payload = payload | request_func_input.prompt
|
|
else:
|
|
payload["input"] = request_func_input.prompt
|
|
|
|
_update_payload_common(payload, request_func_input)
|
|
|
|
headers = _get_headers("application/json")
|
|
_update_headers_common(headers, request_func_input)
|
|
|
|
return await _run_pooling_request(
|
|
session,
|
|
api_url,
|
|
payload=payload,
|
|
headers=headers,
|
|
pbar=pbar,
|
|
num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
|
|
prompt_len=request_func_input.prompt_len,
|
|
)
|
|
|
|
|
|
# TODO: Add more request functions for different API protocols.
|
|
ASYNC_REQUEST_FUNCS: dict[str, RequestFunc] = {
|
|
"vllm": async_request_openai_completions,
|
|
"openai": async_request_openai_completions,
|
|
"openai-chat": async_request_openai_chat_completions,
|
|
"openai-responses": async_request_openai_responses,
|
|
"openai-audio": async_request_openai_audio,
|
|
"openai-embeddings": async_request_openai_embeddings,
|
|
"openai-embeddings-chat": async_request_openai_embeddings_chat,
|
|
"openai-embeddings-clip": async_request_openai_embeddings_clip,
|
|
"openai-embeddings-vlm2vec": async_request_openai_embeddings_vlm2vec,
|
|
# Infinity embedding server: https://github.com/michaelfeil/infinity
|
|
"infinity-embeddings": async_request_infinity_embeddings,
|
|
"infinity-embeddings-clip": async_request_infinity_embeddings_clip,
|
|
# (Infinity embedding server does not support vlm2vec)
|
|
"vllm-pooling": async_request_vllm_pooling,
|
|
"vllm-rerank": async_request_vllm_rerank,
|
|
}
|
|
|
|
POOLING_BACKENDS = {
|
|
"openai-embeddings",
|
|
"openai-embeddings-chat",
|
|
"openai-embeddings-clip",
|
|
"openai-embeddings-vlm2vec",
|
|
"infinity-embeddings",
|
|
"infinity-embeddings-clip",
|
|
"vllm-pooling",
|
|
"vllm-rerank",
|
|
}
|
|
|
|
OPENAI_COMPATIBLE_BACKENDS = [
|
|
k
|
|
for k, v in ASYNC_REQUEST_FUNCS.items()
|
|
if v
|
|
in (
|
|
async_request_openai_completions,
|
|
async_request_openai_chat_completions,
|
|
async_request_openai_responses,
|
|
)
|
|
]
|