[BugFix] Validate --stage-overrides JSON shape with a clear, flag-named error

tests/entrypoints/test_utils.py

@@ -16,10 +16,12 @@
 from vllm_omni.entrypoints.utils import (
     _convert_dataclasses_to_dict,
     _filter_dict_like_object,
+    _validate_stage_overrides,
     coerce_param_message_types,
     filter_dataclass_kwargs,
     filter_stages,
     load_and_resolve_stage_configs,
+    load_stage_configs_from_model,
     load_stage_configs_from_yaml,
     resolve_model_config_path,
 )
@@ -314,6 +316,58 @@ def mock_exists(path):
         assert "glm_image.yaml" in result
 
 
+class TestValidateStageOverrides:
+    """Shape validation for the ``--stage-overrides`` JSON payload."""
+
+    def test_none_and_empty_are_noops(self):
+        # No payload / empty payload must not raise.
+        _validate_stage_overrides(None)
+        _validate_stage_overrides({})
+
+    def test_well_formed_payload_passes(self):
+        _validate_stage_overrides({"0": {"gpu_memory_utilization": 0.8}})
+        # Integer stage ids (programmatic callers) are accepted too.
+        _validate_stage_overrides({0: {"enforce_eager": True}, 2: {}})
+
+    @pytest.mark.parametrize(
+        "payload",
+        [
+            [1, 2],  # top-level list
+            5,  # top-level scalar
+            "hello",  # top-level string
+        ],
+    )
+    def test_non_object_top_level_raises(self, payload):
+        with pytest.raises(ValueError, match="--stage-overrides"):
+            _validate_stage_overrides(payload)
+
+    @pytest.mark.parametrize(
+        "payload",
+        [
+            {"0": 5},  # stage value is a scalar
+            {"0": [1, 2]},  # stage value is a list
+            {"0": "eager"},  # stage value is a string
+        ],
+    )
+    def test_non_dict_stage_value_raises(self, payload):
+        with pytest.raises(ValueError, match="--stage-overrides"):
+            _validate_stage_overrides(payload)
+
+    @pytest.mark.parametrize("bad_key", ["abc", "stage0", "0.5", "-1", True])
+    def test_non_integer_stage_id_raises(self, bad_key):
+        # A non-numeric stage id is silently dropped downstream (the parser only
+        # matches ``stage_<digits>_*``); reject it up-front instead.
+        with pytest.raises(ValueError, match="stage id"):
+            _validate_stage_overrides({bad_key: {"k": 1}})
+
+    def test_loader_rejects_malformed_overrides_before_model_load(self):
+        # The validator runs at the top of load_stage_configs_from_model, so a
+        # malformed payload fails fast with a clear error rather than a cryptic
+        # AttributeError raised deep in config loading.
+        with pytest.raises(ValueError, match="--stage-overrides"):
+            load_stage_configs_from_model("any-model", stage_overrides=[1, 2])
+
+
 class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self):
         """Ensure that dtype survives default stage creation."""

vllm_omni/entrypoints/utils.py

@@ -311,6 +311,50 @@ def resolve_model_config_path(model: str) -> str:
     return str(stage_config_path)
 
 
+def _validate_stage_overrides(stage_overrides: Any) -> None:
+    """Validate the shape of ``--stage-overrides`` before it is consumed.
+
+    ``--stage-overrides`` is documented as a JSON object mapping each
+    ``stage_id`` to a dict of per-stage overrides, e.g.
+    ``{"0": {"gpu_memory_utilization": 0.8}}``. The engine only checks that the
+    value is valid JSON, so a syntactically-valid but malformed payload (a list,
+    a scalar, or a ``stage_id -> scalar`` mapping) would otherwise blow up deep in
+    config loading with a bare ``AttributeError`` that names neither the flag nor
+    the offending value. Stage ids must be integers: the downstream parser
+    (``_STAGE_OVERRIDE_PATTERN = re.compile(r"^stage_(\\d+)_(.+)$")`` in
+    ``config/stage_config.py``) only matches numeric stage ids, so a non-numeric
+    key would be silently dropped rather than applied.
+
+    Raises:
+        ValueError: if ``stage_overrides`` is not a ``stage_id -> {param: value}``
+            mapping with integer stage ids.
+    """
+    if not stage_overrides:
+        return
+
+    example = '{"0": {"gpu_memory_utilization": 0.8}}'
+    if not isinstance(stage_overrides, dict):
+        raise ValueError(
+            "--stage-overrides must be a JSON object mapping stage_id -> "
+            f"{{param: value}}, e.g. '{example}'; got "
+            f"{type(stage_overrides).__name__}."
+        )
+
+    for stage_id, overrides in stage_overrides.items():
+        if isinstance(stage_id, bool) or not (
+            isinstance(stage_id, int) or (isinstance(stage_id, str) and stage_id.isdigit())
+        ):
+            raise ValueError(
+                f'--stage-overrides keys must be integer stage ids (e.g. "0"), got invalid stage id {stage_id!r}.'
+            )
+        if not isinstance(overrides, dict):
+            raise ValueError(
+                f"--stage-overrides[{stage_id!r}] must be a JSON object of "
+                f"{{param: value}}, e.g. '{example}'; got "
+                f"{type(overrides).__name__}."
+            )
+
+
 def load_stage_configs_from_model(
     model: str,
     base_engine_args: dict | None = None,
@@ -337,6 +381,8 @@ def load_stage_configs_from_model(
     if base_engine_args is None:
         base_engine_args = {}
 
+    _validate_stage_overrides(stage_overrides)
+
     cli_overrides = _convert_dataclasses_to_dict(dict(base_engine_args))
     if stage_overrides:
         for stage_id_str, overrides in stage_overrides.items():