Unverified Commit 2c11a29f authored by Ye (Charlotte) Qi's avatar Ye (Charlotte) Qi Committed by GitHub
Browse files

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

parent c76a506b
# 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",
]
\ No newline at end of file
...@@ -3,18 +3,15 @@ ...@@ -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): ...@@ -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
...@@ -4,27 +4,18 @@ import argparse ...@@ -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()]
...@@ -2,51 +2,44 @@ ...@@ -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
......
...@@ -4,27 +4,18 @@ import argparse ...@@ -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()]
...@@ -4,27 +4,18 @@ import argparse ...@@ -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()]
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment