From 2b47faf6012ad3e9c4c2e6668b313a3aea85066b Mon Sep 17 00:00:00 2001 From: Chiquan Mo Date: Thu, 3 Sep 2026 14:19:08 +0800 Subject: [PATCH] [Bugfix] Validate --stage-overrides shape in parse_stage_overrides (drop field allowlist) (#6230) Signed-off-by: Chiquan Mo Signed-off-by: MciG-ggg Co-authored-by: MciG-ggg --- tests/entrypoints/test_serve.py | 72 +++++++++++++++++++++++++++++++++ tests/entrypoints/test_utils.py | 15 +++++-- vllm_omni/entrypoints/utils.py | 56 +++++++++++++++++++++++-- 3 files changed, 137 insertions(+), 6 deletions(-) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index e759e0f15..01d3e7ddb 100644 --- a/tests/entrypoints/test_serve.py +++ b/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__`` 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``, diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 1fc4779a7..6bc6cdc2a 100644 --- a/tests/entrypoints/test_utils.py +++ b/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.""" diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index d1c4c8aba..adc6cbabe 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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__`` 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)