[Refactor] Normalize diffusion request extra arguments in Chat serving (#5171)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>

tests/entrypoints/openai_api/test_diffusion_request_utils.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.entrypoints.openai.diffusion_request_utils import (
+    apply_normalized_diffusion_request_extra_args,
+    normalize_diffusion_request_args,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+REQUEST_SOURCES = (
+    "root.declared",
+    "root.extra_args",
+    "root.extra_params",
+    "nested.declared",
+    "nested.extra_args",
+    "nested.extra_params",
+)
+
+
+def _source_args(source: str, value: object) -> dict[str, object]:
+    scope, container = source.split(".")
+    payload = {"solver": value} if container == "declared" else {container: {"solver": value}}
+    kwargs: dict[str, object] = {scope: payload}
+    if container == "declared":
+        kwargs["registered_extra_fields"] = {"solver"}
+    return kwargs
+
+
+@pytest.mark.parametrize("source", REQUEST_SOURCES)
+@pytest.mark.parametrize(("value", "expected"), [("ddim", "ddim"), (None, "euler")])
+def test_normalizer_routes_values_without_erasing_defaults(
+    source: str,
+    value: object,
+    expected: str,
+) -> None:
+    normalized, _ = normalize_diffusion_request_args(**_source_args(source, value))
+    sampling_params = OmniDiffusionSamplingParams(
+        extra_args={"solver": "euler", "stage_default": True},
+    )
+    apply_normalized_diffusion_request_extra_args(sampling_params, normalized)
+
+    assert sampling_params.extra_args == {"solver": expected, "stage_default": True}
+
+
+def test_non_overlapping_legacy_and_canonical_values_are_merged(mocker) -> None:
+    warning_once = mocker.patch("vllm_omni.entrypoints.openai.diffusion_request_utils.logger.warning_once")
+
+    normalized, _ = normalize_diffusion_request_args(
+        root={"extra_args": {"cfg_text_scale": 7.0}},
+        nested={"extra_params": {"sample_solver": "euler"}},
+    )
+
+    assert normalized == {"cfg_text_scale": 7.0, "sample_solver": "euler"}
+    warning_once.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "kwargs, expected_sources",
+    [
+        (
+            {
+                "root": {"seed": 1, "extra_args": {"seed": 2}},
+                "serving_root_fields": {"seed"},
+            },
+            "request.seed, request.extra_args.seed",
+        ),
+        (
+            {
+                "root": {"cfg_scale": 7.0},
+                "nested": {"true_cfg_scale": 2.0},
+                "serving_root_fields": {"true_cfg_scale"},
+                "root_field_aliases": {"cfg_scale": "true_cfg_scale"},
+            },
+            "request.cfg_scale, request.extra_body.true_cfg_scale",
+        ),
+        (
+            {
+                "root": {
+                    "extra_args": {"solver": "euler"},
+                    "extra_params": {"solver": "ddim"},
+                }
+            },
+            "request.extra_args.solver, request.extra_params.solver",
+        ),
+        (
+            {
+                "root": {"cfg_text_scale": 7.0},
+                "nested": {"extra_args": {"cfg_text_scale": 8.0}},
+                "registered_extra_fields": {"cfg_text_scale"},
+            },
+            "request.cfg_text_scale, request.extra_body.extra_args.cfg_text_scale",
+        ),
+    ],
+)
+def test_duplicate_sources_are_rejected(
+    kwargs: dict[str, object],
+    expected_sources: str,
+    mocker,
+) -> None:
+    warning_once = mocker.patch("vllm_omni.entrypoints.openai.diffusion_request_utils.logger.warning_once")
+
+    with pytest.raises(ValueError, match=expected_sources):
+        normalize_diffusion_request_args(**kwargs)
+
+    warning_once.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("scope", "name"),
+    (
+        ("root", "extra_args"),
+        ("root", "extra_params"),
+        ("nested", "extra_args"),
+        ("nested", "extra_params"),
+    ),
+)
+def test_public_containers_must_be_objects(scope: str, name: str) -> None:
+    with pytest.raises(ValueError, match="must be a JSON object"):
+        normalize_diffusion_request_args(**{scope: {name: ["not", "an", "object"]}})
+
+
+def test_container_keys_must_be_strings() -> None:
+    with pytest.raises(ValueError, match="must be a JSON object with string keys"):
+        normalize_diffusion_request_args(root={"extra_args": {1: "invalid"}})