wehub-resource-sync
453 lines
19 KiB
Python
453 lines
19 KiB
Python
"""MLC LLM bench backends"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
import traceback
|
|
from typing import Optional
|
|
|
|
from typing_extensions import Self
|
|
|
|
from mlc_llm.bench.request_record import Metrics, RequestRecord, ServerMetrics
|
|
from mlc_llm.support import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class APIEndPoint:
|
|
"""Manages the sending of requests to a specified API endpoint and gathers
|
|
inference statistics.
|
|
"""
|
|
|
|
def __init__(self, include_server_metrics: bool = False) -> None:
|
|
self.include_server_metrics = include_server_metrics
|
|
|
|
async def __aenter__(self) -> Self:
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_value, tb) -> None:
|
|
pass
|
|
|
|
async def __call__(self, request: RequestRecord) -> RequestRecord:
|
|
raise NotImplementedError()
|
|
|
|
|
|
class OpenAIChatEndPoint(APIEndPoint):
|
|
"""The backend of sending HTTP requests in OpenAI API through "v1/chat/completions"."""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
timeout: Optional[float] = None,
|
|
include_server_metrics: bool = False,
|
|
) -> None:
|
|
super().__init__(include_server_metrics=include_server_metrics)
|
|
|
|
import aiohttp
|
|
|
|
self.timeout = timeout
|
|
self.client: aiohttp.ClientSession = None
|
|
self.url = f"http://{host}:{port}/v1/chat/completions"
|
|
self.headers = {"Content-Type": "application/json"}
|
|
if os.getenv("MLC_LLM_API_KEY"):
|
|
self.headers["Authorization"] = f"Bearer {os.getenv('MLC_LLM_API_KEY')}"
|
|
|
|
async def __aenter__(self) -> Self:
|
|
import aiohttp
|
|
|
|
self.client = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(self.timeout))
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_value, tb) -> None:
|
|
await self.client.close()
|
|
|
|
async def __call__(self, request_record: RequestRecord) -> RequestRecord:
|
|
payload = request_record.chat_cmpl.model_dump()
|
|
if self.timeout is not None and "timeout" not in payload:
|
|
payload["timeout"] = self.timeout
|
|
if self.include_server_metrics:
|
|
if "stream_options" not in payload or payload["stream_options"] is None:
|
|
payload["stream_options"] = {"include_usage": True}
|
|
else:
|
|
payload["stream_options"]["include_usage"] = True
|
|
if (
|
|
request_record.chat_cmpl.debug_config is not None
|
|
and request_record.chat_cmpl.debug_config.ignore_eos
|
|
):
|
|
payload["ignore_eos"] = True
|
|
|
|
generated_text = ""
|
|
first_chunk_output_str = ""
|
|
time_to_first_token_s = None
|
|
start_time = time.monotonic()
|
|
server_metrics = None
|
|
|
|
try:
|
|
async with self.client.post(self.url, json=payload, headers=self.headers) as response:
|
|
assert response.status == 200, await response.text()
|
|
if payload["stream"]:
|
|
async for chunk in response.content:
|
|
chunk = chunk.strip()
|
|
if not chunk or chunk == b"\n":
|
|
continue
|
|
# Get rid of the prefix "data: " and suffix "\n"
|
|
raw_data = chunk[6:].strip()
|
|
if raw_data == b"[DONE]":
|
|
continue
|
|
data = json.loads(raw_data)
|
|
if not data["choices"]:
|
|
continue
|
|
delta = data["choices"][0]["delta"]
|
|
content = delta.get("content", None)
|
|
if content is not None and not time_to_first_token_s:
|
|
time_to_first_token_s = time.monotonic() - start_time
|
|
first_chunk_output_str = content
|
|
if self.include_server_metrics and data["usage"] is not None:
|
|
# fmt: off
|
|
server_metrics = ServerMetrics(
|
|
input_tokens=data["usage"]["extra"]["prompt_tokens"],
|
|
prefill_tokens=data["usage"]["extra"]["prefill_tokens"],
|
|
output_tokens=data["usage"]["extra"]["completion_tokens"],
|
|
end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"],
|
|
prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"],
|
|
inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"],
|
|
time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"], # noqa: E501
|
|
time_to_first_token_s=data["usage"]["extra"]["ttft_s"],
|
|
)
|
|
# fmt: on
|
|
|
|
if content is not None:
|
|
generated_text += content
|
|
else:
|
|
data = await response.json()
|
|
generated_text = data["choices"][0]["message"]["content"]
|
|
if self.include_server_metrics and data["usage"] is not None:
|
|
# fmt: off
|
|
server_metrics = ServerMetrics(
|
|
input_tokens=data["usage"]["extra"]["prompt_tokens"],
|
|
prefill_tokens=data["usage"]["extra"]["prefill_tokens"],
|
|
output_tokens=data["usage"]["extra"]["completion_tokens"],
|
|
end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"],
|
|
prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"],
|
|
inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"],
|
|
time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"], # noqa: E501
|
|
time_to_first_token_s=data["usage"]["extra"]["ttft_s"],
|
|
)
|
|
# fmt: on
|
|
except Exception:
|
|
error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
|
|
logger.info(error_msg)
|
|
finish_time = time.monotonic()
|
|
request_record.output_str = generated_text
|
|
request_record.first_chunk_output_str = first_chunk_output_str
|
|
request_record.metrics = Metrics(
|
|
success=False,
|
|
start_time=start_time,
|
|
finish_time=finish_time,
|
|
end_to_end_latency_s=finish_time - start_time,
|
|
input_tokens=request_record.metrics.input_tokens,
|
|
time_to_first_token_s=time_to_first_token_s,
|
|
server_metrics=server_metrics,
|
|
exec_feature=request_record.metrics.exec_feature,
|
|
)
|
|
request_record.error_msg = error_msg
|
|
return request_record
|
|
|
|
finish_time = time.monotonic()
|
|
request_record.output_str = generated_text
|
|
request_record.first_chunk_output_str = first_chunk_output_str
|
|
success = True
|
|
error_msg = None
|
|
if len(generated_text) == 0:
|
|
success = False
|
|
error_msg = "Empty generated text."
|
|
request_record.metrics = Metrics(
|
|
success=success,
|
|
start_time=start_time,
|
|
finish_time=finish_time,
|
|
end_to_end_latency_s=finish_time - start_time,
|
|
input_tokens=request_record.metrics.input_tokens,
|
|
time_to_first_token_s=time_to_first_token_s,
|
|
server_metrics=server_metrics,
|
|
exec_feature=request_record.metrics.exec_feature,
|
|
)
|
|
request_record.error_msg = error_msg
|
|
return request_record
|
|
|
|
|
|
class OpenAIEndPoint(APIEndPoint):
|
|
"""The backend of sending HTTP requests in OpenAI API through "v1/completions"."""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
timeout: Optional[float] = None,
|
|
include_server_metrics: bool = False,
|
|
no_debug_config: bool = False,
|
|
) -> None:
|
|
super().__init__(include_server_metrics=include_server_metrics)
|
|
|
|
import aiohttp
|
|
|
|
self.timeout = timeout
|
|
self.client: aiohttp.ClientSession = None
|
|
self.url = f"http://{host}:{port}/v1/completions"
|
|
self.headers = {"Content-Type": "application/json"}
|
|
if os.getenv("MLC_LLM_API_KEY"):
|
|
self.headers["Authorization"] = f"Bearer {os.getenv('MLC_LLM_API_KEY')}"
|
|
assert not include_server_metrics, (
|
|
'"include_server_metrics" only works for "openai-chat" endpoint for now'
|
|
)
|
|
self.no_debug_config = no_debug_config
|
|
|
|
async def __aenter__(self) -> Self:
|
|
import aiohttp
|
|
|
|
self.client = aiohttp.ClientSession()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_value, tb) -> None:
|
|
await self.client.close()
|
|
|
|
async def __call__(self, request_record: RequestRecord) -> RequestRecord:
|
|
assert len(request_record.chat_cmpl.messages) == 1, (
|
|
'Endpoint "openai" does not support system prompt and multi-round conversation.'
|
|
)
|
|
assert isinstance(request_record.chat_cmpl.messages[0].content, str)
|
|
payload = {
|
|
"model": request_record.chat_cmpl.model,
|
|
"prompt": request_record.chat_cmpl.messages[0].content,
|
|
"temperature": request_record.chat_cmpl.temperature,
|
|
"top_p": request_record.chat_cmpl.top_p,
|
|
"max_tokens": request_record.chat_cmpl.max_tokens,
|
|
"stream": True,
|
|
}
|
|
if self.timeout is not None and "timeout" not in payload:
|
|
payload["timeout"] = self.timeout
|
|
if (
|
|
request_record.chat_cmpl.debug_config is not None
|
|
and request_record.chat_cmpl.debug_config.ignore_eos
|
|
):
|
|
payload["ignore_eos"] = True
|
|
if not self.no_debug_config:
|
|
payload["debug_config"] = {"ignore_eos": True}
|
|
|
|
generated_text = ""
|
|
first_chunk_output_str = ""
|
|
time_to_first_token_s = None
|
|
start_time = time.monotonic()
|
|
|
|
try:
|
|
async with self.client.post(
|
|
self.url, json=payload, headers=self.headers, timeout=3600
|
|
) as response:
|
|
assert response.status == 200, await response.text()
|
|
if payload["stream"]:
|
|
async for chunk in response.content:
|
|
chunk = chunk.strip()
|
|
if not chunk or chunk == b"\n":
|
|
continue
|
|
# Get rid of the prefix "data: " and suffix "\n"
|
|
raw_data = chunk[6:].strip()
|
|
if raw_data == b"[DONE]":
|
|
continue
|
|
data = json.loads(raw_data)
|
|
if not data["choices"]:
|
|
continue
|
|
content = data["choices"][0]["text"]
|
|
if content is not None and not time_to_first_token_s:
|
|
time_to_first_token_s = time.monotonic() - start_time
|
|
first_chunk_output_str = content
|
|
if content is not None:
|
|
generated_text += content
|
|
else:
|
|
data = await response.json()
|
|
generated_text = data["choices"][0]["message"]["content"]
|
|
except Exception:
|
|
error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
|
|
logger.info(error_msg)
|
|
finish_time = time.monotonic()
|
|
request_record.output_str = generated_text
|
|
request_record.first_chunk_output_str = first_chunk_output_str
|
|
request_record.metrics = Metrics(
|
|
success=False,
|
|
start_time=start_time,
|
|
finish_time=finish_time,
|
|
end_to_end_latency_s=finish_time - start_time,
|
|
input_tokens=request_record.metrics.input_tokens,
|
|
time_to_first_token_s=time_to_first_token_s,
|
|
server_metrics=None,
|
|
exec_feature=request_record.metrics.exec_feature,
|
|
)
|
|
request_record.error_msg = error_msg
|
|
return request_record
|
|
|
|
finish_time = time.monotonic()
|
|
request_record.output_str = generated_text
|
|
request_record.first_chunk_output_str = first_chunk_output_str
|
|
success = True
|
|
error_msg = None
|
|
if len(generated_text) == 0:
|
|
success = False
|
|
error_msg = "Empty generated text."
|
|
request_record.metrics = Metrics(
|
|
success=success,
|
|
start_time=start_time,
|
|
finish_time=finish_time,
|
|
end_to_end_latency_s=finish_time - start_time,
|
|
input_tokens=request_record.metrics.input_tokens,
|
|
time_to_first_token_s=time_to_first_token_s,
|
|
server_metrics=None,
|
|
exec_feature=request_record.metrics.exec_feature,
|
|
)
|
|
request_record.error_msg = error_msg
|
|
return request_record
|
|
|
|
|
|
class TensorRTLLMEndPoint(APIEndPoint):
|
|
"""The backend of sending HTTP requests in TensorRT-LLM API."""
|
|
|
|
def __init__(self, host: str, port: int, timeout: Optional[float] = None) -> None:
|
|
super().__init__(include_server_metrics=False)
|
|
|
|
import aiohttp
|
|
|
|
self.timeout = timeout
|
|
self.client: aiohttp.ClientSession = None
|
|
self.url_stream = f"http://{host}:{port}/v2/models/ensemble/generate_stream"
|
|
self.url_no_stream = f"http://{host}:{port}/v2/models/ensemble/generate"
|
|
|
|
async def __aenter__(self) -> Self:
|
|
import aiohttp
|
|
|
|
self.client = aiohttp.ClientSession()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_value, tb) -> None:
|
|
await self.client.close()
|
|
|
|
async def __call__(self, request_record: RequestRecord) -> RequestRecord:
|
|
assert len(request_record.chat_cmpl.messages) == 1
|
|
assert isinstance(request_record.chat_cmpl.messages[0].content, str)
|
|
payload = {
|
|
"accumulate_tokens": True,
|
|
"text_input": request_record.chat_cmpl.messages[0].content,
|
|
"temperature": (
|
|
max(request_record.chat_cmpl.temperature, 1e-5)
|
|
if request_record.chat_cmpl.temperature
|
|
else 1
|
|
),
|
|
"top_p": request_record.chat_cmpl.top_p if request_record.chat_cmpl.top_p else 1,
|
|
"max_tokens": request_record.chat_cmpl.max_tokens,
|
|
"stream": request_record.chat_cmpl.stream,
|
|
}
|
|
if (
|
|
request_record.chat_cmpl.debug_config is not None
|
|
and request_record.chat_cmpl.debug_config.ignore_eos
|
|
):
|
|
payload["min_length"] = payload["max_tokens"]
|
|
if self.timeout is not None and "timeout" not in payload:
|
|
payload["timeout"] = self.timeout
|
|
|
|
generated_text = ""
|
|
first_chunk_output_str = ""
|
|
url = self.url_stream if request_record.chat_cmpl.stream else self.url_no_stream
|
|
time_to_first_token_s = None
|
|
start_time = time.monotonic()
|
|
|
|
try:
|
|
async with self.client.post(url, json=payload) as response:
|
|
assert response.status == 200, await response.text()
|
|
if payload["stream"]:
|
|
async for chunk in response.content:
|
|
chunk = chunk.strip()
|
|
if not chunk or chunk == b"\n":
|
|
continue
|
|
# Get rid of the prefix "data:" and suffix "\n"
|
|
raw_data = chunk[5:].strip()
|
|
data = json.loads(raw_data)
|
|
delta = data["text_output"]
|
|
if delta is None:
|
|
continue
|
|
|
|
if not time_to_first_token_s:
|
|
time_to_first_token_s = time.monotonic() - start_time
|
|
first_chunk_output_str = delta
|
|
generated_text += delta
|
|
else:
|
|
data = await response.json()
|
|
generated_text = data["text_output"]
|
|
except Exception:
|
|
error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
|
|
logger.info(error_msg)
|
|
finish_time = time.monotonic()
|
|
request_record.output_str = generated_text
|
|
request_record.first_chunk_output_str = first_chunk_output_str
|
|
request_record.metrics = Metrics(
|
|
success=False,
|
|
start_time=start_time,
|
|
finish_time=finish_time,
|
|
end_to_end_latency_s=finish_time - start_time,
|
|
input_tokens=request_record.metrics.input_tokens,
|
|
time_to_first_token_s=time_to_first_token_s,
|
|
exec_feature=request_record.metrics.exec_feature,
|
|
)
|
|
request_record.error_msg = error_msg
|
|
return request_record
|
|
|
|
finish_time = time.monotonic()
|
|
request_record.output_str = generated_text
|
|
request_record.first_chunk_output_str = first_chunk_output_str
|
|
success = True
|
|
error_msg = None
|
|
if len(generated_text) == 0:
|
|
success = False
|
|
error_msg = "Empty generated text."
|
|
request_record.metrics = Metrics(
|
|
success=success,
|
|
start_time=start_time,
|
|
finish_time=finish_time,
|
|
end_to_end_latency_s=finish_time - start_time,
|
|
input_tokens=request_record.metrics.input_tokens,
|
|
time_to_first_token_s=time_to_first_token_s,
|
|
exec_feature=request_record.metrics.exec_feature,
|
|
)
|
|
request_record.error_msg = error_msg
|
|
return request_record
|
|
|
|
|
|
# Todo: APIEndPoint with AsyncOpenAI Python interface
|
|
# class OpenAIPythonEndPoint(APIEndPoint):
|
|
# pass
|
|
|
|
SUPPORTED_BACKENDS = [
|
|
"openai",
|
|
"openai-chat",
|
|
"mlc",
|
|
"sglang",
|
|
"tensorrt-llm",
|
|
"vllm",
|
|
]
|
|
|
|
|
|
def create_api_endpoint(args: argparse.Namespace) -> APIEndPoint:
|
|
"""Create an API endpoint instance with regard to the specified endpoint kind."""
|
|
if args.api_endpoint in ["openai", "mlc", "sglang"]:
|
|
return OpenAIEndPoint(args.host, args.port, args.timeout, args.include_server_metrics)
|
|
if args.api_endpoint == "vllm":
|
|
return OpenAIEndPoint(
|
|
args.host,
|
|
args.port,
|
|
args.timeout,
|
|
include_server_metrics=False,
|
|
no_debug_config=True,
|
|
)
|
|
if args.api_endpoint == "openai-chat":
|
|
return OpenAIChatEndPoint(args.host, args.port, args.timeout, args.include_server_metrics)
|
|
if args.api_endpoint == "tensorrt-llm":
|
|
return TensorRTLLMEndPoint(args.host, args.port, args.timeout)
|
|
raise ValueError(f'Unrecognized endpoint "{args.api_endpoint}"')
|