Move vllm serve code into vllm.cmd.serve

Signed-off-by: Russell Bryant <[email protected]>

Author: russellb

vllm/cmd/main.py

@@ -7,15 +7,11 @@
 import sys
 from typing import List, Optional
 
-import uvloop
 from openai import OpenAI
 from openai.types.chat import ChatCompletionMessageParam
 
+import vllm.cmd.serve
 import vllm.version
-from vllm.engine.arg_utils import EngineArgs
-from vllm.entrypoints.openai.api_server import run_server
-from vllm.entrypoints.openai.cli_args import (make_arg_parser,
-                                              validate_parsed_serve_args)
 from vllm.logger import init_logger
 from vllm.utils import FlexibleArgumentParser
 
@@ -31,19 +27,6 @@ def signal_handler(sig, frame):
     signal.signal(signal.SIGTSTP, signal_handler)
 
 
-def serve(args: argparse.Namespace) -> None:
-    # The default value of `--model`
-    if args.model != EngineArgs.model:
-        raise ValueError(
-            "With `vllm serve`, you should provide the model as a "
-            "positional argument instead of via the `--model` option.")
-
-    # EngineArgs expects the model name to be passed as --model.
-    args.model = args.model_tag
-
-    uvloop.run(run_server(args))
-
-
 def interactive_cli(args: argparse.Namespace) -> None:
     register_signal_handlers()
 
@@ -150,28 +133,16 @@ def main():
                         '--version',
                         action='version',
                         version=vllm.version.__version__)
-
     subparsers = parser.add_subparsers(required=True, dest="subparser")
 
-    serve_parser = subparsers.add_parser(
-        "serve",
-        help="Start the vLLM OpenAI Compatible API server",
-        usage="vllm serve <model_tag> [options]")
-    serve_parser.add_argument("model_tag",
-                              type=str,
-                              help="The model tag to serve")
-    serve_parser.add_argument(
-        "--config",
-        type=str,
-        default='',
-        required=False,
-        help="Read CLI options from a config file."
-        "Must be a YAML with the following options:"
-        "https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#cli-reference"
-    )
-
-    serve_parser = make_arg_parser(serve_parser)
-    serve_parser.set_defaults(dispatch_function=serve)
+    cmd_modules = [
+        vllm.cmd.serve,
+    ]
+    cmds = {}
+    for cmd_module in cmd_modules:
+        cmd = cmd_module.cmd_init()
+        cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
+        cmds[cmd.name] = cmd
 
     complete_parser = subparsers.add_parser(
         "complete",
@@ -196,8 +167,8 @@ def main():
     chat_parser.set_defaults(dispatch_function=interactive_cli, command="chat")
 
     args = parser.parse_args()
-    if args.subparser == "serve":
-        validate_parsed_serve_args(args)
+    if args.subparser in cmds:
+        cmds[args.subparser].validate(args)
 
     # One of the sub commands should be executed.
     if hasattr(args, "dispatch_function"):

vllm/cmd/serve.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+
+import uvloop
+
+from vllm.cmd.types import CLISubcommand
+from vllm.engine.arg_utils import EngineArgs
+from vllm.entrypoints.openai.api_server import run_server
+from vllm.entrypoints.openai.cli_args import (make_arg_parser,
+                                              validate_parsed_serve_args)
+from vllm.utils import FlexibleArgumentParser
+
+
+class ServeSubcommand(CLISubcommand):
+    """The `serve` subcommand for the vLLM CLI. """
+
+    def __init__(self):
+        self.name = "serve"
+
+    @staticmethod
+    def cmd(args: argparse.Namespace) -> None:
+        # The default value of `--model`
+        if args.model != EngineArgs.model:
+            raise ValueError(
+                "With `vllm serve`, you should provide the model as a "
+                "positional argument instead of via the `--model` option.")
+
+        # EngineArgs expects the model name to be passed as --model.
+        args.model = args.model_tag
+
+        uvloop.run(run_server(args))
+
+    def validate(self, args: argparse.Namespace) -> None:
+        validate_parsed_serve_args(args)
+
+    def subparser_init(
+            self,
+            subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+        serve_parser = subparsers.add_parser(
+            "serve",
+            help="Start the vLLM OpenAI Compatible API server",
+            usage="vllm serve <model_tag> [options]")
+        serve_parser.add_argument("model_tag",
+                                  type=str,
+                                  help="The model tag to serve")
+        serve_parser.add_argument(
+            "--config",
+            type=str,
+            default='',
+            required=False,
+            help="Read CLI options from a config file."
+            "Must be a YAML with the following options:"
+            "https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#cli-reference"
+        )
+
+        return make_arg_parser(serve_parser)
+
+
+def cmd_init() -> CLISubcommand:
+    return ServeSubcommand()

vllm/cmd/types.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+
+from vllm.utils import FlexibleArgumentParser
+
+
+class CLISubcommand:
+    """Base class for CLI argument handlers."""
+
+    name: str
+
+    @staticmethod
+    def cmd(args: argparse.Namespace) -> None:
+        raise NotImplementedError("Subclasses should implement this method")
+
+    def validate(self, args: argparse.Namespace) -> None:
+        # No validation by deafult
+        pass
+
+    def subparser_init(
+            self,
+            subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+        raise NotImplementedError("Subclasses should implement this method")