chore: import upstream snapshot with attribution

python/mlc_llm/bench/__init__.py

@@ -0,0 +1 @@
+"""Subdirectory of bench."""

python/mlc_llm/bench/__main__.py

@@ -0,0 +1,403 @@
+"""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())

python/mlc_llm/bench/api_endpoint.py

@@ -0,0 +1,452 @@
+"""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}"')

python/mlc_llm/bench/dataset.py

@@ -0,0 +1,878 @@
+"""MLC LLM benchmark dataset classes"""
+
+import argparse
+import json
+import random
+from datetime import datetime
+from typing import ClassVar, Dict, List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+import pandas as pd
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+from mlc_llm.bench.request_record import GroupedRequestRecord, Metrics, RequestRecord
+from mlc_llm.protocol.openai_api_protocol import (
+    ChatCompletionMessage,
+    ChatCompletionRequest,
+    DebugConfig,
+)
+
+
+class Dataset:
+    """The dataset base class."""
+
+    # We set a truncation limit of 100k.
+    truncate_length = int(1e5)
+    # For some that datasets (e.g., dataset that has shared common prefix),
+    # we need fake warmup requests to avoid prefilling common prefixes to the engine.
+    require_fake_warmup: bool = False
+    # Whether the dataset contains timestamps already.
+    # If the dataset comes with timestamps, the benchmark can just replay
+    # the requests according to their timestamps.
+    timestamp_available: bool = False
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        """Get the raw unprocessed request records of the dataset."""
+        raise NotImplementedError()
+
+
+class ShareGPTDataset(Dataset):
+    """The dataset class for ShareGPT dataset."""
+
+    _tokenized_dataset: List[Tuple[str, List[int], int]]  # noqa: UP006
+    apply_chat_template: bool
+
+    def __init__(
+        self, dataset_path: str, tokenizer: AutoTokenizer, apply_chat_template: bool
+    ) -> None:
+        self.apply_chat_template = apply_chat_template
+        with open(dataset_path, encoding="utf-8") as f:
+            raw_dataset = json.load(f)
+        # Filter out the conversations with less than 2 turns.
+        _dataset = [
+            (data["conversations"][0]["value"], data["conversations"][1]["value"])
+            for data in raw_dataset
+            if len(data["conversations"]) >= 2 and data["conversations"][0]["from"] == "human"
+        ]
+        # Tokenize the prompts and completions.
+        self.tokenizer = tokenizer
+        prompts = [prompt for prompt, _ in _dataset]
+        if apply_chat_template:
+            assert getattr(tokenizer, "chat_template", None) is not None, (
+                '"--apply-chat-template" is set but the tokenizer does not have chat template.'
+            )
+            prompts = [
+                tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}],
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+                for prompt in prompts
+            ]
+
+        prompt_token_ids = list(
+            tokenizer(
+                prompts,
+                truncation=True,
+                max_length=min(tokenizer.model_max_length, self.truncate_length),
+                add_special_tokens=False,
+            ).input_ids
+        )
+        completions = [completion for _, completion in _dataset]
+        completion_token_ids = tokenizer(
+            completions,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        self._tokenized_dataset: List[Tuple[str, List[int], int]] = []  # noqa: UP006
+        for i in range(len(_dataset)):
+            if (
+                len(prompt_token_ids[i]) < 4
+                or len(completion_token_ids[i]) < 4
+                or len(prompt_token_ids[i]) + len(completion_token_ids[i])
+                >= min(tokenizer.model_max_length, 8192)
+            ):
+                # Filter out sequences that are too short or too long
+                continue
+            self._tokenized_dataset.append(
+                (prompts[i], prompt_token_ids[i], len(completion_token_ids[i]))
+            )
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if self.apply_chat_template:
+            assert input_len is None, (
+                '"--apply-chat-template" is not supported when "--input-len" is specified.'
+            )
+
+        request_records = []
+        for prompt, input_token_ids, output_length in self._tokenized_dataset:
+            input_length = len(input_token_ids)
+            # If the request does not have enough length, discard it.
+            if input_len is not None and input_length < input_len + 4 * input_len_std:
+                continue
+
+            if input_len is not None:
+                input_length = round(
+                    float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0])
+                )
+                input_token_ids = input_token_ids[:input_length]
+                input_truncated = True
+            else:
+                input_truncated = False
+            if output_len is not None:
+                output_length = round(
+                    float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0])
+                )
+            elif output_length <= 1:
+                continue
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[
+                            {
+                                "role": "user",
+                                "content": (
+                                    self.tokenizer.decode(input_token_ids)
+                                    if input_truncated
+                                    else prompt
+                                ),
+                            }
+                        ],
+                        model="",
+                        max_tokens=output_length,
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=len(input_token_ids),
+                    ),
+                )
+            )
+        return request_records
+
+
+class LoogleDataset(Dataset):
+    """The dataset class for Loogle dataset."""
+
+    task2prompt: ClassVar[Dict[str, str]] = {  # noqa: UP006
+        "shortdep_qa": "Please answer the question based on the long texts below. \n{input}\nQuestion: {Q}\nAnswer: ",  # noqa: E501
+        "longdep_qa": "Please answer the question based on the long texts below. \n{input}\nQuestion: {Q}\nAnswer: ",  # noqa: E501
+        "longdep_summarization": "Please generate a summary of the below paper. \n{input}\n Summarization: ",  # noqa: E501
+        "shortdep_cloze": "Please fill in the clozes based on the given long texts below. Each of the placeholder '<mask-n>' in the question could be an entity of Person, Location or Organiocation. The same masks represent the same entity. Output a json format answer, for example: {{'<mask-0>': 'Bob', '<mask-1>': 'Gorrosion Magazine','<mask-2>': 'Bethel Horizon'}}\n{input}\n Question: {Q} What are the masked entities? \nAnswer:",  # noqa: E501
+    }
+    require_fake_warmup: bool = True
+
+    def __init__(self, tokenizer: AutoTokenizer, testset_name: str) -> None:
+        raw_dataset = load_dataset("bigainlco/LooGLE", testset_name, split="test")
+        self.tokenizer = tokenizer
+        self.dataset = []
+        self.prompt_format = self.task2prompt[testset_name]
+        prompts = []
+        generate_lens = []
+        questions = []
+        for data in raw_dataset:
+            prompt = data["input"]
+            prompts.append(prompt)
+            qa_pairs = eval(data["qa_pairs"])
+            questions.append([j["Q"] for j in qa_pairs])
+            generate_lens.append(
+                [len(tokenizer.encode(j["A"], add_special_tokens=False)) for j in qa_pairs]
+            )
+        prompt_token_ids = tokenizer(
+            prompts,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        for prompt, prompt_token_id, question, generate_len in zip(
+            prompts, prompt_token_ids, questions, generate_lens
+        ):
+            self.dataset.append((prompt, prompt_token_id, question, generate_len))
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        request_records = []
+        for prompt, input_token_ids, questions, generate_lens in self.dataset:
+            input_length = round(float(np.random.normal(loc=input_len, scale=input_len_std)))
+            if len(input_token_ids) > input_length:
+                input_token_ids = input_token_ids[:input_length]
+                prompt = self.tokenizer.decode(input_token_ids)
+            grouped_request_records = []
+            for question, generate_len in zip(questions, generate_lens):
+                json_obj = {"input": prompt, "Q": question}
+                full_prompt = self.prompt_format.format(**json_obj)
+
+                output_length = (
+                    round(float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0]))
+                    if output_len is not None
+                    else generate_len
+                )
+                grouped_request_records.append(
+                    RequestRecord(
+                        chat_cmpl=ChatCompletionRequest(
+                            messages=[
+                                {
+                                    "role": "user",
+                                    "content": full_prompt,
+                                }
+                            ],
+                            model="",
+                            max_tokens=output_length,
+                        ),
+                        metrics=Metrics(
+                            success=False,
+                            start_time=0,
+                            finish_time=0,
+                            end_to_end_latency_s=0,
+                            input_tokens=len(input_token_ids),
+                        ),
+                    )
+                )
+            request_records.append(
+                GroupedRequestRecord(
+                    # Create a dummy ChatCompletionRequest.
+                    chat_cmpl=ChatCompletionRequest(messages=[]),
+                    records=grouped_request_records,
+                )
+            )
+        return request_records
+
+
+class LLMPerfDataset(Dataset):
+    """The dataset class for LLMPerf dataset."""
+
+    def __init__(self, dataset_path: str, num_requests: int, tokenizer: AutoTokenizer) -> None:
+        self.tokenizer = tokenizer
+        self.num_requests = num_requests
+
+        with open(dataset_path, encoding="utf-8") as f:
+            untokenized_data = f.readlines()
+        # Tokenize the prompts and completions.
+        tokenized_data = tokenizer(
+            untokenized_data,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        tokenized_data_lengths = [len(tokens) for tokens in tokenized_data]
+        self.dataset: List[Tuple[str, List[int], int]] = list(  # noqa: UP006
+            zip(untokenized_data, tokenized_data, tokenized_data_lengths)
+        )
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int] = None,
+        output_len: Optional[int] = None,
+        input_len_std: float = 250,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if input_len is None or input_len < 40:
+            input_len = 550
+        if output_len is None:
+            output_len = 150
+
+        request_records = []
+        for _ in range(self.num_requests):
+            input_length = round(float(np.random.normal(loc=input_len, scale=input_len_std)))
+            output_length = round(float(np.random.normal(loc=output_len, scale=output_len_std)))
+
+            prompt = (
+                "Randomly stream lines from the following text "
+                f"with {output_length} output tokens. "
+                "Don't generate eos tokens:\n\n"
+            )
+
+            remaining_token_length = input_length - len(
+                self.tokenizer.encode(prompt, add_special_tokens=False)
+            )
+
+            random.shuffle(self.dataset)
+
+            while remaining_token_length > 0:
+                for text, tokens, token_length in self.dataset:
+                    if remaining_token_length < token_length:
+                        prompt += self.tokenizer.decode(tokens[:remaining_token_length])
+                    else:
+                        prompt += text
+
+                    remaining_token_length -= token_length
+                    if remaining_token_length < 0:
+                        break
+
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[{"role": "user", "content": prompt}],
+                        model="",
+                        max_tokens=output_length,
+                        debug_config=DebugConfig(ignore_eos=True),
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=input_length,
+                    ),
+                )
+            )
+        return request_records
+
+
+class JSONModeEvalDataset(Dataset):
+    """The dataset class for JSON dataset."""
+
+    def __init__(self, tokenizer: AutoTokenizer) -> None:
+        raw_dataset = load_dataset("NousResearch/json-mode-eval")
+        self.tokenizer = tokenizer
+        self.dataset = []
+        for data in raw_dataset["train"]:
+            messages = data["prompt"]
+            schema = {
+                "type": "json_object",
+                "schema": data["schema"],
+            }
+            num_tokens = 0
+            for message in messages:
+                num_tokens += len(
+                    self.tokenizer.encode(message["content"], add_special_tokens=False)
+                )
+            self.dataset.append((messages, schema, num_tokens))
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        request_records = []
+        for messages, schema, num_tokens in self.dataset:
+            # If the request does not have enough length, discard it.
+            if input_len is not None and num_tokens < input_len + 4 * input_len_std:
+                continue
+
+            if output_len is not None:
+                output_length = max(
+                    round(np.random.normal(loc=output_len, scale=output_len_std)), 1
+                )
+            else:
+                output_length = None
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[
+                            ChatCompletionMessage(content=message["content"], role=message["role"])
+                            for message in messages
+                        ],
+                        model="",
+                        max_tokens=output_length,
+                        response_format=schema,
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=num_tokens,
+                    ),
+                )
+            )
+        return request_records
+
+
+class ReActDataset(Dataset):
+    """The dataset class for replaying a given ReAct trace for benchmark purpose.
+    It is not an actual ReAct agent implementation.
+    """
+
+    _dataset: List[List[Tuple[str, int, int]]]  # noqa: UP006
+    require_fake_warmup: bool = True
+    prefix: str = """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types:
+(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
+(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage.
+(3) Finish[answer], which returns the answer and finishes the task.
+Here are some examples.
+Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into?
+Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area.
+Action 1: Search[Colorado orogeny]
+Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas.
+Thought 2: It does not mention the eastern sector. So I need to look up eastern sector.
+Action 2: Lookup[eastern sector]
+Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny.
+Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range.
+Action 3: Search[High Plains]
+Observation 3: High Plains refers to one of two distinct land regions:
+Thought 4: I need to instead search High Plains (United States).
+Action 4: Search[High Plains (United States)]
+Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3]
+Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft.
+Action 5: Finish[1,800 to 7,000 ft]
+Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
+Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after.
+Action 1: Search[Milhouse]
+Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
+Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after".
+Action 2: Lookup[named after]
+Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous.
+Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon.
+Action 3: Finish[Richard Nixon]
+Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture?
+Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups.
+Action 1: Search[Adam Clayton Powell]
+Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito'].
+Thought 2: To find the documentary, I can search Adam Clayton Powell (film).
+Action 2: Search[Adam Clayton Powell (film)]
+Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg.
+The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience.
+Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture.
+Action 3: Finish[The Saimaa Gesture]
+Question: What profession does Nicholas Ray and Elia Kazan have in common?
+Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common.
+Action 1: Search[Nicholas Ray]
+Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause.
+Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions.
+Action 2: Search[Elia Kazan]
+Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor.
+Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor.
+Action 3: Finish[director, screenwriter, actor]
+Question: Which magazine was started first Arthur's Magazine or First for Women?
+Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first.
+Action 1: Search[Arthur's Magazine]
+Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century.
+Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next.
+Action 2: Search[First for Women]
+Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989.
+Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first.
+Action 3: Finish[Arthur's Magazine]
+Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
+Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same.
+Action 1: Search[Pavel Urysohn]
+Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
+Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work.
+Action 2: Search[Leonid Levin]
+Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist.
+Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work.
+Action 3: Finish[yes]
+"""  # noqa: E501, RUF001
+
+    def __init__(self, dataset_path: str, tokenizer: AutoTokenizer) -> None:
+        raw_entries: List[Dict] = []  # noqa: UP006
+        with open(dataset_path) as fin:
+            for line in fin:
+                line_content = json.loads(line)
+                raw_entries += list({"question": k, "triplets": v} for k, v in line_content.items())
+
+        self._dataset = []
+        max_rounds = 0
+        for raw_entry in raw_entries:
+            processed_entry = []
+            question = raw_entry["question"]
+            triplets = raw_entry["triplets"]
+            seq = self.prefix + question
+            max_rounds = max(max_rounds, len(triplets) + 1)
+            output_lengths: List[int] = []  # noqa: UP006
+            for i, triplet in enumerate(triplets):
+                output_lengths.append(
+                    len(
+                        tokenizer(
+                            triplet["thought"]
+                            + "\nAction "
+                            + str(i + 1)
+                            + ": "
+                            + triplet["action"]
+                            + "\n",
+                            truncation=True,
+                            max_length=min(tokenizer.model_max_length, self.truncate_length),
+                            add_special_tokens=False,
+                        ).input_ids
+                    )
+                )
+
+            for i in range(1, len(triplets) + 2):
+                seq += "Thought " + str(i) + ":"
+                input_len = len(
+                    tokenizer(
+                        seq,
+                        truncation=True,
+                        max_length=min(tokenizer.model_max_length, self.truncate_length),
+                        add_special_tokens=False,
+                    ).input_ids
+                )
+                output_length = (
+                    output_lengths[i - 1]
+                    if i <= len(triplets)
+                    else int(sum(output_lengths) / len(triplets))
+                )
+                processed_entry.append((seq, input_len, output_length))
+                if i != len(triplets) + 1:
+                    seq += (
+                        triplets[i - 1]["thought"]
+                        + "\nAction "
+                        + str(i)
+                        + ": "
+                        + triplets[i - 1]["action"]
+                        + "\nObservation "
+                        + str(i)
+                        + ": "
+                        + triplets[i - 1]["observation"]
+                        + "\n"
+                    )
+            self._dataset.append(processed_entry)
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if input_len is not None or output_len is not None:
+            raise ValueError("ReAct dataset does not support specifying input/output length.")
+
+        request_records = []
+        for processed_entries in self._dataset:
+            grouped_request_records = []
+            for prompt, input_length, output_length in processed_entries:
+                grouped_request_records.append(
+                    RequestRecord(
+                        chat_cmpl=ChatCompletionRequest(
+                            messages=[{"role": "user", "content": prompt}],
+                            model="",
+                            max_tokens=output_length,
+                        ),
+                        metrics=Metrics(
+                            success=False,
+                            start_time=0,
+                            finish_time=0,
+                            end_to_end_latency_s=0,
+                            input_tokens=input_length,
+                        ),
+                    )
+                )
+            request_records.append(
+                GroupedRequestRecord(
+                    # Create a dummy ChatCompletionRequest.
+                    chat_cmpl=ChatCompletionRequest(messages=[]),
+                    records=grouped_request_records,
+                )
+            )
+        return request_records
+
+
+class WildChatDataset(Dataset):
+    """The dataset class for WildChat dataset."""
+
+    apply_chat_template: bool
+
+    def __init__(self, tokenizer: AutoTokenizer, apply_chat_template: bool) -> None:
+        raw_dataset = load_dataset("allenai/WildChat", split="train")
+        self.tokenizer = tokenizer
+        self.apply_chat_template = apply_chat_template
+
+        # Filter out the conversations with less than 2 turns.
+        _dataset = [
+            (entry["conversation"][0]["content"], entry["conversation"][1]["content"])
+            for entry in raw_dataset
+            if len(entry["conversation"]) >= 2
+            and entry["conversation"][0]["role"] == "user"
+            and entry["conversation"][1]["role"] == "assistant"
+        ]
+
+        prompts = []
+        completions = []
+        for prompt, completion in _dataset:
+            prompts.append(prompt)
+            completions.append(completion)
+        if apply_chat_template:
+            assert getattr(tokenizer, "chat_template", None) is not None, (
+                '"--apply-chat-template" is set but the tokenizer does not have chat template.'
+            )
+            prompts = [
+                tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}],
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+                for prompt in prompts
+            ]
+
+        prompt_token_ids = list(
+            tokenizer(
+                prompts,
+                truncation=True,
+                max_length=min(tokenizer.model_max_length, self.truncate_length),
+                add_special_tokens=False,
+            ).input_ids
+        )
+        completion_token_ids = tokenizer(
+            completions,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        self._tokenized_dataset: List[Tuple[str, List[int], int]] = []  # noqa: UP006
+        for i in range(len(_dataset)):
+            if len(prompt_token_ids[i]) < 4 or len(completion_token_ids[i]) < 4:
+                # Filter out sequences that are too short
+                continue
+            self._tokenized_dataset.append(
+                (prompts[i], prompt_token_ids[i], len(completion_token_ids[i]))
+            )
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if self.apply_chat_template:
+            assert input_len is None, (
+                '"--apply-chat-template" is not supported when "--input-len" is specified.'
+            )
+
+        request_records = []
+        for prompt, input_token_ids, output_length in self._tokenized_dataset:
+            input_length = len(input_token_ids)
+            # If the request does not have enough length, discard it.
+            if input_len is not None and input_length < input_len + 4 * input_len_std:
+                continue
+
+            if input_len is not None:
+                input_length = round(
+                    float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0])
+                )
+                input_token_ids = input_token_ids[:input_length]
+                input_truncated = True
+            else:
+                input_truncated = False
+            if output_len is not None:
+                output_length = round(
+                    float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0])
+                )
+            elif output_length <= 1:
+                continue
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[
+                            {
+                                "role": "user",
+                                "content": (
+                                    self.tokenizer.decode(input_token_ids)
+                                    if input_truncated
+                                    else prompt
+                                ),
+                            }
+                        ],
+                        model="",
+                        max_tokens=output_length,
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=len(input_token_ids),
+                    ),
+                )
+            )
+        return request_records
+
+
+class AzureLLMInferenceDataset(Dataset):
+    """The dataset class for AzureLLMInference dataset.
+    Reference: https://github.com/Azure/AzurePublicDataset
+    """
+
+    timestamp_available: bool = True
+
+    def __init__(self, dataset_path: str, tokenizer: AutoTokenizer) -> None:
+        df = pd.read_csv(dataset_path)
+        self.tokenizer = tokenizer
+
+        # Filter out the conversations with less than 2 turns.
+        self.dataset = [
+            (
+                entry["TIMESTAMP"],
+                min(
+                    entry["ContextTokens"],
+                    tokenizer.model_max_length,
+                    self.truncate_length,
+                ),
+                min(
+                    entry["GeneratedTokens"],
+                    tokenizer.model_max_length,
+                    self.truncate_length,
+                ),
+            )
+            for _, entry in df.iterrows()
+            if entry["ContextTokens"] >= 4 and entry["GeneratedTokens"] >= 4
+        ]
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        time_fmt = "%Y-%m-%d %H:%M:%S.%f"
+        start_time = datetime.strptime(self.dataset[0][0][:-1], time_fmt)
+        request_records = []
+        for timestamp, input_length, output_length in self.dataset:
+            # If the request does not have enough length, discard it.
+            if input_len is not None and input_length < input_len + 4 * input_len_std:
+                continue
+
+            if input_len is not None:
+                input_length = round(
+                    float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0])
+                )
+            if output_len is not None:
+                output_length = round(
+                    float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0])
+                )
+            elif output_length <= 1:
+                continue
+
+            prompt_token_ids = [
+                random.randint(0, self.tokenizer.vocab_size - 1) for _ in range(input_length)
+            ]
+            while True:
+                # Adjust the token ids until the retokenization on the decoded string
+                # matches the required input length.
+                prompt = self.tokenizer.decode(prompt_token_ids)
+                retokenized_token_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
+                if len(retokenized_token_ids) < input_length:
+                    prompt_token_ids = retokenized_token_ids + [
+                        random.randint(0, self.tokenizer.vocab_size - 1)
+                        for _ in range(input_length - len(retokenized_token_ids))
+                    ]
+                elif len(retokenized_token_ids) > input_length:
+                    prompt_token_ids = retokenized_token_ids[:input_length]
+                else:
+                    break
+
+            time_diff = (datetime.strptime(timestamp[:-1], time_fmt) - start_time).total_seconds()
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[{"role": "user", "content": prompt}],
+                        model="",
+                        max_tokens=output_length,
+                    ),
+                    timestamp=time_diff,
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=input_length,
+                    ),
+                )
+            )
+        return request_records
+
+
+SUPPORTED_DATASET = [
+    "sharegpt",
+    "llmperf",
+    "json-mode-eval",
+    "loogle",
+    "react",
+    "wildchat",
+    "azure-llm-inference",
+]
+
+
+def create_dataset(args: argparse.Namespace, tokenizer: AutoTokenizer) -> Dataset:
+    """Create a dataset instance with regard to the specified dataset kind and file path."""
+    if args.dataset_path is not None and not isinstance(args.dataset_path, str):
+        raise TypeError(f"Invalid dataset path {args.dataset_path}. Please use a string.")
+    if args.dataset is None and args.dataset_path is not None:
+        # Auto-detect the dataset kind by looking into the dataset path.
+        if "sharegpt" in args.dataset_path.lower():
+            args.dataset = "sharegpt"
+        else:
+            raise ValueError(
+                f"Unable to detect the dataset kind from dataset path {args.dataset_path}. "
+                'Please specify the dataset kind via "--dataset".'
+            )
+    if args.dataset == "sharegpt":
+        if args.dataset_path is None:
+            raise ValueError(
+                'ShareGPT dataset requires dataset path. Please specify it with "--dataset-path".'
+            )
+        return ShareGPTDataset(args.dataset_path, tokenizer, args.apply_chat_template)
+    if args.dataset == "llmperf":
+        if args.dataset_path is None:
+            raise ValueError(
+                'LLMPerf dataset requires dataset path. Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "LLMPerf dataset does not support applying chat template"
+        )
+        return LLMPerfDataset(
+            args.dataset_path,
+            (args.num_requests + args.num_warmup_requests) * 4,
+            tokenizer,
+        )
+    if args.dataset == "json-mode-eval":
+        assert args.apply_chat_template is False, (
+            "JSON mode evaluation does not support applying chat template"
+        )
+        return JSONModeEvalDataset(tokenizer)
+    if args.dataset == "loogle":
+        if args.dataset_path is None:
+            raise ValueError(
+                'Loogle dataset requires a testset name. Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "Loogle dataset does not support applying chat template"
+        )
+        return LoogleDataset(tokenizer, testset_name=args.dataset_path)
+    if args.dataset == "react":
+        if args.dataset_path is None:
+            raise ValueError(
+                'ReAct dataset requires dataset path. Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "ReAct dataset does not support applying chat template"
+        )
+        return ReActDataset(args.dataset_path, tokenizer)
+    if args.dataset == "wildchat":
+        return WildChatDataset(tokenizer, args.apply_chat_template)
+    if args.dataset == "azure-llm-inference":
+        if args.dataset_path is None:
+            raise ValueError(
+                "AzureLLMInference dataset requires dataset path. "
+                'Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "AzureLLMInference dataset does not support applying chat template"
+        )
+        return AzureLLMInferenceDataset(args.dataset_path, tokenizer)
+    raise ValueError(f"Unrecognized dataset {args.dataset}")

python/mlc_llm/bench/evaluation/gsm8k.py

@@ -0,0 +1,315 @@
+"""Eval GSM8K with MLCEngine."""
+
+import argparse
+import asyncio
+import json
+import random
+import re
+from datetime import datetime
+from pathlib import Path
+from typing import List, Literal, Optional  # noqa: UP035
+
+import tqdm
+
+from mlc_llm import AsyncMLCEngine
+
+DEVICES = ["cuda", "rocm", "metal", "vulkan"]
+ANSWER_TRIGGER = "The answer is"
+INVALID_ANS = "[invalid]"
+
+
+def extract_answer(text: str, regex: re.Pattern, select_index: int) -> str:
+    """Extract the answer from the text."""
+    match_all = regex.findall(text)
+    if len(match_all) == 0:
+        return INVALID_ANS
+    match = match_all[select_index]
+    if isinstance(match, tuple):
+        match = next(m for m in match if m)
+    match_str: str = match.strip()
+    match_str = match_str.lstrip("$").rstrip(".").replace(",", "")
+    return match_str
+
+
+def extract_ground_truth(text: str) -> str:
+    """Extract the ground truth from the text."""
+    return extract_answer(text, re.compile(r"#### (\-?[0-9\.\,]+)"), 0)
+
+
+def strict_extract_answer(text: str) -> str:
+    """Strictly extract the answer from the text."""
+    return extract_answer(text, re.compile(r"The answer is \$?(\-?[0-9\.\,]+)."), 0)
+
+
+def flexible_extract_answer(text: str) -> str:
+    """Extract the last number from the text."""
+    return extract_answer(text, re.compile(r"(-?[$0-9.,]{2,})|(-?[0-9]+)"), -1)
+
+
+def create_few_shot_prompt(n_shot: int, use_cot: bool, random_order=False) -> str:
+    """
+    Create a prompt for the few-shot learning task.
+
+    Note
+    ----
+    The examples are taken from the paper https://arxiv.org/pdf/2201.11903.pdf page 35.
+    """
+    question, chain, answer = [], [], []
+
+    question.append(
+        "There are 15 trees in the grove. "
+        "Grove workers will plant trees in the grove today. "
+        "After they are done, there will be 21 trees. "
+        "How many trees did the grove workers plant today?"
+    )
+    chain.append(
+        "There are 15 trees originally. "
+        "Then there were 21 trees after some more were planted. "
+        "So there must have been 21 - 15 = 6."
+    )
+    answer.append("6")
+
+    question.append(
+        "If there are 3 cars in the parking lot and 2 more cars arrive, "
+        "how many cars are in the parking lot?"
+    )
+    chain.append("There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5.")
+    answer.append("5")
+
+    question.append(
+        "Leah had 32 chocolates and her sister had 42. If they ate 35, "
+        "how many pieces do they have left in total?"
+    )
+    chain.append(
+        "Originally, Leah had 32 chocolates. "
+        "Her sister had 42. So in total they had 32 + 42 = 74. "
+        "After eating 35, they had 74 - 35 = 39."
+    )
+    answer.append("39")
+
+    question.append(
+        "Jason had 20 lollipops. He gave Denny some lollipops. Now Jason "
+        "has 12 lollipops. How many lollipops did Jason give to Denny?"
+    )
+    chain.append(
+        "Jason started with 20 lollipops. Then he had 12 after giving some "
+        "to Denny. So he gave Denny 20 - 12 = 8."
+    )
+    answer.append("8")
+
+    question.append(
+        "Shawn has five toys. For Christmas, he got two toys each from his "
+        "mom and dad. How many toys does he have now?"
+    )
+    chain.append(
+        "Shawn started with 5 toys. If he got 2 toys each from his mom and "
+        "dad, then that is 4 more toys. 5 + 4 = 9."
+    )
+    answer.append("9")
+
+    question.append(
+        "There were nine computers in the server room. Five more computers "
+        "were installed each day, from monday to thursday. "
+        "How many computers are now in the server room?"
+    )
+    chain.append(
+        "There were originally 9 computers. For each of 4 days, 5 more "
+        "computers were added. So 5 * 4 = 20 computers were added. "
+        "9 + 20 is 29."
+    )
+    answer.append("29")
+
+    question.append(
+        "Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On "
+        "wednesday, he lost 2 more. "
+        "How many golf balls did he have at the end of wednesday?"
+    )
+    chain.append(
+        "Michael started with 58 golf balls. After losing 23 on tuesday, "
+        "he had 58 - 23 = 35. After losing 2 more, "
+        "he had 35 - 2 = 33 golf balls."
+    )
+    answer.append("33")
+
+    question.append(
+        "Olivia has $23. She bought five bagels for $3 each. How much money does she have left?"
+    )
+    chain.append(
+        "Olivia had 23 dollars. "
+        "5 bagels for 3 dollars each will be 5 x 3 = 15 dollars. "
+        "So she has 23 - 15 dollars left. 23 - 15 is 8."
+    )
+    answer.append("8")
+
+    index_list = list(range(len(question)))
+    if random_order:
+        random.shuffle(index_list)
+
+    prompt = ""
+    for i in index_list[:n_shot]:
+        if use_cot:
+            prompt += f"Q: {question[i]}\nA: {chain[i]} {ANSWER_TRIGGER} {answer[i]}.\n\n"
+        else:
+            prompt += f"Question: {question[i]}\nAnswer: {ANSWER_TRIGGER} {answer[i]}.\n\n"
+    return prompt
+
+
+def create_prompt(question: str, n_shot: int, use_cot: bool, random_order: bool = False) -> str:
+    """Create a prompt for the few-shot learning task."""
+    prompt = create_few_shot_prompt(n_shot, use_cot, random_order)
+    if use_cot:
+        prompt += f"Q: {question}\nA:"
+    else:
+        prompt += f"Question: {question}\nAnswer:"
+    return prompt
+
+
+def parse_args():
+    """Parse command line arguments."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument(
+        "--dataset", type=Path, required=True, help="Path to GSM8K test dataset home."
+    )
+    parser.add_argument("--device", type=str, choices=["auto", *DEVICES], default="auto")
+    parser.add_argument("--model-lib", type=str, default=None)
+    parser.add_argument("--n-shot", type=int, default=8)
+    parser.add_argument("--disable_cot", action="store_true", default=False)
+    parser.add_argument("-bs", "--batch-size", type=int, default=16)
+    parser.add_argument("--log-dir", type=Path, default=None)
+    return parser.parse_args()
+
+
+async def send_request(
+    async_engine: AsyncMLCEngine,
+    prompts: List[str],  # noqa: UP006
+    semaphore: asyncio.Semaphore,
+):
+    """Send the calibration requests to the engine."""
+    tasks = []
+
+    async def generate_task(prompt):
+        async with semaphore:
+            return await async_engine.completions.create(
+                prompt=prompt,
+                stream=False,
+                max_tokens=512,
+                stop=["Q:", "Question:"],
+                temperature=0.0,
+            )
+
+    for prompt in prompts:
+        task = asyncio.create_task(generate_task(prompt))
+        tasks.append(task)
+
+    return await tqdm.asyncio.tqdm.gather(*tasks)
+
+
+async def evaluate(
+    model: str,
+    device: str,
+    dataset: Path,
+    model_lib: Optional[str],
+    n_shot: int,
+    use_cot: bool,
+    batch_size: int,
+    log_dir: Optional[Path],
+):
+    """Evaluate GSM8K for the model."""
+    mode: Literal["local", "interactive", "server"] = (
+        "server" if batch_size > 4 else "interactive" if batch_size == 1 else "local"
+    )
+    async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode=mode)
+
+    with open(dataset / "test.jsonl", encoding="utf-8") as file:
+        tests = [json.loads(line) for line in file]
+
+    prompts = [create_prompt(test["question"], n_shot, use_cot) for test in tests]
+    responses = await send_request(async_engine, prompts, asyncio.Semaphore(batch_size))
+    assert len(responses) == len(tests)
+
+    num_strict_correct, num_flexible_correct = 0, 0
+    num_tests = len(tests)
+    logs = []
+
+    for response, test in zip(responses, tests):
+        response_text = response.choices[0].text.strip()
+        gt_answer = extract_ground_truth(test["answer"])
+        assert gt_answer != INVALID_ANS
+        strict_answer = strict_extract_answer(response_text)
+        flexible_answer = flexible_extract_answer(response_text)
+
+        if gt_answer == strict_extract_answer(response_text):
+            # If the answer is exactly the same as the response, then it is correct
+            num_strict_correct += 1
+            num_flexible_correct += 1
+
+        elif gt_answer == flexible_extract_answer(response_text):
+            # Try flexible extract if the strict match fails
+            num_flexible_correct += 1
+
+        logs.append(
+            {
+                "question": test["question"],
+                "response": response_text,
+                "ground_truth": gt_answer,
+                "strict_answer": strict_answer,
+                "flexible_answer": flexible_answer,
+                "strict_match": gt_answer == strict_answer,
+                "flexible_match": gt_answer == flexible_answer,
+            }
+        )
+
+    results = {
+        "config": {
+            "model": model,
+            "device": device,
+            "model_lib": model_lib,
+            "n_shot": n_shot,
+            "use_cot": use_cot,
+        },
+        "results": {
+            "strict_match": num_strict_correct,
+            "flexible_match": num_flexible_correct,
+            "total": num_tests,
+        },
+    }
+    print(
+        f"Strict Matching Accuracy: {num_strict_correct} / {num_tests} = "
+        f"{num_strict_correct / num_tests * 100:.2f}%"
+    )
+    print(
+        f"Flexible Matching Accuracy: {num_flexible_correct} / {num_tests} = "
+        f"{num_flexible_correct / num_tests * 100:.2f}%"
+    )
+
+    if log_dir:
+        with open(log_dir / "summary.json", "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2)
+        with open(log_dir / "logs.json", "w", encoding="utf-8") as f:
+            json.dump(logs, f, indent=2)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    start_time = datetime.now()
+    log_dir: Optional[Path] = None
+    if args.log_dir is not None:
+        time_dir = start_time.strftime("%Y-%m-%d_%H-%M-%S")
+        log_dir = args.log_dir / time_dir
+        log_dir.mkdir(parents=True, exist_ok=True)
+    asyncio.run(
+        evaluate(
+            model=args.model,
+            device=args.device,
+            dataset=args.dataset,
+            model_lib=args.model_lib,
+            n_shot=args.n_shot,
+            use_cot=not args.disable_cot,
+            batch_size=args.batch_size,
+            log_dir=log_dir,
+        )
+    )
+    end_time = datetime.now()
+    print(f"Time used: {end_time - start_time}")

python/mlc_llm/bench/evaluation/mmlu.py

@@ -0,0 +1,245 @@
+"""Eval MMLU with MLCEngine."""
+
+import argparse
+import asyncio
+import csv
+import json
+import string
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional  # noqa: UP035
+
+import numpy as np
+import tqdm
+
+from mlc_llm import AsyncMLCEngine
+
+SUBJECTS = [
+    "abstract_algebra",
+    "anatomy",
+    "astronomy",
+    "business_ethics",
+    "clinical_knowledge",
+    "college_biology",
+    "college_chemistry",
+    "college_computer_science",
+    "college_mathematics",
+    "college_medicine",
+    "college_physics",
+    "computer_security",
+    "conceptual_physics",
+    "econometrics",
+    "electrical_engineering",
+    "elementary_mathematics",
+    "formal_logic",
+    "global_facts",
+    "high_school_biology",
+    "high_school_chemistry",
+    "high_school_computer_science",
+    "high_school_european_history",
+    "high_school_geography",
+    "high_school_government_and_politics",
+    "high_school_macroeconomics",
+    "high_school_mathematics",
+    "high_school_microeconomics",
+    "high_school_physics",
+    "high_school_psychology",
+    "high_school_statistics",
+    "high_school_us_history",
+    "high_school_world_history",
+    "human_aging",
+    "human_sexuality",
+    "international_law",
+    "jurisprudence",
+    "logical_fallacies",
+    "machine_learning",
+    "management",
+    "marketing",
+    "medical_genetics",
+    "miscellaneous",
+    "moral_disputes",
+    "moral_scenarios",
+    "nutrition",
+    "philosophy",
+    "prehistory",
+    "professional_accounting",
+    "professional_law",
+    "professional_medicine",
+    "professional_psychology",
+    "public_relations",
+    "security_studies",
+    "sociology",
+    "us_foreign_policy",
+    "virology",
+    "world_religions",
+]
+PADDING_LEN = max(len(subject) for subject in SUBJECTS)
+DEVICES = ["cuda", "rocm", "metal", "vulkan"]
+PROMPT_TEMPLATE = string.Template("$Q\nA. $A\nB. $B\nC. $C\nD. $D\nAnswer:")
+
+
+def parse_args():
+    """Parse command line arguments."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument(
+        "--dataset", type=Path, required=True, help="Path to MMLU test dataset home."
+    )
+    parser.add_argument("--device", type=str, choices=["auto", *DEVICES], default="auto")
+    parser.add_argument("--model-lib", type=str, default=None)
+    parser.add_argument("-s", "--subject", nargs="+", type=str, choices=SUBJECTS, default=SUBJECTS)
+    parser.add_argument("-bs", "--batch-size", type=int, default=16)
+    parser.add_argument("--log-dir", type=Path, default=None)
+    return parser.parse_args()
+
+
+async def send_request(
+    async_engine: AsyncMLCEngine,
+    prompts: List[str],  # noqa: UP006
+    semaphore: asyncio.Semaphore,
+    subject: str,
+):
+    """Send the calibration requests to the engine."""
+    tasks = []
+
+    async def generate_task(prompt):
+        async with semaphore:
+            return await async_engine.completions.create(
+                prompt=prompt,
+                stream=False,
+                max_tokens=1,
+                temperature=1.0,
+                logprobs=True,
+                top_logprobs=5,
+            )
+
+    for prompt in prompts:
+        task = asyncio.create_task(generate_task(prompt))
+        tasks.append(task)
+
+    return await tqdm.asyncio.tqdm.gather(
+        *tasks,
+        desc=f"Running {subject.ljust(PADDING_LEN)}",
+        bar_format="{desc} {percentage:3.0f}%|{bar}{r_bar}",
+    )
+
+
+async def evaluate(
+    model: str,
+    device: str,
+    dataset: Path,
+    model_lib: Optional[str],
+    subjects: List[str],  # noqa: UP006
+    semaphore: asyncio.Semaphore,
+    log_dir: Optional[Path],
+):
+    """Evaluate MMLU for the model."""
+    async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode="server")
+
+    results: Dict[str, Any] = {}  # noqa: UP006
+    for subject in subjects:
+        with open(dataset / "test" / f"{subject}_test.csv", encoding="utf-8") as csvfile:
+            tests = list(csv.reader(csvfile, delimiter=",", quotechar='"'))
+            assert all(len(test) == 6 for test in tests)
+
+        logs = []
+        num_correct = 0
+        prompts = [
+            PROMPT_TEMPLATE.substitute(Q=test[0], A=test[1], B=test[2], C=test[3], D=test[4])
+            for test in tests
+        ]
+        responses = await send_request(async_engine, prompts, semaphore, subject)
+
+        assert len(responses) == len(tests)
+        for response, test in zip(responses, tests):
+            token_logprobs = {}
+            logprobs = response.choices[0].logprobs.content[0].top_logprobs
+            for logprob in logprobs:
+                if logprob.token not in token_logprobs:
+                    token_logprobs[logprob.token] = logprob.logprob
+
+            abcd_logprobs = {}
+            for choice in ["A", "B", "C", "D"]:
+                abcd_logprobs[choice] = token_logprobs[choice] if choice in token_logprobs else -100
+
+            pred = {0: "A", 1: "B", 2: "C", 3: "D"}[int(np.argmax(list(abcd_logprobs.values())))]
+            num_correct += pred == test[5]
+
+            logs.append(
+                {
+                    "Question": {
+                        "Q": test[0],
+                        "A": test[1],
+                        "B": test[2],
+                        "C": test[3],
+                        "D": test[4],
+                    },
+                    "Answer": test[5],
+                    "Response": {
+                        "pred": pred,
+                        "logprobs": list(abcd_logprobs.values()),
+                    },
+                }
+            )
+
+        results[subject] = {
+            "correct": num_correct,
+            "total": len(tests),
+            "accuracy": num_correct / len(tests),
+        }
+
+        if log_dir:
+            with open(log_dir / "subjects" / f"{subject}.json", "w", encoding="utf-8") as f:
+                json.dump(logs, f, indent=2)
+
+    total_correct, total_tests = 0, 0
+    for subject, v in results.items():
+        num_correct, num_tests, accuracy = v["correct"], v["total"], v["accuracy"]
+        print(f"{subject}: {num_correct} / {num_tests} = {accuracy * 100:.2f}%")
+        total_correct += num_correct
+        total_tests += num_tests
+
+    total_accuracy = total_correct / total_tests
+    results["total"] = {
+        "correct": total_correct,
+        "total": total_tests,
+        "accuracy": total_accuracy,
+    }
+    print(f"Total accuracy: {total_correct} / {total_tests} = {total_accuracy * 100:.2f}%")
+
+    if log_dir:
+        results = {
+            "config": {
+                "model": model,
+                "device": device,
+                "model_lib": model_lib,
+                "subjects": subjects,
+            },
+            "results": results,
+        }
+        with open(log_dir / "summary.json", "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    start_time = datetime.now()
+    log_dir: Optional[Path] = None
+    if args.log_dir is not None:
+        time_dir = start_time.strftime("%Y-%m-%d_%H-%M-%S")
+        log_dir = args.log_dir / time_dir
+        (log_dir / "subjects").mkdir(parents=True, exist_ok=True)
+    asyncio.run(
+        evaluate(
+            model=args.model,
+            device=args.device,
+            dataset=args.dataset,
+            model_lib=args.model_lib,
+            subjects=args.subject,
+            semaphore=asyncio.Semaphore(args.batch_size),
+            log_dir=log_dir,
+        )
+    )
+    end_time = datetime.now()
+    print(f"Time used: {end_time - start_time}")

python/mlc_llm/bench/request_processor.py

@@ -0,0 +1,740 @@
+"""MLC LLM Bench Request"""
+
+import argparse
+import asyncio
+import concurrent.futures
+import copy
+import os
+import random
+import time
+from typing import Any, Callable, Dict, List, Optional  # noqa: UP035
+
+import numpy as np
+import requests
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+from mlc_llm.bench.api_endpoint import APIEndPoint
+from mlc_llm.bench.dataset import Dataset
+from mlc_llm.bench.request_record import GroupedRequestRecord, RequestRecord
+from mlc_llm.protocol.openai_api_protocol import (
+    ChatCompletionMessage,
+    ChatCompletionRequest,
+    DebugConfig,
+)
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+class RequestProcessor:
+    """The request processor base class.
+    Each processor can take a list of RequestRecord, applying the process,
+    and returning the processed RequestRecord in the end.
+    """
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        raise NotImplementedError()
+
+
+class LogMessage(RequestProcessor):
+    """The processor that prints the logger message."""
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        logger.info(self.message)
+        return request_records
+
+
+class SampleRequests(RequestProcessor):
+    """The processor that samples requests out from the given request list."""
+
+    def __init__(self, num_requests: int, take_first_x_requests: bool = False) -> None:
+        self.num_requests = num_requests
+        # If `take_first_x_requests` is True, the first `num_requests` requests
+        # are returned and sampling will not happen.
+        self.take_first_x_requests = take_first_x_requests
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        assert len(request_records) > 0, "Empty input request record."
+
+        # We expect the input request records to be all grouped or all plain.
+        if isinstance(request_records[0], GroupedRequestRecord):
+            assert all(isinstance(record, GroupedRequestRecord) for record in request_records)
+            return self._sample_from_grouped_request_records(request_records)
+
+        assert all(not isinstance(record, GroupedRequestRecord) for record in request_records)
+        return self._sample_from_plain_request_records(request_records)
+
+    def _sample_from_plain_request_records(
+        self,
+        request_records: List[RequestRecord],  # noqa: UP006
+    ) -> List[RequestRecord]:  # noqa: UP006
+        samples: List[RequestRecord] = []  # noqa: UP006
+        if self.take_first_x_requests:
+            if len(request_records) < self.num_requests:
+                raise ValueError(
+                    f"Insufficient requests. Requiring {self.num_requests} requests "
+                    f"but only {len(request_records)} are available."
+                )
+            samples = copy.deepcopy(list(request_records[: self.num_requests]))
+        else:
+            while len(samples) < self.num_requests:
+                # Create a new list so that the in-place shuffle does not mutate the input list.
+                records = list(request_records)
+                random.shuffle(records)
+                samples += copy.deepcopy(records)
+            samples = samples[: self.num_requests]
+        for i, record in enumerate(samples):
+            record.request_id = i
+        return samples
+
+    def _sample_from_grouped_request_records(
+        self,
+        grouped_request_records: List[GroupedRequestRecord],  # noqa: UP006
+    ) -> List[RequestRecord]:  # noqa: UP006
+        num_total_available_requests = sum(
+            len(record.records) for record in grouped_request_records
+        )
+        if self.num_requests > num_total_available_requests:
+            raise ValueError(
+                "Due to the existence of shared common prefixes, we do not allow "
+                "benchmarking with requests more than the available requests in the dataset. "
+                f"The required number of requests {self.num_requests} exceeds the "
+                f"number of total available requests {num_total_available_requests}."
+            )
+
+        # Create a new list so that the in-place shuffle does not mutate the input list.
+        records = list(grouped_request_records)
+        if not self.take_first_x_requests:
+            random.shuffle(records)
+        remaining = self.num_requests
+        samples: List[RequestRecord] = []  # noqa: UP006
+        for grouped_request_record in grouped_request_records:
+            num_used_requests = min(len(grouped_request_record.records), remaining)
+            samples += grouped_request_record.records[:num_used_requests]
+            remaining -= num_used_requests
+            if remaining == 0:
+                break
+        for i, record in enumerate(samples):
+            record.request_id = i
+        return samples
+
+
+class AttachModelName(RequestProcessor):
+    """The processor that attaches model name to requests."""
+
+    def __init__(self, model: str) -> None:
+        self.model = model
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            request_record.chat_cmpl.model = self.model
+        return request_records
+
+
+class AttachRequestRateTimestamp(RequestProcessor):
+    """The processor that applies timestamps to the requests."""
+
+    def __init__(self, request_rate: np.float32) -> None:
+        self.request_rate = request_rate
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        timestamp = 0.0
+        for request_record in request_records:
+            assert request_record.timestamp is None, "The request record already has a timestamp"
+            request_record.timestamp = timestamp
+            timestamp += float(np.random.exponential(1.0 / self.request_rate))
+        return request_records
+
+
+class AttachExecutionFeature(RequestProcessor):
+    """The processor that attaches execution features to all requests"""
+
+    def __init__(self, exec_feature: Dict[str, Any]) -> None:  # noqa: UP006
+        self.exec_feature = exec_feature
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            assert request_record.metrics is not None
+            request_record.metrics.exec_feature = self.exec_feature
+        return request_records
+
+
+class AttachStreamFlag(RequestProcessor):
+    """The processor that attaches the stream flag to the requests."""
+
+    def __init__(self, stream: Optional[bool]) -> None:
+        self.stream = stream
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        if self.stream is None:
+            return request_records
+        for request_record in request_records:
+            request_record.chat_cmpl.stream = self.stream
+        return request_records
+
+
+class AttachSamplingOptions(RequestProcessor):
+    """The processor that attaches the stream flag to the requests."""
+
+    def __init__(self, temperature: float, top_p: float, ignore_eos: bool) -> None:
+        self.temperature = temperature
+        self.top_p = top_p
+        self.ignore_eos = ignore_eos
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            request_record.chat_cmpl.temperature = self.temperature
+            request_record.chat_cmpl.top_p = self.top_p
+            request_record.chat_cmpl.frequency_penalty = 0.0
+            request_record.chat_cmpl.presence_penalty = 0.0
+            request_record.chat_cmpl.tool_choice = "none"
+            if self.ignore_eos:
+                request_record.chat_cmpl.debug_config = DebugConfig(ignore_eos=True)
+        return request_records
+
+
+class ScaleTimestamp(RequestProcessor):
+    """Scale the timestamp of requests by the given scale factor."""
+
+    def __init__(self, timestamp_scale: float):
+        self.timestamp_scale = timestamp_scale
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            if request_record.timestamp is None:
+                raise ValueError(
+                    f"The timestamp of request {request_record} has not been initialized."
+                )
+            request_record.timestamp *= self.timestamp_scale
+        return request_records
+
+
+class MetricAnalyzer(RequestProcessor):
+    """The processor that analyzes the raw benchmark results and computes more detailed metrics."""
+
+    def __init__(self, tokenizer: AutoTokenizer) -> None:
+        self.tokenizer = tokenizer
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        updated_records = []
+        for request_record in request_records:
+            metrics = request_record.metrics
+            if not metrics.success:
+                assert request_record.error_msg is not None
+                continue
+
+            metrics.output_tokens = len(
+                self.tokenizer.encode(request_record.output_str, add_special_tokens=False)
+            )
+            first_chunk_output_tokens = len(
+                self.tokenizer.encode(
+                    request_record.first_chunk_output_str, add_special_tokens=False
+                )
+            )
+            if metrics.output_tokens <= first_chunk_output_tokens:
+                metrics.success = False
+                request_record.error_msg = (
+                    f"Total output token num ({metrics.output_tokens}) equals "
+                    f'the first chunk output token. Output text "{request_record.output_str}", '
+                    f'first chunk output text "{request_record.first_chunk_output_str}"'
+                )
+                continue
+            assert metrics.input_tokens > 0, "Invalid prompt tokens"
+            metrics.inter_token_latency_s = metrics.end_to_end_latency_s / metrics.output_tokens
+            if metrics.time_to_first_token_s is None:
+                metrics.time_to_first_token_s = 0
+            metrics.time_per_output_token_s = (
+                metrics.end_to_end_latency_s - metrics.time_to_first_token_s
+            ) / (metrics.output_tokens - first_chunk_output_tokens)
+            updated_records.append(request_record)
+        return updated_records
+
+
+class WarmupAndRun(RequestProcessor):
+    """The processor that runs warmup first and then runs the benchmark with the given pipeline."""
+
+    def __init__(
+        self,
+        num_warmup_requests: int,
+        num_benchmark_requests: int,
+        pipeline: RequestProcessor,
+        cuda_profile_url: Optional[str],
+        fake_warmup: bool = False,
+    ) -> None:
+        self.num_warmup_requests = num_warmup_requests
+        self.num_benchmark_requests = num_benchmark_requests
+        self.pipeline = pipeline
+        self.cuda_profile_url = cuda_profile_url
+        self.fake_warmup = fake_warmup
+
+    def generate_fake_warmup_requests(
+        self, num_warmup_requests: int, example_request: RequestRecord
+    ) -> List[RequestRecord]:  # noqa: UP006
+        records = []
+        for _ in range(num_warmup_requests):
+            record = copy.deepcopy(example_request)
+            record.chat_cmpl = ChatCompletionRequest(
+                messages=[
+                    {
+                        "role": "user",
+                        "content": "Please output arbitrary coherent sentences. Do not output eos token.",  # noqa: E501
+                    }
+                ],
+                model="",
+                max_tokens=128,
+            )
+            records.append(record)
+        return records
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        # Warmup
+        if self.fake_warmup:
+            assert len(request_records) == self.num_benchmark_requests
+            benchmark_requests = request_records
+            example_request = benchmark_requests[0]
+            warmup_requests = self.generate_fake_warmup_requests(
+                self.num_warmup_requests, example_request=example_request
+            )
+        else:
+            assert len(request_records) == self.num_warmup_requests + self.num_benchmark_requests
+            benchmark_requests = request_records[: -self.num_warmup_requests]
+            warmup_requests = request_records[-self.num_warmup_requests :]
+        for request_record in warmup_requests:
+            request_record.timestamp = 0 if request_record.timestamp is not None else None
+        warmup_requests = self._process_warmup_requests(warmup_requests)
+        logger.info("Warmup with %d request(s)...", self.num_warmup_requests)
+        self.pipeline(warmup_requests)
+
+        # Then run benchmark
+        if self.cuda_profile_url is not None:
+            cuda_profiler_start_url = self.cuda_profile_url + "/debug/cuda_profiler_start"
+            cuda_profiler_start_response = requests.post(cuda_profiler_start_url, timeout=60)
+            assert cuda_profiler_start_response.status_code == 200
+        logger.info("Warmup finished. Start benchmarking...")
+        updated_request_records = self.pipeline(benchmark_requests)
+        if self.cuda_profile_url is not None:
+            cuda_profiler_stop_url = self.cuda_profile_url + "/debug/cuda_profiler_stop"
+            cuda_profiler_stop_response = requests.post(cuda_profiler_stop_url, timeout=60)
+            assert cuda_profiler_stop_response.status_code == 200
+
+        return updated_request_records
+
+    def _process_warmup_requests(self, warmup_requests: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        if len(warmup_requests) == 0:
+            return warmup_requests
+        # NOTE: to warm up the server for as more different batch sizes as possible,
+        # we usese 128 output tokens for the first request and use two more tokens
+        # for every followup request.
+        # Setting a high temperature and top-p to avoid early stop as much as possible.
+        warmup_requests[0].chat_cmpl.max_tokens = 128
+        for i in range(1, len(warmup_requests)):
+            warmup_requests[i].chat_cmpl.max_tokens = (
+                warmup_requests[i - 1].chat_cmpl.max_tokens + 1
+            )
+            warmup_requests[i].chat_cmpl.temperature = 2.0
+            warmup_requests[i].chat_cmpl.top_p = 1.0
+        return warmup_requests
+
+
+class SequentialProcessor(RequestProcessor):
+    """The processor that sequentially applies a list of processors in order."""
+
+    processors: List[RequestProcessor]  # noqa: UP006
+
+    def __init__(self, *processors: RequestProcessor) -> None:
+        self.processors = list(processors)
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for processor in self.processors:
+            request_records = processor(request_records)
+        return request_records
+
+
+class Executor(RequestProcessor):
+    """The executor base class, denoting the kind of benchmark mode."""
+
+    def __init__(
+        self,
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        num_processes: int,
+        disable_tqdm: bool,
+    ) -> None:
+        self.f_create_api_endpoint = f_create_api_endpoint
+        self.disable_tqdm = disable_tqdm
+        self.num_processes = num_processes
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        raise NotImplementedError()
+
+
+class FixedConcurrentRequestExecutor(Executor):
+    """The benchmark executor of fixing the number of concurrent requests."""
+
+    def __init__(
+        self,
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        num_processes: Optional[int],
+        disable_tqdm: bool,
+        num_concurrent_requests: int,
+        multi_round: bool,
+    ) -> None:
+        if num_processes is None:
+            # We assign each process at most 32 concurrent requests to send
+            # so that the asyncio pressure will not be too much.
+            num_processes = min((num_concurrent_requests + 31) // 32, 10)
+        super().__init__(f_create_api_endpoint, num_processes, disable_tqdm)
+        self.num_concurrent_requests = num_concurrent_requests
+        self.multi_round = multi_round
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        partitions: List[List[RequestRecord]] = [  # noqa: UP006
+            request_records[slice(i, len(request_records), self.num_processes)]
+            for i in range(self.num_processes)
+        ]
+        # Package "tokenizers" reports warnings with multiprocessing.
+        # We disable "TOKENIZERS_PARALLELISM" to depress the warnings.
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+        pbar = None if self.disable_tqdm else tqdm(total=len(request_records))
+        with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_processes) as pool:
+            futures = [
+                pool.submit(
+                    FixedConcurrentRequestExecutor._process_task,
+                    self.f_create_api_endpoint,
+                    partition,
+                    self.num_concurrent_requests // self.num_processes
+                    + int(i < self.num_concurrent_requests % self.num_processes),
+                    self.multi_round,
+                )
+                for i, partition in enumerate(partitions)
+            ]
+            results: List[RequestRecord] = []  # noqa: UP006
+            for i, future in enumerate(concurrent.futures.as_completed(futures)):
+                results.extend(future.result())
+                if pbar is not None:
+                    pbar.update(len(partitions[i]))
+
+        return results
+
+    @staticmethod
+    def _process_task(
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        request_records: List[RequestRecord],  # noqa: UP006
+        num_concurrent_requests: int,
+        multi_round: bool,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if len(request_records) == 0:
+            return []
+        chat_history: List[List[ChatCompletionMessage]] = [  # noqa: UP006
+            [] for _ in range(num_concurrent_requests)
+        ]
+
+        async def process_task_impl(
+            f_create_api_endpoint: Callable[[], APIEndPoint],
+            request_records: List[RequestRecord],  # noqa: UP006
+            num_concurrent_requests: int,
+            multi_round: bool,
+        ) -> List[RequestRecord]:  # noqa: UP006
+            api_endpoint = f_create_api_endpoint()
+            updated_request_records: List[RequestRecord] = [None for _ in request_records]  # noqa: UP006
+            async with api_endpoint:
+                num_sent_request = 0
+
+                async def _task(i: int) -> None:
+                    nonlocal num_sent_request
+                    while True:
+                        if num_sent_request == len(request_records):
+                            break
+                        idx = num_sent_request
+                        num_sent_request += 1
+                        request = request_records[idx]
+
+                        if multi_round:
+                            request.chat_cmpl.messages = (
+                                chat_history[i] + request.chat_cmpl.messages
+                            )
+
+                        updated_request_records[idx] = await api_endpoint(request)
+
+                        if multi_round:
+                            chat_history[i] = [
+                                *updated_request_records[idx].chat_cmpl.messages,
+                                ChatCompletionMessage(
+                                    content=updated_request_records[idx].output_str,
+                                    role="assistant",
+                                ),
+                            ]
+
+                tasks = [asyncio.create_task(_task(i)) for i in range(num_concurrent_requests)]
+                await asyncio.gather(*tasks)
+
+            return updated_request_records
+
+        return asyncio.run(
+            process_task_impl(
+                f_create_api_endpoint,
+                request_records,
+                num_concurrent_requests,
+                multi_round,
+            )
+        )
+
+
+class FixTimestampExecutor(Executor):
+    """The benchmark executor of fixing the timestamps of sending requests."""
+
+    def __init__(
+        self,
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        num_processes: Optional[int],
+        disable_tqdm: bool,
+        max_schedule_gap: float,
+        num_requests: int,
+    ) -> None:
+        if num_processes is None:
+            # We assign each process at most 32 requests to send
+            # so that the asyncio pressure will not be too much.
+            num_processes = min((num_requests + 31) // 32, 10)
+        super().__init__(f_create_api_endpoint, num_processes, disable_tqdm)
+        self.max_schedule_gap = max_schedule_gap
+        self.num_requests = num_requests
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        assert len(request_records) > 0
+        assert all(request_record.timestamp is not None for request_record in request_records)
+        # Sort the request records in timestamp ascending order before partitioning.
+        request_records.sort(key=lambda request_record: request_record.timestamp)
+        base_timestamp = request_records[0].timestamp
+        partitions: List[List[RequestRecord]] = [  # noqa: UP006
+            request_records[slice(i, len(request_records), self.num_processes)]
+            for i in range(self.num_processes)
+        ]
+        base_sys_time = time.time()
+        # Package "tokenizers" reports warnings with multiprocessing.
+        # We disable "TOKENIZERS_PARALLELISM" to depress the warnings.
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+        pbar = None if self.disable_tqdm else tqdm(total=len(request_records))
+        with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_processes) as pool:
+            futures = [
+                pool.submit(
+                    FixTimestampExecutor._process_task,
+                    self.f_create_api_endpoint,
+                    partition,
+                    base_timestamp,
+                    base_sys_time,
+                    self.max_schedule_gap,
+                )
+                for partition in partitions
+            ]
+            results: List[RequestRecord] = []  # noqa: UP006
+            for i, future in enumerate(concurrent.futures.as_completed(futures)):
+                results.extend(future.result())
+                if pbar is not None:
+                    pbar.update(len(partitions[i]))
+
+        return results
+
+    @staticmethod
+    def _process_task(
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        request_records: List[RequestRecord],  # noqa: UP006
+        base_timestamp: float,
+        base_sys_time: float,
+        max_schedule_gap: float,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if len(request_records) == 0:
+            return []
+
+        async def process_task_impl(
+            f_create_api_endpoint: Callable[[], APIEndPoint],
+            request_records: List[RequestRecord],  # noqa: UP006
+            base_timestamp: float,
+            base_sys_time: float,
+            max_schedule_gap: float,
+        ) -> List[RequestRecord]:  # noqa: UP006
+            api_endpoint = f_create_api_endpoint()
+            loop = asyncio.get_running_loop()
+            # Get the delta time to convert system time to the loop time.
+            # We must use the system time `time.time()` which is consistent across processes.
+            loop_sys_delta_time = loop.time() - time.time()
+            updated_request_records: List[RequestRecord] = []  # noqa: UP006
+            async with api_endpoint:
+
+                async def _task(request_record: RequestRecord) -> None:
+                    updated_request_records.append(await api_endpoint(request_record))
+
+                tasks = []
+                for request_record in request_records:
+                    launch_time = (
+                        (request_record.timestamp - base_timestamp)
+                        + (base_sys_time + max_schedule_gap)
+                        + loop_sys_delta_time
+                    )
+                    loop.call_at(
+                        launch_time,
+                        lambda record: tasks.append(asyncio.create_task(_task(record))),
+                        request_record,
+                    )
+                    # Sleep to allow runs of other scheduled tasks if any.
+                    await asyncio.sleep(max(launch_time - loop.time() - max_schedule_gap, 0))
+
+                # Sleep until all the tasks are launched.
+                await asyncio.sleep(launch_time - loop.time() + max_schedule_gap)
+                # Wait for all tasks to be scheduled
+                assert len(tasks) == len(request_records)
+                await asyncio.gather(*tasks)
+
+            assert len(updated_request_records) == len(request_records)
+            return updated_request_records
+
+        return asyncio.run(
+            process_task_impl(
+                f_create_api_endpoint,
+                request_records,
+                base_timestamp,
+                base_sys_time,
+                max_schedule_gap,
+            )
+        )
+
+
+def create_pipelines(
+    args: argparse.Namespace,
+    f_create_api_endpoint: Callable[[], APIEndPoint],
+    dataset: Dataset,
+) -> List[RequestProcessor]:  # noqa: UP006
+    """Creating request processing pipelines with regard to the specified args."""
+    cuda_profile_url = f"http://{args.host}:{args.port}" if args.cuda_profile else None
+    pipelines: List[RequestProcessor] = []  # noqa: UP006
+    if args.num_concurrent_requests is not None:
+        if args.request_rate is not None:
+            raise ValueError(
+                'Both "num_concurrent_requests" and "request_rate" are specified. '
+                "Please specify only one of them."
+            )
+        if args.replay_timestamp_scale is not None:
+            raise ValueError(
+                "Dataset replay is unsupported when fixing number of concurrent requests."
+            )
+        for num_concurrent_requests in args.num_concurrent_requests:
+            num_warmup_requests = (
+                args.num_warmup_requests
+                if args.num_warmup_requests is not None
+                else num_concurrent_requests
+            )
+            pipelines.append(
+                SequentialProcessor(
+                    LogMessage(f"Fixing number of concurrent requests: {num_concurrent_requests}"),
+                    SampleRequests(args.num_requests + num_warmup_requests),
+                    AttachModelName(args.tokenizer),
+                    AttachStreamFlag(args.stream),
+                    AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos),
+                    AttachExecutionFeature({"num_concurrent_requests": num_concurrent_requests}),
+                    WarmupAndRun(
+                        num_warmup_requests=num_warmup_requests,
+                        num_benchmark_requests=args.num_requests,
+                        pipeline=FixedConcurrentRequestExecutor(
+                            f_create_api_endpoint,
+                            args.num_process_workers,
+                            args.disable_tqdm,
+                            num_concurrent_requests,
+                            args.multi_round,
+                        ),
+                        cuda_profile_url=cuda_profile_url,
+                        fake_warmup=dataset.require_fake_warmup,
+                    ),
+                )
+            )
+        return pipelines
+    if args.request_rate is not None:
+        if args.num_warmup_requests is None:
+            raise ValueError(
+                "Please specify the number of warmup requests via "
+                '"--num-warmup-requests" when fixing request rate.'
+            )
+        if args.replay_timestamp_scale is not None:
+            raise ValueError("Dataset replay is unsupported when fixing request rates.")
+        num_total_requests = int(
+            args.num_requests if not args.per_gpu_workload else args.num_requests * args.num_gpus
+        )
+        if dataset.require_fake_warmup:
+            num_samples = num_total_requests
+        else:
+            num_samples = num_total_requests + args.num_warmup_requests
+        return [
+            SequentialProcessor(
+                LogMessage(f"Fixing request rate: {request_rate}"),
+                SampleRequests(num_samples),
+                AttachModelName(args.tokenizer),
+                AttachRequestRateTimestamp(
+                    request_rate if not args.per_gpu_workload else request_rate * args.num_gpus
+                ),
+                AttachStreamFlag(args.stream),
+                AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos),
+                AttachExecutionFeature({"request_rate": float(request_rate)}),
+                WarmupAndRun(
+                    num_warmup_requests=args.num_warmup_requests,
+                    num_benchmark_requests=num_total_requests,
+                    pipeline=FixTimestampExecutor(
+                        f_create_api_endpoint,
+                        args.num_process_workers,
+                        args.disable_tqdm,
+                        args.max_schedule_gap,
+                        args.num_requests,
+                    ),
+                    cuda_profile_url=cuda_profile_url,
+                    fake_warmup=dataset.require_fake_warmup,
+                ),
+            )
+            for request_rate in args.request_rate
+        ]
+
+    # Default: dataset replay mode
+    # The dataset must come with timestamps.
+    if not dataset.timestamp_available:
+        raise ValueError(
+            "The dataset does not have timestamps, so dataset replay is unsupported. "
+            'Please specify one of "num_concurrent_requests" '
+            'and "request_rate".'
+        )
+    if args.per_gpu_workload:
+        raise ValueError("Fixing per-GPU workload is not compatible with dataset replay.")
+    if args.num_warmup_requests is None:
+        raise ValueError(
+            "Please specify the number of warmup requests via "
+            '"--num-warmup-requests" for dataset replay.'
+        )
+    timestamp_scale = args.replay_timestamp_scale or 1.0
+    if dataset.require_fake_warmup:
+        num_samples = args.num_requests
+    else:
+        num_samples = args.num_requests + args.num_warmup_requests
+    return [
+        SequentialProcessor(
+            LogMessage(f"Dataset replay with time scaling of {timestamp_scale}"),
+            SampleRequests(num_samples, take_first_x_requests=True),
+            AttachModelName(args.tokenizer),
+            ScaleTimestamp(timestamp_scale),
+            AttachStreamFlag(args.stream),
+            AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos),
+            AttachExecutionFeature({"timestamp_scale": timestamp_scale}),
+            WarmupAndRun(
+                num_warmup_requests=args.num_warmup_requests,
+                num_benchmark_requests=args.num_requests,
+                pipeline=FixTimestampExecutor(
+                    f_create_api_endpoint,
+                    args.num_process_workers,
+                    args.disable_tqdm,
+                    args.max_schedule_gap,
+                    args.num_requests,
+                ),
+                cuda_profile_url=cuda_profile_url,
+                fake_warmup=dataset.require_fake_warmup,
+            ),
+        )
+    ]

python/mlc_llm/bench/request_record.py

@@ -0,0 +1,278 @@
+"""MLC LLM Bench Request"""
+
+from typing import Any, Dict, List, Optional, Tuple, Union  # noqa: UP035
+
+import pandas as pd
+from pydantic import BaseModel
+
+from mlc_llm.protocol.openai_api_protocol import ChatCompletionRequest
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ServerMetrics(BaseModel):
+    """The metrics from the server side."""
+
+    input_tokens: int
+    prefill_tokens: int
+    output_tokens: int
+    end_to_end_latency_s: float
+    prefill_tokens_per_s: float
+    inter_token_latency_s: float
+    time_per_output_token_s: float
+    time_to_first_token_s: Optional[float] = None
+
+
+class Metrics(BaseModel):
+    """The list of metric keys"""
+
+    success: bool
+    start_time: float
+    finish_time: float
+    end_to_end_latency_s: float
+
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    inter_token_latency_s: Optional[float] = None
+    time_per_output_token_s: Optional[float] = None
+    time_to_first_token_s: Optional[float] = None
+    server_metrics: Optional[ServerMetrics] = None
+
+    exec_feature: Optional[Dict[str, Any]] = None  # noqa: UP006
+
+
+class RequestRecord(BaseModel):
+    """The request records collected from LLM inference requests."""
+
+    request_id: Optional[int] = None
+    chat_cmpl: ChatCompletionRequest
+    output_str: Optional[str] = None
+    first_chunk_output_str: str = ""
+    timestamp: Optional[float] = None
+    metrics: Optional[Metrics] = None
+    error_msg: Optional[str] = None
+
+
+class GroupedRequestRecord(RequestRecord):
+    """The data structure for request record groups.
+    For datasets that have common prefix sharing, the request records
+    that share a same common prefix will be wrapped in a GroupedRequestRecord
+    at the beginning.
+    """
+
+    records: List[RequestRecord]  # noqa: UP006
+
+
+def generate_metrics_summary(
+    request_records: List[RequestRecord],  # noqa: UP006
+    num_total_requests: int,
+    num_gpus: int,
+) -> Dict[str, Any]:  # noqa: UP006
+    """Computes summary statistics across all metrics collected.
+    Return a dictionary as the report.
+    """
+    num_completed_requests = len(request_records)
+    assert num_completed_requests <= num_total_requests
+    request_metrics = [record.metrics for record in request_records]
+    duration = (
+        max(metrics.finish_time for metrics in request_metrics)
+        - min(metrics.start_time for metrics in request_metrics)
+        if num_completed_requests > 0
+        else 1e-5
+    )
+
+    report = _compute_metrics_statistics(request_metrics)
+    report["num_gpus"] = num_gpus
+    report["duration"] = duration
+    report["num_total_requests"] = num_total_requests
+    report["num_completed_requests"] = num_completed_requests
+    report["request_throughput"] = num_completed_requests / duration
+
+    total_input_tokens = sum(metric.input_tokens for metric in request_metrics)
+    total_output_tokens = sum(metric.output_tokens for metric in request_metrics)
+    report["total_input_tokens"] = total_input_tokens
+    report["total_output_tokens"] = total_output_tokens
+    report["input_token_throughput"] = total_input_tokens / duration
+    report["input_token_throughput_per_gpu"] = report["input_token_throughput"] / num_gpus
+    report["output_token_throughput"] = total_output_tokens / duration
+    report["output_token_throughput_per_gpu"] = report["output_token_throughput"] / num_gpus
+
+    # Generate the server metrics statistics
+    server_metrics = [metric.server_metrics for metric in request_metrics if metric.server_metrics]
+    server_report = _compute_metrics_statistics(server_metrics)
+    if server_report is not None and len(server_report) > 0:
+        report["server_metrics"] = server_report
+
+    report = {
+        "exec_feature": (
+            request_records[0].metrics.exec_feature if num_completed_requests > 0 else None
+        ),
+        **report,
+    }
+    return report
+
+
+def _compute_metrics_statistics(
+    metrics: List[Union[Metrics, ServerMetrics]],  # noqa: UP006
+) -> Dict[str, Any]:  # noqa: UP006
+    """
+    Compute the statistics of the metrics.
+
+    Parameters
+    ----------
+    metrics : List[Union[Metrics, ServerMetrics]]
+        The list of metrics to get the statistics.
+
+    Returns
+    -------
+    report : Dict
+        The statistics of the metrics.
+    """
+    if not metrics:
+        return {}
+
+    report: Dict = {}  # noqa: UP006
+    df = pd.DataFrame([metric.model_dump() for metric in metrics])
+    for key, _ in metrics[0].model_fields.items():
+        if key in [
+            "success",
+            "start_time",
+            "finish_time",
+            "server_metrics",
+            "exec_feature",
+        ]:
+            continue
+        if key in df.columns:
+            series = df[key].dropna()
+            report[key] = {
+                "quantiles": {
+                    f"p{int(q * 100)}": v
+                    for q, v in series.quantile([0.25, 0.5, 0.75, 0.9, 0.95, 0.99]).items()
+                },
+                "mean": series.mean(),
+                "min": series.min(),
+                "max": series.max(),
+                "stddev": series.std(),
+            }
+    return report
+
+
+def convert_reports_to_df(reports: List[Dict[str, Any]]) -> pd.DataFrame:  # noqa: UP006
+    """Convert benchmark reports to pandas DataFrame."""
+
+    def _flatten_dict(d: Dict[str, Any], parent_key: str = "") -> Dict[str, Any]:  # noqa: UP006
+        items: List[Tuple[str, Any]] = []  # noqa: UP006
+        for key, value in d.items():
+            new_key = f"{parent_key}.{key}" if parent_key != "" else key
+            if isinstance(value, dict):
+                items.extend(_flatten_dict(value, new_key).items())
+            else:
+                items.append((new_key, value))
+        return dict(items)
+
+    return pd.DataFrame([_flatten_dict(report) for report in reports])
+
+
+def pretty_print_report(report: Dict[str, Any]) -> None:  # noqa: UP006
+    """Pretty print the metrics report."""
+
+    def _print(report: Dict[str, Any], server_metrics: bool):  # noqa: UP006
+        # fmt: off
+        title = "Benchmark Result"
+        if server_metrics:
+            title += " (server side)"
+        print(f" {title} ".center(50, "="))
+        if not server_metrics:
+            print(f"{'Total requests:':<40} {report['num_total_requests']:<10}")
+            print(f"{'Completed requests:':<40} {report['num_completed_requests']:<10}")
+            print(f"{'Duration (s):':<40} {report['duration']:<10.2f}")
+            print(f"{'Num GPUs:':<40} {report['num_gpus']:<10}")
+            print(f"{'Total input tokens:':<40} {report['total_input_tokens']:<10}")
+            print(f"{'Total output tokens:':<40} {report['total_output_tokens']:<10}")
+            print(f"{'Request throughput (req/s):':<40} {report['request_throughput']:<10.2f}")
+            print(f"{'Input token throughput (tok/s):':<40} {report['input_token_throughput']:<10.2f}")  # noqa: E501
+            print(f"{'Input token throughput per GPU (tok/s):':<40} {report['input_token_throughput_per_gpu']:<10.2f}")  # noqa: E501
+            print(f"{'Output token throughput (tok/s):':<40} {report['output_token_throughput']:<10.2f}")  # noqa: E501
+            print(f"{'Output token throughput per GPU (tok/s):':<40} {report['output_token_throughput_per_gpu']:<10.2f}")  # noqa: E501
+
+        if report["num_completed_requests"] == 0:
+            return
+        ttft = report["time_to_first_token_s"]
+        print(" Time to First Token (TTFT, ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {ttft['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {ttft['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {ttft['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {ttft['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {ttft['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {ttft['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {ttft['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {ttft['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {ttft['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {ttft['max'] * 1000:<10.2f}")
+
+        tpot = report["time_per_output_token_s"]
+        print(" Time per Output Token (TPOT, ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {tpot['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {tpot['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {tpot['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {tpot['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {tpot['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {tpot['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {tpot['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {tpot['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {tpot['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {tpot['max'] * 1000:<10.2f}")
+
+        itl = report["inter_token_latency_s"]
+        print(" Inter-Token Latency (ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {itl['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {itl['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {itl['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {itl['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {itl['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {itl['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {itl['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {itl['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {itl['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {itl['max'] * 1000:<10.2f}")
+
+        e2e_latency = report["end_to_end_latency_s"]
+        print(" End-to-End Latency (ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {e2e_latency['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {e2e_latency['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {e2e_latency['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {e2e_latency['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {e2e_latency['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {e2e_latency['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {e2e_latency['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {e2e_latency['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {e2e_latency['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {e2e_latency['max'] * 1000:<10.2f}")
+
+        input_tokens = report["input_tokens"]
+        print(" Input Tokens ".center(50, "-"))
+        print(f"{'Mean:':<40} {input_tokens['mean']:<1}")
+        print(f"{'Stddev:':<40} {input_tokens['stddev']:<1}")
+        print(f"{'P25:':<40} {input_tokens['quantiles']['p25']:<1}")
+        print(f"{'P50:':<40} {input_tokens['quantiles']['p50']:<1}")
+        print(f"{'P95:':<40} {input_tokens['quantiles']['p95']:<1}")
+        print(f"{'Min:':<40} {input_tokens['min']:<1}")
+        print(f"{'Max:':<40} {input_tokens['max']:<1}")
+
+        output_tokens = report["output_tokens"]
+        print(" Output Tokens ".center(50, "-"))
+        print(f"{'Mean:':<40} {output_tokens['mean']:<1}")
+        print(f"{'Stddev:':<40} {output_tokens['stddev']:<1}")
+        print(f"{'P25:':<40} {output_tokens['quantiles']['p25']:<1}")
+        print(f"{'P50:':<40} {output_tokens['quantiles']['p50']:<1}")
+        print(f"{'P95:':<40} {output_tokens['quantiles']['p95']:<1}")
+        print(f"{'Min:':<40} {output_tokens['min']:<1}")
+        print(f"{'Max:':<40} {output_tokens['max']:<1}")
+
+        print("=" * 50)
+
+    # fmt: on
+    _print(report, server_metrics=False)
+    if "server_metrics" in report:
+        _print(report["server_metrics"], server_metrics=True)