sglang/benchmark/simulator/bench_serving.py

"""SGLang serving benchmark adapter for simulator traffic.

This script deliberately reuses SGLang's benchmark implementation and dataset
loaders.  It only owns the simulator-specific parts of the protocol:

* convert request-rate or trace timestamps into logical arrival timestamps;
* inject the internal ``sampling_params.custom_params.simulation`` metadata;
* avoid client-side pacing in OFFLINE mode; and
* display backend-produced simulator metrics when they are locally available.

User datasets must not contain simulator metadata.
"""

import argparse
import contextlib
import json
import os
import re
import sys
from dataclasses import fields
from pathlib import Path
from typing import AsyncGenerator, List, Optional

import aiohttp
import numpy as np
from sglang_simulator.compat import validate_benchmark_runtime
from sglang_simulator.dataset.autobench import register_autobench_dataset

register_autobench_dataset()

from sglang.benchmark import serving
from sglang.benchmark.datasets.common import DatasetRow

_ORIGINAL_AIOHTTP_REQUEST = None
_ORIGINAL_CALCULATE_METRICS = serving.calculate_metrics
_ORIGINAL_GET_REQUEST = serving.get_request
_ORIGINAL_RUN_BENCHMARK = serving.run_benchmark
_SIMULATOR_MODE = "offline"
_USE_TRACE_TIMESTAMPS = False


def _metrics_path() -> Path:
    output_dir = Path(
        os.getenv("SGLANG_SIMULATOR_OUTPUT_DIR", "/tmp/sglang_simulator/output")
    )
    return output_dir / "metrics.json"


def _load_backend_metrics() -> Optional[dict]:
    metrics_path = _metrics_path()
    if not metrics_path.is_file():
        return None
    return json.loads(metrics_path.read_text(encoding="utf-8"))


class _DurationReplacingStream:
    """Keep SGLang's output format but print the simulated duration."""

    def __init__(self, target):
        self.target = target

    def write(self, text):
        if "Benchmark duration (s):" in text:
            metrics = _load_backend_metrics()
            if metrics is not None and "duration" in metrics:
                text = "{:<40} {:<10.2f}".format(
                    "Benchmark duration (s):", metrics["duration"]
                )
        return self.target.write(text)

    def flush(self):
        return self.target.flush()


def _set_simulation_metadata(
    request: DatasetRow, *, created_time_ms: float, total_request: int
) -> None:
    """Attach transient metadata without replacing dataset-specific parameters."""
    extra_request_body = dict(request.extra_request_body or {})
    extra_request_body["simulation"] = {
        "created_time_ms": created_time_ms,
        "total_request": total_request,
    }
    request.extra_request_body = extra_request_body


async def simulator_get_request(
    input_requests: List[DatasetRow],
    request_rate: float,
    use_trace_timestamps: bool = False,
    slowdown_factor: float = 1.0,
) -> AsyncGenerator[DatasetRow, None]:
    """Generate simulator traffic while retaining official BLOCKING pacing."""
    # The benchmark may not forward --use-trace-timestamps to get_request(),
    # so preserve the parsed value in this adapter.
    use_trace_timestamps = use_trace_timestamps or _USE_TRACE_TIMESTAMPS
    if _SIMULATOR_MODE == "blocking":
        async for request in _ORIGINAL_GET_REQUEST(
            input_requests,
            request_rate,
            use_trace_timestamps=use_trace_timestamps,
            slowdown_factor=slowdown_factor,
        ):
            yield request
        return

    total_request = len(input_requests)
    if use_trace_timestamps:
        if any(request.timestamp is None for request in input_requests):
            raise ValueError(
                "--use-trace-timestamps requires every request to have timestamp"
            )
        input_requests.sort(key=lambda request: request.timestamp)
        trace_start_time_ms = input_requests[0].timestamp if input_requests else 0.0
        for request in input_requests:
            created_time_ms = (
                float(request.timestamp) - float(trace_start_time_ms)
            ) * slowdown_factor
            _set_simulation_metadata(
                request,
                created_time_ms=created_time_ms,
                total_request=total_request,
            )
            yield request
        return

    created_time_ms = 0.0
    for request in input_requests:
        _set_simulation_metadata(
            request,
            created_time_ms=created_time_ms,
            total_request=total_request,
        )
        yield request
        if request_rate != float("inf"):
            created_time_ms += np.random.exponential(1.0 / request_rate) * 1000.0


def install_aiohttp_json_hijack(
    *, hijack_url_regex: Optional[str] = r"/generate(?:\?.*)?$"
) -> None:
    """Move transient metadata into the already-built sampling parameters."""
    global _ORIGINAL_AIOHTTP_REQUEST
    if _ORIGINAL_AIOHTTP_REQUEST is not None:
        return

    pattern = re.compile(hijack_url_regex) if hijack_url_regex else None
    _ORIGINAL_AIOHTTP_REQUEST = aiohttp.ClientSession._request

    async def patched_request(self, method, url, **kwargs):
        if pattern is None or pattern.search(str(url)):
            payload = kwargs.get("json")
            if isinstance(payload, dict) and "simulation" in payload:
                simulation = payload.pop("simulation")
                sampling_params = payload.setdefault("sampling_params", {})
                custom_params = sampling_params.setdefault("custom_params", {})
                custom_params["simulation"] = simulation
                kwargs["json"] = payload
        return await _ORIGINAL_AIOHTTP_REQUEST(self, method, url, **kwargs)

    aiohttp.ClientSession._request = patched_request


def simulator_calculate_metrics(*args, **kwargs):
    """Use simulator metrics; mark unsupported client-only fields with -1."""
    client_metrics, output_lens = _ORIGINAL_CALCULATE_METRICS(*args, **kwargs)
    backend_metrics = _load_backend_metrics()
    if backend_metrics is None:
        print(
            f"Simulator metrics are not available at {_metrics_path()}; "
            "showing client-side benchmark metrics."
        )
        return client_metrics, output_lens

    metric_names = {field.name for field in fields(serving.BenchmarkMetrics)}
    values = {name: backend_metrics.get(name, -1) for name in metric_names}
    return serving.BenchmarkMetrics(**values), output_lens


def _replace_output_file_duration(
    args: argparse.Namespace, simulated_duration: float
) -> None:
    output_file = getattr(args, "output_file", None)
    if not output_file:
        return
    path = Path(output_file)
    if not path.is_file():
        return
    lines = path.read_text(encoding="utf-8").splitlines()
    if not lines:
        return
    last_result = json.loads(lines[-1])
    last_result["duration"] = simulated_duration
    lines[-1] = json.dumps(last_result)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def simulator_run_benchmark(args: argparse.Namespace):
    global _USE_TRACE_TIMESTAMPS
    if args.backend != "sglang":
        raise ValueError(
            "benchmark/simulator/bench_serving.py requires --backend sglang"
        )
    if args.dataset_name == "mooncake":
        raise ValueError(
            "Mooncake's multi-round scheduler is not supported by the simulator "
            "benchmark adapter"
        )
    _USE_TRACE_TIMESTAMPS = getattr(args, "use_trace_timestamps", False)
    args.profile = True
    with contextlib.redirect_stdout(_DurationReplacingStream(sys.stdout)):
        result = _ORIGINAL_RUN_BENCHMARK(args)

    backend_metrics = _load_backend_metrics()
    if backend_metrics is not None and "duration" in backend_metrics:
        simulated_duration = backend_metrics["duration"]
        if isinstance(result, dict):
            result["duration"] = simulated_duration
        _replace_output_file_duration(args, simulated_duration)
    return result


def _extract_simulator_args(argv: list[str]) -> tuple[str, list[str]]:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument(
        "--simulator-mode",
        choices=("offline", "blocking"),
        default="offline",
        help=argparse.SUPPRESS,
    )
    args, remaining = parser.parse_known_args(argv)
    return args.simulator_mode, remaining


def _simulator_argument_parser(base_parser):
    """Include simulator-owned datasets in SGLang's hard-coded CLI choices."""

    class SimulatorArgumentParser(base_parser):
        def add_argument(self, *name_or_flags, **kwargs):
            choices = kwargs.get("choices")
            if (
                "--dataset-name" in name_or_flags
                and choices is not None
                and "autobench" not in choices
            ):
                kwargs["choices"] = [*choices, "autobench"]
            if "--warmup-requests" in name_or_flags:
                kwargs["default"] = 0
            return super().add_argument(*name_or_flags, **kwargs)

    return SimulatorArgumentParser


def cli_main() -> None:
    global _SIMULATOR_MODE
    validate_benchmark_runtime()
    if any(argument in ("-h", "--help") for argument in sys.argv[1:]):
        print(
            "SGLang Simulator option: "
            "--simulator-mode {offline,blocking} (default: offline)\n"
        )
    _SIMULATOR_MODE, remaining = _extract_simulator_args(sys.argv[1:])
    sys.argv = [sys.argv[0], *remaining]

    serving.get_request = simulator_get_request
    serving.calculate_metrics = simulator_calculate_metrics
    serving.run_benchmark = simulator_run_benchmark
    install_aiohttp_json_hijack()

    print(f"SGLang Simulator benchmark mode: {_SIMULATOR_MODE.upper()}")
    original_parser = serving.ArgumentParser
    serving.ArgumentParser = _simulator_argument_parser(original_parser)
    try:
        serving.cli_main()
    finally:
        serving.ArgumentParser = original_parser


if __name__ == "__main__":
    cli_main()