[Misc] Simplify vllm bench cli subcommand implementation (#19948)

This commit is contained in:
Ye (Charlotte) Qi 2025-06-22 09:34:48 -07:00 committed by GitHub
parent c76a506bd6
commit 2c11a29f0b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 44 additions and 80 deletions

View File

@ -0,0 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.entrypoints.cli.benchmark.latency import BenchmarkLatencySubcommand
from vllm.entrypoints.cli.benchmark.serve import BenchmarkServingSubcommand
from vllm.entrypoints.cli.benchmark.throughput import (
BenchmarkThroughputSubcommand)
__all__: list[str] = [
"BenchmarkLatencySubcommand",
"BenchmarkServingSubcommand",
"BenchmarkThroughputSubcommand",
]

View File

@ -3,18 +3,15 @@
import argparse import argparse
from vllm.entrypoints.cli.types import CLISubcommand from vllm.entrypoints.cli.types import CLISubcommand
from vllm.utils import FlexibleArgumentParser
class BenchmarkSubcommandBase(CLISubcommand): class BenchmarkSubcommandBase(CLISubcommand):
""" The base class of subcommands for vllm bench. """ """ The base class of subcommands for vllm bench. """
@property help: str
def help(self) -> str:
"""The help message of the subcommand."""
raise NotImplementedError
def add_cli_args(self, parser: argparse.ArgumentParser) -> None: @classmethod
def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
"""Add the CLI arguments to the parser.""" """Add the CLI arguments to the parser."""
raise NotImplementedError raise NotImplementedError
@ -26,14 +23,3 @@ class BenchmarkSubcommandBase(CLISubcommand):
args: The arguments to the command. args: The arguments to the command.
""" """
raise NotImplementedError raise NotImplementedError
def subparser_init(
self,
subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
parser = subparsers.add_parser(
self.name,
help=self.help,
description=self.help,
usage=f"vllm bench {self.name} [options]")
self.add_cli_args(parser)
return parser

View File

@ -4,27 +4,18 @@ import argparse
from vllm.benchmarks.latency import add_cli_args, main from vllm.benchmarks.latency import add_cli_args, main
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
from vllm.entrypoints.cli.types import CLISubcommand
class BenchmarkLatencySubcommand(BenchmarkSubcommandBase): class BenchmarkLatencySubcommand(BenchmarkSubcommandBase):
""" The `latency` subcommand for vllm bench. """ """ The `latency` subcommand for vllm bench. """
def __init__(self): name = "latency"
self.name = "latency" help = "Benchmark the latency of a single batch of requests."
super().__init__()
@property @classmethod
def help(self) -> str: def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
return "Benchmark the latency of a single batch of requests."
def add_cli_args(self, parser: argparse.ArgumentParser) -> None:
add_cli_args(parser) add_cli_args(parser)
@staticmethod @staticmethod
def cmd(args: argparse.Namespace) -> None: def cmd(args: argparse.Namespace) -> None:
main(args) main(args)
def cmd_init() -> list[CLISubcommand]:
return [BenchmarkLatencySubcommand()]

View File

@ -2,51 +2,44 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse import argparse
import vllm.entrypoints.cli.benchmark.latency from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
import vllm.entrypoints.cli.benchmark.serve
import vllm.entrypoints.cli.benchmark.throughput
from vllm.entrypoints.cli.types import CLISubcommand from vllm.entrypoints.cli.types import CLISubcommand
from vllm.utils import FlexibleArgumentParser from vllm.utils import FlexibleArgumentParser
BENCHMARK_CMD_MODULES = [
vllm.entrypoints.cli.benchmark.latency,
vllm.entrypoints.cli.benchmark.serve,
vllm.entrypoints.cli.benchmark.throughput,
]
class BenchmarkSubcommand(CLISubcommand): class BenchmarkSubcommand(CLISubcommand):
""" The `bench` subcommand for the vLLM CLI. """ """ The `bench` subcommand for the vLLM CLI. """
def __init__(self): name = "bench"
self.name = "bench" help = "vLLM bench subcommand."
super().__init__()
@staticmethod @staticmethod
def cmd(args: argparse.Namespace) -> None: def cmd(args: argparse.Namespace) -> None:
args.dispatch_function(args) args.dispatch_function(args)
def validate(self, args: argparse.Namespace) -> None: def validate(self, args: argparse.Namespace) -> None:
if args.bench_type in self.cmds: pass
self.cmds[args.bench_type].validate(args)
def subparser_init( def subparser_init(
self, self,
subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
bench_parser = subparsers.add_parser( bench_parser = subparsers.add_parser(
"bench", self.name,
help="vLLM bench subcommand.", help=self.help,
description="vLLM bench subcommand.", description=self.help,
usage="vllm bench <bench_type> [options]") usage="vllm bench <bench_type> [options]")
bench_subparsers = bench_parser.add_subparsers(required=True, bench_subparsers = bench_parser.add_subparsers(required=True,
dest="bench_type") dest="bench_type")
self.cmds = {}
for cmd_module in BENCHMARK_CMD_MODULES: for cmd_cls in BenchmarkSubcommandBase.__subclasses__():
new_cmds = cmd_module.cmd_init() cmd_subparser = bench_subparsers.add_parser(
for cmd in new_cmds: cmd_cls.name,
cmd.subparser_init(bench_subparsers).set_defaults( help=cmd_cls.help,
dispatch_function=cmd.cmd) description=cmd_cls.help,
self.cmds[cmd.name] = cmd )
cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd)
cmd_cls.add_cli_args(cmd_subparser)
return bench_parser return bench_parser

View File

@ -4,27 +4,18 @@ import argparse
from vllm.benchmarks.serve import add_cli_args, main from vllm.benchmarks.serve import add_cli_args, main
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
from vllm.entrypoints.cli.types import CLISubcommand
class BenchmarkServingSubcommand(BenchmarkSubcommandBase): class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
""" The `serve` subcommand for vllm bench. """ """ The `serve` subcommand for vllm bench. """
def __init__(self): name = "serve"
self.name = "serve" help = "Benchmark the online serving throughput."
super().__init__()
@property @classmethod
def help(self) -> str: def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
return "Benchmark the online serving throughput."
def add_cli_args(self, parser: argparse.ArgumentParser) -> None:
add_cli_args(parser) add_cli_args(parser)
@staticmethod @staticmethod
def cmd(args: argparse.Namespace) -> None: def cmd(args: argparse.Namespace) -> None:
main(args) main(args)
def cmd_init() -> list[CLISubcommand]:
return [BenchmarkServingSubcommand()]

View File

@ -4,27 +4,18 @@ import argparse
from vllm.benchmarks.throughput import add_cli_args, main from vllm.benchmarks.throughput import add_cli_args, main
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
from vllm.entrypoints.cli.types import CLISubcommand
class BenchmarkThroughputSubcommand(BenchmarkSubcommandBase): class BenchmarkThroughputSubcommand(BenchmarkSubcommandBase):
""" The `throughput` subcommand for vllm bench. """ """ The `throughput` subcommand for vllm bench. """
def __init__(self): name = "throughput"
self.name = "throughput" help = "Benchmark offline inference throughput."
super().__init__()
@property @classmethod
def help(self) -> str: def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
return "Benchmark offline inference throughput."
def add_cli_args(self, parser: argparse.ArgumentParser) -> None:
add_cli_args(parser) add_cli_args(parser)
@staticmethod @staticmethod
def cmd(args: argparse.Namespace) -> None: def cmd(args: argparse.Namespace) -> None:
main(args) main(args)
def cmd_init() -> list[CLISubcommand]:
return [BenchmarkThroughputSubcommand()]