[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

tests/entrypoints/test_serve.py

@@ -377,7 +377,7 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
         return_value=(None, None, None),
     )
     mocker.patch(
-        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        "vllm_omni.engine.stage_init_utils.extract_legacy_stage_metadata",
         return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
     )
     mock_inject = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
@@ -457,7 +457,7 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
         return_value=(None, None, None),
     )
     mocker.patch(
-        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        "vllm_omni.engine.stage_init_utils.extract_legacy_stage_metadata",
         return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")

vllm_omni/engine/stage_engine_startup.py

@@ -1347,7 +1347,10 @@ def launch_headless_diffusion_replicas(
         stage_id,
     )
 
-    metadata = stage_init_utils.extract_stage_metadata(stage_cfg)
+    # Headless diffusion startup and its downstream helpers still consume the
+    # legacy StageConfig shape; switch this with the coordinated RFC #4021
+    # stage-init cutover.
+    metadata = stage_init_utils.extract_legacy_stage_metadata(stage_cfg)
     if omni_conn_cfg:
         inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
     # Headless single-stage launch must still infer cross-stage TP topology

vllm_omni/engine/stage_init_utils.py

@@ -25,6 +25,8 @@ from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine.input_processor import InputProcessor
 from vllm.v1.executor import Executor
 
+from vllm_omni.config.omni_config import BaseVllmOmniStageConfig
+from vllm_omni.config.stage_config import StageType
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
@@ -346,8 +348,12 @@ class StageMetadata:
     replica_id: int = 0
 
 
-def extract_stage_metadata(stage_config: Any) -> StageMetadata:
-    """Pure data extraction from a stage_config object."""
+def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
+    """Extract metadata through the active production legacy path.
+
+    Keep production callers on this path until RFC #4021 migrates the
+    engine-argument and stage-init consumers together.
+    """
     stage_id: int = stage_config.stage_id
     stage_type: Literal["llm", "diffusion"] = _get_attr_or_item(stage_config, "stage_type", "llm")
     engine_args = stage_config.engine_args
@@ -432,6 +438,68 @@ def extract_stage_metadata(stage_config: Any) -> StageMetadata:
     )
 
 
+def extract_stage_metadata(stage_config: Any) -> StageMetadata:
+    """Preserve the legacy one-argument API for external callers."""
+    return extract_legacy_stage_metadata(stage_config)
+
+
+def _resolve_omni_metadata_hook(path: str | None) -> Callable | None:
+    if not path:
+        return None
+    module_path, function_name = path.rsplit(".", 1)
+    return getattr(importlib.import_module(module_path), function_name)
+
+
+def extract_stage_metadata_from_omni_stage_config(
+    stage_config: BaseVllmOmniStageConfig,
+) -> StageMetadata:
+    """Project one typed stage config into metadata for a future cutover.
+
+    This projection is not used by production startup yet. Current replica
+    layout, engine-argument, remote-diffusion, and platform setup paths still
+    require the legacy StageConfig/OmegaConf shape.
+    """
+    stage_type: Literal["llm", "diffusion"] = "diffusion" if stage_config.stage_type == StageType.DIFFUSION else "llm"
+    sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
+    sampling_params: OmniSamplingParams = sampling_params_cls(
+        **(stage_config.model_config.default_sampling_params or {})
+    )
+    custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func)
+
+    if stage_type == "diffusion":
+        return StageMetadata(
+            stage_id=stage_config.stage_id,
+            stage_type="diffusion",
+            engine_output_type=None,
+            is_comprehension=False,
+            requires_multimodal_data=False,
+            engine_input_source=stage_config.input_sources,
+            final_output=stage_config.final_output,
+            final_output_type=stage_config.final_output_type,
+            default_sampling_params=sampling_params,
+            custom_process_input_func=custom_process_input_func,
+            model_stage=stage_config.model_stage,
+            runtime_cfg=stage_config.runtime_config,
+            cfg_kv_collect_func=_resolve_omni_metadata_hook(stage_config.cfg_kv_collect_func),
+        )
+
+    return StageMetadata(
+        stage_id=stage_config.stage_id,
+        stage_type="llm",
+        engine_output_type=stage_config.engine_output_type,
+        is_comprehension=stage_config.is_comprehension,
+        requires_multimodal_data=stage_config.requires_multimodal_data,
+        engine_input_source=stage_config.input_sources,
+        final_output=stage_config.final_output,
+        final_output_type=stage_config.final_output_type,
+        default_sampling_params=sampling_params,
+        custom_process_input_func=custom_process_input_func,
+        model_stage=stage_config.model_stage,
+        runtime_cfg=stage_config.runtime_config,
+        prompt_expand_func=_resolve_omni_metadata_hook(stage_config.prompt_expand_func),
+    )
+
+
 def prepare_engine_environment() -> None:
     """One-time global setup: load plugins, set multiprocessing spawn method."""
     from vllm_omni.plugins import load_omni_general_plugins

vllm_omni/engine/stage_runtime.py

@@ -51,7 +51,7 @@ from vllm_omni.engine.stage_init_utils import (
     build_llm_stage_output_processor,
     build_vllm_config,
     compute_replica_layout,
-    extract_stage_metadata,
+    extract_legacy_stage_metadata,
     get_stage_connector_spec,
     inject_kv_stage_info,
     inject_omni_kv_connector_config,
@@ -339,8 +339,12 @@ class StageRuntime:
         """Build startup plans for every logical stage and replica."""
         stage_plans: list[LogicalStageInitPlan] = []
 
+        # RFC #4021 transition boundary: stage planning still relies on legacy
+        # StageConfig.runtime and StageConfig.engine_args for replica and engine
+        # setup. Keep metadata extraction on the legacy path until the
+        # coordinated stage-init cutover.
         for stage_idx, stage_cfg in enumerate(self._stage_configs):
-            base_metadata = extract_stage_metadata(stage_cfg)
+            base_metadata = extract_legacy_stage_metadata(stage_cfg)
             stage_id = int(base_metadata.stage_id)
             if stage_id != stage_idx:
                 raise ValueError(
@@ -392,7 +396,7 @@ class StageRuntime:
                 if stage_idx in replica_devices_map:
                     replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id]
 
-                replica_metadata = extract_stage_metadata(replica_cfg)
+                replica_metadata = extract_legacy_stage_metadata(replica_cfg)
                 replica_metadata.replica_id = replica_id
                 if launch_mode == "remote" and replica_metadata.stage_type != "diffusion":
                     replica_metadata.runtime_cfg = None
@@ -865,8 +869,10 @@ class DistStageRuntime(StageRuntime):
         if registered_stage_cfg is None:
             raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config")
 
+        # Remote diffusion registration still transports the legacy mapping
+        # shape. Reconstruct and project that shape until its RFC #4021 cutover.
         metadata = (
-            extract_stage_metadata(OmegaConf.create(registered_stage_cfg))
+            extract_legacy_stage_metadata(OmegaConf.create(registered_stage_cfg))
             if plan.metadata.stage_type == "diffusion"
             else copy.deepcopy(plan.metadata)
         )

vllm_omni/platforms/rocm/platform.py

@@ -30,7 +30,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
     when the selected_backend is not specified.
 
     So the behaviour of the attention backend overriding logic currently lives in
-    extract_stage_metadata in `vllm_omni/engine/stage_init_utils.py`
+    extract_legacy_stage_metadata in `vllm_omni/engine/stage_init_utils.py`
 
     ```
     if current_omni_platform.is_rocm():