| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | |
| 5 | from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG |
| 6 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 7 | |
| 8 | from .plot import SweepPlotArgs |
| 9 | from .plot import main as plot_main |
| 10 | from .plot_pareto import SweepPlotParetoArgs |
| 11 | from .plot_pareto import main as plot_pareto_main |
| 12 | from .serve import SweepServeArgs |
| 13 | from .serve import main as serve_main |
| 14 | from .serve_workload import SweepServeWorkloadArgs |
| 15 | from .serve_workload import main as serve_workload_main |
| 16 | from .startup import SweepStartupArgs |
| 17 | from .startup import main as startup_main |
| 18 | |
| 19 | SUBCOMMANDS = ( |
| 20 | (SweepServeArgs, serve_main), |
| 21 | (SweepServeWorkloadArgs, serve_workload_main), |
| 22 | (SweepStartupArgs, startup_main), |
| 23 | (SweepPlotArgs, plot_main), |
| 24 | (SweepPlotParetoArgs, plot_pareto_main), |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | def add_cli_args(parser: FlexibleArgumentParser): |
| 29 | subparsers = parser.add_subparsers(required=True, dest="sweep_type") |
| 30 | |
| 31 | for cmd, entrypoint in SUBCOMMANDS: |
| 32 | cmd_subparser = subparsers.add_parser( |
| 33 | cmd.parser_name, |
| 34 | description=cmd.parser_help, |
| 35 | usage=f"vllm bench sweep {cmd.parser_name} [options]", |
| 36 | ) |
| 37 | cmd_subparser.set_defaults(dispatch_function=entrypoint) |
| 38 | cmd.add_cli_args(cmd_subparser) |
| 39 | cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format( |
| 40 | subcmd=f"sweep {cmd.parser_name}" |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | def main(args: argparse.Namespace): |
| 45 | args.dispatch_function(args) |