[Refactor] Extract maybe_run_omni() and defer CLI imports past omni early-return (#56103)
Signed-off-by: Jeff Zvier <zvier20@gmail.com>
zvier committed
Sep 10, 2026 at 11:21 UTC
8c34a372323b3becc92c183257d195562882ef1a
1 file changed
+63
-54
vllm/entrypoints/cli/main.py
+63
-54
@@ -14,7 +14,32 @@ from vllm.logger import init_logger
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
@@ -38,65 +63,49 @@ def main():
63
64
cli_env_setup()
65
41
- # If `--omni` arg is passed to the CLI, delegate to vLLM Omni's entrypoint handling
42
- if "--omni" in sys.argv:
43
- # NOTE: Check the spec instead of importing directly here, since things could
44
- # fail with ImportError due to mismatched versions if things are moved around.
45
- spec = find_spec("vllm_omni")
46
- if spec is None:
47
- logger.error(
48
- "--omni flag requires a valid instance of vllm-omni to be installed."
49
- )
50
- sys.exit(1)
66
+ vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench()
67
52
- from vllm_omni.entrypoints.cli.main import main as omni_main
68
+ # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default
69
+ if len(sys.argv) > 1 and sys.argv[1] == "bench":
70
+ logger.debug(
71
+ "Bench command detected, must ensure current platform is not "
72
+ "UnspecifiedPlatform to avoid device type inference error"
73
+ )
74
+ from vllm import platforms
75
54
- logger.info("Delegating entrypoint handling to vllm-omni")
55
- omni_main()
56
- else:
57
- vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench()
76
+ if platforms.current_platform.is_unspecified():
77
+ from vllm.platforms.cpu import CpuPlatform
78
59
- # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default
60
- if len(sys.argv) > 1 and sys.argv[1] == "bench":
61
- logger.debug(
62
- "Bench command detected, must ensure current platform is not "
63
- "UnspecifiedPlatform to avoid device type inference error"
79
+ platforms.current_platform = CpuPlatform()
80
+ logger.info(
81
+ "Unspecified platform detected, switching to CPU Platform instead."
82
)
65
- from vllm import platforms
66
-
67
- if platforms.current_platform.is_unspecified():
68
- from vllm.platforms.cpu import CpuPlatform
83
70
- platforms.current_platform = CpuPlatform()
71
- logger.info(
72
- "Unspecified platform detected, switching to CPU Platform instead."
73
- )
74
-
75
- parser = FlexibleArgumentParser(
76
- description="vLLM CLI",
77
- epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
78
- )
79
- parser.add_argument(
80
- "-v",
81
- "--version",
82
- action="version",
83
- version=importlib.metadata.version("vllm"),
84
- )
85
- subparsers = parser.add_subparsers(required=False, dest="subparser")
86
- cmds = {}
87
- for cmd_module in CMD_MODULES:
88
- new_cmds = cmd_module.cmd_init()
89
- for cmd in new_cmds:
90
- cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
91
- cmds[cmd.name] = cmd
92
- args = parser.parse_args()
93
- if args.subparser in cmds:
94
- cmds[args.subparser].validate(args)
95
-
96
- if hasattr(args, "dispatch_function"):
97
- args.dispatch_function(args)
98
- else:
99
- parser.print_help()
84
+ parser = FlexibleArgumentParser(
85
+ description="vLLM CLI",
86
+ epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
87
+ )
88
+ parser.add_argument(
89
+ "-v",
90
+ "--version",
91
+ action="version",
92
+ version=importlib.metadata.version("vllm"),
93
+ )
94
+ subparsers = parser.add_subparsers(required=False, dest="subparser")
95
+ cmds = {}
96
+ for cmd_module in CMD_MODULES:
97
+ new_cmds = cmd_module.cmd_init()
98
+ for cmd in new_cmds:
99
+ cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
100
+ cmds[cmd.name] = cmd
101
+ args = parser.parse_args()
102
+ if args.subparser in cmds:
103
+ cmds[args.subparser].validate(args)
104
+
105
+ if hasattr(args, "dispatch_function"):
106
+ args.dispatch_function(args)
107
+ else:
108
+ parser.print_help()
109
110
111
if __name__ == "__main__":