chore: import upstream snapshot with attribution

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,
+            ),
+        )
+    ]