wehub-resource-sync
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import argparse
|
|
|
|
from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
|
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
|
|
|
from .plot import SweepPlotArgs
|
|
from .plot import main as plot_main
|
|
from .plot_pareto import SweepPlotParetoArgs
|
|
from .plot_pareto import main as plot_pareto_main
|
|
from .serve import SweepServeArgs
|
|
from .serve import main as serve_main
|
|
from .serve_workload import SweepServeWorkloadArgs
|
|
from .serve_workload import main as serve_workload_main
|
|
from .startup import SweepStartupArgs
|
|
from .startup import main as startup_main
|
|
|
|
SUBCOMMANDS = (
|
|
(SweepServeArgs, serve_main),
|
|
(SweepServeWorkloadArgs, serve_workload_main),
|
|
(SweepStartupArgs, startup_main),
|
|
(SweepPlotArgs, plot_main),
|
|
(SweepPlotParetoArgs, plot_pareto_main),
|
|
)
|
|
|
|
|
|
def add_cli_args(parser: FlexibleArgumentParser):
|
|
subparsers = parser.add_subparsers(required=True, dest="sweep_type")
|
|
|
|
for cmd, entrypoint in SUBCOMMANDS:
|
|
cmd_subparser = subparsers.add_parser(
|
|
cmd.parser_name,
|
|
description=cmd.parser_help,
|
|
usage=f"vllm bench sweep {cmd.parser_name} [options]",
|
|
)
|
|
cmd_subparser.set_defaults(dispatch_function=entrypoint)
|
|
cmd.add_cli_args(cmd_subparser)
|
|
cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(
|
|
subcmd=f"sweep {cmd.parser_name}"
|
|
)
|
|
|
|
|
|
def main(args: argparse.Namespace):
|
|
args.dispatch_function(args)
|