Upgrade to vllm 0.17.0 corex v4.1 overlay
This commit is contained in:
328
vllm/benchmarks/sweep/serve_workload.py
Normal file
328
vllm/benchmarks/sweep/serve_workload.py
Normal file
@@ -0,0 +1,328 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import argparse
|
||||
import math
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import ClassVar, Literal, get_args
|
||||
|
||||
import numpy as np
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from vllm.benchmarks.datasets import DEFAULT_NUM_PROMPTS
|
||||
from vllm.utils.import_utils import PlaceholderModule
|
||||
|
||||
from .param_sweep import ParameterSweep, ParameterSweepItem
|
||||
from .serve import (
|
||||
SweepServeArgs,
|
||||
_get_comb_base_path,
|
||||
run_comb,
|
||||
server_ctx,
|
||||
)
|
||||
from .server import ServerProcess
|
||||
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
pd = PlaceholderModule("pandas")
|
||||
|
||||
|
||||
WorkloadVariable = Literal["request_rate", "max_concurrency"]
|
||||
|
||||
|
||||
def _estimate_workload_value(
|
||||
run_data: dict[str, object],
|
||||
workload_var: WorkloadVariable,
|
||||
):
|
||||
request_throughput = float(run_data["request_throughput"]) # type: ignore
|
||||
if workload_var == "request_rate":
|
||||
return request_throughput
|
||||
if workload_var == "max_concurrency":
|
||||
mean_latency_ms = float(run_data["mean_e2el_ms"]) # type: ignore
|
||||
return request_throughput * mean_latency_ms / 1000
|
||||
|
||||
assert_never(workload_var)
|
||||
|
||||
|
||||
def _estimate_workload_avg(
|
||||
runs: list[dict[str, object]],
|
||||
workload_var: WorkloadVariable,
|
||||
):
|
||||
total = sum(_estimate_workload_value(run, workload_var) for run in runs)
|
||||
return total / len(runs)
|
||||
|
||||
|
||||
def run_comb_workload(
|
||||
server: ServerProcess | None,
|
||||
bench_cmd: list[str],
|
||||
*,
|
||||
serve_comb: ParameterSweepItem,
|
||||
bench_comb: ParameterSweepItem,
|
||||
link_vars: list[tuple[str, str]],
|
||||
experiment_dir: Path,
|
||||
num_runs: int,
|
||||
dry_run: bool,
|
||||
workload_var: WorkloadVariable,
|
||||
workload_value: int,
|
||||
) -> list[dict[str, object]] | None:
|
||||
bench_comb_workload = bench_comb | {workload_var: workload_value}
|
||||
|
||||
return run_comb(
|
||||
server,
|
||||
bench_cmd,
|
||||
serve_comb=serve_comb,
|
||||
bench_comb=bench_comb_workload,
|
||||
link_vars=link_vars,
|
||||
base_path=_get_comb_base_path(
|
||||
experiment_dir,
|
||||
serve_comb,
|
||||
bench_comb,
|
||||
extra_parts=("WL-", f"{workload_var}={workload_value}"),
|
||||
),
|
||||
num_runs=num_runs,
|
||||
dry_run=dry_run,
|
||||
)
|
||||
|
||||
|
||||
def explore_comb_workloads(
|
||||
server: ServerProcess | None,
|
||||
bench_cmd: list[str],
|
||||
*,
|
||||
serve_comb: ParameterSweepItem,
|
||||
bench_comb: ParameterSweepItem,
|
||||
link_vars: list[tuple[str, str]],
|
||||
workload_var: WorkloadVariable,
|
||||
workload_iters: int,
|
||||
experiment_dir: Path,
|
||||
num_runs: int,
|
||||
dry_run: bool,
|
||||
):
|
||||
print("[WL START]")
|
||||
print(f"Serve parameters: {serve_comb.as_text() or '(None)'}")
|
||||
print(f"Bench parameters: {bench_comb.as_text() or '(None)'}")
|
||||
print(f"Number of workload iterations: {workload_iters}")
|
||||
|
||||
if workload_iters < 2:
|
||||
raise ValueError("`workload_iters` should be at least 2")
|
||||
|
||||
dataset_size = DEFAULT_NUM_PROMPTS
|
||||
if "num_prompts" in bench_comb:
|
||||
dataset_size = int(bench_comb["num_prompts"]) # type: ignore
|
||||
else:
|
||||
for i, arg in enumerate(bench_cmd):
|
||||
if arg == "--num-prompts" and i + 1 < len(bench_cmd):
|
||||
dataset_size = int(bench_cmd[i + 1])
|
||||
break
|
||||
elif arg.startswith("--num-prompts="):
|
||||
dataset_size = int(arg.split("=", 1)[1])
|
||||
break
|
||||
|
||||
print(f"Dataset size: {dataset_size}")
|
||||
|
||||
serial_workload_data = run_comb_workload(
|
||||
server,
|
||||
bench_cmd,
|
||||
serve_comb=serve_comb,
|
||||
bench_comb=bench_comb | {"max_concurrency": 1},
|
||||
link_vars=link_vars,
|
||||
experiment_dir=experiment_dir,
|
||||
num_runs=num_runs,
|
||||
dry_run=dry_run,
|
||||
workload_var=workload_var,
|
||||
workload_value=1,
|
||||
)
|
||||
batch_workload_data = run_comb_workload(
|
||||
server,
|
||||
bench_cmd,
|
||||
serve_comb=serve_comb,
|
||||
bench_comb=bench_comb | {"max_concurrency": dataset_size},
|
||||
link_vars=link_vars,
|
||||
experiment_dir=experiment_dir,
|
||||
num_runs=num_runs,
|
||||
dry_run=dry_run,
|
||||
workload_var=workload_var,
|
||||
workload_value=dataset_size,
|
||||
)
|
||||
|
||||
if serial_workload_data is None or batch_workload_data is None:
|
||||
if dry_run:
|
||||
print("Omitting intermediate Workload iterations.")
|
||||
print("[WL END]")
|
||||
|
||||
return
|
||||
|
||||
serial_workload_value = math.ceil(
|
||||
_estimate_workload_avg(serial_workload_data, workload_var)
|
||||
)
|
||||
print(f"Serial inference: {workload_var}={serial_workload_value}")
|
||||
|
||||
batch_workload_value = math.floor(
|
||||
_estimate_workload_avg(batch_workload_data, workload_var)
|
||||
)
|
||||
print(f"Batch inference: {workload_var}={batch_workload_value}")
|
||||
|
||||
# Avoid duplicated runs for intermediate values if the range between
|
||||
# `serial_workload_value` and `batch_workload_value` is small
|
||||
inter_workload_values = np.linspace(
|
||||
serial_workload_value, batch_workload_value, workload_iters
|
||||
)[1:-1]
|
||||
inter_workload_values = sorted(set(map(round, inter_workload_values)))
|
||||
|
||||
inter_workloads_data: list[dict[str, object]] = []
|
||||
for inter_workload_value in inter_workload_values:
|
||||
print(f"Exploring: {workload_var}={inter_workload_value}")
|
||||
inter_workload_data = run_comb_workload(
|
||||
server,
|
||||
bench_cmd,
|
||||
serve_comb=serve_comb,
|
||||
bench_comb=bench_comb,
|
||||
link_vars=link_vars,
|
||||
experiment_dir=experiment_dir,
|
||||
num_runs=num_runs,
|
||||
dry_run=dry_run,
|
||||
workload_var=workload_var,
|
||||
workload_value=inter_workload_value,
|
||||
)
|
||||
if inter_workload_data is not None:
|
||||
inter_workloads_data.extend(inter_workload_data)
|
||||
|
||||
print("[WL END]")
|
||||
|
||||
return serial_workload_data + inter_workloads_data + batch_workload_data
|
||||
|
||||
|
||||
def explore_combs_workloads(
|
||||
serve_cmd: list[str],
|
||||
bench_cmd: list[str],
|
||||
after_bench_cmd: list[str],
|
||||
*,
|
||||
show_stdout: bool,
|
||||
server_ready_timeout: int,
|
||||
serve_params: ParameterSweep,
|
||||
bench_params: ParameterSweep,
|
||||
link_vars: list[tuple[str, str]],
|
||||
workload_var: WorkloadVariable,
|
||||
workload_iters: int,
|
||||
experiment_dir: Path,
|
||||
num_runs: int,
|
||||
dry_run: bool,
|
||||
):
|
||||
if any(bench_comb.has_param(workload_var) for bench_comb in bench_params):
|
||||
raise ValueError(
|
||||
f"You should not override `{workload_var}` in `bench_params` "
|
||||
"since it is supposed to be explored automatically."
|
||||
)
|
||||
|
||||
all_data = list[dict[str, object]]()
|
||||
for serve_comb in serve_params:
|
||||
with server_ctx(
|
||||
serve_cmd,
|
||||
after_bench_cmd,
|
||||
show_stdout=show_stdout,
|
||||
server_ready_timeout=server_ready_timeout,
|
||||
serve_comb=serve_comb,
|
||||
bench_params=bench_params,
|
||||
experiment_dir=experiment_dir,
|
||||
dry_run=dry_run,
|
||||
) as server:
|
||||
for bench_comb in bench_params:
|
||||
comb_data = explore_comb_workloads(
|
||||
server,
|
||||
bench_cmd,
|
||||
serve_comb=serve_comb,
|
||||
bench_comb=bench_comb,
|
||||
link_vars=link_vars,
|
||||
workload_var=workload_var,
|
||||
workload_iters=workload_iters,
|
||||
experiment_dir=experiment_dir,
|
||||
num_runs=num_runs,
|
||||
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(experiment_dir / "summary.csv")
|
||||
|
||||
return combined_df
|
||||
|
||||
|
||||
@dataclass
|
||||
class SweepServeWorkloadArgs(SweepServeArgs):
|
||||
workload_var: WorkloadVariable
|
||||
workload_iters: int
|
||||
|
||||
parser_name: ClassVar[str] = "serve_workload"
|
||||
parser_help: ClassVar[str] = (
|
||||
"Explore the latency-throughput tradeoff for different workload levels."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_cli_args(cls, args: argparse.Namespace):
|
||||
# NOTE: Don't use super() as `from_cli_args` calls `cls()`
|
||||
base_args = SweepServeArgs.from_cli_args(args)
|
||||
|
||||
return cls(
|
||||
**asdict(base_args),
|
||||
workload_var=args.workload_var,
|
||||
workload_iters=args.workload_iters,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
|
||||
parser = super().add_cli_args(parser)
|
||||
|
||||
workload_group = parser.add_argument_group("workload options")
|
||||
workload_group.add_argument(
|
||||
"--workload-var",
|
||||
type=str,
|
||||
choices=get_args(WorkloadVariable),
|
||||
default="request_rate",
|
||||
help="The variable to adjust in each iteration.",
|
||||
)
|
||||
workload_group.add_argument(
|
||||
"--workload-iters",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Number of workload levels to explore. "
|
||||
"This includes the first two iterations used to interpolate the value of "
|
||||
"`workload_var` for remaining iterations.",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def run_main(args: SweepServeWorkloadArgs):
|
||||
experiment_dir = args.resolve_experiment_dir()
|
||||
|
||||
with args.run_ctx(experiment_dir):
|
||||
return explore_combs_workloads(
|
||||
serve_cmd=args.serve_cmd,
|
||||
bench_cmd=args.bench_cmd,
|
||||
after_bench_cmd=args.after_bench_cmd,
|
||||
show_stdout=args.show_stdout,
|
||||
server_ready_timeout=args.server_ready_timeout,
|
||||
serve_params=args.serve_params,
|
||||
bench_params=args.bench_params,
|
||||
link_vars=args.link_vars,
|
||||
workload_var=args.workload_var,
|
||||
workload_iters=args.workload_iters,
|
||||
experiment_dir=experiment_dir,
|
||||
num_runs=args.num_runs,
|
||||
dry_run=args.dry_run,
|
||||
)
|
||||
|
||||
|
||||
def main(args: argparse.Namespace):
|
||||
run_main(SweepServeWorkloadArgs.from_cli_args(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=SweepServeWorkloadArgs.parser_help)
|
||||
SweepServeWorkloadArgs.add_cli_args(parser)
|
||||
|
||||
main(parser.parse_args())
|
||||
Reference in New Issue
Block a user