|
| 1 | +"""MLC LLM benchmark main entrance""" |
| 2 | + |
| 3 | +import asyncio |
| 4 | +import random |
| 5 | +from typing import Any, Dict, List, Optional |
| 6 | + |
| 7 | +import numpy as np |
| 8 | +from transformers import AutoTokenizer # pylint: disable=import-error |
| 9 | + |
| 10 | +import mlc_llm |
| 11 | +from mlc_llm.bench.api_endpoint import SUPPORTED_BACKENDS, create_api_endpoint |
| 12 | +from mlc_llm.bench.dataset import SUPPORTED_DATASET, Dataset, create_dataset |
| 13 | +from mlc_llm.bench.executor import Executor, create_executors |
| 14 | +from mlc_llm.bench.request_processor import ( |
| 15 | + AttachStreamFlag, |
| 16 | + MetricAnalyzer, |
| 17 | + SampleRequests, |
| 18 | + SequentialProcessor, |
| 19 | +) |
| 20 | +from mlc_llm.bench.request_record import convert_reports_to_df, generate_metrics_summary |
| 21 | +from mlc_llm.cli.serve import EngineConfigOverride |
| 22 | +from mlc_llm.serve import EngineConfig |
| 23 | +from mlc_llm.support import argparse, logging |
| 24 | + |
| 25 | +logging.enable_logging() |
| 26 | +logger = logging.getLogger(__name__) |
| 27 | + |
| 28 | + |
| 29 | +def _parse_num_concurrent_requests(num_str: Optional[str]) -> Optional[List[int]]: |
| 30 | + if num_str is None: |
| 31 | + return None |
| 32 | + numbers = num_str.split(",") |
| 33 | + if any(not number.isdigit() for number in numbers): |
| 34 | + raise ValueError(f"Unrecognized num_concurrent_requests list: {numbers}") |
| 35 | + return list(int(number) for number in numbers) |
| 36 | + |
| 37 | + |
| 38 | +def _parse_mlc_engine_config(config_str: Optional[str]) -> EngineConfig: |
| 39 | + if config_str is None: |
| 40 | + return None |
| 41 | + engine_config_override = EngineConfigOverride.from_str(config_str) |
| 42 | + return EngineConfig( |
| 43 | + tensor_parallel_shards=engine_config_override.tensor_parallel_shards, |
| 44 | + max_num_sequence=engine_config_override.max_num_sequence, |
| 45 | + max_total_sequence_length=engine_config_override.max_total_seq_length, |
| 46 | + prefill_chunk_size=engine_config_override.prefill_chunk_size, |
| 47 | + sliding_window_size=engine_config_override.sliding_window_size, |
| 48 | + attention_sink_size=engine_config_override.attention_sink_size, |
| 49 | + max_history_size=engine_config_override.max_history_size, |
| 50 | + gpu_memory_utilization=engine_config_override.gpu_memory_utilization, |
| 51 | + spec_draft_length=engine_config_override.spec_draft_length, |
| 52 | + ) |
| 53 | + |
| 54 | + |
| 55 | +def _launch_mlc_server(args: argparse.argparse.Namespace): |
| 56 | + return mlc_llm.serve.PopenServer( |
| 57 | + model=args.tokenizer, |
| 58 | + mode="server", |
| 59 | + model_lib=args.mlc_model_lib, |
| 60 | + enable_tracing=False, |
| 61 | + host=args.host, |
| 62 | + port=args.port, |
| 63 | + engine_config=args.mlc_engine_config, |
| 64 | + ) |
| 65 | + |
| 66 | + |
| 67 | +def run_executor( |
| 68 | + executor: Executor, |
| 69 | + dataset: Dataset, |
| 70 | + tokenizer: AutoTokenizer, |
| 71 | + args: argparse.argparse.Namespace, |
| 72 | +) -> Dict[str, Any]: |
| 73 | + """Run the executor with the given dataset and args. Return the benchmark report dict.""" |
| 74 | + # Pre-process |
| 75 | + num_warmup_requests = executor.get_num_warmup_requests() |
| 76 | + pre_processor = SequentialProcessor( |
| 77 | + SampleRequests(args.num_requests + num_warmup_requests), |
| 78 | + AttachStreamFlag(args.stream), |
| 79 | + ) |
| 80 | + request_records = dataset.generate_request_records( |
| 81 | + args.input_len, |
| 82 | + args.output_len, |
| 83 | + args.input_len_std, |
| 84 | + args.output_len_std, |
| 85 | + ) |
| 86 | + request_records = pre_processor(request_records) |
| 87 | + assert len(request_records) == args.num_requests + num_warmup_requests |
| 88 | + warmup_requests = request_records[:num_warmup_requests] |
| 89 | + request_records = request_records[num_warmup_requests:] |
| 90 | + |
| 91 | + # Warmup and run |
| 92 | + logger.info( |
| 93 | + "Executor %s created for %s dataset at %s", |
| 94 | + type(executor).__name__, |
| 95 | + args.dataset, |
| 96 | + args.dataset_path, |
| 97 | + ) |
| 98 | + logger.info("Warmup with %d request(s)...", len(warmup_requests)) |
| 99 | + asyncio.run(executor.warmup(warmup_requests)) |
| 100 | + logger.info("Warmup finished. Start benchmarking...") |
| 101 | + request_records, duration = asyncio.run(executor.run_benchmark(request_records)) |
| 102 | + |
| 103 | + # Post-process |
| 104 | + request_records = MetricAnalyzer(tokenizer)(request_records) |
| 105 | + report = generate_metrics_summary(request_records, duration, args.num_requests, args.num_gpus) |
| 106 | + report = {**report, **executor.get_executor_feature_dict()} |
| 107 | + return report |
| 108 | + |
| 109 | + |
| 110 | +def main(args: argparse.argparse.Namespace): |
| 111 | + """Main benchmark entrance.""" |
| 112 | + random.seed(args.seed) |
| 113 | + np.random.seed(args.seed) |
| 114 | + |
| 115 | + mlc_server = None |
| 116 | + if args.mlc_model_lib: |
| 117 | + mlc_server = _launch_mlc_server(args) |
| 118 | + |
| 119 | + def _main(): |
| 120 | + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) |
| 121 | + dataset = create_dataset(args, tokenizer) |
| 122 | + api_endpoint = create_api_endpoint(args) |
| 123 | + executors = create_executors(args, api_endpoint) |
| 124 | + reports = [] |
| 125 | + for executor in executors: |
| 126 | + reports.append(run_executor(executor, dataset, tokenizer, args)) |
| 127 | + |
| 128 | + # Construct data frame |
| 129 | + df = convert_reports_to_df(reports) |
| 130 | + print(df) |
| 131 | + df.to_csv(args.output) |
| 132 | + logger.info("Benchmark results dumped to file %s", args.output) |
| 133 | + |
| 134 | + if mlc_server is not None: |
| 135 | + with mlc_server: |
| 136 | + _main() |
| 137 | + else: |
| 138 | + _main() |
| 139 | + |
| 140 | + |
| 141 | +if __name__ == "__main__": |
| 142 | + parser = argparse.ArgumentParser("MLC LLM benchmark") |
| 143 | + |
| 144 | + parser.add_argument( |
| 145 | + "--dataset", |
| 146 | + type=str, |
| 147 | + choices=SUPPORTED_DATASET, |
| 148 | + help=f"The benchmark dataset kind. Supporting {SUPPORTED_DATASET}", |
| 149 | + ) |
| 150 | + parser.add_argument( |
| 151 | + "--dataset-path", |
| 152 | + type=str, |
| 153 | + required=True, |
| 154 | + help="The dataset file path.", |
| 155 | + ) |
| 156 | + parser.add_argument( |
| 157 | + "--api-endpoint", |
| 158 | + type=str, |
| 159 | + choices=SUPPORTED_BACKENDS, |
| 160 | + default="openai", |
| 161 | + help="The API endpoint API for benchmarking.", |
| 162 | + ) |
| 163 | + parser.add_argument( |
| 164 | + "--tokenizer", |
| 165 | + type=str, |
| 166 | + required=True, |
| 167 | + help="The path of the tokenizer directory.", |
| 168 | + ) |
| 169 | + parser.add_argument( |
| 170 | + "--num-gpus", |
| 171 | + type=int, |
| 172 | + required=True, |
| 173 | + help="The number of GPUs used by the server. " |
| 174 | + "We need this to better analyze the throughput per GPU.", |
| 175 | + ) |
| 176 | + parser.add_argument( |
| 177 | + "--num-requests", |
| 178 | + type=int, |
| 179 | + required=True, |
| 180 | + help="The number of requests for benchmark.", |
| 181 | + ) |
| 182 | + parser.add_argument( |
| 183 | + "--num-concurrent-requests", |
| 184 | + type=_parse_num_concurrent_requests, |
| 185 | + help="The number(s) of concurrent requests to benchmark. " |
| 186 | + 'It can be either one integer or a list of integer separated by commas(","). ' |
| 187 | + "When specified, for each integer, the benchmark keeps these many consistent " |
| 188 | + "number of concurrently running requests.", |
| 189 | + ) |
| 190 | + parser.add_argument( |
| 191 | + "--request-rate", |
| 192 | + type=int, |
| 193 | + help="The request rate, denoting the number of new requests each second. " |
| 194 | + "When specified, the benchmark sends these many new requests each second.", |
| 195 | + ) |
| 196 | + parser.add_argument( |
| 197 | + "--input-len", |
| 198 | + type=int, |
| 199 | + help="The benchmark request average input length. Default to None, " |
| 200 | + "which means the request input length depends on the dataset being used.", |
| 201 | + ) |
| 202 | + parser.add_argument( |
| 203 | + "--input-len-std", |
| 204 | + type=float, |
| 205 | + default=0, |
| 206 | + help="The benchmark request input length standard deviation. Default to 0.", |
| 207 | + ) |
| 208 | + parser.add_argument( |
| 209 | + "--output-len", |
| 210 | + type=int, |
| 211 | + help="The benchmark request average output length. Default to None, " |
| 212 | + "which means the request output length depends on the dataset being used.", |
| 213 | + ) |
| 214 | + parser.add_argument( |
| 215 | + "--output-len-std", |
| 216 | + type=float, |
| 217 | + default=0, |
| 218 | + help="The benchmark request output length standard deviation. Default to 0.", |
| 219 | + ) |
| 220 | + parser.add_argument( |
| 221 | + "--stream", |
| 222 | + type=bool, |
| 223 | + default=True, |
| 224 | + help="Whether to benchmark stream responses. " |
| 225 | + "When not enabled, metrics such as time-to-first-token (TTFT) will not be available. " |
| 226 | + "Default to True.", |
| 227 | + ) |
| 228 | + parser.add_argument( |
| 229 | + # NOTE: The current implementation of server metrics still has some issues that need fixes, |
| 230 | + # which makes it not work to include server metrics. |
| 231 | + "--include-server-metrics", |
| 232 | + action="store_true", |
| 233 | + help="Whether to also benchmark the server side request metrics. " |
| 234 | + "This option is only available when benchmarking MLC server.", |
| 235 | + ) |
| 236 | + parser.add_argument( |
| 237 | + "--host", |
| 238 | + type=str, |
| 239 | + required=True, |
| 240 | + help="The host address of the backend API.", |
| 241 | + ) |
| 242 | + parser.add_argument( |
| 243 | + "--port", |
| 244 | + type=int, |
| 245 | + required=True, |
| 246 | + help="The port of the backend API.", |
| 247 | + ) |
| 248 | + parser.add_argument( |
| 249 | + "--timeout", |
| 250 | + type=float, |
| 251 | + help="The timeout limit of each request.", |
| 252 | + ) |
| 253 | + parser.add_argument( |
| 254 | + "--seed", |
| 255 | + type=int, |
| 256 | + default=0, |
| 257 | + help="The random number seed. Default to 0.", |
| 258 | + ) |
| 259 | + parser.add_argument( |
| 260 | + "--disable-tqdm", |
| 261 | + action="store_true", |
| 262 | + help="Whether to disable showing progress bar with tqdm during benchmarking.", |
| 263 | + ) |
| 264 | + parser.add_argument( |
| 265 | + "--mlc-model-lib", |
| 266 | + type=str, |
| 267 | + help="The model lib path when benchmarking MLC serve. " |
| 268 | + "When specified, the server is automatic launched and no external server launch is needed.", |
| 269 | + ) |
| 270 | + parser.add_argument( |
| 271 | + "--mlc-engine-config", |
| 272 | + type=_parse_mlc_engine_config, |
| 273 | + help="The engine config used when launch MLC server.", |
| 274 | + ) |
| 275 | + parser.add_argument( |
| 276 | + "--output", |
| 277 | + "-o", |
| 278 | + type=str, |
| 279 | + default="mlc_benchmark.csv", |
| 280 | + help="The path of the output file where to dump the benchmark results.", |
| 281 | + ) |
| 282 | + |
| 283 | + main(parser.parse_args()) |
0 commit comments