# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
r"""Benchmark online serving throughput.

On the server side, run one of the following commands
to launch the vLLM OpenAI API server:
    vllm serve <your_model> <engine arguments>

On the client side, run:
    vllm bench serve \
        --backend <backend or endpoint type. Default 'openai'> \
        --label <benchmark result label. Default using backend> \
        --model <your_model. Optional, defaults to first model from server> \
        --dataset-name <dataset_name. Default 'random'> \
        --input-len <general input length. Optional, maps to dataset-specific args> \
        --output-len <general output length. Optional, maps to dataset-specific args> \
        --request-rate <request_rate. Default inf> \
        --num-prompts <num_prompts. Default 1000>
"""

import argparse
import asyncio
import contextlib
import importlib.util
import json
import math
import os
import random
import shutil
import ssl
import time
import uuid
import warnings
from collections.abc import AsyncGenerator, Iterable, Iterator
from dataclasses import dataclass, replace
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Literal

import aiohttp
import numpy as np
from tqdm.asyncio import tqdm

from vllm.benchmarks.datasets import SampleRequest, add_dataset_parser, get_samples
from vllm.benchmarks.lib.endpoint_request_func import (
    ASYNC_REQUEST_FUNCS,
    OPENAI_COMPATIBLE_BACKENDS,
    POOLING_BACKENDS,
    RequestFuncInput,
    RequestFuncOutput,
)
from vllm.benchmarks.lib.ready_checker import wait_for_endpoint
from vllm.benchmarks.lib.utils import (
    convert_to_pytorch_benchmark_format,
    redact_sensitive_namespace,
    write_to_json,
)
from vllm.tokenizers import TokenizerLike, get_tokenizer
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.utils.gc_utils import freeze_gc_heap
from vllm.utils.network_utils import join_host_port

MILLISECONDS_TO_SECONDS_CONVERSION = 1000
_SENSITIVE_ARG_FIELDS = ("header",)


def _merge_overrides(base: dict | None, override: dict | None) -> dict | None:
    """Shallow merge; per-request wins. Returns None if both are empty."""
    if not base and not override:
        return None
    return {**(base or {}), **(override or {})}


TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) and (
    shutil.which("gnuplot") is not None
)


async def _align_prompts_to_server_tokenizer(
    base_url: str,
    model_id: str,
    input_requests: list[SampleRequest],
    ssl_context: ssl.SSLContext | bool | None = None,
) -> list[SampleRequest]:
    """Re-align prompts if local/server tokenizers disagree."""
    if not input_requests or not isinstance(input_requests[0].prompt, str):
        return input_requests

    tok_url = f"{base_url}/tokenize"
    detok_url = f"{base_url}/detokenize"
    connector = aiohttp.TCPConnector(ssl=ssl_context)

    async with aiohttp.ClientSession(connector=connector) as session:
        sem = asyncio.Semaphore(64)

        async def _tokenize(prompt: str) -> list[int]:
            async with (
                sem,
                session.post(
                    tok_url,
                    json={
                        "model": model_id,
                        "prompt": prompt,
                        "add_special_tokens": False,
                    },
                ) as r,
            ):
                r.raise_for_status()
                return (await r.json())["tokens"]

        async def _detokenize(tokens: list[int]) -> str:
            async with (
                sem,
                session.post(
                    detok_url, json={"model": model_id, "tokens": tokens}
                ) as r,
            ):
                r.raise_for_status()
                return (await r.json())["prompt"]

        try:
            first_tokens = await _tokenize(input_requests[0].prompt)
        except Exception:
            print("WARNING: /tokenize unavailable, skipping alignment.")
            return input_requests

        expected = input_requests[0].prompt_len
        if len(first_tokens) == expected:
            return input_requests

        print(
            f"WARNING: tokenizer mismatch "
            f"(server={len(first_tokens)}, expected={expected}), "
            f"re-aligning prompts."
        )

        async def _fix_one(req: SampleRequest) -> SampleRequest:
            assert isinstance(req.prompt, str)
            tokens = await _tokenize(req.prompt)
            if len(tokens) <= req.prompt_len:
                return req
            corrected = await _detokenize(tokens[: req.prompt_len])
            return replace(req, prompt=corrected, prompt_len=req.prompt_len)

        results = await asyncio.gather(
            *[_fix_one(r) for r in input_requests], return_exceptions=True
        )
        return [
            res if not isinstance(res, BaseException) else orig
            for orig, res in zip(input_requests, results)
        ]


async def get_first_model_from_server(
    base_url: str,
    headers: dict | None = None,
    ssl_context: ssl.SSLContext | bool | None = None,
) -> tuple[str, str]:
    """Fetch the first model from the server's /v1/models endpoint."""
    models_url = f"{base_url}/v1/models"
    connector = aiohttp.TCPConnector(ssl=ssl_context)
    async with aiohttp.ClientSession(connector=connector) as session:
        try:
            async with session.get(models_url, headers=headers) as response:
                response.raise_for_status()
                data = await response.json()
                if "data" in data and len(data["data"]) > 0:
                    return data["data"][0]["id"], data["data"][0]["root"]
                else:
                    raise ValueError(
                        f"No models found on the server at {base_url}. "
                        "Make sure the server is running and has models loaded."
                    )
        except (aiohttp.ClientError, json.JSONDecodeError) as e:
            raise RuntimeError(
                f"Failed to fetch models from server at {models_url}. "
                "Check that:\n"
                "1. The server is running\n"
                "2. The server URL is correct\n"
                f"Error: {e}"
            ) from e


@dataclass
class SpecDecodeMetrics:
    """Speculative decoding metrics from the server's Prometheus endpoint."""

    num_drafts: int
    num_draft_tokens: int
    num_accepted_tokens: int
    accepted_per_pos: dict[int, int]


async def fetch_spec_decode_metrics(
    base_url: str, session: aiohttp.ClientSession
) -> SpecDecodeMetrics | None:
    """Fetch speculative decoding metrics from the server's Prometheus endpoint.

    Returns None if speculative decoding is not enabled or metrics are not available.
    """
    metrics_url = f"{base_url}/metrics"
    try:
        async with session.get(metrics_url) as response:
            if response.status != 200:
                return None
            text = await response.text()

            num_drafts = 0
            num_draft_tokens = 0
            num_accepted_tokens = 0
            accepted_per_pos: dict[int, int] = {}
            found_spec_decode = False

            for line in text.split("\n"):
                line = line.strip()
                if not line or line.startswith("#"):
                    continue

                if line.startswith("vllm:spec_decode"):
                    # Extract metric name (before labels) to avoid matching
                    # substrings inside label values.
                    parts = line.split(None, 1)
                    metric_name = parts[0].split("{")[0]
                    if not metric_name.endswith("_total"):
                        continue
                    found_spec_decode = True
                    with contextlib.suppress(ValueError):
                        if "num_drafts" in metric_name:
                            num_drafts += int(float(parts[-1]))
                        elif "num_draft_tokens" in metric_name:
                            num_draft_tokens += int(float(parts[-1]))
                        elif "num_accepted_tokens_per_pos" in metric_name:
                            pos_label = 'position="'
                            if pos_label in line:
                                start = line.index(pos_label) + len(pos_label)
                                end = line.index('"', start)
                                pos = int(line[start:end])
                                val = int(float(parts[-1]))
                                accepted_per_pos[pos] = (
                                    accepted_per_pos.get(pos, 0) + val
                                )
                        elif "num_accepted_tokens" in metric_name:
                            num_accepted_tokens += int(float(parts[-1]))

            if not found_spec_decode:
                return None

            return SpecDecodeMetrics(
                num_drafts=num_drafts,
                num_draft_tokens=num_draft_tokens,
                num_accepted_tokens=num_accepted_tokens,
                accepted_per_pos=accepted_per_pos,
            )
    except (aiohttp.ClientError, asyncio.TimeoutError):
        return None


@dataclass
class DiffusionMetrics:
    """Diffusion (dLLM) decoding metrics from the server's Prometheus endpoint."""

    num_denoising_steps: int
    num_canvas_positions: int
    num_committed_tokens: int


async def fetch_diffusion_metrics(
    base_url: str, session: aiohttp.ClientSession
) -> DiffusionMetrics | None:
    """Fetch diffusion decoding metrics from the server's Prometheus endpoint.

    Returns None if the model is not a diffusion model or metrics are not
    available.
    """
    metrics_url = f"{base_url}/metrics"
    try:
        async with session.get(metrics_url) as response:
            if response.status != 200:
                return None
            text = await response.text()

            num_denoising_steps = 0
            num_canvas_positions = 0
            num_committed_tokens = 0
            found_diffusion = False

            for line in text.split("\n"):
                line = line.strip()
                if not line or line.startswith("#"):
                    continue

                if line.startswith("vllm:diffusion"):
                    # Extract metric name (before labels) to avoid matching
                    # substrings inside label values.
                    parts = line.split(None, 1)
                    metric_name = parts[0].split("{")[0]
                    if not metric_name.endswith("_total"):
                        continue
                    found_diffusion = True
                    with contextlib.suppress(ValueError):
                        if "num_denoising_steps" in metric_name:
                            num_denoising_steps += int(float(parts[-1]))
                        elif "num_canvas_positions" in metric_name:
                            num_canvas_positions += int(float(parts[-1]))
                        elif "num_committed_tokens" in metric_name:
                            num_committed_tokens += int(float(parts[-1]))

            if not found_diffusion:
                return None

            return DiffusionMetrics(
                num_denoising_steps=num_denoising_steps,
                num_canvas_positions=num_canvas_positions,
                num_committed_tokens=num_committed_tokens,
            )
    except (aiohttp.ClientError, asyncio.TimeoutError):
        return None


class TaskType(Enum):
    GENERATION = "generation"
    POOLING = "pooling"


@dataclass
class BenchmarkMetrics:
    completed: int
    failed: int
    total_input: int
    total_output: int
    request_throughput: float
    request_goodput: float
    output_throughput: float
    total_token_throughput: float
    mean_ttft_ms: float
    median_ttft_ms: float
    std_ttft_ms: float
    percentiles_ttft_ms: list[tuple[float, float]]
    mean_tpot_ms: float
    median_tpot_ms: float
    std_tpot_ms: float
    percentiles_tpot_ms: list[tuple[float, float]]
    mean_itl_ms: float
    median_itl_ms: float
    std_itl_ms: float
    percentiles_itl_ms: list[tuple[float, float]]
    # E2EL stands for end-to-end latency per request.
    # It is the time taken on the client side from sending
    # a request to receiving a complete response.
    mean_e2el_ms: float
    median_e2el_ms: float
    std_e2el_ms: float
    percentiles_e2el_ms: list[tuple[float, float]]
    # Max output tokens per second and concurrent requests at that peak
    max_output_tokens_per_s: float
    max_concurrent_requests: int
    rtfx: float = 0.0  # Inverse Real-Time Factor for ASR benchmarks


@dataclass
class EmbedBenchmarkMetrics:
    completed: int
    failed: int
    total_input: int
    total_input_sequences: int
    request_throughput: float
    input_sequence_throughput: float
    total_token_throughput: float
    mean_e2el_ms: float
    std_e2el_ms: float
    median_e2el_ms: float
    percentiles_e2el_ms: list[tuple[float, float]]


def _get_current_request_rate(
    ramp_up_strategy: Literal["linear", "exponential"] | None,
    ramp_up_start_rps: int | None,
    ramp_up_end_rps: int | None,
    request_index: int,
    total_requests: int,
    request_rate: float,
) -> float:
    if (
        ramp_up_strategy
        and ramp_up_start_rps is not None
        and ramp_up_end_rps is not None
    ):
        progress = request_index / max(total_requests - 1, 1)
        if ramp_up_strategy == "linear":
            increase = (ramp_up_end_rps - ramp_up_start_rps) * progress
            return ramp_up_start_rps + increase
        elif ramp_up_strategy == "exponential":
            ratio = ramp_up_end_rps / ramp_up_start_rps
            return ramp_up_start_rps * (ratio**progress)
        else:
            raise ValueError(f"Unknown ramp-up strategy: {ramp_up_strategy}")
    return request_rate


async def get_request(
    input_requests: list[SampleRequest],
    request_rate: float,
    burstiness: float = 1.0,
    ramp_up_strategy: Literal["linear", "exponential"] | None = None,
    ramp_up_start_rps: int | None = None,
    ramp_up_end_rps: int | None = None,
    self_timed: bool = False,
) -> AsyncGenerator[tuple[SampleRequest, float], None]:
    """Asynchronously generates requests at a specified rate
    with OPTIONAL burstiness and OPTIONAL ramp-up strategy.

    Args:
        input_requests:
            A list of input requests, each represented as a SampleRequest.
        self_timed:
            If True, the requests carry their own arrival timing and no
            request rate, burstiness or ramp-up is applied.
        request_rate:
            The rate at which requests are generated (requests/s).
        burstiness (optional):
            The burstiness factor of the request generation.
            Only takes effect when request_rate is not inf.
            Default value is 1, which follows a Poisson process.
            Otherwise, the request intervals follow a gamma distribution.
            A lower burstiness value (0 < burstiness < 1) results
            in more bursty requests, while a higher burstiness value
            (burstiness > 1) results in a more uniform arrival of requests.
        ramp_up_strategy (optional):
            The ramp-up strategy. Can be "linear" or "exponential".
            If None, uses constant request rate (specified by request_rate).
        ramp_up_start_rps (optional):
            The starting request rate for ramp-up.
        ramp_up_end_rps (optional):
            The ending request rate for ramp-up.

    """
    assert burstiness > 0, (
        f"A positive burstiness factor is expected, but given {burstiness}."
    )
    # Convert to list to get length for ramp-up calculations
    if isinstance(input_requests, Iterable) and not isinstance(input_requests, list):
        input_requests = list(input_requests)

    total_requests = len(input_requests)
    assert total_requests > 0, "No requests provided."

    # Precompute delays among requests to minimize request send laggings
    request_rates: list[float] = []
    delay_ts: list[float] = []

    # if the traces have timing info then:
    if not self_timed:
        for request_index, request in enumerate(input_requests):
            current_request_rate = _get_current_request_rate(
                ramp_up_strategy,
                ramp_up_start_rps,
                ramp_up_end_rps,
                request_index,
                total_requests,
                request_rate,
            )
            assert current_request_rate > 0.0, (
                f"Obtained non-positive request rate {current_request_rate}."
            )
            request_rates.append(current_request_rate)
            if current_request_rate == float("inf"):
                delay_ts.append(0)
            elif burstiness == float("inf"):
                # when burstiness tends to infinity, the delay time becomes constant
                # and tends to the inverse of the request rate
                delay_ts.append(1.0 / current_request_rate)
            else:
                theta = 1.0 / (current_request_rate * burstiness)

                # Sample the request interval from the gamma distribution.
                # If burstiness is 1, it follows exponential distribution.
                delay_ts.append(np.random.gamma(shape=burstiness, scale=theta))

        # Calculate the cumulative delay time from the first sent out requests.
        for i in range(1, len(delay_ts)):
            delay_ts[i] += delay_ts[i - 1]
        if ramp_up_strategy is None and delay_ts[-1] != 0:
            # When ramp_up_strategy is not set, we assume the request rate is fixed
            # and all requests should be sent in target_total_delay_s, the following
            # logic would re-scale delay time to ensure the final delay_ts
            # align with target_total_delay_s.
            #
            # NOTE: If we simply accumulate the random delta values
            # from the gamma distribution, their sum would have 1-2% gap
            # from target_total_delay_s. The purpose of the following logic is to
            # close the gap for stabilizing the throughput data
            # from different random seeds.
            target_total_delay_s = total_requests / request_rate
            normalize_factor = target_total_delay_s / delay_ts[-1]
            delay_ts = [delay * normalize_factor for delay in delay_ts]
    else:
        for request_index, request in enumerate(input_requests):
            # this is cumulative running ts, from which sleep is calculated later
            if request.timestamp is not None:
                delay_ts.append(request.timestamp)
            else:
                delay_ts.append(0.0)
            # TODO: there is no notion of RPS here, may be we can calculate
            # from the trace.
            request_rates.append(0.0)

    start_ts = time.time()
    for request_index, request in enumerate(input_requests):
        if delay_ts[request_index] > 0:
            current_ts = time.time()
            sleep_interval_s = start_ts + delay_ts[request_index] - current_ts
            if sleep_interval_s > 0:
                await asyncio.sleep(sleep_interval_s)
        yield request, request_rates[request_index]


def calculate_metrics_for_embeddings(
    outputs: list[RequestFuncOutput],
    dur_s: float,
    selected_percentiles: list[float],
) -> EmbedBenchmarkMetrics:
    """Calculate the metrics for the embedding requests.

    Args:
        outputs: The outputs of the requests.
        dur_s: The duration of the benchmark.
        selected_percentiles: The percentiles to select.

    Returns:
        The calculated benchmark metrics.

    """
    total_input = 0
    total_input_sequences = 0
    completed = 0
    failed = 0
    e2els: list[float] = []
    for i in range(len(outputs)):
        if outputs[i].success:
            e2els.append(outputs[i].latency)
            completed += 1
            total_input += outputs[i].prompt_len
            total_input_sequences += outputs[i].num_input_sequences
        else:
            failed += 1

    if completed == 0:
        warnings.warn(
            "All requests failed. This is likely due to a misconfiguration "
            "on the benchmark arguments.",
            stacklevel=2,
        )
    metrics = EmbedBenchmarkMetrics(
        completed=completed,
        failed=failed,
        total_input=total_input,
        total_input_sequences=total_input_sequences,
        request_throughput=completed / dur_s,
        input_sequence_throughput=total_input_sequences / dur_s,
        total_token_throughput=total_input / dur_s,
        mean_e2el_ms=np.mean(e2els or 0) * 1000,
        std_e2el_ms=np.std(e2els or 0) * 1000,
        median_e2el_ms=np.median(e2els or 0) * 1000,
        percentiles_e2el_ms=[
            (p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles
        ],
    )
    return metrics


def calculate_metrics(
    input_requests: list[SampleRequest],
    outputs: list[RequestFuncOutput],
    dur_s: float,
    tokenizer: TokenizerLike | None,
    selected_percentiles: list[float],
    goodput_config_dict: dict[str, float],
) -> tuple[BenchmarkMetrics, list[int]]:
    """Calculate the metrics for the benchmark.

    Args:
        input_requests: The input requests.
        outputs: The outputs of the requests.
        dur_s: The duration of the benchmark.
        tokenizer: The tokenizer to use.
        selected_percentiles: The percentiles to select.
        goodput_config_dict: The goodput configuration.

    Returns:
        A tuple of the benchmark metrics and the actual output lengths.

    """
    actual_output_lens: list[int] = []
    total_input = 0
    completed = 0
    good_completed = 0
    itls: list[float] = []
    tpots: list[float] = []
    all_tpots: list[float] = []
    ttfts: list[float] = []
    e2els: list[float] = []
    input_audio_duration = 0.0
    for i in range(len(outputs)):
        if outputs[i].success:
            output_len = outputs[i].output_tokens

            if not output_len:
                if tokenizer is None:
                    output_len = 1
                else:
                    # We use the tokenizer to count the number of output tokens
                    # for some serving backends instead of looking at
                    # len(outputs[i].itl) since multiple output tokens may be
                    # bundled together
                    # Note : this may inflate the output token count slightly
                    output_len = len(
                        tokenizer(
                            outputs[i].generated_text, add_special_tokens=False
                        ).input_ids
                    )
            actual_output_lens.append(output_len)
            total_input += outputs[i].prompt_len
            tpot = 0.0
            if output_len > 1:
                latency_minus_ttft = outputs[i].latency - outputs[i].ttft
                tpot = latency_minus_ttft / (output_len - 1)
                tpots.append(tpot)
            # Note: if output_len <= 1, we regard tpot as 0 for goodput
            all_tpots.append(tpot)
            itls += outputs[i].itl
            ttfts.append(outputs[i].ttft)
            e2els.append(outputs[i].latency)
            input_audio_duration += outputs[i].input_audio_duration
            completed += 1
        else:
            actual_output_lens.append(0)

    if goodput_config_dict:
        valid_metrics = []
        slo_values = []

        if "ttft" in goodput_config_dict:
            valid_metrics.append(ttfts)
            slo_values.append(
                goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION
            )
        if "tpot" in goodput_config_dict:
            valid_metrics.append(all_tpots)
            slo_values.append(
                goodput_config_dict["tpot"] / MILLISECONDS_TO_SECONDS_CONVERSION
            )
        if "e2el" in goodput_config_dict:
            valid_metrics.append(e2els)
            slo_values.append(
                goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION
            )

        for req_metric in zip(*valid_metrics):
            is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])
            if is_good_req:
                good_completed += 1

    if completed == 0:
        warnings.warn(
            "All requests failed. This is likely due to a misconfiguration "
            "on the benchmark arguments.",
            stacklevel=2,
        )

    # Calculate max output tokens per second metric
    max_output_tokens_per_s = 0.0
    max_concurrent_requests = 0

    # Find the time range across all successful requests
    successful_outputs = [output for output in outputs if output.success]
    failed_outputs = [output for output in outputs if not output.success]

    if len(failed_outputs) > 0:
        print("Failed requests during benchmark run detected (capping to 10):")
        for i, err in enumerate(failed_outputs[:10]):
            print(f"Error {i}: {err.error}")

    if successful_outputs:
        min_start_time = min(output.start_time for output in successful_outputs)
        max_end_time = max(
            output.start_time + output.latency for output in successful_outputs
        )

        # Create second buckets (ceiling to ensure we capture all time)
        duration_seconds = int(np.ceil(max_end_time - min_start_time)) + 1
        tokens_per_second = np.zeros(duration_seconds)
        concurrent_requests_per_second = np.zeros(duration_seconds)

        for i, output in enumerate(successful_outputs):
            # Calculate token generation timestamp using
            # start_time, ttft, and itl
            token_times = [output.start_time + output.ttft]
            current_time = token_times[0]
            for itl_value in output.itl:
                current_time += itl_value
                token_times.append(current_time)

            # Add tokens to second buckets
            for token_time in token_times:
                second_bucket = int(token_time - min_start_time)
                if 0 <= second_bucket < duration_seconds:
                    tokens_per_second[second_bucket] += 1

            # Track concurrent requests for each second this request was active
            request_start_second = int(output.start_time - min_start_time)
            request_end_second = int(
                (output.start_time + output.latency) - min_start_time
            )

            for second in range(request_start_second, request_end_second + 1):
                concurrent_requests_per_second[second] += 1

        # Find the maximum tokens per second and corresponding
        # concurrent requests
        if len(tokens_per_second) > 0:
            max_output_tokens_per_s = float(np.max(tokens_per_second))
            max_concurrent_requests = int(np.max(concurrent_requests_per_second))

        if TERM_PLOTLIB_AVAILABLE:
            import termplotlib as tpl

            fig = tpl.figure()
            fig.plot(
                np.arange(len(tokens_per_second)),
                tokens_per_second,
                title="Output tokens per second",
            )
            fig.plot(
                np.arange(len(concurrent_requests_per_second)),
                concurrent_requests_per_second,
                title="Concurrent requests per second",
            )
            fig.show()
        else:
            print("tip: install termplotlib and gnuplot to plot the metrics")

    metrics = BenchmarkMetrics(
        completed=completed,
        failed=len(failed_outputs),
        total_input=total_input,
        total_output=sum(actual_output_lens),
        request_throughput=completed / dur_s,
        request_goodput=good_completed / dur_s,
        output_throughput=sum(actual_output_lens) / dur_s,
        total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
        mean_ttft_ms=np.mean(ttfts or 0)
        * 1000,  # ttfts is empty if streaming is not supported by the endpoint
        std_ttft_ms=np.std(ttfts or 0) * 1000,
        median_ttft_ms=np.median(ttfts or 0) * 1000,
        percentiles_ttft_ms=[
            (p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles
        ],
        mean_tpot_ms=np.mean(tpots or 0) * 1000,
        std_tpot_ms=np.std(tpots or 0) * 1000,
        median_tpot_ms=np.median(tpots or 0) * 1000,
        percentiles_tpot_ms=[
            (p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles
        ],
        mean_itl_ms=np.mean(itls or 0) * 1000,
        std_itl_ms=np.std(itls or 0) * 1000,
        median_itl_ms=np.median(itls or 0) * 1000,
        percentiles_itl_ms=[
            (p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles
        ],
        mean_e2el_ms=np.mean(e2els or 0) * 1000,
        std_e2el_ms=np.std(e2els or 0) * 1000,
        median_e2el_ms=np.median(e2els or 0) * 1000,
        percentiles_e2el_ms=[
            (p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles
        ],
        max_output_tokens_per_s=max_output_tokens_per_s,
        max_concurrent_requests=max_concurrent_requests,
        rtfx=input_audio_duration / dur_s,
    )

    return metrics, actual_output_lens


async def benchmark(
    task_type: TaskType,
    endpoint_type: str,
    api_url: str,
    base_url: str,
    model_id: str,
    model_name: str,
    tokenizer: TokenizerLike | None,
    input_requests: list[SampleRequest],
    logprobs: int | None,
    request_rate: float,
    burstiness: float,
    disable_tqdm: bool,
    num_warmups: int,
    profile: bool,
    selected_percentile_metrics: list[str],
    selected_percentiles: list[float],
    ignore_eos: bool,
    goodput_config_dict: dict[str, float],
    max_concurrency: int | None,
    lora_modules: Iterable[str] | None,
    extra_headers: dict | None,
    extra_body: dict | None,
    lora_assignment: Literal["random", "round-robin"] = "random",
    ramp_up_strategy: Literal["linear", "exponential"] | None = None,
    ramp_up_start_rps: int | None = None,
    ramp_up_end_rps: int | None = None,
    ready_check_timeout_sec: int = 600,
    ssl_context: ssl.SSLContext | bool | None = None,
    self_timed: bool = False,
    probe_request_rate: float = 0.0,
):
    try:
        request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
    except KeyError:
        raise ValueError(f"Unknown backend: {endpoint_type}") from None

    # Reuses connections across requests to reduce TLS handshake overhead.
    # Use ssl_context if provided, otherwise default to True for https URLs
    ssl_setting = ssl_context if ssl_context is not None else ("https://" in api_url)
    connector = aiohttp.TCPConnector(
        limit=max_concurrency or 0,
        limit_per_host=max_concurrency or 0,
        ttl_dns_cache=300,
        use_dns_cache=True,
        keepalive_timeout=60,
        enable_cleanup_closed=True,
        force_close=False,
        ssl=ssl_setting,
    )

    session = aiohttp.ClientSession(
        connector=connector,
        trust_env=True,
        timeout=aiohttp.ClientTimeout(total=6 * 60 * 60),
    )

    print("Starting initial single prompt test run...")
    test_prompt, test_prompt_len, test_output_len, test_mm_content = (
        input_requests[0].prompt,
        input_requests[0].prompt_len,
        input_requests[0].expected_output_len,
        input_requests[0].multi_modal_data,
    )
    test_extra_body = _merge_overrides(extra_body, input_requests[0].request_overrides)
    test_chat_messages = input_requests[0].chat_messages

    assert (
        test_mm_content is None
        or isinstance(test_mm_content, dict)
        or (
            isinstance(test_mm_content, list)
            and all(isinstance(item, dict) for item in test_mm_content)
        )
    ), "multi_modal_data must be a dict or list[dict]"
    test_input = RequestFuncInput(
        model=model_id,
        model_name=model_name,
        prompt=test_prompt,
        api_url=api_url,
        prompt_len=test_prompt_len,
        output_len=test_output_len,
        logprobs=logprobs,
        multi_modal_content=test_mm_content,
        ignore_eos=ignore_eos,
        extra_headers=extra_headers,
        extra_body=test_extra_body,
        chat_messages=test_chat_messages,
    )

    if ready_check_timeout_sec > 0:
        test_output = await wait_for_endpoint(
            request_func,
            test_input,
            session,
            timeout_seconds=ready_check_timeout_sec,
        )
        if not test_output.success:
            raise ValueError(
                "Initial test run failed - Please make sure benchmark "
                "arguments are correctly specified. "
                f"Error: {test_output.error}"
            )
        else:
            print("Initial test run completed.")
    else:
        print("Skipping endpoint ready check.")

    if num_warmups > 0:
        print(f"Warming up with {num_warmups} requests...")
        warmup_pbar = None if disable_tqdm else tqdm(total=num_warmups)
        warmup_semaphore = (
            asyncio.Semaphore(max_concurrency)
            if max_concurrency
            else contextlib.nullcontext()
        )
        warmup_tasks = []

        async def warmup_limited_request_func():
            async with warmup_semaphore:
                return await request_func(
                    request_func_input=test_input, session=session, pbar=warmup_pbar
                )

        for _ in range(num_warmups):
            request_task = asyncio.create_task(warmup_limited_request_func())
            warmup_tasks.append(request_task)
        _ = await asyncio.gather(*warmup_tasks)

        if warmup_pbar is not None:
            warmup_pbar.close()
        print("Warmup run completed.")

    print("Starting main benchmark run...")

    lora_modules_iter: Iterator[str] | None = None
    if lora_modules:
        lora_modules_list = list(lora_modules)
        if lora_assignment == "round-robin":
            # Deterministic round-robin assignment across requests.
            lora_modules_iter = iter(
                [
                    lora_modules_list[i % len(lora_modules_list)]
                    for i in range(len(input_requests))
                ]
            )
        else:
            # For each input request, choose a LoRA module at random.
            lora_modules_iter = iter(
                [random.choice(lora_modules_list) for _ in range(len(input_requests))]
            )

    if profile:
        print("Starting profiler...")
        profile_input = RequestFuncInput(
            model=model_id,
            model_name=model_name,
            prompt=test_prompt,
            api_url=base_url + "/start_profile",
            prompt_len=test_prompt_len,
            output_len=test_output_len,
            logprobs=logprobs,
            multi_modal_content=test_mm_content,
            ignore_eos=ignore_eos,
            extra_headers=extra_headers,
            extra_body=test_extra_body,
            chat_messages=test_chat_messages,
        )
        profile_output = await request_func(
            request_func_input=profile_input, session=session
        )
        if profile_output.success:
            print("Profiler started")

    distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution"
    if not self_timed:
        if ramp_up_strategy is not None:
            print(f"Traffic ramp-up strategy: {ramp_up_strategy}.")
            print(
                f"Will increase RPS from {ramp_up_start_rps} to "
                f"{ramp_up_end_rps} RPS over the duration of the benchmark."
            )
        else:
            print(f"Traffic request rate: {request_rate}")

        print(f"Burstiness factor: {burstiness} ({distribution})")
        print(f"Maximum request concurrency: {max_concurrency}")
    else:
        print("Self timing is set, using the timestamps from the trace file.")

    spec_decode_metrics_before = await fetch_spec_decode_metrics(base_url, session)
    diffusion_metrics_before = await fetch_diffusion_metrics(base_url, session)

    pbar = None if disable_tqdm else tqdm(total=len(input_requests))

    semaphore = (
        asyncio.Semaphore(max_concurrency)
        if max_concurrency
        else contextlib.nullcontext()
    )

    async def limited_request_func(
        request_func_input, session, pbar, request_arrival_time
    ):
        async with semaphore:
            output = await request_func(
                request_func_input=request_func_input, session=session, pbar=pbar
            )
        # Preserve start_time as the time the request was sent. It is used by
        # throughput calculations; record client-side semaphore delay separately.
        output.client_queue_time = output.start_time - request_arrival_time
        return output

    probe_outputs: list[RequestFuncOutput] = []
    probe_stop = asyncio.Event()

    async def probe_loop():
        probe_input = replace(
            test_input,
            prompt="Hi",
            prompt_len=1,
            output_len=1,
            multi_modal_content=None,
            chat_messages=None,
        )
        interval = 1 / probe_request_rate
        while not probe_stop.is_set():
            probe_outputs.append(
                await request_func(request_func_input=probe_input, session=session)
            )
            await asyncio.sleep(interval)

    probe_task: asyncio.Task | None = None
    if probe_request_rate > 0:
        print(f"Probe request rate: {probe_request_rate} req/s")
        probe_task = asyncio.create_task(probe_loop())

    benchmark_start_time = time.perf_counter()
    tasks: list[asyncio.Task] = []

    rps_change_events = []
    last_int_rps = -1
    if ramp_up_strategy is not None and ramp_up_start_rps is not None:
        last_int_rps = ramp_up_start_rps
        rps_change_events.append(
            {
                "rps": last_int_rps,
                "timestamp": datetime.now().isoformat(),
            }
        )

    async for request, current_request_rate in get_request(
        input_requests,
        request_rate,
        burstiness,
        ramp_up_strategy,
        ramp_up_start_rps,
        ramp_up_end_rps,
        self_timed,
    ):
        if ramp_up_strategy is not None:
            current_int_rps = int(current_request_rate)
            if current_int_rps > last_int_rps:
                timestamp = datetime.now().isoformat()
                for rps_val in range(last_int_rps + 1, current_int_rps + 1):
                    rps_change_events.append({"rps": rps_val, "timestamp": timestamp})
                last_int_rps = current_int_rps
        prompt, prompt_len, output_len, mm_content, request_id = (
            request.prompt,
            request.prompt_len,
            request.expected_output_len,
            request.multi_modal_data,
            request.request_id,
        )
        per_request_extra_body = _merge_overrides(extra_body, request.request_overrides)
        req_model_id, req_model_name = model_id, model_name
        if lora_modules_iter:
            req_lora_module = next(lora_modules_iter)
            req_model_id, req_model_name = req_lora_module, req_lora_module

        mm_content_typed: dict[str, Any] | list[dict[str, Any]] | None = None
        if isinstance(mm_content, (dict, list)):
            mm_content_typed = mm_content

        request_func_input = RequestFuncInput(
            model=req_model_id,
            model_name=req_model_name,
            prompt=prompt,
            api_url=api_url,
            prompt_len=prompt_len,
            output_len=output_len or 0,
            logprobs=logprobs,
            multi_modal_content=mm_content_typed,
            ignore_eos=ignore_eos,
            extra_headers=extra_headers,
            extra_body=per_request_extra_body,
            request_id=request_id,
            chat_messages=request.chat_messages,
        )
        request_arrival_time = time.perf_counter()
        tasks.append(
            asyncio.create_task(
                limited_request_func(
                    request_func_input, session, pbar, request_arrival_time
                )
            )
        )
    outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks)

    if probe_task is not None:
        probe_stop.set()
        await probe_task

    if pbar is not None:
        pbar.close()

    benchmark_duration = time.perf_counter() - benchmark_start_time

    spec_decode_metrics_after = await fetch_spec_decode_metrics(base_url, session)
    spec_decode_stats: dict[str, Any] | None = None
    if spec_decode_metrics_before is not None and spec_decode_metrics_after is not None:
        delta_drafts = (
            spec_decode_metrics_after.num_drafts - spec_decode_metrics_before.num_drafts
        )
        delta_draft_tokens = (
            spec_decode_metrics_after.num_draft_tokens
            - spec_decode_metrics_before.num_draft_tokens
        )
        delta_accepted = (
            spec_decode_metrics_after.num_accepted_tokens
            - spec_decode_metrics_before.num_accepted_tokens
        )
        per_pos_rates: list[float] = []
        if delta_drafts > 0:
            positions = sorted(
                set(spec_decode_metrics_before.accepted_per_pos.keys())
                | set(spec_decode_metrics_after.accepted_per_pos.keys())
            )
            for pos in positions:
                before_val = spec_decode_metrics_before.accepted_per_pos.get(pos, 0)
                after_val = spec_decode_metrics_after.accepted_per_pos.get(
                    pos, before_val
                )
                delta_pos = after_val - before_val
                per_pos_rates.append(delta_pos / delta_drafts)

        if delta_draft_tokens > 0:
            acceptance_rate = (delta_accepted / delta_draft_tokens) * 100
            acceptance_length = (
                1 + delta_accepted / delta_drafts if delta_drafts > 0 else 0.0
            )
            spec_decode_stats = {
                "num_drafts": delta_drafts,
                "draft_tokens": delta_draft_tokens,
                "accepted_tokens": delta_accepted,
                "acceptance_rate": acceptance_rate,
                "acceptance_length": acceptance_length,
                "per_position_acceptance_rates": per_pos_rates,
            }

    diffusion_metrics_after = await fetch_diffusion_metrics(base_url, session)
    diffusion_stats: dict[str, Any] | None = None
    if diffusion_metrics_before is not None and diffusion_metrics_after is not None:
        delta_steps = (
            diffusion_metrics_after.num_denoising_steps
            - diffusion_metrics_before.num_denoising_steps
        )
        delta_positions = (
            diffusion_metrics_after.num_canvas_positions
            - diffusion_metrics_before.num_canvas_positions
        )
        delta_committed = (
            diffusion_metrics_after.num_committed_tokens
            - diffusion_metrics_before.num_committed_tokens
        )
        if delta_steps > 0 and delta_committed > 0:
            block_size = delta_positions / delta_steps  # canvas length (CL)
            num_canvases = delta_committed / block_size  # = number of commit steps
            denoising_steps = delta_steps - num_canvases  # exclude commit steps
            diffusion_stats = {
                "denoising_steps": denoising_steps,
                "canvas_positions": delta_positions,
                "committed_tokens": delta_committed,
                "committed_throughput": delta_committed / benchmark_duration,
                "steps_per_canvas": denoising_steps / num_canvases,
                "committed_per_step": delta_committed / denoising_steps,
            }

    metrics: BenchmarkMetrics | EmbedBenchmarkMetrics
    actual_output_lens: list[int] | int
    if task_type == TaskType.GENERATION:
        metrics, actual_output_lens = calculate_metrics(
            input_requests=input_requests,
            outputs=outputs,
            dur_s=benchmark_duration,
            tokenizer=tokenizer,
            selected_percentiles=selected_percentiles,
            goodput_config_dict=goodput_config_dict,
        )
    else:
        metrics = calculate_metrics_for_embeddings(
            outputs=outputs,
            dur_s=benchmark_duration,
            selected_percentiles=selected_percentiles,
        )
        actual_output_lens = 0

    print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
    print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
    print("{:<40} {:<10}".format("Failed requests:", metrics.failed))
    if max_concurrency is not None:
        print("{:<40} {:<10}".format("Maximum request concurrency:", max_concurrency))
    if request_rate != float("inf"):
        print("{:<40} {:<10.2f}".format("Request rate configured (RPS):", request_rate))
    print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))
    print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
    if isinstance(metrics, BenchmarkMetrics) and tokenizer:
        print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
    print(
        "{:<40} {:<10.2f}".format(
            "Request throughput (req/s):", metrics.request_throughput
        )
    )
    if isinstance(metrics, EmbedBenchmarkMetrics):
        print(
            "{:<40} {:<10.2f}".format(
                "Input throughput (inputs/s):", metrics.input_sequence_throughput
            )
        )
    if goodput_config_dict and isinstance(metrics, BenchmarkMetrics):
        print(
            "{:<40} {:<10.2f}".format(
                "Request goodput (req/s):", metrics.request_goodput
            )
        )
    if isinstance(metrics, BenchmarkMetrics):
        if tokenizer:
            print(
                "{:<40} {:<10.2f}".format(
                    "Output token throughput (tok/s):", metrics.output_throughput
                )
            )
            print(
                "{:<40} {:<10.2f}".format(
                    "Peak output token throughput (tok/s):",
                    metrics.max_output_tokens_per_s,
                )
            )
        print(
            "{:<40} {:<10.2f}".format(
                "Peak concurrent requests:", metrics.max_concurrent_requests
            )
        )
        if metrics.rtfx > 0.0:
            print(
                "{:<40} {:<10.2f}".format(
                    "RTFx (Inverse Real-Time Factor):", metrics.rtfx
                )
            )
    if tokenizer:
        print(
            "{:<40} {:<10.2f}".format(
                "Total token throughput (tok/s):", metrics.total_token_throughput
            )
        )

    probe_stats: dict[str, Any] | None = None
    if probe_task is not None:
        probe_lats = [o.latency for o in probe_outputs if o.success]
        if probe_lats:
            probe_stats = {
                "probe_completed": len(probe_lats),
                "probe_failed": len(probe_outputs) - len(probe_lats),
                "probe_median_e2el_ms": float(np.median(probe_lats)) * 1000,
                "probe_p99_e2el_ms": float(np.percentile(probe_lats, 99)) * 1000,
                "probe_max_e2el_ms": float(max(probe_lats)) * 1000,
            }
            print("{s:{c}^{n}}".format(s="Probe Requests", n=50, c="-"))
            print(
                "{:<40} {:<10}".format(
                    "Probe requests completed:", probe_stats["probe_completed"]
                )
            )
            print(
                "{:<40} {:<10}".format(
                    "Probe requests failed:", probe_stats["probe_failed"]
                )
            )
            print(
                "{:<40} {:<10.2f}".format(
                    "Median probe E2EL (ms):", probe_stats["probe_median_e2el_ms"]
                )
            )
            print(
                "{:<40} {:<10.2f}".format(
                    "P99 probe E2EL (ms):", probe_stats["probe_p99_e2el_ms"]
                )
            )
            print(
                "{:<40} {:<10.2f}".format(
                    "Max probe E2EL (ms):", probe_stats["probe_max_e2el_ms"]
                )
            )

    result: dict[str, Any]
    if isinstance(metrics, BenchmarkMetrics):
        result = {
            "duration": benchmark_duration,
            "completed": metrics.completed,
            "failed": metrics.failed,
            "total_input_tokens": metrics.total_input,
            "total_output_tokens": metrics.total_output,
            "request_throughput": metrics.request_throughput,
            "request_goodput": metrics.request_goodput if goodput_config_dict else None,
            "output_throughput": metrics.output_throughput,
            "total_token_throughput": metrics.total_token_throughput,
            "input_lens": [output.prompt_len for output in outputs],
            "output_lens": actual_output_lens,
            "ttfts": [output.ttft for output in outputs],
            "itls": [output.itl for output in outputs],
            "latencies": [output.latency for output in outputs],
            "start_times": [output.start_time for output in outputs],
            "queue_times": [output.client_queue_time for output in outputs],
            "generated_texts": [output.generated_text for output in outputs],
            "errors": [output.error for output in outputs],
            "max_output_tokens_per_s": metrics.max_output_tokens_per_s,
            "max_concurrent_requests": metrics.max_concurrent_requests,
            "rtfx": metrics.rtfx,
        }
    else:
        result = {
            "duration": benchmark_duration,
            "completed": metrics.completed,
            "total_input_tokens": metrics.total_input,
            "total_input_sequences": metrics.total_input_sequences,
            "request_throughput": metrics.request_throughput,
            "input_sequence_throughput": metrics.input_sequence_throughput,
            "total_token_throughput": metrics.total_token_throughput,
            "input_lens": [output.prompt_len for output in outputs],
            "latencies": [output.latency for output in outputs],
            "queue_times": [output.client_queue_time for output in outputs],
            "errors": [output.error for output in outputs],
        }

    queue_times: list[float] | None = None
    e2els_including_queue: list[float] | None = None
    if max_concurrency is not None:
        queue_times = [output.client_queue_time for output in outputs if output.success]
        if not math.isinf(request_rate):
            e2els_including_queue = [
                output.latency + output.client_queue_time
                for output in outputs
                if output.success
            ]

    if probe_stats is not None:
        result.update(probe_stats)

    if rps_change_events:
        result["rps_change_events"] = rps_change_events

    if spec_decode_stats is not None:
        result["spec_decode_acceptance_rate"] = spec_decode_stats["acceptance_rate"]
        result["spec_decode_acceptance_length"] = spec_decode_stats["acceptance_length"]
        result["spec_decode_num_drafts"] = int(spec_decode_stats["num_drafts"])
        result["spec_decode_draft_tokens"] = int(spec_decode_stats["draft_tokens"])
        result["spec_decode_accepted_tokens"] = int(
            spec_decode_stats["accepted_tokens"]
        )
        result["spec_decode_per_position_acceptance_rates"] = spec_decode_stats.get(
            "per_position_acceptance_rates", []
        )

    if diffusion_stats is not None:
        result["diffusion_committed_throughput"] = diffusion_stats[
            "committed_throughput"
        ]
        result["diffusion_steps_per_canvas"] = diffusion_stats["steps_per_canvas"]
        result["diffusion_committed_per_step"] = diffusion_stats["committed_per_step"]
        result["diffusion_committed_tokens"] = int(diffusion_stats["committed_tokens"])
        result["diffusion_denoising_steps"] = int(diffusion_stats["denoising_steps"])
        result["diffusion_canvas_positions"] = int(diffusion_stats["canvas_positions"])

    def process_one_metric(
        # E.g., "ttft"
        metric_attribute_name: str,
        # E.g., "TTFT"
        metric_name: str,
        # E.g., "Time to First Token"
        metric_header: str,
        values: list[float] | None = None,
    ):
        # This function prints and adds statistics of the specified
        # metric.
        if metric_attribute_name not in selected_percentile_metrics:
            return
        if values is None:
            mean = getattr(metrics, f"mean_{metric_attribute_name}_ms")
            median = getattr(metrics, f"median_{metric_attribute_name}_ms")
            std = getattr(metrics, f"std_{metric_attribute_name}_ms")
            percentiles = getattr(metrics, f"percentiles_{metric_attribute_name}_ms")
        else:
            mean = np.mean(values or 0) * 1000
            median = np.median(values or 0) * 1000
            std = np.std(values or 0) * 1000
            percentiles = [
                (p, np.percentile(values or 0, p) * 1000) for p in selected_percentiles
            ]
        print("{s:{c}^{n}}".format(s=metric_header, n=50, c="-"))
        print("{:<40} {:<10.2f}".format(f"Mean {metric_name} (ms):", mean))
        print("{:<40} {:<10.2f}".format(f"Median {metric_name} (ms):", median))
        result[f"mean_{metric_attribute_name}_ms"] = mean
        result[f"median_{metric_attribute_name}_ms"] = median
        result[f"std_{metric_attribute_name}_ms"] = std
        for p, value in percentiles:
            p_word = str(int(p)) if int(p) == p else str(p)
            print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", value))
            result[f"p{p_word}_{metric_attribute_name}_ms"] = value

    if task_type == TaskType.GENERATION and tokenizer:
        process_one_metric("ttft", "TTFT", "Time to First Token")
        process_one_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)")
        process_one_metric("itl", "ITL", "Inter-token Latency")
    process_one_metric("e2el", "E2EL", "End-to-end Latency")
    if queue_times is not None:
        process_one_metric(
            "client_queue_time",
            "Client Queue Time",
            "Client-side Queueing",
            queue_times,
        )
    if e2els_including_queue is not None:
        process_one_metric(
            "e2el_including_client_queue",
            "E2EL incl. Client Queue",
            "Queue-inclusive End-to-end Latency",
            e2els_including_queue,
        )

    if diffusion_stats is not None:
        print("{s:{c}^{n}}".format(s="Diffusion Decoding", n=50, c="-"))
        for label, key, value_fmt in (
            ("Committed throughput (tok/s):", "committed_throughput", "{:<10.2f}"),
            ("Denoising steps per canvas:", "steps_per_canvas", "{:<10.2f}"),
            ("Committed per denoising step:", "committed_per_step", "{:<10.2f}"),
            ("Committed tokens:", "committed_tokens", "{:<10d}"),
            ("Denoising steps:", "denoising_steps", "{:<10d}"),
            ("Canvas positions evaluated:", "canvas_positions", "{:<10d}"),
        ):
            value = diffusion_stats[key]
            if value_fmt.endswith("d}"):
                value = int(value)
            print("{:<40} ".format(label) + value_fmt.format(value))

    if spec_decode_stats is not None and diffusion_stats is None:
        print("{s:{c}^{n}}".format(s="Speculative Decoding", n=50, c="-"))
        print(
            "{:<40} {:<10.2f}".format(
                "Acceptance rate (%):", spec_decode_stats["acceptance_rate"]
            )
        )
        print(
            "{:<40} {:<10.2f}".format(
                "Acceptance length:", spec_decode_stats["acceptance_length"]
            )
        )
        print("{:<40} {:<10}".format("Drafts:", int(spec_decode_stats["num_drafts"])))
        print(
            "{:<40} {:<10}".format(
                "Draft tokens:", int(spec_decode_stats["draft_tokens"])
            )
        )
        print(
            "{:<40} {:<10}".format(
                "Accepted tokens:", int(spec_decode_stats["accepted_tokens"])
            )
        )
        per_pos = spec_decode_stats.get("per_position_acceptance_rates", [])
        if per_pos:
            print("Per-position acceptance (%):")
            for i, rate in enumerate(per_pos):
                print("{:<40} {:<10.2f}".format(f"  Position {i}:", rate * 100))

    print("=" * 50)

    if profile:
        print("Stopping profiler...")
        profile_input = RequestFuncInput(
            model=model_id,
            prompt=test_prompt,
            api_url=base_url + "/stop_profile",
            prompt_len=test_prompt_len,
            output_len=test_output_len,
            logprobs=logprobs,
        )
        profile_output = await request_func(
            request_func_input=profile_input, session=session
        )
        if profile_output.success:
            print("Profiler stopped")

    await session.close()
    return result


def check_goodput_args(args):
    # Check and parse goodput arguments
    goodput_config_dict = {}
    VALID_NAMES = ["ttft", "tpot", "e2el"]
    if args.goodput:
        goodput_config_dict = parse_goodput(args.goodput)
        for slo_name, slo_val in goodput_config_dict.items():
            if slo_name not in VALID_NAMES:
                raise ValueError(
                    f"Invalid metric name found, {slo_name}: {slo_val}. "
                    "The service level objective name should be one of "
                    f"{str(VALID_NAMES)}. "
                )
            if slo_val < 0:
                raise ValueError(
                    f"Invalid value found, {slo_name}: {slo_val}. "
                    "The service level objective value should be "
                    "non-negative."
                )
    return goodput_config_dict


def parse_goodput(slo_pairs):
    goodput_config_dict = {}
    try:
        for slo_pair in slo_pairs:
            slo_name, slo_val = slo_pair.split(":")
            goodput_config_dict[slo_name] = float(slo_val)
    except ValueError as err:
        raise argparse.ArgumentTypeError(
            "Invalid format found for service level objectives. "
            'Specify service level objectives for goodput as "KEY:VALUE" '
            "pairs, where the key is a metric name, and the value is a "
            "number in milliseconds."
        ) from err
    return goodput_config_dict


def save_to_pytorch_benchmark_format(
    args: argparse.Namespace, results: dict[str, Any], file_name: str
) -> None:
    metrics = [
        "median_ttft_ms",
        "mean_ttft_ms",
        "std_ttft_ms",
        "p99_ttft_ms",
        "mean_tpot_ms",
        "median_tpot_ms",
        "std_tpot_ms",
        "p99_tpot_ms",
        "median_itl_ms",
        "mean_itl_ms",
        "std_itl_ms",
        "p99_itl_ms",
    ]
    # These raw data might be useful, but they are rather big. They can be added
    # later if needed
    ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"]
    pt_records = convert_to_pytorch_benchmark_format(
        args=redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS),
        metrics={k: [results[k]] for k in metrics if k in results},
        extra_info={
            k: results[k]
            for k in results
            if k not in metrics and k not in ignored_metrics
        },
    )
    if pt_records:
        # Don't use json suffix here as we don't want CI to pick it up
        pt_file = f"{os.path.splitext(file_name)[0]}.pytorch.json"
        write_to_json(pt_file, pt_records)


def compute_result_filename(
    args: argparse.Namespace,
    model_id: str,
    label: str,
    current_dt: str,
) -> str | None:
    """Compute the result filename based on benchmark configuration.

    Args:
        args: Command line arguments containing result configuration
        model_id: The model identifier
        label: The benchmark label
        current_dt: Current datetime string

    Returns:
        The computed filename path or None if no result saving is requested

    """
    if not (args.plot_timeline or args.save_result or args.append_result):
        return None

    base_model_id = model_id.split("/")[-1]
    max_concurrency_str = (
        f"-concurrency{args.max_concurrency}"
        if args.max_concurrency is not None
        else ""
    )
    label = label or args.backend

    if args.ramp_up_strategy is not None:
        file_name = f"{label}-ramp-up-{args.ramp_up_strategy}-{args.ramp_up_start_rps}qps-{args.ramp_up_end_rps}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"  # noqa
    else:
        file_name = f"{label}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"  # noqa

    if args.result_filename:
        file_name = args.result_filename

    if args.result_dir:
        os.makedirs(args.result_dir, exist_ok=True)
        file_name = os.path.join(args.result_dir, file_name)

    return file_name


def add_cli_args(parser: FlexibleArgumentParser):
    add_dataset_parser(parser)
    parser.add_argument(
        "--label",
        type=str,
        default=None,
        help="The label (prefix) of the benchmark results. If not specified, "
        "the value of '--backend' will be used as the label.",
    )
    parser.add_argument(
        "--backend",
        type=str,
        default="openai",
        choices=list(ASYNC_REQUEST_FUNCS.keys()),
        help="The type of backend or endpoint to use for the benchmark.",
    )
    parser.add_argument(
        "--base-url",
        type=str,
        default=None,
        help="Server or API base url if not using http host and port.",
    )
    # Use 127.0.0.1 here instead of localhost to force the use of ipv4
    parser.add_argument("--host", type=str, default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument(
        "--endpoint",
        type=str,
        default="/v1/completions",
        help="API endpoint.",
    )
    parser.add_argument(
        "--header",
        metavar="KEY=VALUE",
        nargs="*",
        help="Key-value pairs (e.g, --header x-additional-info=0.3.3) "
        "for headers to be passed with each request. These headers override "
        "per backend constants and values set via environment variable, and "
        "will be overridden by other arguments (such as request ids).",
    )
    parser.add_argument(
        "--max-concurrency",
        type=int,
        default=None,
        help="Maximum number of concurrent requests. This can be used "
        "to help simulate an environment where a higher level component "
        "is enforcing a maximum number of concurrent requests. While the "
        "--request-rate argument controls the rate at which requests are "
        "initiated, this argument will control how many are actually allowed "
        "to execute at a time. This means that when used in combination, the "
        "actual request rate may be lower than specified with --request-rate, "
        "if the server is not processing requests fast enough to keep up.",
    )

    parser.add_argument(
        "--model",
        type=str,
        required=False,
        default=None,
        help="Name of the model. If not specified, will fetch the first model "
        "from the server's /v1/models endpoint.",
    )
    parser.add_argument(
        "--input-len",
        type=int,
        default=None,
        help="General input length for datasets. Maps to dataset-specific "
        "input length arguments (e.g., --random-input-len, --sonnet-input-len). "
        "If not specified, uses dataset defaults.",
    )
    parser.add_argument(
        "--output-len",
        type=int,
        default=None,
        help="General output length for datasets. Maps to dataset-specific "
        "output length arguments (e.g., --random-output-len, --sonnet-output-len). "
        "If not specified, uses dataset defaults.",
    )
    parser.add_argument(
        "--tokenizer",
        type=str,
        help="Name or path of the tokenizer, if not using the default tokenizer.",  # noqa: E501
    )
    parser.add_argument(
        "--tokenizer-mode",
        type=str,
        default="auto",
        help="""Tokenizer mode:\n
        - "auto" will use the tokenizer from `mistral_common` for Mistral models
        if available, otherwise it will use the "hf" tokenizer.\n
        - "hf" will use the fast tokenizer if available.\n
        - "mistral" will always use the tokenizer from `mistral_common`.\n
        - "deepseek_v32" will always use the tokenizer from `deepseek_v32`.\n
        - Other custom values can be supported via plugins.""",
    )
    parser.add_argument("--use-beam-search", action="store_true")
    parser.add_argument(
        "--logprobs",
        type=int,
        default=None,
        help=(
            "Number of logprobs-per-token to compute & return as part of "
            "the request. If unspecified, then either (1) if beam search "
            "is disabled, no logprobs are computed & a single dummy "
            "logprob is returned for each token; or (2) if beam search "
            "is enabled 1 logprob per token is computed"
        ),
    )
    parser.add_argument(
        "--request-rate",
        type=float,
        default=float("inf"),
        help="Number of requests per second. If this is inf, "
        "then all the requests are sent at time 0. "
        "Otherwise, we use Poisson process or gamma distribution "
        "to synthesize the request arrival times.",
    )
    parser.add_argument(
        "--burstiness",
        type=float,
        default=1.0,
        help="Burstiness factor of the request generation. "
        "Only take effect when request_rate is not inf. "
        "Default value is 1, which follows Poisson process. "
        "Otherwise, the request intervals follow a gamma distribution. "
        "A lower burstiness value (0 < burstiness < 1) results in more "
        "bursty requests. A higher burstiness value (burstiness > 1) "
        "results in a more uniform arrival of requests.",
    )
    parser.add_argument(
        "--probe-request-rate",
        type=float,
        default=0.0,
        help="If positive, send single-token text-only probe requests at "
        "this rate (req/s) alongside the main workload, bypassing "
        "--max-concurrency, and report their latency separately. Useful "
        "for measuring how the main workload stalls unrelated requests.",
    )
    parser.add_argument(
        "--disable-tqdm",
        action="store_true",
        help="Specify to disable tqdm progress bar.",
    )
    parser.add_argument(
        "--num-warmups",
        type=int,
        default=0,
        help="Number of warmup requests.",
    )
    parser.add_argument(
        "--profile",
        action="store_true",
        help="Use vLLM Profiling. --profiler-config must be provided on the server.",
    )
    parser.add_argument(
        "--save-result",
        action="store_true",
        help="Specify to save benchmark results to a json file",
    )
    parser.add_argument(
        "--save-detailed",
        action="store_true",
        help="When saving the results, whether to include per request "
        "information such as response, error, ttfts, tpots, etc.",
    )
    parser.add_argument(
        "--append-result",
        action="store_true",
        help="Append the benchmark result to the existing json file.",
    )
    parser.add_argument(
        "--metadata",
        metavar="KEY=VALUE",
        nargs="*",
        help="Key-value pairs (e.g, --metadata version=0.3.3 tp=1) "
        "for metadata of this run to be saved in the result JSON file "
        "for record keeping purposes.",
    )
    parser.add_argument(
        "--result-dir",
        type=str,
        default=None,
        help="Specify directory to save benchmark json results."
        "If not specified, results are saved in the current directory.",
    )
    parser.add_argument(
        "--result-filename",
        type=str,
        default=None,
        help="Specify the filename to save benchmark json results."
        "If not specified, results will be saved in "
        "{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"  # noqa
        " format.",
    )
    parser.add_argument(
        "--ignore-eos",
        action="store_true",
        help="Set ignore_eos flag when sending the benchmark request."
        "Warning: ignore_eos is not supported in deepspeed_mii and tgi.",
    )
    parser.add_argument(
        "--self-timed",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="Use timing information from the traces instead of the configuration. "
        "This is useful when replaying traces faithfully based on their timestamps. "
        "When unset, defaults to False, except for --dataset-name=timed_trace where "
        "it defaults to True. Use --no-self-timed to force off. When off, user "
        "defined generation rates are used and in trace timing info is ignored.",
    )
    parser.add_argument(
        "--percentile-metrics",
        type=str,
        default=None,
        help="Comma-separated list of selected metrics to report percentiles. "
        "This argument specifies the metrics to report percentiles. "
        'Allowed metric names are "ttft", "tpot", "itl", "e2el", '
        '"client_queue_time", "e2el_including_client_queue". '
        'If not specified, defaults to "ttft,tpot,itl" for generative models '
        'and "e2el" for pooling models.',
    )
    parser.add_argument(
        "--metric-percentiles",
        type=str,
        default="99",
        help="Comma-separated list of percentiles for selected metrics. "
        'To report 25-th, 50-th, and 75-th percentiles, use "25,50,75". '
        'Default value is "99".'
        'Use "--percentile-metrics" to select metrics.',
    )
    parser.add_argument(
        "--goodput",
        nargs="+",
        required=False,
        help='Specify service level objectives for goodput as "KEY:VALUE" '
        "pairs, where the key is a metric name, and the value is in "
        'milliseconds. Multiple "KEY:VALUE" pairs can be provided, '
        "separated by spaces. Allowed request level metric names are "
        '"ttft", "tpot", "e2el". For more context on the definition of '
        "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 "
        "and the blog: https://hao-ai-lab.github.io/blogs/distserve",
    )
    parser.add_argument(
        "--request-id-prefix",
        type=str,
        required=False,
        default=f"bench-{uuid.uuid4().hex[:8]}-",
        help="Specify the prefix of request id.",
    )

    sampling_group = parser.add_argument_group("sampling parameters")
    sampling_group.add_argument(
        "--top-p",
        type=float,
        default=None,
        help="Top-p sampling parameter. Only has effect on openai-compatible backends.",
    )
    sampling_group.add_argument(
        "--top-k",
        type=int,
        default=None,
        help="Top-k sampling parameter. Only has effect on openai-compatible backends.",
    )
    sampling_group.add_argument(
        "--min-p",
        type=float,
        default=None,
        help="Min-p sampling parameter. Only has effect on openai-compatible backends.",
    )
    sampling_group.add_argument(
        "--temperature",
        type=float,
        default=None,
        help="Temperature sampling parameter. Only has effect on "
        "openai-compatible backends.",
    )
    sampling_group.add_argument(
        "--frequency-penalty",
        type=float,
        default=None,
        help="Frequency penalty sampling parameter. Only has effect on "
        "openai-compatible backends.",
    )
    sampling_group.add_argument(
        "--presence-penalty",
        type=float,
        default=None,
        help="Presence penalty sampling parameter. Only has effect on "
        "openai-compatible backends.",
    )
    sampling_group.add_argument(
        "--repetition-penalty",
        type=float,
        default=None,
        help="Repetition penalty sampling parameter. Only has effect on "
        "openai-compatible backends.",
    )

    parser.add_argument(
        "--served-model-name",
        type=str,
        default=None,
        help="The model name used in the API. "
        "If not specified, the model name will be the "
        "same as the `--model` argument. ",
    )

    parser.add_argument(
        "--lora-modules",
        nargs="+",
        default=None,
        help="A subset of LoRA module names passed in when "
        "launching the server. For each request, the "
        "script chooses a LoRA module at random by default. "
        "Use --lora-assignment to control selection strategy.",
    )

    parser.add_argument(
        "--lora-assignment",
        type=str,
        default="random",
        choices=["random", "round-robin"],
        help="Strategy for assigning LoRA modules to requests. "
        "'random' (default) selects a LoRA at random for each request. "
        "'round-robin' cycles through LoRA modules deterministically.",
    )

    parser.add_argument(
        "--ramp-up-strategy",
        type=str,
        default=None,
        choices=["linear", "exponential"],
        help="The ramp-up strategy. This would be used to "
        "ramp up the request rate from initial RPS to final "
        "RPS rate (specified by --ramp-up-start-rps and "
        "--ramp-up-end-rps.) over the duration of the benchmark.",
    )
    parser.add_argument(
        "--ramp-up-start-rps",
        type=int,
        default=None,
        help="The starting request rate for ramp-up (RPS). "
        "Needs to be specified when --ramp-up-strategy is used.",
    )
    parser.add_argument(
        "--ramp-up-end-rps",
        type=int,
        default=None,
        help="The ending request rate for ramp-up (RPS). "
        "Needs to be specified when --ramp-up-strategy is used.",
    )
    parser.add_argument(
        "--ready-check-timeout-sec",
        type=int,
        default=0,
        help="Maximum time to wait for the endpoint to become ready "
        "in seconds. Ready check will be skipped by default.",
    )

    parser.add_argument(
        "--chat-template-kwargs",
        type=json.loads,
        default=None,
        help="A JSON string of kwargs forwarded to the tokenizer's "
        "apply_chat_template when a dataset renders prompts client-side "
        "(e.g. custom / speed_bench). "
        "Example: '{\"thinking\": true}' to enable reasoning models.",
    )
    parser.add_argument(
        "--extra-body",
        help="A JSON string representing extra body parameters to include "
        "in each request."
        'Example: \'{"chat_template_kwargs":{"enable_thinking":false}}\'',
        type=json.loads,
        default=None,
    )
    parser.add_argument(
        "--skip-tokenizer-init",
        action="store_true",
        default=False,
        help="Skip initialization of tokenizer and detokenizer",
    )

    parser.add_argument(
        "--insecure",
        action="store_true",
        default=False,
        help="Disable SSL certificate verification. Use this option when "
        "connecting to servers with self-signed certificates.",
    )

    parser.add_argument(
        "--plot-timeline",
        action="store_true",
        help="Generate an HTML timeline plot showing request execution. "
        "The plot will be saved alongside the results JSON file.",
    )
    parser.add_argument(
        "--timeline-itl-thresholds",
        type=str,
        default="25,50",
        help="ITL thresholds in milliseconds for timeline plot coloring. "
        "Specify two comma-separated values to categorize inter-token "
        "latencies into three groups: below first threshold (green), "
        "between thresholds (orange), and above second threshold (red).",
    )
    parser.add_argument(
        "--plot-dataset-stats",
        action="store_true",
        help="Generate a matplotlib figure with dataset statistics showing "
        "prompt tokens, output tokens, and combined token distributions.",
    )


def main(args: argparse.Namespace) -> dict[str, Any]:
    return asyncio.run(main_async(args))


async def main_async(args: argparse.Namespace) -> dict[str, Any]:
    print(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS))
    if args.max_concurrency is not None and args.max_concurrency <= 0:
        raise ValueError("--max-concurrency must be greater than 0")

    random.seed(args.seed)
    np.random.seed(args.seed)

    # Validate timeline ITL thresholds
    if args.plot_timeline:
        try:
            itl_thresholds = [
                float(t.strip()) for t in args.timeline_itl_thresholds.split(",")
            ]
            if len(itl_thresholds) != 2:
                raise ValueError(
                    f"Expected 2 ITL threshold values, got {len(itl_thresholds)}"
                )
        except ValueError as e:
            raise ValueError(f"Invalid --timeline-itl-thresholds format: {e}") from e

    # Validate ramp-up arguments
    if args.ramp_up_strategy is not None:
        if args.request_rate != float("inf"):
            raise ValueError(
                "When using ramp-up, do not specify --request-rate. "
                "The request rate will be controlled by ramp-up parameters. "
                "Please remove the --request-rate argument."
            )
        if args.ramp_up_start_rps is None or args.ramp_up_end_rps is None:
            raise ValueError(
                "When using --ramp-up-strategy, both --ramp-up-start-rps and "
                "--ramp-up-end-rps must be specified"
            )
        if args.ramp_up_start_rps < 0 or args.ramp_up_end_rps < 0:
            raise ValueError("Ramp-up start and end RPS must be non-negative")
        if args.ramp_up_start_rps > args.ramp_up_end_rps:
            raise ValueError("Ramp-up start RPS must be less than end RPS")
        if args.ramp_up_strategy == "exponential" and args.ramp_up_start_rps == 0:
            raise ValueError("For exponential ramp-up, the start RPS cannot be 0.")

    label = args.label

    if args.base_url is not None:
        api_url = f"{args.base_url}{args.endpoint}"
        base_url = f"{args.base_url}"
    else:
        host_port = join_host_port(args.host, args.port)
        api_url = f"http://{host_port}{args.endpoint}"
        base_url = f"http://{host_port}"

    # Headers
    headers = None
    if args.header:
        headers = {}
        for item in args.header:
            if "=" in item:
                kvstring = item.split("=", 1)
                headers[kvstring[0].strip()] = kvstring[1].strip()
            else:
                raise ValueError("Invalid header format. Please use KEY=VALUE format.")

    # SSL context configuration
    ssl_context: ssl.SSLContext | bool | None = None
    if args.insecure:
        # Disable SSL certificate verification
        ssl_context = False
    elif "https://" in base_url:
        # Use default SSL context for HTTPS
        ssl_context = True

    # Fetch model from server if not specified
    if args.model is None:
        print("Model not specified, fetching first model from server...")
        model_name, model_id = await get_first_model_from_server(
            base_url, headers, ssl_context
        )
        print(f"First model name: {model_name}, first model id: {model_id}")
    else:
        model_name = args.served_model_name
        model_id = args.model

    if args.skip_tokenizer_init:
        tokenizer_id = None
        tokenizer_mode = None
        tokenizer = None
    else:
        tokenizer_id = args.tokenizer if args.tokenizer is not None else model_id
        tokenizer_mode = args.tokenizer_mode
        tokenizer = get_tokenizer(
            tokenizer_id,
            tokenizer_mode=tokenizer_mode,
            trust_remote_code=args.trust_remote_code,
        )

    # Validate dataset name/path
    if args.dataset_name is None:
        raise ValueError(
            "Please specify '--dataset-name' and the corresponding "
            "'--dataset-path' if required."
        )

    if (
        args.dataset_name
        in ["random", "random-mm", "random-rerank", "prefix_repetition"]
        and args.dataset_path is not None
    ):
        raise ValueError(
            f"Cannot use '{args.dataset_name}' dataset with --dataset-path. "
            "Please specify the appropriate --dataset-name (e.g., "
            "'sharegpt', 'custom', 'sonnet') for your dataset file: "
            f"{args.dataset_path}"
        )

    # Map general --input-len and --output-len to all dataset-specific arguments
    if args.input_len is not None:
        args.random_input_len = args.input_len
        args.sonnet_input_len = args.input_len

    if args.output_len is not None:
        args.random_output_len = args.output_len
        args.sonnet_output_len = args.output_len
        args.sharegpt_output_len = args.output_len
        args.custom_output_len = args.output_len
        args.hf_output_len = args.output_len
        args.spec_bench_output_len = args.output_len
        args.prefix_repetition_output_len = args.output_len

    # when using random datasets, default to ignoring EOS
    # so generation runs to the requested length
    if (
        args.dataset_name in ("random", "random-mm")
        and args.backend in OPENAI_COMPATIBLE_BACKENDS
    ):
        args.ignore_eos = True

    if args.dataset_name == "timed_trace":
        if args.backend not in ("vllm", "openai"):
            raise ValueError(
                "timed_trace dataset passes pre-tokenized prompts (list[int])"
                " and requires a completions backend ('vllm' or 'openai')."
            )
        # timed_trace carries per-request timestamps;
        # ignore EOS so generation runs to the trace's specified output length,
        # and default to using those timestamps for scheduling unless the user
        # opted out.
        args.ignore_eos = True
        if args.self_timed is None:
            args.self_timed = True
    else:
        # if this is set for anything else, it is an error
        if args.self_timed is not None:
            raise ValueError(
                "--self-timed/--no-self-timed is only supported with "
                "--dataset-name=timed_trace"
            )
        # for any non self-timed trace, this is False
        args.self_timed = False

    # Load the dataset.
    input_requests = get_samples(args, tokenizer)

    if args.dataset_name in ("random", "prefix_repetition"):
        input_requests = await _align_prompts_to_server_tokenizer(
            base_url, model_id, input_requests, ssl_context
        )

    goodput_config_dict = check_goodput_args(args)

    backend = args.backend
    task_type = TaskType.POOLING if backend in POOLING_BACKENDS else TaskType.GENERATION

    # Collect the sampling parameters.
    if task_type == TaskType.GENERATION:
        sampling_params = {
            k: v
            for k, v in {
                "top_p": args.top_p,
                "top_k": args.top_k,
                "min_p": args.min_p,
                "temperature": args.temperature,
                "frequency_penalty": args.frequency_penalty,
                "presence_penalty": args.presence_penalty,
                "repetition_penalty": args.repetition_penalty,
            }.items()
            if v is not None
        }

        # Sampling parameters are only supported by openai-compatible backend.
        if sampling_params and args.backend not in OPENAI_COMPATIBLE_BACKENDS:
            raise ValueError(
                "Sampling parameters are only supported by openai-compatible backends."
            )

        # The Responses API accepts every sampling parameter above except
        # min_p, which it would silently drop as an unknown field.
        if args.backend == "openai-responses" and "min_p" in sampling_params:
            raise ValueError("--min-p is not supported by the Responses API.")

        if "temperature" not in sampling_params:
            print(
                "WARNING: vllm bench serve no longer sets temperature==0 (greedy) "
                "in requests by default. The default will be determined on the "
                "server side and can be model/API specific. "
                "For the old behavior, include --temperature=0."
            )

        default_percentile_metrics = "ttft,tpot,itl"
    else:
        sampling_params = {}
        default_percentile_metrics = "e2el"

    extra_body = args.extra_body or {}
    extra_body = {**sampling_params, **extra_body}

    percentile_metrics: str = args.percentile_metrics or default_percentile_metrics

    # Avoid GC processing "static" data - reduce pause times.
    freeze_gc_heap()

    benchmark_result = await benchmark(
        task_type=task_type,
        endpoint_type=backend,
        api_url=api_url,
        base_url=base_url,
        model_id=model_id,
        model_name=model_name,
        tokenizer=tokenizer,
        input_requests=input_requests,
        logprobs=args.logprobs,
        request_rate=args.request_rate,
        burstiness=args.burstiness,
        disable_tqdm=args.disable_tqdm,
        num_warmups=args.num_warmups,
        profile=args.profile,
        selected_percentile_metrics=percentile_metrics.split(","),
        selected_percentiles=[float(p) for p in args.metric_percentiles.split(",")],
        ignore_eos=args.ignore_eos,
        goodput_config_dict=goodput_config_dict,
        max_concurrency=args.max_concurrency,
        lora_modules=args.lora_modules,
        lora_assignment=args.lora_assignment,
        extra_headers=headers,
        extra_body=extra_body,
        ramp_up_strategy=args.ramp_up_strategy,
        ramp_up_start_rps=args.ramp_up_start_rps,
        ramp_up_end_rps=args.ramp_up_end_rps,
        ready_check_timeout_sec=args.ready_check_timeout_sec,
        ssl_context=ssl_context,
        self_timed=args.self_timed,
        probe_request_rate=args.probe_request_rate,
    )

    # Save config and results to json
    result_json: dict[str, Any] = {}

    # Setup
    current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
    result_json["date"] = current_dt
    result_json["endpoint_type"] = args.backend  # for backward compatibility
    result_json["backend"] = args.backend
    result_json["label"] = label
    result_json["model_id"] = model_id
    result_json["tokenizer_id"] = tokenizer_id
    result_json["num_prompts"] = args.num_prompts

    # Metadata
    if args.metadata:
        for item in args.metadata:
            if "=" in item:
                kvstring = item.split("=", 1)
                result_json[kvstring[0].strip()] = kvstring[1].strip()
            else:
                raise ValueError(
                    "Invalid metadata format. Please use KEY=VALUE format."
                )

    # Traffic
    result_json["request_rate"] = (
        args.request_rate if args.request_rate < float("inf") else "inf"
    )
    result_json["burstiness"] = args.burstiness
    result_json["max_concurrency"] = args.max_concurrency

    if args.ramp_up_strategy is not None:
        result_json["ramp_up_strategy"] = args.ramp_up_strategy
        result_json["ramp_up_start_rps"] = args.ramp_up_start_rps
        result_json["ramp_up_end_rps"] = args.ramp_up_end_rps

    # Merge with benchmark result
    result_json = {**result_json, **benchmark_result}

    # Compute file_name once before using it for plots or saving results
    file_name = compute_result_filename(args, model_id, label, current_dt)

    # Generate timeline plot if requested
    if args.plot_timeline:
        assert file_name is not None, (
            "file_name must be set when plot_timeline is enabled"
        )
        try:
            from vllm.benchmarks.plot import generate_timeline_plot

            # Prepare per-request data for timeline
            per_request_data = []
            start_times = benchmark_result.get("start_times", [])
            ttfts = benchmark_result.get("ttfts", [])
            itls = benchmark_result.get("itls", [])
            input_lens = benchmark_result.get("input_lens", [])
            output_lens = benchmark_result.get("output_lens", [])

            if start_times and ttfts and itls:
                for i in range(len(start_times)):
                    # Calculate latency as ttft + sum of all itls
                    latency = ttfts[i] + sum(itls[i]) if itls[i] else ttfts[i]

                    per_request_data.append(
                        {
                            "start_time": start_times[i],
                            "ttft": ttfts[i],
                            "itl": itls[i],
                            "latency": latency,
                            "prompt_len": input_lens[i],
                            "output_tokens": output_lens[i],
                        }
                    )

                timeline_path = Path(file_name).with_suffix(".timeline.html")
                # Convert thresholds from milliseconds to seconds
                itl_thresholds_sec = [
                    float(t) / 1000.0 for t in args.timeline_itl_thresholds.split(",")
                ]
                generate_timeline_plot(
                    per_request_data, timeline_path, itl_thresholds=itl_thresholds_sec
                )
            else:
                warnings.warn(
                    "Timeline plot requires detailed metrics. "
                    "Ensure the benchmark completed successfully.",
                    stacklevel=2,
                )
        except Exception as e:
            warnings.warn(f"Failed to generate timeline plot: {e}", stacklevel=2)

    # Generate dataset statistics plot if requested
    if args.plot_dataset_stats:
        assert file_name is not None, (
            "file_name must be set when plot_dataset_stats is enabled"
        )
        try:
            from vllm.benchmarks.plot import generate_dataset_stats_plot

            # Prepare per-request data for dataset stats
            per_request_data = []
            input_lens = benchmark_result.get("input_lens", [])
            output_lens = benchmark_result.get("output_lens", [])

            if input_lens and output_lens:
                for req_input_len, req_output_len in zip(input_lens, output_lens):
                    per_request_data.append(
                        {
                            "prompt_len": req_input_len,
                            "output_tokens": req_output_len,
                        }
                    )

                stats_path = Path(file_name).with_suffix(".dataset_stats.png")
                generate_dataset_stats_plot(per_request_data, stats_path)
            else:
                warnings.warn(
                    "Dataset statistics plot requires input and "
                    "output length data. Ensure the benchmark completed "
                    "successfully.",
                    stacklevel=2,
                )
        except Exception as e:
            warnings.warn(
                f"Failed to generate dataset statistics plot: {e}", stacklevel=2
            )

    if not args.save_detailed:
        # Remove fields with too many data points
        for field in [
            "input_lens",
            "output_lens",
            "start_times",
            "ttfts",
            "itls",
            "generated_texts",
            "errors",
        ]:
            if field in result_json:
                del result_json[field]
            if field in benchmark_result:
                del benchmark_result[field]

    # Save to file
    if args.save_result or args.append_result:
        assert file_name is not None, (
            "file_name must be set when save_result or append_result is enabled"
        )
        with open(
            file_name, mode="a+" if args.append_result else "w", encoding="utf-8"
        ) as outfile:
            # Append a newline.
            if args.append_result and outfile.tell() != 0:
                outfile.write("\n")
            json.dump(result_json, outfile)
        save_to_pytorch_benchmark_format(args, result_json, file_name)

    return result_json
