startup.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import shlex
import subprocess
from dataclasses import dataclass
from datetime import datetime
from functools import lru_cache
from pathlib import Path
from typing import ClassVar

from vllm.benchmarks.startup import add_cli_args as add_startup_cli_args
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.utils.import_utils import PlaceholderModule

from .param_sweep import ParameterSweep, ParameterSweepItem
from .utils import sanitize_filename

try:
    import pandas as pd
except ImportError:
    pd = PlaceholderModule("pandas")


@lru_cache(maxsize=1)
def _get_supported_startup_keys() -> set[str]:
    parser = FlexibleArgumentParser(add_help=False)
    add_startup_cli_args(parser)

    supported: set[str] = {"config"}
    for action in parser._actions:
        if action.dest and action.dest is not argparse.SUPPRESS:
            supported.add(action.dest)
        for option in action.option_strings:
            if option.startswith("--"):
                supported.add(option.lstrip("-").replace("-", "_"))

    return supported


def _is_supported_param(param_key: str, supported: set[str]) -> bool:
    if param_key == "_benchmark_name":
        return True
    prefix = param_key.split(".", 1)[0]
    normalized = prefix.replace("-", "_")
    return normalized in supported


def _filter_params(
    params: ParameterSweep, *, supported: set[str], strict: bool
) -> ParameterSweep:
    filtered = []
    for item in params:
        kept: dict[str, object] = {}
        dropped: list[str] = []
        for key, value in item.items():
            if _is_supported_param(key, supported):
                kept[key] = value
            else:
                dropped.append(key)

        if dropped:
            label = item.get("_benchmark_name") or item.as_text()
            message = (
                "Ignoring unsupported startup params"
                f"{' for ' + str(label) if label else ''}: "
                f"{', '.join(sorted(dropped))}"
            )
            if strict:
                raise ValueError(message)
            print(message)

        filtered.append(ParameterSweepItem.from_record(kept))

    return ParameterSweep(filtered)


def _update_run_data(
    run_data: dict[str, object],
    serve_overrides: ParameterSweepItem,
    startup_overrides: ParameterSweepItem,
    run_number: int,
) -> dict[str, object]:
    run_data["run_number"] = run_number
    run_data.update(serve_overrides)
    run_data.update(startup_overrides)
    return run_data


def _strip_arg(cmd: list[str], keys: tuple[str, ...]) -> list[str]:
    stripped: list[str] = []
    skip_next = False
    for arg in cmd:
        if skip_next:
            skip_next = False
            continue
        if arg in keys:
            skip_next = True
            continue
        if any(arg.startswith(f"{key}=") for key in keys):
            continue
        stripped.append(arg)
    return stripped


def _apply_output_json(cmd: list[str], output_path: Path) -> list[str]:
    keys = ("--output-json", "--output_json")
    cmd = _strip_arg(cmd, keys)
    return [*cmd, keys[0], str(output_path)]


def _get_comb_base_path(
    output_dir: Path,
    serve_comb: ParameterSweepItem,
    startup_comb: ParameterSweepItem,
) -> Path:
    parts = list[str]()
    if serve_comb:
        parts.extend(("SERVE-", serve_comb.name))
    if startup_comb:
        parts.extend(("STARTUP-", startup_comb.name))
    return output_dir / sanitize_filename("-".join(parts))


def _get_comb_run_path(base_path: Path, run_number: int | None) -> Path:
    if run_number is None:
        return base_path / "summary.json"
    return base_path / f"run={run_number}.json"


def run_benchmark(
    startup_cmd: list[str],
    *,
    serve_overrides: ParameterSweepItem,
    startup_overrides: ParameterSweepItem,
    run_number: int,
    output_path: Path,
    show_stdout: bool,
    dry_run: bool,
) -> dict[str, object] | None:
    cmd = serve_overrides.apply_to_cmd(startup_cmd)
    cmd = startup_overrides.apply_to_cmd(cmd)
    cmd = _apply_output_json(cmd, output_path)

    print("[BEGIN BENCHMARK]")
    print(f"Serve overrides: {serve_overrides}")
    print(f"Startup overrides: {startup_overrides}")
    print(f"Run Number: {run_number}")
    print(f"Benchmark command: {cmd}")
    print(f"Output file: {output_path}")

    if output_path.exists():
        print("Found existing results. Skipping.")

        with output_path.open("r", encoding="utf-8") as f:
            run_data = json.load(f)
            return _update_run_data(
                run_data, serve_overrides, startup_overrides, run_number
            )

    if dry_run:
        print("[END BENCHMARK]")
        return None

    output_path.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(
        cmd,
        stdout=None if show_stdout else subprocess.DEVNULL,
        check=True,
    )

    with output_path.open("r", encoding="utf-8") as f:
        run_data = json.load(f)

    run_data = _update_run_data(
        run_data, serve_overrides, startup_overrides, run_number
    )

    with output_path.open("w", encoding="utf-8") as f:
        json.dump(run_data, f, indent=4)

    print("[END BENCHMARK]")
    return run_data


def run_comb(
    startup_cmd: list[str],
    *,
    serve_comb: ParameterSweepItem,
    startup_comb: ParameterSweepItem,
    base_path: Path,
    num_runs: int,
    show_stdout: bool,
    dry_run: bool,
) -> list[dict[str, object]] | None:
    comb_data = list[dict[str, object]]()
    for run_number in range(num_runs):
        run_data = run_benchmark(
            startup_cmd,
            serve_overrides=serve_comb,
            startup_overrides=startup_comb,
            run_number=run_number,
            output_path=_get_comb_run_path(base_path, run_number),
            show_stdout=show_stdout,
            dry_run=dry_run,
        )
        if run_data is not None:
            comb_data.append(run_data)

    if dry_run:
        return None

    with _get_comb_run_path(base_path, run_number=None).open(
        "w", encoding="utf-8"
    ) as f:
        json.dump(comb_data, f, indent=4)

    return comb_data


def run_combs(
    startup_cmd: list[str],
    *,
    serve_params: ParameterSweep,
    startup_params: ParameterSweep,
    output_dir: Path,
    num_runs: int,
    show_stdout: bool,
    dry_run: bool,
) -> "pd.DataFrame | None":
    all_data = list[dict[str, object]]()
    for serve_comb in serve_params:
        for startup_comb in startup_params:
            base_path = _get_comb_base_path(output_dir, serve_comb, startup_comb)
            comb_data = run_comb(
                startup_cmd,
                serve_comb=serve_comb,
                startup_comb=startup_comb,
                base_path=base_path,
                num_runs=num_runs,
                show_stdout=show_stdout,
                dry_run=dry_run,
            )
            if comb_data is not None:
                all_data.extend(comb_data)

    if dry_run:
        return None

    combined_df = pd.DataFrame.from_records(all_data)
    combined_df.to_csv(output_dir / "summary.csv")
    return combined_df


@dataclass
class SweepStartupArgs:
    startup_cmd: list[str]
    serve_params: ParameterSweep
    startup_params: ParameterSweep
    output_dir: Path
    num_runs: int
    show_stdout: bool
    dry_run: bool
    resume: str | None
    strict_params: bool

    parser_name: ClassVar[str] = "startup"
    parser_help: ClassVar[str] = (
        "Benchmark vLLM startup time over parameter combinations."
    )

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace):
        startup_cmd = shlex.split(args.startup_cmd)

        if args.serve_params:
            serve_params = ParameterSweep.read_json(args.serve_params)
        else:
            serve_params = ParameterSweep.from_records([{}])

        if args.startup_params:
            startup_params = ParameterSweep.read_json(args.startup_params)
        else:
            startup_params = ParameterSweep.from_records([{}])

        supported = _get_supported_startup_keys()
        serve_params = _filter_params(
            serve_params, supported=supported, strict=args.strict_params
        )
        startup_params = _filter_params(
            startup_params, supported=supported, strict=args.strict_params
        )

        if args.num_runs < 1:
            raise ValueError("`num_runs` should be at least 1.")

        return cls(
            startup_cmd=startup_cmd,
            serve_params=serve_params,
            startup_params=startup_params,
            output_dir=Path(args.output_dir),
            num_runs=args.num_runs,
            show_stdout=args.show_stdout,
            dry_run=args.dry_run,
            resume=args.resume,
            strict_params=args.strict_params,
        )

    @classmethod
    def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
        parser.add_argument(
            "--startup-cmd",
            type=str,
            default="vllm bench startup",
            help="The command used to run the startup benchmark.",
        )
        parser.add_argument(
            "--serve-params",
            type=str,
            default=None,
            help="Path to JSON file containing parameter combinations "
            "for the `vllm serve` command. Only parameters supported by "
            "`vllm bench startup` will be applied.",
        )
        parser.add_argument(
            "--startup-params",
            type=str,
            default=None,
            help="Path to JSON file containing parameter combinations "
            "for the `vllm bench startup` command.",
        )
        parser.add_argument(
            "-o",
            "--output-dir",
            type=str,
            default="results",
            help="The directory to which results are written.",
        )
        parser.add_argument(
            "--num-runs",
            type=int,
            default=1,
            help="Number of runs per parameter combination.",
        )
        parser.add_argument(
            "--show-stdout",
            action="store_true",
            help="If set, logs the standard output of subcommands.",
        )
        parser.add_argument(
            "--dry-run",
            action="store_true",
            help="If set, prints the commands to run, "
            "then exits without executing them.",
        )
        parser.add_argument(
            "--resume",
            type=str,
            default=None,
            help="Set this to the name of a directory under `output_dir` (which is a "
            "timestamp) to resume a previous execution of this script, i.e., only run "
            "parameter combinations for which there are still no output files.",
        )
        parser.add_argument(
            "--strict-params",
            action="store_true",
            help="If set, unknown parameters in sweep files raise an error "
            "instead of being ignored.",
        )
        return parser


def run_main(args: SweepStartupArgs):
    timestamp = args.resume or datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = args.output_dir / timestamp

    if args.resume and not output_dir.exists():
        raise ValueError(f"Cannot resume from non-existent directory ({output_dir})")

    try:
        return run_combs(
            startup_cmd=args.startup_cmd,
            serve_params=args.serve_params,
            startup_params=args.startup_params,
            output_dir=output_dir,
            num_runs=args.num_runs,
            show_stdout=args.show_stdout,
            dry_run=args.dry_run,
        )
    except BaseException as exc:
        raise RuntimeError(
            f"The script was terminated early. Use `--resume {timestamp}` "
            f"to continue the script from its last checkpoint."
        ) from exc


def main(args: argparse.Namespace):
    run_main(SweepStartupArgs.from_cli_args(args))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=SweepStartupArgs.parser_help)
    SweepStartupArgs.add_cli_args(parser)
    main(parser.parse_args())