Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 74 additions & 0 deletions tests/entrypoints/test_serve.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
44 changes: 44 additions & 0 deletions vllm_omni/entrypoints/cli/serve.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,50 @@ def validate(self, args: argparse.Namespace) -> None:
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.
Expand Down
Loading