A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 16:15 UTC
[Frontend] Delegate to vLLM Omni When `--omni` Passed (#40744)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Alex Brooks committed
Apr 24, 2026 at 10:30 UTC
5e11b403657ebd5507e07200c2ba2b8f186d07da
1 file changed
+54
-37
vllm/entrypoints/cli/main.py
+54
-37
@@ -7,6 +7,7 @@ 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 init_logger
13
@@ -34,47 +35,63 @@ def main():
35
36
cli_env_setup()
37
37
- # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default
38
- if len(sys.argv) > 1 and sys.argv[1] == "bench":
39
- logger.debug(
40
- "Bench command detected, must ensure current platform is not "
41
- "UnspecifiedPlatform to avoid device type inference error"
42
- )
43
- from vllm import platforms
38
+ # If `--omni` arg is passed to the CLI, delegate to vLLM Omni's entrypoint handling
39
+ if "--omni" in sys.argv:
40
+ # NOTE: Check the spec instead of importing directly here, since things could
41
+ # fail with ImportError due to mismatched versions if things are moved around.
42
+ spec = find_spec("vllm_omni")
43
+ if spec is None:
44
+ logger.error(
45
+ "--omni flag requires a valid instance of vllm-omni to be installed."
46
+ )
47
+ sys.exit(1)
48
45
- if platforms.current_platform.is_unspecified():
46
- from vllm.platforms.cpu import CpuPlatform
49
+ from vllm_omni.entrypoints.cli.main import main as omni_main
50
48
- platforms.current_platform = CpuPlatform()
49
- logger.info(
50
- "Unspecified platform detected, switching to CPU Platform instead."
51
+ logger.info("Delegating entrypoint handling to vllm-omni")
52
+ omni_main()
53
+ else:
54
+ # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default
55
+ if len(sys.argv) > 1 and sys.argv[1] == "bench":
56
+ logger.debug(
57
+ "Bench command detected, must ensure current platform is not "
58
+ "UnspecifiedPlatform to avoid device type inference error"
59
)
60
+ from vllm import platforms
61
53
- parser = FlexibleArgumentParser(
54
- description="vLLM CLI",
55
- epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
56
- )
57
- parser.add_argument(
58
- "-v",
59
- "--version",
60
- action="version",
61
- version=importlib.metadata.version("vllm"),
62
- )
63
- subparsers = parser.add_subparsers(required=False, dest="subparser")
64
- cmds = {}
65
- for cmd_module in CMD_MODULES:
66
- new_cmds = cmd_module.cmd_init()
67
- for cmd in new_cmds:
68
- cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
69
- cmds[cmd.name] = cmd
70
- args = parser.parse_args()
71
- if args.subparser in cmds:
72
- cmds[args.subparser].validate(args)
73
-
74
- if hasattr(args, "dispatch_function"):
75
- args.dispatch_function(args)
76
- else:
77
- parser.print_help()
62
+ if platforms.current_platform.is_unspecified():
63
+ from vllm.platforms.cpu import CpuPlatform
64
+
65
+ platforms.current_platform = CpuPlatform()
66
+ logger.info(
67
+ "Unspecified platform detected, switching to CPU Platform instead."
68
+ )
69
+
70
+ parser = FlexibleArgumentParser(
71
+ description="vLLM CLI",
72
+ epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
73
+ )
74
+ parser.add_argument(
75
+ "-v",
76
+ "--version",
77
+ action="version",
78
+ version=importlib.metadata.version("vllm"),
79
+ )
80
+ subparsers = parser.add_subparsers(required=False, dest="subparser")
81
+ cmds = {}
82
+ for cmd_module in CMD_MODULES:
83
+ new_cmds = cmd_module.cmd_init()
84
+ for cmd in new_cmds:
85
+ cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
86
+ cmds[cmd.name] = cmd
87
+ args = parser.parse_args()
88
+ if args.subparser in cmds:
89
+ cmds[args.subparser].validate(args)
90
+
91
+ if hasattr(args, "dispatch_function"):
92
+ args.dispatch_function(args)
93
+ else:
94
+ parser.print_help()
95
96
97
if __name__ == "__main__":