diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index db7d53b79..87edc594b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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 diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index dc06c0c8d..7cf4c9d75 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/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", diff --git a/tests/engine/test_stage_metadata.py b/tests/engine/test_stage_metadata.py new file mode 100644 index 000000000..f56bbc2a9 --- /dev/null +++ b/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 diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 61ffed3d5..bbed8fc08 100644 --- a/tests/entrypoints/test_serve.py +++ b/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") diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 7a55bb93f..d23d9f3e0 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/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 diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 3b1bb93af..2310daebe 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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 diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 4ee4ccc61..1f2024fd2 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/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) ) diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 38a14c325..f44bb86fa 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/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():