86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import argparse
|
|
import importlib
|
|
import logging
|
|
import typing
|
|
|
|
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
|
|
from vllm.entrypoints.cli.types import CLISubcommand
|
|
from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
|
else:
|
|
FlexibleArgumentParser = argparse.ArgumentParser
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _load_benchmark_subcommands() -> None:
|
|
modules = [
|
|
"vllm.entrypoints.cli.benchmark.latency",
|
|
"vllm.entrypoints.cli.benchmark.mm_processor",
|
|
"vllm.entrypoints.cli.benchmark.serve",
|
|
"vllm.entrypoints.cli.benchmark.startup",
|
|
"vllm.entrypoints.cli.benchmark.sweep",
|
|
"vllm.entrypoints.cli.benchmark.throughput",
|
|
]
|
|
|
|
for module_name in modules:
|
|
try:
|
|
importlib.import_module(module_name)
|
|
except ModuleNotFoundError as e:
|
|
logger.warning(
|
|
"Skipping benchmark subcommand module %s because an optional "
|
|
"dependency could not be imported: %r",
|
|
module_name,
|
|
e,
|
|
)
|
|
|
|
|
|
class BenchmarkSubcommand(CLISubcommand):
|
|
"""The `bench` subcommand for the vLLM CLI."""
|
|
|
|
name = "bench"
|
|
help = "vLLM bench subcommand."
|
|
|
|
@staticmethod
|
|
def cmd(args: argparse.Namespace) -> None:
|
|
args.dispatch_function(args)
|
|
|
|
def validate(self, args: argparse.Namespace) -> None:
|
|
pass
|
|
|
|
def subparser_init(
|
|
self, subparsers: argparse._SubParsersAction
|
|
) -> FlexibleArgumentParser:
|
|
bench_parser = subparsers.add_parser(
|
|
self.name,
|
|
help=self.help,
|
|
description=self.help,
|
|
usage=f"vllm {self.name} <bench_type> [options]",
|
|
)
|
|
bench_subparsers = bench_parser.add_subparsers(required=True, dest="bench_type")
|
|
|
|
_load_benchmark_subcommands()
|
|
|
|
for cmd_cls in BenchmarkSubcommandBase.__subclasses__():
|
|
cmd_subparser = bench_subparsers.add_parser(
|
|
cmd_cls.name,
|
|
help=cmd_cls.help,
|
|
description=cmd_cls.help,
|
|
usage=f"vllm {self.name} {cmd_cls.name} [options]",
|
|
)
|
|
cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd)
|
|
cmd_cls.add_cli_args(cmd_subparser)
|
|
cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(
|
|
subcmd=f"{self.name} {cmd_cls.name}"
|
|
)
|
|
return bench_parser
|
|
|
|
|
|
def cmd_init() -> list[CLISubcommand]:
|
|
return [BenchmarkSubcommand()]
|