[Core] Integrate TrackingArgumentParser (#3369)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

vllm_omni/entrypoints/cli/benchmark/main.py

@@ -9,7 +9,7 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase
 
 if typing.TYPE_CHECKING:
-    from vllm.utils.argparse_utils import FlexibleArgumentParser
+    from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 
 class OmniBenchmarkSubcommand(CLISubcommand):
@@ -25,7 +25,7 @@ class OmniBenchmarkSubcommand(CLISubcommand):
     def validate(self, args: argparse.Namespace) -> None:
         pass
 
-    def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+    def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgumentParser:
         bench_parser = subparsers.add_parser(
             self.name, description=self.help, usage=f"vllm {self.name} <bench_type> [options]"
         )

vllm_omni/entrypoints/cli/main.py

@@ -24,10 +24,10 @@ def main():
             os.environ["VLLM_LOGGING_COLOR"] = "1"
 
         from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
-        from vllm.utils.argparse_utils import FlexibleArgumentParser
 
         import vllm_omni.entrypoints.cli.benchmark.main
         import vllm_omni.entrypoints.cli.serve
+        from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
         CMD_MODULES = [
             vllm_omni.entrypoints.cli.serve,
@@ -40,7 +40,7 @@ def main():
 
         _ensure_vllm_platform()
 
-        parser = FlexibleArgumentParser(
+        parser = TrackingArgumentParser(
             description="vLLM OMNI CLI",
             epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
         )

vllm_omni/entrypoints/cli/serve.py

@@ -19,11 +19,10 @@ from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.cli.logo import log_logo
 from vllm_omni.entrypoints.openai.api_server import omni_run_server
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 logger = init_logger(__name__)
 
@@ -84,10 +83,10 @@ class OmniServeCommand(CLISubcommand):
     name = "serve"
     # Parser stashed at subparser_init so ``cmd`` can resolve each user-typed
     # flag to its real ``dest`` via the parser's action table.
-    _parser: FlexibleArgumentParser | None = None
+    _parser: TrackingArgumentParser
 
     @staticmethod
-    def cmd(args: argparse.Namespace) -> None:
+    def cmd(args: TrackingNamespace) -> None:
         if not os.environ.get("VLLM_DISABLE_LOG_LOGO"):
             os.environ["VLLM_DISABLE_LOG_LOGO"] = "1"
             log_logo()
@@ -176,7 +175,7 @@ class OmniServeCommand(CLISubcommand):
             return
         validate_parsed_serve_args(args)
 
-    def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+    def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgumentParser:
         serve_parser = subparsers.add_parser(
             self.name,
             description=DESCRIPTION,
@@ -682,22 +681,10 @@ class OmniServeCommand(CLISubcommand):
         # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError.
         type(self)._parser = serve_parser
 
-        nullify_stage_engine_defaults(serve_parser)
         return serve_parser
 
 
-def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[str, Any]]:
-    """Create default diffusion stage configuration.
-
-    Uses AsyncOmniEngine's implementation which doesn't have OmegaConf
-    compatibility issues.
-    """
-    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-
-    return AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))
-
-
-def run_headless(args: argparse.Namespace) -> None:
+def run_headless(args: TrackingNamespace) -> None:
     """Run a single stage in headless mode.
 
     Honors ``--omni-dp-size-local``: launches that many replicas locally for
@@ -738,43 +725,39 @@ def run_headless(args: argparse.Namespace) -> None:
     stage_id: int | None = args.stage_id
     omni_master_address: str | None = args.omni_master_address
     omni_master_port: int | None = args.omni_master_port
+    worker_backend: str | None = args.worker_backend
+    stage_configs_path: str | None = args.stage_configs_path
     omni_replica_address: str | None = getattr(args, "omni_replica_address", None)
     omni_dp_size_local: int = max(1, int(getattr(args, "omni_dp_size_local", 1) or 1))
 
+    if not model:
+        raise ValueError("Failed to pass model from kwargs")
     if stage_id is None:
         raise ValueError("--stage-id is required in headless mode")
+    if omni_master_address is None or omni_master_port is None:
+        raise ValueError("--omni-master-address and --omni-master-port are required in headless mode")
+    if worker_backend != "multi_process":
+        raise ValueError("headless mode requires worker_backend=multi_process")
+
+    # Filter down to a dict of things explicitly requested by the user
+    args_dict = args.get_explicit_kwargs_dict()
 
     # ``--replica-id`` is deprecated and ignored — replica ids are
     # auto-assigned by ``OmniMasterServer`` so headless processes carry
     # no knowledge of their per-replica id at launch time. Warn (don't
     # error) when the operator still supplies it so existing launchers
     # keep working with a single log line.
-    explicit_cli_keys: set[str] = getattr(args, "_cli_explicit_keys", set()) or set()
-    if "replica_id" in explicit_cli_keys:
+    if "replica_id" in args_dict:
         logger.warning(
             "[Headless] --replica-id is deprecated and ignored "
             "(supplied value: %s). Replica ids are auto-assigned by the "
             "master server.",
             args.replica_id,
         )
-    if omni_master_address is None or omni_master_port is None:
-        raise ValueError("--omni-master-address and --omni-master-port are required in headless mode")
-    api_server_count = args.api_server_count or 0
-    if api_server_count > 1:
-        raise ValueError("api_server_count can't be set in headless mode")
-    if args.worker_backend != "multi_process":
-        raise ValueError("headless mode requires worker_backend=multi_process")
 
-    args_dict = vars(args).copy()
-    args_dict.pop("_cli_explicit_keys", None)
-    # Forward ``--deploy-config`` so the headless reads the same YAML the
-    # head was launched with — otherwise ``load_and_resolve_stage_configs``
-    # falls back to ``vllm_omni/deploy/<model>.yaml`` and the headless's
-    # view of ``stage.runtime.devices`` diverges from the head's, breaking
-    # the per-replica device split.
     config_path, stage_configs = load_and_resolve_stage_configs(
         model,
-        args_dict.get("stage_configs_path"),
+        stage_configs_path,
         args_dict,
         deploy_config_path=args_dict.get("deploy_config"),
     )