[Bugfix] Validate --stage-overrides shape in parse_stage_overrides (drop field allowlist) (#6230)

Signed-off-by: Chiquan Mo <mcigggg399026@gmail.com>
Signed-off-by: MciG-ggg <MciG-ggg@users.noreply.github.com>
Co-authored-by: MciG-ggg <MciG-ggg@users.noreply.github.com>

tests/entrypoints/test_serve.py

@@ -179,6 +179,78 @@ def test_parse_stage_overrides_invalid_json_raises() -> None:
     assert f"Got: {bad!r}" in message
 
 
+def test_parse_stage_overrides_rejects_non_dict_top_level() -> None:
+    """Top-level must be a JSON object (dict). A list, scalar, or non-dict
+    mapping is rejected with a ValueError naming ``--stage-overrides`` and
+    pointing at the bad shape. Without this guard, ``json.loads`` happily
+    returns a list/scalar and the override silently never applies."""
+    for bad in ("[1, 2, 3]", '"oops"', "42"):
+        with pytest.raises(ValueError, match="must be a JSON object"):
+            parse_stage_overrides(bad)
+
+
+def test_parse_stage_overrides_rejects_non_integer_stage_id() -> None:
+    """Stage-id keys must be non-negative ASCII integer strings. Letters,
+    signs, floats, Unicode digit classes, and integer (non-string) keys all
+    fail. ``str.isdigit() and str.isascii()`` is the minimal check: it
+    rejects ``"-1"``, ``"abc"``, ``"1.5"``, fullwidth ``"０"``, and bare ``1``.
+
+    Note: ``json.loads`` normalizes integer object keys (``{"1": {}}``) into
+    the string ``"1"`` and would pass our digit check, so the integer-key case
+    is exercised via the already-parsed-dict code path (``parse_stage_overrides({1: {}})``)."""
+    bad_string_keys = ('"abc"', '"-1"', '"1.5"', '"\uff10"')
+    for bad_key in bad_string_keys:
+        with pytest.raises(ValueError, match="non-negative integer stage ids"):
+            parse_stage_overrides("{" + bad_key + ": {}}")
+    # Integer (non-string) key: must reach the structural check, not the JSON
+    # parser, so pass an already-parsed mapping directly.
+    with pytest.raises(ValueError, match="non-negative integer stage ids"):
+        parse_stage_overrides({1: {}})
+
+
+def test_parse_stage_overrides_accepts_stage_merge_extras_and_engine_args() -> None:
+    """Three classes of keys pass through the shape-only parser:
+
+    - ``extras`` is read by the default-diffusion fallback
+      (``async_omni_engine.py``); registered pipelines carry it on
+      ``StagePipelineConfig.extras`` directly.
+    - Engine arguments (``kv_cache_dtype``, ``stage_connector_spec``, ...)
+      are forwarded as ``stage_<id>_<key>`` and applied via
+      ``OmniEngineArgs``.
+    - Unknown keys parse through and are dropped with a warning at
+      ``filter_dataclass_kwargs`` (see
+      ``tests/entrypoints/test_utils.py::TestFilterDataclassKwargs``).
+
+    ``typo_field_xyz`` in the payload exercises the third case.
+    """
+    parsed = parse_stage_overrides(
+        '{"0": {"extras": {"ltx2_use_conv_vae": true},'
+        ' "kv_cache_dtype": "fp8", "seed": 42,'
+        ' "stage_connector_spec": {"name": "SharedMemoryConnector", "extra": {}},'
+        ' "typo_field_xyz": 1}}'
+    )
+    assert parsed == {
+        "0": {
+            "extras": {"ltx2_use_conv_vae": True},
+            "kv_cache_dtype": "fp8",
+            "seed": 42,
+            "stage_connector_spec": {"name": "SharedMemoryConnector", "extra": {}},
+            "typo_field_xyz": 1,
+        },
+    }
+    # Already-parsed mapping path carries the same trust.
+    assert parse_stage_overrides(dict(parsed)) == parsed
+
+
+def test_parse_stage_overrides_accepts_empty_inner_dict() -> None:
+    """Per-stage overrides may be empty (``{}``): shape stays valid, every
+    stage simply receives no per-stage tweaks. Locks in the
+    ``not parsed`` -> ``None`` short-circuit's twin: an outer
+    non-empty dict with empty inner dicts is a valid no-op pass-through."""
+    parsed = parse_stage_overrides('{"0": {}, "1": {}}')
+    assert parsed == {"0": {}, "1": {}}
+
+
 def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None:
     """Regression: the headless path must parse ``--stage-overrides`` (a JSON
     string) and forward the parsed dict to ``load_and_resolve_stage_configs``,

tests/entrypoints/test_utils.py

@@ -3,6 +3,7 @@
 
 """Unit tests for vllm_omni.entrypoints.utils module."""
 
+import logging
 import os
 from collections import Counter
 from dataclasses import dataclass
@@ -255,8 +256,12 @@ class TestFilterDataclassKwargs:
         with pytest.raises(ValueError, match="kwargs must be a dictionary"):
             filter_dataclass_kwargs(SimpleConfig, "invalid")
 
-    def test_filters_omni_engine_args_unknown_fields(self):
-        """Test that OmniEngineArgs kwargs are filtered to valid fields only."""
+    def test_filters_omni_engine_args_unknown_fields(self, caplog):
+        """Test that OmniEngineArgs kwargs are filtered to valid fields only,
+        and that the WARNING contract fires for every drop — the affordance
+        that lets a ``--stage-overrides`` typo (``{"0":{"kv_cache_dtpye":...}}``)
+        surface in the log rather than being silently misapplied downstream.
+        """
         kwargs = {
             "model": "dummy",
             "stage_id": 1,
@@ -264,12 +269,16 @@ class TestFilterDataclassKwargs:
             "unknown_field": "drop_me",
         }
 
-        result = filter_dataclass_kwargs(OmniEngineArgs, kwargs)
+        with caplog.at_level(logging.WARNING, logger="vllm_omni.entrypoints.utils"):
+            result = filter_dataclass_kwargs(OmniEngineArgs, kwargs)
 
         assert "model" in result
         assert "stage_id" in result
         assert "engine_output_type" in result
         assert "unknown_field" not in result
+        assert any(rec.levelno == logging.WARNING and "unknown_field" in rec.message for rec in caplog.records), (
+            f"expected WARNING naming 'unknown_field'; got {[r.message for r in caplog.records]}"
+        )
 
     def test_filters_omni_diffusion_config_union_dataclass(self):
         """Test that OmniDiffusionConfig filters nested dataclass in Union fields."""

vllm_omni/entrypoints/utils.py

@@ -544,17 +544,62 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
     ``value`` may be a raw JSON string (as supplied on the CLI) or an
     already-parsed mapping. Returns ``None`` when no overrides are given.
 
+    Only the **shape** of the override mapping is validated here:
+
+    - top-level must be a dict
+    - keys must be non-negative ASCII integer strings (stage ids)
+    - values must be dicts
+
+    Field-name / type / range validation is intentionally not done at this
+    layer. The downstream resolver forwards every surviving key as
+    ``stage_<id>_<key>`` into ``cli_overrides`` (see
+    ``load_stage_configs_from_model``); field semantics live in
+    ``StageConfigFactory`` for registered pipelines, and the
+    default-diffusion fallback in ``async_omni_engine.py`` reads
+    ``stage_zero_overrides["extras"]`` for unregistered ones. Unknown
+    engine-arg keys that don't match ``OmniEngineArgs`` are dropped with
+    a warning at ``filter_dataclass_kwargs`` rather than failing here,
+    so a typo surfaces in the log instead of being silently lost.
+
     Raises:
-        ValueError: when ``value`` is a string that is not valid JSON.
+        ValueError: when ``value`` is not a valid per-stage override
+            mapping (invalid JSON, non-dict top level, non-ASCII-digit
+            stage id, non-dict value).
     """
     if not value:
         return None
     if isinstance(value, str):
         try:
-            return json.loads(value)
+            parsed = json.loads(value)
         except json.JSONDecodeError as exc:
             raise ValueError(f"--stage-overrides is not valid JSON: {exc}. Got: {value!r}") from exc
-    return value
+    else:
+        parsed = value
+
+    if not isinstance(parsed, dict):
+        raise ValueError(
+            f"--stage-overrides must be a JSON object mapping stage_id -> overrides, "
+            f"got {type(parsed).__name__}: {parsed!r}"
+        )
+    if not parsed:
+        return None
+
+    for stage_id_str, overrides in parsed.items():
+        # ``str.isdigit()`` accepts Unicode digit classes (e.g. fullwidth
+        # ``"０"``); downstream's ``_STAGE_OVERRIDE_PATTERN`` is ASCII-only
+        # (``^stage_(\d+)_(.+)$``), so a non-ASCII-digit key would parse
+        # here and silently misroute as a global key. ``isascii()`` closes
+        # that hole.
+        if not isinstance(stage_id_str, str) or not stage_id_str.isdigit() or not stage_id_str.isascii():
+            raise ValueError(
+                f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id_str!r}"
+            )
+        if not isinstance(overrides, dict):
+            raise ValueError(
+                f"--stage-overrides[{stage_id_str!r}] must be an object, got {type(overrides).__name__}: {overrides!r}"
+            )
+
+    return parsed
 
 
 def load_and_resolve_stage_configs(
@@ -721,6 +766,11 @@ def filter_dataclass_kwargs(cls: Any, kwargs: dict) -> dict:
     filtered_kwargs = {}
     for k, v in kwargs.items():
         if k not in valid_fields:
+            logger.warning(
+                "Dropping unknown %s field %r (not declared on the dataclass)",
+                cls.__name__,
+                k,
+            )
             continue
         field = valid_fields[k]
         filtered_kwargs[k] = _filter_value(v, field.type)