Files
mlc-ai--mlc-llm/python/mlc_llm/bench/__main__.py
T
wehub-resource-sync 770d92cb1f
Lint / lint (push) Waiting to run
Windows CI / Windows (push) Waiting to run
Build Docs / Deploy Docs (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:23:58 +08:00

404 lines
14 KiB
Python

"""MLC LLM benchmark main entrance"""
import functools
import json
import random
from typing import Any, Dict, List, Optional, Tuple # noqa: UP035
import numpy as np
import requests
from transformers import AutoTokenizer
import mlc_llm
from mlc_llm.bench.api_endpoint import SUPPORTED_BACKENDS, create_api_endpoint
from mlc_llm.bench.dataset import SUPPORTED_DATASET, Dataset, create_dataset
from mlc_llm.bench.request_processor import (
MetricAnalyzer,
RequestProcessor,
create_pipelines,
)
from mlc_llm.bench.request_record import (
RequestRecord,
convert_reports_to_df,
generate_metrics_summary,
pretty_print_report,
)
from mlc_llm.cli.serve import EngineConfigOverride
from mlc_llm.serve import EngineConfig
from mlc_llm.support import argparse, logging
logging.enable_logging()
logger = logging.getLogger(__name__)
def _parse_num_concurrent_requests(num_str: Optional[str]) -> Optional[List[int]]: # noqa: UP006
if num_str is None:
return None
numbers = num_str.split(",")
if any(not number.isdigit() for number in numbers):
raise ValueError(f"Unrecognized num_concurrent_requests list: {numbers}")
return list(int(number) for number in numbers)
def _parse_request_rate(request_rate_str: Optional[str]) -> Optional[List[np.float32]]: # noqa: UP006
if request_rate_str is None:
return None
request_rates = request_rate_str.split(",")
results = []
for rate_str in request_rates:
request_rate = float(rate_str)
if request_rate <= 0:
raise ValueError(f"Invalid request rate {request_rate}")
results.append(np.float32(request_rate))
return results
def _parse_mlc_engine_config(config_str: Optional[str]) -> EngineConfig:
if config_str is None:
return None
engine_config_override = EngineConfigOverride.from_str(config_str)
return EngineConfig(
tensor_parallel_shards=engine_config_override.tensor_parallel_shards,
max_num_sequence=engine_config_override.max_num_sequence,
max_total_sequence_length=engine_config_override.max_total_seq_length,
prefill_chunk_size=engine_config_override.prefill_chunk_size,
sliding_window_size=engine_config_override.sliding_window_size,
attention_sink_size=engine_config_override.attention_sink_size,
max_history_size=engine_config_override.max_history_size,
gpu_memory_utilization=engine_config_override.gpu_memory_utilization,
spec_draft_length=engine_config_override.spec_draft_length,
prefill_mode=engine_config_override.prefill_mode,
prefix_cache_max_num_recycling_seqs=engine_config_override.prefix_cache_max_num_recycling_seqs,
prefix_cache_mode=engine_config_override.prefix_cache_mode,
)
def _launch_mlc_server(args: argparse.argparse.Namespace):
return mlc_llm.serve.PopenServer(
model=args.tokenizer,
mode="server",
model_lib=args.mlc_model_lib,
enable_tracing=False,
host=args.host,
port=args.port,
engine_config=args.mlc_engine_config,
)
def run_pipeline(
pipeline: RequestProcessor,
dataset: Dataset,
tokenizer: AutoTokenizer,
args: argparse.argparse.Namespace,
) -> Tuple[Dict[str, Any], List[RequestRecord]]: # noqa: UP006
"""Run the pipeline with the given dataset and args. Return the benchmark report dict."""
random.seed(args.seed)
np.random.seed(args.seed)
request_records = dataset.generate_request_records(
args.input_len,
args.output_len,
args.input_len_std,
args.output_len_std,
)
request_records = pipeline(request_records)
num_total_requests = (
args.num_requests if not args.per_gpu_workload else args.num_requests * args.num_gpus
)
assert len(request_records) == num_total_requests
sorted_requests: List[RequestRecord] = [None] * num_total_requests # noqa: UP006
for request_record in request_records:
assert request_record.request_id is not None
assert sorted_requests[request_record.request_id] is None
sorted_requests[request_record.request_id] = request_record
request_records = MetricAnalyzer(tokenizer)(request_records)
report = generate_metrics_summary(request_records, num_total_requests, args.num_gpus)
return report, sorted_requests
def query_mlc_server_metrics(host: str, port: int):
"""Try to get the MLC server metrics whenever it exists."""
try:
r = requests.post(f"http://{host}:{port}/debug/dump_engine_metrics", json={}, timeout=10)
if r.status_code == 200:
print(f"MLC server metrics: {r.json()}")
except Exception:
pass
def main(args: argparse.argparse.Namespace):
"""Main benchmark entrance."""
mlc_server = None
if args.mlc_model_lib:
mlc_server = _launch_mlc_server(args)
if args.num_requests <= 0:
raise ValueError("Number of requests to benchmark must be positive.")
def _main():
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
dataset = create_dataset(args, tokenizer)
f_create_api_endpoint = functools.partial(create_api_endpoint, args)
pipelines = create_pipelines(args, f_create_api_endpoint, dataset)
reports = []
alltime_records = {}
for i, pipeline in enumerate(pipelines):
report, request_records = run_pipeline(pipeline, dataset, tokenizer, args)
exec_feature = (
json.dumps(report["exec_feature"])
if report["exec_feature"] is not None
else f"pipeline{i}"
)
alltime_records[exec_feature] = [
request_record.model_dump() for request_record in request_records
]
reports.append(report)
pretty_print_report(report)
query_mlc_server_metrics(args.host, args.port)
# Construct data frame
df = convert_reports_to_df(reports)
print(df)
df.to_csv(args.output, index=False)
logger.info("Benchmark results dumped to file %s", args.output)
if args.debug_dump:
debug_dump_filepath = (
args.output[:-4] if args.output.endswith(".csv") else args.output
) + "_debug_dump.log"
with open(debug_dump_filepath, "w", encoding="utf-8") as file:
json.dump(alltime_records, file, indent=4)
logger.info("Debug log dumped to file %s", debug_dump_filepath)
if mlc_server is not None:
with mlc_server:
_main()
else:
_main()
if __name__ == "__main__":
parser = argparse.ArgumentParser("MLC LLM benchmark")
parser.add_argument(
"--dataset",
type=str,
choices=SUPPORTED_DATASET,
help=f"The benchmark dataset kind. Supporting {SUPPORTED_DATASET}",
)
parser.add_argument(
"--dataset-path",
type=str,
help="The dataset file path.",
)
parser.add_argument(
"--api-endpoint",
type=str,
choices=SUPPORTED_BACKENDS,
default="openai",
help="The API endpoint API for benchmarking.",
)
parser.add_argument(
"--tokenizer",
type=str,
required=True,
help="The path of the tokenizer directory.",
)
parser.add_argument(
"--num-gpus",
type=int,
required=True,
help="The number of GPUs used by the server. "
"We need this to better analyze the throughput per GPU.",
)
parser.add_argument(
"--num-requests",
type=int,
required=True,
help="The number of requests for benchmark.",
)
parser.add_argument(
"--num-warmup-requests",
type=int,
help="The number of requests for warmup. "
"It is optional when fixing the number of concurrent requests, and is required otherwise.",
)
parser.add_argument(
"--per-gpu-workload",
default=False,
action="store_true",
help='When set to True, the specified "num_concurrent_requests"/"request_rate" '
"denote the workload **per GPU**, which means that the real values of "
'"num_concurrent_requests"/"request_rate" used in benchmark'
'will be multiplied by "num_gpus".',
)
parser.add_argument(
"--num-concurrent-requests",
type=_parse_num_concurrent_requests,
help="The number(s) of concurrent requests to benchmark. "
'It can be either one integer or a list of integer separated by commas(","). '
"When specified, for each integer, the benchmark keeps these many consistent "
"number of concurrently running requests.",
)
parser.add_argument(
"--request-rate",
type=_parse_request_rate,
help="The request rate(s) denoting the number of new requests each second. "
'It can be either one float number (or "inf") or a list of numbers separated '
'by commas(","). '
"When specified, the benchmark sends these many new requests each second. "
'If it is "inf", all requests will be sent together at once.',
)
parser.add_argument(
"--replay-timestamp-scale",
type=float,
help="The timestamp scale when replaying the timestamps in a dataset. "
'The dataset replay mode is enabled when neither "--num-concurrent-requests" and '
'"--request-rate" is specified. '
"The scale is 1 by default in the replay mode.",
)
parser.add_argument(
"--input-len",
type=int,
help="The benchmark request average input length. Default to None, "
"which means the request input length depends on the dataset being used.",
)
parser.add_argument(
"--input-len-std",
type=float,
default=0,
help="The benchmark request input length standard deviation. Default to 0.",
)
parser.add_argument(
"--output-len",
type=int,
help="The benchmark request average output length. Default to None, "
"which means the request output length depends on the dataset being used.",
)
parser.add_argument(
"--output-len-std",
type=float,
default=0,
help="The benchmark request output length standard deviation. Default to 0.",
)
parser.add_argument(
"--stream",
type=bool,
default=True,
help="Whether to benchmark stream responses. "
"When not enabled, metrics such as time-to-first-token (TTFT) will not be available. "
"Default to True.",
)
parser.add_argument(
# NOTE: The current implementation of server metrics still has some issues that need fixes,
# which makes it not work to include server metrics.
"--include-server-metrics",
action="store_true",
help="Whether to also benchmark the server side request metrics. "
"This option is only available when benchmarking MLC server.",
)
parser.add_argument(
"--host",
type=str,
required=True,
help="The host address of the backend API.",
)
parser.add_argument(
"--port",
type=int,
required=True,
help="The port of the backend API.",
)
parser.add_argument(
"--timeout",
type=float,
default=3 * 60 * 60,
help="The timeout limit of each request.",
)
parser.add_argument(
"--seed",
type=int,
default=0,
help="The random number seed. Default to 0.",
)
parser.add_argument(
"--temperature",
type=float,
default=1.0,
help="The temperature value for logit adjustment. Default to 1.",
)
parser.add_argument(
"--top-p",
type=float,
default=1.0,
help="The top-p value for sampling. Default to 1.",
)
parser.add_argument(
"--ignore-eos",
default=False,
action="store_true",
help='Whether to set the "ignore_eos" field.',
)
parser.add_argument(
"--apply-chat-template",
default=False,
action="store_true",
help="Whether to apply chat template to the request input text. "
'It is not supported when "--input-len" is specified.',
)
parser.add_argument(
"--num-process-workers",
type=int,
help="The number of parallel process workers to send the requests.",
)
parser.add_argument(
"--disable-tqdm",
action="store_true",
help="Whether to disable showing progress bar with tqdm during benchmarking.",
)
parser.add_argument(
"--max-schedule-gap",
type=float,
default=0.5,
help="The maximum allowed delay between the scheduled time in seconds.",
)
parser.add_argument(
"--mlc-model-lib",
type=str,
help="The model lib path when benchmarking MLC serve. "
"When specified, the server is automatic launched and no external server launch is needed.",
)
parser.add_argument(
"--mlc-engine-config",
type=_parse_mlc_engine_config,
help="The engine config used when launch MLC server.",
)
parser.add_argument(
"--cuda-profile",
default=False,
action="store_true",
help="Whether to enable cuda profile on server. "
"The --mlc-model-lib path should be provided when enabling this option.",
)
parser.add_argument(
"--debug-dump",
default=False,
action="store_true",
help="Whether to dump all request record raw data to file.",
)
parser.add_argument(
"--multi-round",
default=False,
action="store_true",
help="Whether to chat like multi round conversion with history log each request. "
"Only enabled when benchmarked with fixed concurrent request mode."
"The --num-concurrent-requests should be provided when enabling this option.",
)
parser.add_argument(
"--output",
"-o",
type=str,
default="mlc_benchmark.csv",
help="The path of the output file where to dump the benchmark results.",
)
main(parser.parse_args())