[Config] Read stage metadata from VllmOmniConfig (#5343)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -342,7 +342,7 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
 
     monkeypatch.setattr(
         runtime_mod,
-        "extract_stage_metadata",
+        "extract_legacy_stage_metadata",
         lambda cfg: types.SimpleNamespace(**metadata_by_stage[cfg.stage_id].__dict__),
     )
     monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {})
@@ -836,9 +836,9 @@ def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stag
     assert not hasattr(stage_configs[0].engine_args, "diffusion_attention_config")
 
 
-def test_extract_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
+def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
     """ROCm default attention logic only applies to LLM stages, not diffusion."""
-    from vllm_omni.engine.stage_init_utils import extract_stage_metadata
+    from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata
 
     monkeypatch.setattr("vllm_omni.engine.stage_init_utils.current_omni_platform.is_rocm", lambda: True)
 
@@ -852,7 +852,7 @@ def test_extract_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeyp
         final_output_type=None,
     )
 
-    metadata = extract_stage_metadata(stage_cfg)
+    metadata = extract_legacy_stage_metadata(stage_cfg)
 
     assert metadata.stage_type == "diffusion"
     assert "diffusion_attention_config" not in stage_cfg.engine_args

tests/engine/test_single_stage_mode.py

@@ -450,7 +450,7 @@ class TestSingleStageInitialization:
         monkeypatch = pytest.MonkeyPatch()
         monkeypatch.setattr(
             runtime_mod,
-            "extract_stage_metadata",
+            "extract_legacy_stage_metadata",
             lambda cfg: SimpleNamespace(
                 stage_id=cfg.stage_id,
                 stage_type=getattr(cfg, "stage_type", "llm"),
@@ -477,7 +477,7 @@ class TestSingleStageInitialization:
         monkeypatch = pytest.MonkeyPatch()
         monkeypatch.setattr(
             runtime_mod,
-            "extract_stage_metadata",
+            "extract_legacy_stage_metadata",
             lambda cfg: SimpleNamespace(
                 stage_id=cfg.stage_id,
                 stage_type=getattr(cfg, "stage_type", "llm"),
@@ -550,7 +550,7 @@ class TestSingleStageInitialization:
         monkeypatch = pytest.MonkeyPatch()
         monkeypatch.setattr(
             runtime_mod,
-            "extract_stage_metadata",
+            "extract_legacy_stage_metadata",
             lambda cfg: SimpleNamespace(
                 stage_id=cfg.stage_id,
                 stage_type="llm",
@@ -609,7 +609,7 @@ class TestSingleStageInitialization:
         monkeypatch = pytest.MonkeyPatch()
         monkeypatch.setattr(
             runtime_mod,
-            "extract_stage_metadata",
+            "extract_legacy_stage_metadata",
             lambda cfg: SimpleNamespace(
                 stage_id=cfg.stage_id,
                 stage_type="diffusion",
@@ -943,7 +943,7 @@ class TestSingleStageReplicaInitialization:
         plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0]
         sentinel_client = SimpleNamespace()
 
-        mocker.patch.object(runtime_mod, "extract_stage_metadata", return_value=remote_metadata)
+        mocker.patch.object(runtime_mod, "extract_legacy_stage_metadata", return_value=remote_metadata)
         mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect)
         mock_from_addresses = mocker.patch(
             "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses",

tests/engine/test_stage_metadata.py

@@ -0,0 +1,160 @@
+import copy
+import operator
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.stage_config import (
+    DeployConfig,
+    PipelineConfig,
+    StageDeployConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+    merge_pipeline_deploy,
+)
+from vllm_omni.engine.stage_init_utils import (
+    extract_legacy_stage_metadata,
+    extract_stage_metadata,
+    extract_stage_metadata_from_omni_stage_config,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _metadata_inputs() -> tuple[PipelineConfig, DeployConfig]:
+    pipeline = PipelineConfig(
+        model_type="metadata-test",
+        stages=(
+            StagePipelineConfig(
+                stage_id=0,
+                model_stage="thinker",
+                execution_type=StageExecutionType.LLM_AR,
+                owns_tokenizer=True,
+                requires_multimodal_data=True,
+                engine_output_type="token_ids",
+                custom_process_input_func="operator.add",
+                prompt_expand_func="operator.mul",
+            ),
+            StagePipelineConfig(
+                stage_id=1,
+                model_stage="talker",
+                execution_type=StageExecutionType.LLM_GENERATION,
+                input_sources=(0,),
+                engine_output_type="audio_tokens",
+                custom_process_input_func="operator.sub",
+                sync_process_input_func="operator.floordiv",
+            ),
+            StagePipelineConfig(
+                stage_id=2,
+                model_stage="dit",
+                execution_type=StageExecutionType.DIFFUSION,
+                input_sources=(1,),
+                final_output=True,
+                final_output_type="image",
+                custom_process_input_func="operator.neg",
+                cfg_kv_collect_func="operator.concat",
+            ),
+        ),
+    )
+    deploy = DeployConfig(
+        async_chunk=False,
+        stages=[
+            StageDeployConfig(
+                stage_id=0,
+                devices="0",
+                env={"STAGE": "thinker"},
+                default_sampling_params={"temperature": 0.25},
+            ),
+            StageDeployConfig(
+                stage_id=1,
+                devices="1,2",
+                num_replicas=2,
+                env={"STAGE": "talker"},
+                default_sampling_params={"temperature": 0.75},
+            ),
+            StageDeployConfig(
+                stage_id=2,
+                devices="3",
+                env={"STAGE": "dit"},
+                default_sampling_params={"seed": 7},
+            ),
+        ],
+    )
+    return pipeline, deploy
+
+
+def test_extract_stage_metadata_matches_legacy_projection():
+    pipeline, deploy = _metadata_inputs()
+    omni_config = VllmOmniConfig.from_pipeline_config(
+        pipeline,
+        user_deploy_config=copy.deepcopy(deploy),
+    )
+    legacy_configs = [
+        stage.to_omegaconf()
+        for stage in merge_pipeline_deploy(
+            pipeline,
+            copy.deepcopy(deploy),
+        )
+    ]
+
+    parity_fields = (
+        "stage_id",
+        "stage_type",
+        "engine_output_type",
+        "is_comprehension",
+        "requires_multimodal_data",
+        "engine_input_source",
+        "final_output",
+        "final_output_type",
+        "custom_process_input_func",
+        "model_stage",
+        "prompt_expand_func",
+        "cfg_kv_collect_func",
+    )
+    for legacy_config in legacy_configs:
+        stage_id = legacy_config.stage_id
+        omni_stage_config = omni_config.stage_by_id(stage_id)
+        structured = extract_stage_metadata_from_omni_stage_config(omni_stage_config)
+        legacy = extract_legacy_stage_metadata(legacy_config)
+
+        assert {field: getattr(structured, field) for field in parity_fields} == {
+            field: getattr(legacy, field) for field in parity_fields
+        }
+        assert type(structured.default_sampling_params) is type(legacy.default_sampling_params)
+        assert structured.runtime_cfg is omni_stage_config.runtime_config
+        for runtime_field in ("devices", "num_replicas", "env"):
+            assert getattr(structured.runtime_cfg, runtime_field) == getattr(legacy.runtime_cfg, runtime_field)
+
+    thinker = extract_stage_metadata_from_omni_stage_config(omni_config.stage_by_id(0))
+    talker = extract_stage_metadata_from_omni_stage_config(omni_config.stage_by_id(1))
+    diffusion = extract_stage_metadata_from_omni_stage_config(omni_config.stage_by_id(2))
+
+    assert isinstance(thinker.default_sampling_params, SamplingParams)
+    assert thinker.default_sampling_params.temperature == 0.25
+    assert thinker.custom_process_input_func is operator.add
+    assert thinker.prompt_expand_func is operator.mul
+
+    assert isinstance(talker.default_sampling_params, SamplingParams)
+    assert talker.default_sampling_params.temperature == 0.75
+    assert talker.custom_process_input_func is operator.floordiv
+
+    assert isinstance(diffusion.default_sampling_params, OmniDiffusionSamplingParams)
+    assert diffusion.default_sampling_params.seed == 7
+    assert diffusion.custom_process_input_func is operator.neg
+    assert diffusion.cfg_kv_collect_func is operator.concat
+
+
+def test_extract_stage_metadata_preserves_legacy_one_argument_api():
+    pipeline, deploy = _metadata_inputs()
+    legacy_config = merge_pipeline_deploy(
+        pipeline,
+        deploy,
+    )[0].to_omegaconf()
+
+    metadata = extract_stage_metadata(legacy_config)
+
+    assert metadata.stage_id == 0
+    assert metadata.model_stage == "thinker"
+    assert metadata.custom_process_input_func is operator.add