@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 124 lines 3.82 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """The CLI entrypoints of vLLM.
4
5 Note that all future modules must be lazily loaded within main
6 to avoid certain eager import breakage."""
7
8 import importlib.metadata
9 import sys
10 from importlib.util import find_spec
11
12 from vllm.logger import configure_logging_from_args, init_logger
13
14 logger = init_logger(__name__)
15
16
17 def maybe_run_omni() -> bool:
18 # If `--omni` arg is passed to the CLI, delegate to vLLM Omni's entrypoint handling
19 if "--omni" not in sys.argv:
20 return False
21
22 # NOTE: Check the spec instead of importing directly here, since things could
23 # fail with ImportError due to mismatched versions if things are moved around.
24 spec = find_spec("vllm_omni")
25 if spec is None:
26 logger.error(
27 "--omni flag requires a valid instance of vllm-omni to be installed."
28 )
29 sys.exit(1)
30
31 from vllm_omni.entrypoints.cli.main import main as omni_main
32
33 logger.info("Delegating entrypoint handling to vllm-omni")
34 omni_main()
35
36 return True
37
38
39 def main():
40 if maybe_run_omni():
41 return
42
43 import vllm.entrypoints.cli.benchmark.main
44 import vllm.entrypoints.cli.collect_env
45 import vllm.entrypoints.cli.launch
46 import vllm.entrypoints.cli.openai
47 import vllm.entrypoints.cli.preload
48 import vllm.entrypoints.cli.run_batch
49 import vllm.entrypoints.cli.serve
50 import vllm.entrypoints.cli.snapshot
51 from vllm.entrypoints.serve.utils.api_utils import (
52 VLLM_SUBCMD_PARSER_EPILOG,
53 cli_env_setup,
54 )
55 from vllm.utils.argparse_utils import FlexibleArgumentParser
56
57 CMD_MODULES = [
58 vllm.entrypoints.cli.openai,
59 vllm.entrypoints.cli.serve,
60 vllm.entrypoints.cli.launch,
61 vllm.entrypoints.cli.benchmark.main,
62 vllm.entrypoints.cli.collect_env,
63 vllm.entrypoints.cli.preload,
64 vllm.entrypoints.cli.run_batch,
65 vllm.entrypoints.cli.snapshot,
66 ]
67
68 if sys.argv[1:2] != ["snapshot"]:
69 cli_env_setup()
70
71 vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench()
72
73 # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default
74 if len(sys.argv) > 1 and sys.argv[1] == "bench":
75 logger.debug(
76 "Bench command detected, must ensure current platform is not "
77 "UnspecifiedPlatform to avoid device type inference error"
78 )
79 from vllm import platforms
80
81 if platforms.current_platform.is_unspecified():
82 from vllm.platforms.cpu import CpuPlatform
83
84 platforms.current_platform = CpuPlatform()
85 logger.info(
86 "Unspecified platform detected, switching to CPU Platform instead."
87 )
88
89 parser = FlexibleArgumentParser(
90 description="vLLM CLI",
91 epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
92 )
93 parser.add_argument(
94 "-v",
95 "--version",
96 action="version",
97 version=importlib.metadata.version("vllm"),
98 )
99 subparsers = parser.add_subparsers(required=False, dest="subparser")
100 cmds = {}
101 for cmd_module in CMD_MODULES:
102 if cmd_module is vllm.entrypoints.cli.snapshot:
103 new_cmds = cmd_module.cmd_init(
104 create_requested=sys.argv[1:3] == ["snapshot", "create"]
105 )
106 else:
107 new_cmds = cmd_module.cmd_init()
108 for cmd in new_cmds:
109 cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
110 cmds[cmd.name] = cmd
111 args = parser.parse_args()
112 if args.subparser in cmds:
113 cmd = cmds[args.subparser]
114 configure_logging_from_args(args)
115 cmd.validate(args)
116
117 if hasattr(args, "dispatch_function"):
118 args.dispatch_function(args)
119 else:
120 parser.print_help()
121
122
123 if __name__ == "__main__":
124 main()