[Bugfix] Require a model for vllm serve --omni (fixes #4158) (#4167)

Signed-off-by: abinggo <107740309+abinggo@users.noreply.github.com>

tests/entrypoints/test_serve.py

@@ -44,6 +44,80 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
     assert not explicit["async_chunk"]
 
 
+def _parse_serve_args(argv: list[str]) -> TrackingNamespace:
+    """Parse a ``serve`` argv through the real Omni parser, returning the
+    TrackingNamespace (with ``explicit_keys``) that ``validate`` receives."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+    return parser.parse_args(argv)
+
+
+def test_omni_serve_requires_model_when_none_provided() -> None:
+    """Regression for https://github.com/vllm-project/vllm-omni/issues/4158:
+    ``vllm serve --omni`` with no model must fail fast instead of silently
+    falling back to vLLM's default model and crashing in the diffusion worker.
+    """
+    args = _parse_serve_args(["serve", "--omni"])
+    cmd = OmniServeCommand()
+    with pytest.raises(ValueError, match="requires an explicit model"):
+        cmd.validate(args)
+
+
+@pytest.mark.parametrize(
+    "argv",
+    [
+        # A deploy YAML supplied *without* an explicit model must still trip
+        # the guard. It carries per-stage engine args only -- the
+        # checkpoint is always threaded in from ``args.model`` -- so its mere
+        # presence does not establish that a model was provided. Without an
+        # explicit model ``args.model`` stays at vLLM's default and the launch
+        # reproduces the issue-4158 diffusion-registry crash.
+        ["serve", "--omni", "--deploy-config", "deploy.yaml"],
+        # An empty/whitespace model must not satisfy the guard either -- this is
+        # the ``--model "$MODEL"``-with-unset-var footgun. Otherwise the empty
+        # value flows downstream and crashes with the same confusing error.
+        ["serve", "", "--omni"],  # empty positional
+        ["serve", "--omni", "--model", ""],  # empty --model
+        ["serve", "--omni", "--model", "   "],  # whitespace-only --model
+    ],
+)
+def test_omni_serve_rejects_missing_or_empty_model(argv: list[str]) -> None:
+    """Regression for review feedback on PR #4167: a deploy YAML is not a
+    model source, so ``--deploy-config`` alone must not bypass the
+    require-a-model guard; neither may an empty/whitespace model value."""
+    args = _parse_serve_args(argv)
+    cmd = OmniServeCommand()
+    with pytest.raises(ValueError, match="requires an explicit model"):
+        cmd.validate(args)
+
+
+@pytest.mark.parametrize(
+    "argv",
+    [
+        ["serve", "fake-model", "--omni"],  # positional model
+        ["serve", "--omni", "--model", "fake-model"],  # --model flag
+        # An explicit model with a deploy YAML is the legitimate multi-stage
+        # path and must pass the guard.
+        ["serve", "fake-model", "--omni", "--deploy-config", "deploy.yaml"],
+    ],
+)
+def test_omni_serve_accepts_explicit_model(argv: list[str], mocker: MockerFixture) -> None:
+    """A model supplied positionally or via ``--model`` -- with or without a
+    deploy YAML -- must NOT trip the require-a-model guard. Downstream
+    validation is mocked so only the new guard is exercised."""
+    mocker.patch(
+        "vllm_omni.diffusion.utils.hf_utils.is_diffusion_model",
+        return_value=False,
+    )
+    mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
+
+    args = _parse_serve_args(argv)
+    cmd = OmniServeCommand()
+    # Must not raise the "requires a model" error for any of these.
+    cmd.validate(args)
+
+
 def test_serve_parser_accepts_strategy_config() -> None:
     """``--strategy-config`` must parse onto the ``strategy_config`` dest and be
     forwarded as an explicit kwarg so the engine can overlay the strategy."""

vllm_omni/entrypoints/cli/serve.py

@@ -148,6 +148,50 @@ class OmniServeCommand(CLISubcommand):
         if args.stage_id is not None and (args.omni_master_address is None or args.omni_master_port is None):
             raise ValueError("--stage-id requires both --omni-master-address and --omni-master-port to be set")
 
+        # Require an explicit model under --omni. ``args.model`` always carries
+        # vLLM's ModelConfig default (``Qwen/Qwen3-0.6B``), so an omit is silent:
+        # the default text LLM is routed into the diffusion stage and startup
+        # crashes deep in the diffusion worker with a confusing registry error.
+        # Fail fast instead.
+        #
+        # The model must come from the CLI -- positionally
+        # (``vllm serve <model> --omni``) or via ``--model``. A deploy YAML
+        # (``--deploy-config``) is NOT a model source: it carries per-stage
+        # engine args only, and the checkpoint is always threaded in from
+        # ``args.model`` (see
+        # ``load_and_resolve_stage_configs``, which takes ``model`` as its first
+        # argument). Treating its mere presence as "model provided" let
+        # ``vllm serve --omni --deploy-config <cfg>`` slip through with the
+        # default model and reproduce the very crash this guard prevents. See
+        # https://github.com/vllm-project/vllm-omni/issues/4158.
+        if getattr(args, "omni", False):
+            explicit_keys = getattr(args, "explicit_keys", None) or frozenset()
+            # Resolve the model the user actually supplied on the CLI. A
+            # positional ``model_tag`` takes precedence (``cmd`` later copies it
+            # onto ``args.model``); otherwise ``--model`` counts only when it was
+            # explicitly passed. An empty/whitespace value (e.g. ``--model
+            # "$MODEL"`` with ``MODEL`` unset) is treated as "not provided" so it
+            # fails here with a clear message instead of the same confusing
+            # downstream crash.
+            model_tag = getattr(args, "model_tag", None)
+            if model_tag is not None:
+                explicit_model = model_tag
+            elif "model" in explicit_keys:
+                explicit_model = getattr(args, "model", None)
+            else:
+                explicit_model = None
+            model_provided = explicit_model is not None and str(explicit_model).strip() != ""
+            if not model_provided:
+                raise ValueError(
+                    "`vllm serve --omni` requires an explicit model. Pass it "
+                    "positionally (`vllm serve <model> --omni`) or via `--model`. "
+                    "`--deploy-config` carries per-stage engine args only and "
+                    "does not supply a model; without an "
+                    "explicit model, vLLM's default (Qwen/Qwen3-0.6B) is selected "
+                    "and routed into the diffusion stage, which fails with a "
+                    "confusing diffusion-registry error."
+                )
+
         # --omni-replica-address is only consulted in run_headless(); reject it
         # on the head so a misconfigured launch fails loudly instead of being
         # silently ignored.