[FEAT] support multi-stage deployment (#2396)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Chenguang Zheng <645327136@qq.com>

tests/engine/test_async_omni_engine_input.py

@@ -59,7 +59,7 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
     assert request.additional_information is not None
     assert request.additional_information.entries["text"].list_data == ["hello world"]
     assert request.additional_information.entries["speaker"].list_data == ["vivian"]
-    output_processor.add_request.assert_called_once()
+    output_processor.add_request.assert_not_called()
 
 
 def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture):

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,15 +1,138 @@
 import importlib
 import os
 import threading
+import time
 import types
 
 import pytest
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.stage_init_utils import (
+    LogicalStageInitPlan,
+    ReplicaInitPlan,
+    build_stage0_input_processor,
+    compute_replica_layout,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _make_llm_metadata(
+    stage_id: int,
+    *,
+    replica_id: int = 0,
+    final_output: bool = False,
+    final_output_type: str | None = None,
+    is_comprehension: bool = False,
+):
+    return types.SimpleNamespace(
+        stage_id=stage_id,
+        stage_type="llm",
+        runtime_cfg={},
+        prompt_expand_func=None,
+        final_output=final_output,
+        final_output_type=final_output_type,
+        default_sampling_params=types.SimpleNamespace(name=f"sp-{stage_id}-{replica_id}"),
+        custom_process_input_func=None,
+        engine_input_source=[] if stage_id == 0 else [stage_id - 1],
+        engine_output_type="token_ids",
+        replica_id=replica_id,
+        is_comprehension=is_comprehension,
+    )
+
+
+def _make_diffusion_metadata(stage_id: int, *, replica_id: int = 0, final_output_type: str = "image"):
+    return types.SimpleNamespace(
+        stage_id=stage_id,
+        stage_type="diffusion",
+        runtime_cfg={"devices": str(replica_id)},
+        prompt_expand_func=None,
+        final_output=True,
+        final_output_type=final_output_type,
+        default_sampling_params=types.SimpleNamespace(name=f"dsp-{stage_id}-{replica_id}"),
+        custom_process_input_func=None,
+        engine_input_source=[],
+        cfg_kv_collect_func=None,
+        replica_id=replica_id,
+    )
+
+
+def _make_llm_plan(
+    stage_idx: int,
+    *,
+    configured_stage_id: int,
+    vllm_config: object,
+    num_replicas: int = 1,
+    final_output: bool = False,
+    final_output_type: str | None = None,
+    is_comprehension: bool = False,
+):
+    replicas: list[ReplicaInitPlan] = []
+    for replica_id in range(num_replicas):
+        stage_cfg = types.SimpleNamespace(
+            stage_id=configured_stage_id,
+            stage_type="llm",
+            runtime=types.SimpleNamespace(devices=str(replica_id)),
+            engine_args={},
+        )
+        replicas.append(
+            ReplicaInitPlan(
+                replica_id=replica_id,
+                num_replicas=num_replicas,
+                launch_mode="local",
+                stage_cfg=stage_cfg,
+                metadata=_make_llm_metadata(
+                    configured_stage_id,
+                    replica_id=replica_id,
+                    final_output=final_output,
+                    final_output_type=final_output_type,
+                    is_comprehension=is_comprehension and replica_id == 0,
+                ),
+                stage_connector_spec={},
+                omni_kv_connector=(None, None, None),
+                stage_vllm_config=vllm_config,
+                executor_class=object,
+            )
+        )
+    return LogicalStageInitPlan(
+        stage_idx=stage_idx,
+        configured_stage_id=configured_stage_id,
+        replicas=replicas,
+    )
+
+
+def _make_diffusion_plan(
+    stage_idx: int,
+    *,
+    configured_stage_id: int,
+    num_replicas: int = 1,
+):
+    replicas: list[ReplicaInitPlan] = []
+    for replica_id in range(num_replicas):
+        stage_cfg = types.SimpleNamespace(
+            stage_id=configured_stage_id,
+            stage_type="diffusion",
+            runtime=types.SimpleNamespace(devices=str(replica_id)),
+            engine_args={},
+        )
+        replicas.append(
+            ReplicaInitPlan(
+                replica_id=replica_id,
+                num_replicas=num_replicas,
+                launch_mode="local",
+                stage_cfg=stage_cfg,
+                metadata=_make_diffusion_metadata(configured_stage_id, replica_id=replica_id),
+                stage_connector_spec={},
+                omni_kv_connector=(None, None, None),
+            )
+        )
+    return LogicalStageInitPlan(
+        stage_idx=stage_idx,
+        configured_stage_id=configured_stage_id,
+        replicas=replicas,
+    )
+
+
 def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     """Regression test for forward references in make_async_mp_client."""
     import vllm_omni.engine.stage_engine_core_client as client_mod
@@ -21,15 +144,279 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     )
 
 
-def test_initialize_stages_restores_device_visibility_after_diffusion_init(monkeypatch):
-    """Regression test for stage device env leakage across stage init.
+def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_args={"parallel_config": {"tensor_parallel_size": 2}},
+        runtime={"devices": "0,1,2,3", "num_replicas": 2},
+    )
 
-    Diffusion init mutates process-level CUDA visibility. Ensure AsyncOmniEngine
-    restores the previous value after diffusion stage setup.
-    """
+    replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg])
+
+    assert replicas_per_stage == [2]
+    assert replica_devices_map == {0: ["0,1", "2,3"]}
+
+
+def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
+    shared = types.SimpleNamespace(name="shared")
+    extra = types.SimpleNamespace(name="extra")
+
+    cleanup_clients = AsyncOmniEngine._collect_initialized_clients_for_cleanup(
+        stage_pools=[types.SimpleNamespace(clients=[shared, None])],
+        initialized_clients_by_stage={0: [shared], 1: [extra]},
+    )
+
+    assert cleanup_clients == [shared, extra]
+
+
+def test_initialize_stages_rejects_non_diffusion_replicas_in_single_stage_mode():
+    engine = object.__new__(AsyncOmniEngine)
+    engine.single_stage_mode = True
+    engine.stage_configs = [types.SimpleNamespace(stage_id=0, runtime={"num_replicas": 2})]
+
+    with pytest.raises(
+        ValueError,
+        match="single_stage_mode only supports num_replicas > 1 for diffusion stages",
+    ):
+        engine._validate_single_stage_mode_replica_constraints()
+
+
+def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
     from vllm_omni.platforms import current_omni_platform
 
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.num_stages = 1
+    engine.diffusion_batch_size = 1
+    engine.single_stage_mode = False
+    engine._omni_master_server = None
+    engine.stage_configs = []
+
+    plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
+
+    env_var = current_omni_platform.device_control_env_var
+    old_env = os.environ.get(env_var)
+    os.environ[env_var] = "0,1"
+
+    def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
+        current_omni_platform.set_device_control_env_var("1")
+
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
+    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace())
+
+    try:
+        engine._initialize_diffusion_replica(plan, stage_init_timeout=1, stage_launch_lock=threading.Lock())
+        assert os.environ.get(env_var) == "0,1"
+    finally:
+        if old_env is None:
+            os.environ.pop(env_var, None)
+        else:
+            os.environ[env_var] = old_env
+
+
+def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.num_stages = 1
+    engine.diffusion_batch_size = 4
+    engine.single_stage_mode = False
+    engine._omni_master_server = None
+    engine.stage_configs = []
+
+    plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
+
+    captured: dict[str, object] = {}
+
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
+
+    def _capture_initialize_diffusion_stage(
+        stage_id, _model, _stage_cfg, _metadata, *, stage_init_timeout, batch_size, use_inline
+    ):
+        captured["stage_id"] = stage_id
+        captured["stage_init_timeout"] = stage_init_timeout
+        captured["batch_size"] = batch_size
+        captured["use_inline"] = use_inline
+        return types.SimpleNamespace()
+
+    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _capture_initialize_diffusion_stage)
+
+    engine._initialize_diffusion_replica(plan, stage_init_timeout=302, stage_launch_lock=threading.Lock())
+
+    assert captured == {
+        "stage_id": 0,
+        "stage_init_timeout": 302,
+        "batch_size": 4,
+        "use_inline": True,
+    }
+
+
+def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_input_processor(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.num_stages = 2
+    engine.async_chunk = False
+    engine.diffusion_batch_size = 1
+    engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
+    engine._omni_master_server = None
+    engine.stage_configs = [types.SimpleNamespace(), types.SimpleNamespace()]
+
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [
+        _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True),
+        _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, final_output=True),
+    ]
+
+    stage0_client_r0 = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=True,
+        final_output=False,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(name="sp0"),
+    )
+    stage0_client_r1 = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=False,
+        final_output=False,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(name="sp0r1"),
+    )
+    stage1_client_r0 = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(name="sp1"),
+    )
+    initialized_clients = {
+        0: [stage0_client_r0, stage0_client_r1],
+        1: [stage1_client_r0],
+    }
+
+    stage0_output_processor = object()
+    stage1_output_processor = object()
+    top_level_input_processor = object()
+
+    monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
+    monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2, 1], {}))
+    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
+    monkeypatch.setattr(engine, "_initialize_stage_replicas", lambda *_: initialized_clients)
+    monkeypatch.setattr(
+        engine_mod,
+        "build_llm_stage_output_processor",
+        lambda plan, _cfg: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor,
+    )
+    monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor)
+
+    engine._initialize_stages(stage_init_timeout=1)
+
+    assert len(engine.stage_pools) == 2
+    assert engine.input_processor is top_level_input_processor
+    assert engine.stage_clients == [stage0_client_r0, stage1_client_r0]
+    assert engine.stage_vllm_configs == [cfg0, cfg1]
+    assert engine.output_processors == [stage0_output_processor, stage1_output_processor]
+    assert engine.default_sampling_params_list == [
+        stage0_client_r0.default_sampling_params,
+        stage1_client_r0.default_sampling_params,
+    ]
+    assert engine.stage_metadata == [
+        {"final_output": False, "final_output_type": None, "stage_type": "llm"},
+        {"final_output": True, "final_output_type": None, "stage_type": "llm"},
+    ]
+
+
+def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.async_chunk = False
+    engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
+    engine.stage_configs = [
+        types.SimpleNamespace(stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0")),
+        types.SimpleNamespace(
+            stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3")
+        ),
+    ]
+
+    metadata_by_stage = {
+        0: _make_llm_metadata(0),
+        1: _make_llm_metadata(1),
+    }
+
+    monkeypatch.setattr(
+        engine_mod,
+        "extract_stage_metadata",
+        lambda cfg: types.SimpleNamespace(**metadata_by_stage[cfg.stage_id].__dict__),
+    )
+    monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
+    monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
+    monkeypatch.setattr(
+        engine_mod,
+        "build_vllm_config",
+        lambda stage_cfg, *_args, **_kwargs: (types.SimpleNamespace(tag=f"cfg-{stage_cfg.stage_id}"), object),
+    )
+
+    stage_plans, prompt_expand_func = engine._build_logical_stage_init_plans(
+        omni_transfer_config=None,
+        replicas_per_stage=[1, 3],
+        replica_devices_map={1: ["1", "2", "3"]},
+    )
+
+    assert prompt_expand_func is None
+    assert [plan.configured_stage_id for plan in stage_plans] == [0, 1]
+    assert [replica.stage_cfg.runtime.devices for replica in stage_plans[1].replicas] == ["1", "2", "3"]
+    assert [replica.replica_id for replica in stage_plans[1].replicas] == [0, 1, 2]
+    assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas)
+
+
+def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
+    engine = object.__new__(AsyncOmniEngine)
+
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [
+        _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2),
+        _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, num_replicas=2),
+    ]
+
+    clients = {
+        (0, 0): types.SimpleNamespace(name="stage0-replica0"),
+        (0, 1): types.SimpleNamespace(name="stage0-replica1"),
+        (1, 0): types.SimpleNamespace(name="stage1-replica0"),
+        (1, 1): types.SimpleNamespace(name="stage1-replica1"),
+    }
+
+    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+        time.sleep(0.02 * (3 - plan.metadata.stage_id - plan.replica_id))
+        return clients[(plan.metadata.stage_id, plan.replica_id)]
+
+    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
+
+    initialized_clients = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=123)
+
+    assert initialized_clients == {
+        0: [clients[(0, 0)], clients[(0, 1)]],
+        1: [clients[(1, 0)], clients[(1, 1)]],
+    }
+
+
+def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
     engine = object.__new__(AsyncOmniEngine)
     engine.model = "dummy-model"
     engine.config_path = "dummy-config"
@@ -39,158 +426,111 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
     engine.single_stage_mode = False
     engine._single_stage_id_filter = None
     engine._omni_master_server = None
-    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")]
+    engine.stage_configs = [types.SimpleNamespace()]
 
-    env_var = current_omni_platform.device_control_env_var
-    old_env = os.environ.get(env_var)
-    os.environ[env_var] = "0,1"
-
-    diffusion_client = types.SimpleNamespace(is_comprehension=False)
-
-    metadata = types.SimpleNamespace(
-        stage_id=0,
-        stage_type="diffusion",
-        runtime_cfg={"devices": "1"},
-        prompt_expand_func=None,
-    )
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)]
+    initialized_client = types.SimpleNamespace(shutdown=lambda: None)
 
     monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
     monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
-    monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
-    monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {}))
+    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
 
-    def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
-        # Simulate diffusion setup mutating process-global visibility.
-        current_omni_platform.set_device_control_env_var("1")
+    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+        if plan.replica_id == 0:
+            return initialized_client
+        time.sleep(0.05)
+        raise RuntimeError("replica launch failed")
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
-    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: diffusion_client)
-    monkeypatch.setattr(
-        engine_mod,
-        "finalize_initialized_stages",
-        lambda stage_clients, _input_processor: (
-            stage_clients,
-            [types.SimpleNamespace()],
-            [{"final_output_type": "image"}],
-        ),
-    )
+    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
 
-    try:
+    captured_cleanup: list[list[object]] = []
+
+    def _capture_shutdown(clients):
+        captured_cleanup.append(list(clients))
+
+    monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
+
+    with pytest.raises(RuntimeError, match="replica launch failed"):
         engine._initialize_stages(stage_init_timeout=1)
-        assert os.environ.get(env_var) == "0,1"
-    finally:
-        if old_env is None:
-            os.environ.pop(env_var, None)
-        else:
-            os.environ[env_var] = old_env
+
+    assert captured_cleanup == [[initialized_client]]
 
 
-def test_initialize_stages_passes_stage_init_timeout_to_diffusion_handshake(monkeypatch):
-    """Regression test for stage_init_timeout passing to complete_diffusion_handshake
-    in the diffusion stage path.
-    """
-    import vllm_omni.diffusion.data as diffusion_data_mod
-    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_replica_failure(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
-    from vllm_omni.platforms import current_omni_platform
 
     engine = object.__new__(AsyncOmniEngine)
-    engine.log_stats = False
     engine.model = "dummy-model"
     engine.config_path = "dummy-config"
-    engine.num_stages = 2
+    engine.num_stages = 1
     engine.async_chunk = False
     engine.diffusion_batch_size = 1
     engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
     engine._omni_master_server = None
-    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion", engine_args={})]
+    engine.stage_configs = [types.SimpleNamespace()]
 
-    metadata = types.SimpleNamespace(
-        stage_id=0,
-        stage_type="diffusion",
-        runtime_cfg={"devices": "0"},
-        prompt_expand_func=None,
-        final_output=True,
-        final_output_type="image",
-        default_sampling_params=None,
-        custom_process_input_func=None,
-        engine_input_source=None,
-        cfg_kv_collect_func=None,
-    )
-
-    captured_timeout = None
-    device_env_var = current_omni_platform.device_control_env_var
-    prev_device_env = os.environ.get(device_env_var)
-    os.environ[device_env_var] = "0"
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)]
+    initialized_client = types.SimpleNamespace(shutdown=lambda: None)
+    replica1_started = threading.Event()
 
     monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
     monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
-    monkeypatch.setattr(
-        engine_mod,
-        "finalize_initialized_stages",
-        lambda stage_clients, _input_processor: (
-            stage_clients,
-            [types.SimpleNamespace()],
-            [{"final_output_type": "image"}],
-        ),
-    )
-    monkeypatch.setattr(
-        diffusion_data_mod.OmniDiffusionConfig,
-        "from_kwargs",
-        classmethod(lambda cls, **kwargs: types.SimpleNamespace(parallel_config=types.SimpleNamespace(world_size=1))),
-    )
-    monkeypatch.setattr(
-        client_mod,
-        "spawn_diffusion_proc",
-        lambda model, od_cfg: (object(), "ipc://handshake", "ipc://request", "ipc://response"),
-    )
+    monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {}))
+    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
 
-    def _capture_handshake_timeout(proc, handshake_address, handshake_timeout):
-        nonlocal captured_timeout
-        captured_timeout = handshake_timeout
+    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+        if plan.replica_id == 0:
+            assert replica1_started.wait(timeout=1)
+            raise RuntimeError("replica launch failed")
+        replica1_started.set()
+        time.sleep(0.05)
+        return initialized_client
 
-    monkeypatch.setattr(client_mod, "complete_diffusion_handshake", _capture_handshake_timeout)
-    monkeypatch.setattr(
-        client_mod.zmq,
-        "Context",
-        lambda: types.SimpleNamespace(socket=lambda _: types.SimpleNamespace(connect=lambda _: None)),
-    )
+    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
 
-    try:
-        engine._initialize_stages(stage_init_timeout=302)
-    finally:
-        if prev_device_env is None:
-            os.environ.pop(device_env_var, None)
-        else:
-            os.environ[device_env_var] = prev_device_env
+    captured_cleanup: list[list[object]] = []
 
-    assert captured_timeout == 302
+    def _capture_shutdown(clients):
+        captured_cleanup.append(list(clients))
+
+    monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
+
+    with pytest.raises(RuntimeError, match="replica launch failed"):
+        engine._initialize_stages(stage_init_timeout=1)
+
+    assert captured_cleanup == [[initialized_client]]
 
 
-def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
-    """Regression test for stage_init_timeout reaching complete_stage_handshake
-    in the LLM stage path.
-    """
+def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
     from vllm_omni.platforms import current_omni_platform
 
     engine = object.__new__(AsyncOmniEngine)
-    engine.log_stats = False
     engine.model = "dummy-model"
     engine.single_stage_mode = False
     engine._omni_master_server = None
     engine.stage_configs = []
 
-    metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"})
     fake_vllm_config = types.SimpleNamespace()
-    fake_addresses = types.SimpleNamespace()
+    fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None)
     fake_proc = types.SimpleNamespace()
+    captured_timeout: int | None = None
 
-    captured_timeout = None
+    plan = ReplicaInitPlan(
+        replica_id=0,
+        num_replicas=1,
+        launch_mode="local",
+        stage_cfg=types.SimpleNamespace(engine_args={}, runtime=types.SimpleNamespace(devices="0")),
+        metadata=types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}),
+        stage_connector_spec={},
+        omni_kv_connector=(None, None, None),
+        stage_vllm_config=fake_vllm_config,
+        executor_class=object,
+    )
 
     device_env_var = current_omni_platform.device_control_env_var
     prev_device_env = os.environ.get(device_env_var)
@@ -198,28 +538,22 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(
 
     monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
     monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
-    monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object))
     monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
-    monkeypatch.setattr(
-        engine_mod,
-        "spawn_stage_core",
-        lambda **_: (fake_addresses, fake_proc, "ipc://handshake"),
-    )
+    monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake"))
 
     def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout):
         nonlocal captured_timeout
         captured_timeout = handshake_timeout
 
     monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout)
+    monkeypatch.setattr(
+        engine_mod.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        staticmethod(lambda **_: types.SimpleNamespace(shutdown=lambda: None)),
+    )
 
     try:
-        engine._launch_llm_stage(
-            stage_cfg=types.SimpleNamespace(engine_args={}),
-            metadata=metadata,
-            stage_connector_spec={},
-            stage_init_timeout=302,
-            llm_stage_launch_lock=threading.Lock(),
-        )
+        engine._initialize_llm_replica(plan, 302, threading.Lock())
     finally:
         if prev_device_env is None:
             os.environ.pop(device_env_var, None)
@@ -229,155 +563,27 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(
     assert captured_timeout == 302
 
 
-def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(monkeypatch):
-    """Regression test for parallel LLM stage startup during handshake wait."""
-    import vllm_omni.engine.async_omni_engine as engine_mod
-    from vllm_omni.platforms import current_omni_platform
-
-    engine = object.__new__(AsyncOmniEngine)
-    engine.log_stats = False
-    engine.model = "dummy-model"
-    engine.single_stage_mode = False
-    engine._omni_master_server = None
-    engine.stage_configs = []
-
-    fake_vllm_config = types.SimpleNamespace()
-    fake_addresses = types.SimpleNamespace()
-    shared_launch_lock = threading.Lock()
-    counter_lock = threading.Lock()
-    first_handshake_started = threading.Event()
-    second_stage_spawned = threading.Event()
-    allow_first_handshake_to_finish = threading.Event()
-    launch_errors: list[BaseException] = []
-    spawn_count = 0
-
-    device_env_var = current_omni_platform.device_control_env_var
-    prev_device_env = os.environ.get(device_env_var)
-    os.environ[device_env_var] = "0"
-
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
-    monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object))
-    monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
-
-    def _spawn_stage_core(**_):
-        nonlocal spawn_count
-        with counter_lock:
-            spawn_count += 1
-            call_idx = spawn_count
-        if call_idx == 2:
-            second_stage_spawned.set()
-        return fake_addresses, types.SimpleNamespace(), f"ipc://handshake-{call_idx}"
-
-    def _complete_stage_handshake(_proc, handshake_address, _addresses, _vllm_cfg, _timeout):
-        if handshake_address == "ipc://handshake-1":
-            first_handshake_started.set()
-            assert second_stage_spawned.wait(timeout=1), (
-                "second stage did not reach spawn_stage_core while first stage waited in handshake"
-            )
-            assert allow_first_handshake_to_finish.wait(timeout=1), (
-                "second stage did not enter handshake while first stage was still waiting"
-            )
-        else:
-            allow_first_handshake_to_finish.set()
-
-    monkeypatch.setattr(engine_mod, "spawn_stage_core", _spawn_stage_core)
-    monkeypatch.setattr(engine_mod, "complete_stage_handshake", _complete_stage_handshake)
-
-    def _launch_stage(stage_id: int) -> None:
-        metadata = types.SimpleNamespace(stage_id=stage_id, runtime_cfg={"devices": str(stage_id)})
-        try:
-            engine._launch_llm_stage(
-                stage_cfg=types.SimpleNamespace(engine_args={}),
-                metadata=metadata,
-                stage_connector_spec={},
-                stage_init_timeout=302,
-                llm_stage_launch_lock=shared_launch_lock,
-            )
-        except BaseException as exc:  # pragma: no cover - surfaced through assertion below
-            launch_errors.append(exc)
-
-    try:
-        first_thread = threading.Thread(target=_launch_stage, args=(0,))
-        first_thread.start()
-        assert first_handshake_started.wait(timeout=1), "first stage never entered handshake"
-
-        second_thread = threading.Thread(target=_launch_stage, args=(1,))
-        second_thread.start()
-
-        first_thread.join(timeout=3)
-        second_thread.join(timeout=3)
-    finally:
-        if prev_device_env is None:
-            os.environ.pop(device_env_var, None)
-        else:
-            os.environ[device_env_var] = prev_device_env
-
-    assert not first_thread.is_alive()
-    assert not second_thread.is_alive()
-    assert second_stage_spawned.is_set()
-    assert not launch_errors
-
-
-def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
-    """Regression test for GLM-Image t2i preprocessing path.
-
-    Stage-0 InputProcessor must use OmniInputPreprocessor so text prompts with
-    mm_processor_kwargs go through multimodal preprocessing.
-    """
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    class DummyStageEngineCoreClient:
-        def __init__(self, **kwargs):
-            self.kwargs = kwargs
-
-        def shutdown(self):
-            return None
-
-    class DummyOutputProcessor:
-        def __init__(self, **kwargs):
-            self.kwargs = kwargs
+def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
+    import vllm_omni.engine.stage_init_utils as init_mod
 
     class DummyInputProcessor:
         def __init__(self, vllm_config):
             self.vllm_config = vllm_config
             self.renderer = object()
-            self.input_preprocessor = object()
+            self.input_preprocessor = None
 
     class DummyOmniInputPreprocessor:
         def __init__(self, vllm_config, renderer=None):
             self.vllm_config = vllm_config
             self.renderer = renderer
 
-    monkeypatch.setattr(
-        engine_mod.StageEngineCoreClientBase,
-        "make_async_mp_client",
-        staticmethod(lambda **kwargs: DummyStageEngineCoreClient(**kwargs)),
-    )
-    monkeypatch.setattr(engine_mod, "MultimodalOutputProcessor", DummyOutputProcessor)
-    monkeypatch.setattr(engine_mod, "InputProcessor", DummyInputProcessor)
-    monkeypatch.setattr(engine_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)
+    monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
+    monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)
 
-    started = types.SimpleNamespace(
-        stage_id=0,
-        metadata=types.SimpleNamespace(stage_id=0, engine_output_type="token_ids"),
-        vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(skip_tokenizer_init=True)),
-        executor_class=object,
-        engine_manager=object(),
-        coordinator=object(),
-        proc=None,
-        addresses=types.SimpleNamespace(
-            inputs=["inproc://input"],
-            outputs=["inproc://output"],
-            frontend_stats_publish_address=None,
-        ),
+    input_processor = build_stage0_input_processor(
+        types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {}))
     )
 
-    engine = object.__new__(AsyncOmniEngine)
-
-    _stage_client, _out_proc, _vllm_cfg, input_processor = engine._attach_llm_stage(started)
-
-    assert input_processor is not None
     assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor)
     assert input_processor.input_preprocessor.renderer is input_processor.renderer
 

tests/engine/test_cfg_companion_tracker.py

@@ -1,7 +1,6 @@
 import pytest
 
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -25,18 +24,6 @@ def test_register_companion_and_cleanup():
     assert tracker.get_companion_request_ids("req1") == {}
 
 
-def test_attach_cfg_request_ids_clones_diffusion_params():
-    tracker = CfgCompanionTracker()
-    tracker.register_companion("req1", "cfg_text", "req1__cfg_text")
-
-    params = OmniDiffusionSamplingParams()
-    updated = tracker.attach_cfg_request_ids("req1", params)
-
-    assert updated is not params
-    assert params.cfg_kv_request_ids is None
-    assert updated.cfg_kv_request_ids == {"cfg_text": "req1__cfg_text"}
-
-
 def test_abort_parent_expands_to_companions_and_cleans_up_deferred_parent():
     tracker = CfgCompanionTracker()
     tracker.register_companion("req1", "cfg_text", "req1__cfg_text")

tests/engine/test_orchestrator.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import logging
 import queue
 import threading
 import time
@@ -16,8 +17,9 @@ from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.core_client import AsyncMPClient
 
-from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -73,7 +75,7 @@ class FakeStageClient:
     def set_engine_outputs(self, outputs) -> None:
         return None
 
-    def process_engine_inputs(self, stage_list, prompt=None, streaming_context=None):
+    def process_engine_inputs(self, source_outputs, prompt=None, streaming_context=None):
         return list(self.next_inputs)
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
@@ -90,12 +92,33 @@ class FakeStageClient:
         self._diffusion_outputs.put_nowait(output)
 
 
+class FakeCollectiveRpcStageClient(FakeStageClient):
+    def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.rpc_result = rpc_result
+        self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = []
+
+    async def collective_rpc_async(
+        self,
+        *,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        normalized_kwargs = dict(kwargs or {})
+        self.collective_rpc_calls.append((method, timeout, args, normalized_kwargs))
+        return self.rpc_result
+
+
 class FakeOutputProcessor:
     def __init__(self, *, request_outputs: list[object] | None = None) -> None:
         self.request_outputs = list(request_outputs or [])
+        self.add_request_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
         self.abort_calls: list[list[str]] = []
 
-    def add_request(self, *_args, **_kwargs) -> None:
+    def add_request(self, *args, **kwargs) -> None:
+        self.add_request_calls.append((args, kwargs))
         return None
 
     def process_outputs(self, *_args, **_kwargs):
@@ -185,17 +208,60 @@ def _build_request_output(
     )
 
 
+def _build_stage_pools(
+    stage_clients: list[list[FakeStageClient]],
+    *,
+    output_processors: list[FakeOutputProcessor] | None = None,
+    stage_vllm_configs: list[object] | None = None,
+) -> list[StagePool]:
+    """Build StagePool list from per-stage replica lists.
+
+    ``stage_clients[i]`` is the list of FakeStageClient replicas for stage i.
+    """
+    num_stages = len(stage_clients)
+    if output_processors is None:
+        output_processors = [FakeOutputProcessor() for _ in stage_clients]
+    if stage_vllm_configs is None:
+        stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients]
+
+    pools: list[StagePool] = []
+    for stage_id in range(num_stages):
+        clients = stage_clients[stage_id]
+        if clients[0].stage_type == "diffusion":
+            pools.append(StagePool(stage_id, clients[0]))
+        else:
+            pools.append(
+                StagePool(
+                    stage_id,
+                    clients,
+                    output_processor=output_processors[stage_id],
+                    stage_vllm_config=stage_vllm_configs[stage_id],
+                )
+            )
+    return pools
+
+
 def _build_harness(
     stage_clients: list[object],
     *,
     output_processors: list[object] | None = None,
     stage_vllm_configs: list[object] | None = None,
     async_chunk: bool = False,
+    stage_pools: list[StagePool] | None = None,
 ) -> OrchestratorFixture:
-    if output_processors is None:
-        output_processors = [FakeOutputProcessor() for _ in stage_clients]
-    if stage_vllm_configs is None:
-        stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients]
+    """Build an Orchestrator test harness.
+
+    Accepts either pre-built ``stage_pools`` or flat lists of single-replica
+    clients/processors.
+    """
+    if stage_pools is None:
+        # Wrap flat lists into per-stage single-replica lists.
+        nested_clients = [[c] for c in stage_clients]
+        stage_pools = _build_stage_pools(
+            nested_clients,
+            output_processors=output_processors,
+            stage_vllm_configs=stage_vllm_configs,
+        )
 
     ready_future: concurrent.futures.Future[tuple[Orchestrator, janus.Queue, janus.Queue, janus.Queue]] = (
         concurrent.futures.Future()
@@ -214,9 +280,7 @@ def _build_harness(
                 request_async_queue=request_queue.async_q,
                 output_async_queue=output_queue.async_q,
                 rpc_async_queue=rpc_queue.async_q,
-                stage_clients=stage_clients,
-                output_processors=output_processors,
-                stage_vllm_configs=stage_vllm_configs,
+                stage_pools=stage_pools,
                 async_chunk=async_chunk,
             )
             ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
@@ -283,6 +347,18 @@ async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, time
             return msg
 
 
+async def _get_rpc_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+    deadline = time.monotonic() + timeout
+    rpc_sync_q = orchestrator_fixture.queues[2].sync_q
+    while True:
+        if time.monotonic() >= deadline:
+            raise AssertionError("Timed out waiting for orchestrator rpc output")
+        try:
+            return rpc_sync_q.get_nowait()
+        except queue.Empty:
+            await asyncio.sleep(0.01)
+
+
 async def _enqueue_add_request(
     orchestrator_fixture: OrchestratorFixture,
     *,
@@ -387,6 +463,11 @@ def orchestrator_factory():
             q.close()
 
 
+# ---------------------------------------------------------------------------
+# Existing single-replica tests (adapted to StagePool interface)
+# ---------------------------------------------------------------------------
+
+
 @pytest.mark.asyncio
 async def test_run_two_stage_llm(orchestrator_factory) -> None:
     stage0 = FakeStageClient(stage_type="llm", final_output=False)
@@ -600,12 +681,425 @@ async def test_run_abort(orchestrator_factory) -> None:
         await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
 
         await _enqueue_abort_request(orchestrator_fixture, ["req-abort"])
-        await _wait_for(lambda: all(stage.abort_calls for stage in stages))
+        await _wait_for(lambda: bool(stages[0].abort_calls))
 
-        for stage in stages:
-            assert stage.abort_calls == [["req-abort"]]
-        for proc in processors:
-            assert proc.abort_calls == [["req-abort"]]
+        assert stages[0].abort_calls == [["req-abort"]]
+        assert stages[1].abort_calls == []
         assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
+
+
+# ---------------------------------------------------------------------------
+# Multi-replica tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_round_robin_distribution(orchestrator_factory) -> None:
+    """Two replicas at stage-0, single replica at stage-1.
+
+    Send two requests — they should land on different stage-0 replicas
+    (round-robin), then both forward to the single stage-1 replica.
+    """
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(
+        stage_type="llm",
+        final_output=True,
+        next_inputs=[{"prompt_token_ids": [7, 8]}],
+    )
+
+    proc0 = FakeOutputProcessor(request_outputs=[_build_request_output("req-0", token_ids=[3], finished=True)])
+    proc1 = FakeOutputProcessor(request_outputs=[_build_request_output("req-0", token_ids=[10], finished=True)])
+
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1], [stage1]],
+        output_processors=[proc0, proc1],
+        stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg],
+    )
+
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        # Request 0 → should land on replica 0 (RR starts at 0)
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-0",
+            prompt=SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]),
+            original_prompt={"prompt": "hello 0"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1)
+        assert len(stage0_r1.add_request_calls) == 0
+
+        # Request 1 → should land on replica 1 (RR advances)
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-1",
+            prompt=SimpleNamespace(request_id="req-1", prompt_token_ids=[5, 6]),
+            original_prompt={"prompt": "hello 1"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1)
+        assert len(stage0_r0.add_request_calls) == 1  # unchanged
+
+        # Complete req-0 at stage-0 replica-0 → should forward to stage-1
+        stage0_r0.push_engine_core_outputs(_engine_core_outputs("s0r0-raw", 1.0))
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        assert stage1.add_request_calls[0][0].request_id == "req-0"
+
+        # Complete req-0 at stage-1 → final output
+        proc1.request_outputs = [_build_request_output("req-0", token_ids=[10], finished=True)]
+        stage1.push_engine_core_outputs(_engine_core_outputs("s1-raw", 2.0))
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-0"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert "req-0" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_abort_broadcasts_to_all_replicas(orchestrator_factory) -> None:
+    """Abort must be sent to every replica across all stages."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="llm", final_output=True)
+
+    proc0 = FakeOutputProcessor()
+    proc1 = FakeOutputProcessor()
+
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1], [stage1]],
+        output_processors=[proc0, proc1],
+        stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-abort-mr",
+            prompt=SimpleNamespace(request_id="req-abort-mr", prompt_token_ids=[1]),
+            original_prompt={"prompt": "cancel"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1)
+
+        await _enqueue_abort_request(orchestrator_fixture, ["req-abort-mr"])
+        await _wait_for(lambda: bool(stage0_r0.abort_calls))
+
+        assert stage0_r0.abort_calls == [["req-abort-mr"]]
+        assert stage0_r1.abort_calls == []
+        assert stage1.abort_calls == []
+        assert "req-abort-mr" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_shutdown_all_replicas(orchestrator_factory) -> None:
+    """Shutdown must shut down every replica across all stages."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="llm", final_output=True)
+
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1], [stage1]],
+        stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    await _shutdown_orchestrator(orchestrator_fixture)
+
+    assert not orchestrator_fixture.thread.is_alive()
+    for client in [stage0_r0, stage0_r1, stage1]:
+        assert client.shutdown_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_submit_update_reuses_existing_binding() -> None:
+    """A request admitted to one replica must keep using that replica on updates."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    pool = StagePool(
+        0,
+        [stage0_r0, stage0_r1],
+        output_processor=FakeOutputProcessor(),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+
+    req0_state = OrchestratorRequestState(
+        request_id="req-0",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+    req1_state = OrchestratorRequestState(
+        request_id="req-1",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+
+    await pool.submit_initial("req-0", req0_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]))
+    await pool.submit_update("req-0", req0_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[3]))
+    await pool.submit_initial("req-1", req1_state, SimpleNamespace(request_id="req-1", prompt_token_ids=[4, 5]))
+    await pool.submit_update("req-1", req1_state, SimpleNamespace(request_id="req-1", prompt_token_ids=[6]))
+
+    assert pool.get_bound_replica_id("req-0") == 0
+    assert pool.get_bound_replica_id("req-1") == 1
+    assert len(stage0_r0.add_request_calls) == 2
+    assert len(stage0_r1.add_request_calls) == 2
+    assert stage0_r0.add_request_calls[0][0].request_id == "req-0"
+    assert stage0_r0.add_request_calls[1][0].request_id == "req-0"
+    assert stage0_r1.add_request_calls[0][0].request_id == "req-1"
+    assert stage0_r1.add_request_calls[1][0].request_id == "req-1"
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_submit_update_refreshes_output_processor_state() -> None:
+    output_processor = FakeOutputProcessor()
+
+    class AssertingStageClient(FakeStageClient):
+        async def add_request_async(self, *args, **kwargs) -> None:
+            if len(self.add_request_calls) == 1:
+                prompts = [call_kwargs["prompt"] for _, call_kwargs in output_processor.add_request_calls]
+                assert prompts == ["seg-1", "seg-2"]
+            await super().add_request_async(*args, **kwargs)
+
+    stage0 = AssertingStageClient(stage_type="llm", final_output=False)
+    pool = StagePool(
+        0,
+        [stage0],
+        output_processor=output_processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    req_state = OrchestratorRequestState(
+        request_id="req-0",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+
+    await pool.submit_initial(
+        "req-0",
+        req_state,
+        SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]),
+        prompt_text="seg-1",
+    )
+    await pool.submit_update(
+        "req-0",
+        req_state,
+        SimpleNamespace(request_id="req-0", prompt_token_ids=[3], resumable=True),
+        prompt_text="seg-2",
+    )
+
+    assert len(output_processor.add_request_calls) == 2
+    assert output_processor.add_request_calls[1][1]["prompt"] == "seg-2"
+
+
+@pytest.mark.asyncio
+async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
+    class RecordingPool:
+        def __init__(self) -> None:
+            self.calls: list[tuple[str, Any]] = []
+
+        async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int:
+            self.calls.append((request_id, prompt_text))
+            return 0
+
+    pool = RecordingPool()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.request_states = {
+        "req-stream": OrchestratorRequestState(
+            request_id="req-stream",
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+    }
+    orchestrator.stage_pools = [pool]
+
+    await orchestrator._handle_streaming_update(
+        {
+            "request_id": "req-stream",
+            "prompt": SimpleNamespace(request_id="req-stream", prompt_token_ids=[1], resumable=True),
+            "sampling_params_list": [_sampling_params()],
+            "output_prompt_text": "segment-2",
+        }
+    )
+
+    assert pool.calls == [("req-stream", "segment-2")]
+    assert orchestrator.request_states["req-stream"].streaming.enabled is True
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None:
+    class FailingStageClient(FakeStageClient):
+        async def add_request_async(self, *args, **_kwargs) -> None:
+            raise RuntimeError("submit failed")
+
+    class TrackingOutputProcessor(FakeOutputProcessor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.added_request_ids: list[str] = []
+            self.removed_request_ids: list[str] = []
+
+        def add_request(self, request, *_args, **_kwargs) -> None:
+            self.added_request_ids.append(request.request_id)
+
+        def remove_request(self, request_id: str) -> None:
+            self.removed_request_ids.append(request_id)
+
+    client = FailingStageClient(stage_type="llm", final_output=False)
+    output_processor = TrackingOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=output_processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    req_state = OrchestratorRequestState(
+        request_id="req-0",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+
+    with pytest.raises(RuntimeError, match="submit failed"):
+        await pool.submit_initial("req-0", req_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]))
+
+    assert output_processor.added_request_ids == ["req-0"]
+    assert output_processor.removed_request_ids == ["req-0"]
+    assert pool.get_bound_replica_id("req-0") is None
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_abort_requests_logs_when_binding_is_missing(caplog) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    pool = StagePool(
+        0,
+        [stage0],
+        output_processor=FakeOutputProcessor(),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+
+    target_logger = logging.getLogger("vllm_omni.engine.stage_pool")
+    target_logger.addHandler(caplog.handler)
+    prev_level = target_logger.level
+    target_logger.setLevel(logging.DEBUG)
+    try:
+        await pool.abort_requests(["missing-req"])
+    finally:
+        target_logger.removeHandler(caplog.handler)
+        target_logger.setLevel(prev_level)
+
+    assert not stage0.abort_calls
+    assert "abort: no binding for req=missing-req in stage-0" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, caplog) -> None:
+    stage0 = FakeCollectiveRpcStageClient(stage_type="llm", final_output=True, rpc_result={"stage": 0})
+    stage1 = FakeCollectiveRpcStageClient(stage_type="llm", final_output=True, rpc_result={"stage": 1})
+    stage_pools = _build_stage_pools(
+        [[stage0], [stage1]],
+        output_processors=[FakeOutputProcessor(), FakeOutputProcessor()],
+        stage_vllm_configs=[
+            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+        ],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        target_logger = logging.getLogger("vllm_omni.engine.orchestrator")
+        target_logger.addHandler(caplog.handler)
+        prev_level = target_logger.level
+        target_logger.setLevel(logging.WARNING)
+        try:
+            orchestrator_fixture.request_sync_q.put_nowait(
+                {
+                    "type": "collective_rpc",
+                    "rpc_id": "rpc-1",
+                    "method": "list_loras",
+                    "stage_ids": [99, 1],
+                }
+            )
+
+            msg = await _get_rpc_message(orchestrator_fixture)
+        finally:
+            target_logger.removeHandler(caplog.handler)
+            target_logger.setLevel(prev_level)
+
+        assert msg["type"] == "collective_rpc_result"
+        assert msg["rpc_id"] == "rpc-1"
+        assert msg["stage_ids"] == [1]
+        assert msg["results"] == [{"stage": 1}]
+        assert not stage0.collective_rpc_calls
+        assert len(stage1.collective_rpc_calls) == 1
+        assert "collective_rpc: ignoring invalid stage_id 99" in caplog.text
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator_factory) -> None:
+    """CFG companions should be routed to the same stage-0 replica as their parent."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1]],
+        output_processors=[FakeOutputProcessor()],
+        stage_vllm_configs=[default_vllm_cfg],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        # Consume replica-0 first so the parent request binds to replica-1.
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="warmup",
+            prompt=SimpleNamespace(request_id="warmup", prompt_token_ids=[0]),
+            original_prompt={"prompt": "warmup"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+        await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1)
+
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="parent",
+            prompt=SimpleNamespace(request_id="parent", prompt_token_ids=[1, 2]),
+            original_prompt={"prompt": "parent"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+        await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1)
+
+        orchestrator_fixture.request_sync_q.put_nowait(
+            {
+                "type": "add_companion_request",
+                "companion_id": "parent-neg",
+                "parent_id": "parent",
+                "role": "negative",
+                "prompt": SimpleNamespace(request_id="parent-neg", prompt_token_ids=[9]),
+                "companion_prompt_text": {"prompt": "negative"},
+                "sampling_params_list": [_sampling_params()],
+            }
+        )
+        await _wait_for(lambda: len(stage0_r1.add_request_calls) == 2)
+
+        assert stage_pools[0].get_bound_replica_id("parent") == 1
+        assert stage_pools[0].get_bound_replica_id("parent-neg") == 1
+        assert len(stage0_r0.add_request_calls) == 1
+        assert stage0_r1.add_request_calls[0][0].request_id == "parent"
+        assert stage0_r1.add_request_calls[1][0].request_id == "parent-neg"
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)

tests/engine/test_orchestrator_error_handling.py

@@ -124,8 +124,8 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
 @pytest.mark.asyncio
 async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -> None:
     """When a diffusion stage returns an OmniRequestOutput with a non-None
-    error, the orchestrator must route it as a finished output message and
-    clean up the request state.
+    error, the orchestrator must route it as an error message and clean up
+    the request state.
     """
     stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
     orchestrator_fixture = orchestrator_factory([stage0])
@@ -148,10 +148,10 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -
 
         msg = await _get_any_output_message(orchestrator_fixture)
 
-        assert msg["type"] == "output"
+        assert msg["type"] == "error"
         assert msg["request_id"] == "req-err"
-        assert msg["finished"] is True
-        assert msg["engine_outputs"].error == "gpu fault"
+        assert msg["stage_id"] == 0
+        assert msg["error"] == "gpu fault"
 
         # Request state should be cleaned up.
         await _wait_for(lambda: "req-err" not in orchestrator_fixture.orchestrator.request_states)

tests/engine/test_orchestrator_kv_sender_info.py

@@ -7,18 +7,18 @@ from vllm import SamplingParams
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummySenderStage:
+    stage_type = "llm"
+    final_output = False
+
     def __init__(self, sender_info):
         self._sender_info = sender_info
-        self.engine_outputs = None
-
-    def set_engine_outputs(self, outputs):
-        self.engine_outputs = outputs
 
     def get_kv_sender_info(self):
         return self._sender_info
@@ -26,6 +26,7 @@ class _DummySenderStage:
 
 class _DummyDiffusionStage:
     stage_type = "diffusion"
+    final_output = True
     custom_process_input_func = None
 
     def __init__(self, engine_input_source=None):
@@ -43,6 +44,15 @@ class _DummyDiffusionStage:
         )
 
 
+def _build_sender_pool(stage_id: int, sender_info: dict[str, object]) -> StagePool:
+    return StagePool(
+        stage_id,
+        _DummySenderStage(sender_info),
+        output_processor=object(),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+
+
 def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
     client = object.__new__(StageEngineCoreClient)
     client.stage_id = 0
@@ -126,14 +136,13 @@ def test_stage_engine_core_client_preserves_explicit_loopback_sender_host():
 
 def test_forward_to_diffusion_attaches_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
-    sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+    sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
+    diffusion_pool = StagePool(1, diffusion_stage)
 
     orchestrator.num_stages = 2
-    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.stage_pools = [sender_pool, diffusion_pool]
     orchestrator._cfg_tracker = CfgCompanionTracker()
-    orchestrator.stage_vllm_configs = [None, None]
-    orchestrator.output_processors = [None, None]
 
     params = OmniDiffusionSamplingParams()
     req_state = OrchestratorRequestState(
@@ -144,9 +153,8 @@ def test_forward_to_diffusion_attaches_kv_sender_info():
     )
 
     output = SimpleNamespace(request_id="req-1", finished=True)
-    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state))
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", sender_pool.stage_id, output, req_state))
 
-    assert sender_stage.engine_outputs == [output]
     assert diffusion_stage.calls[0]["request_id"] == "req-1"
     assert diffusion_stage.calls[0]["kv_sender_info"] == {
         0: {"host": "10.0.0.2", "zmq_port": 50151},
@@ -156,15 +164,14 @@ def test_forward_to_diffusion_attaches_kv_sender_info():
 
 def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
-    source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
-    previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+    source_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
+    previous_pool = _build_sender_pool(1, {"host": "10.0.0.9", "zmq_port": 59999})
+    diffusion_pool = StagePool(2, diffusion_stage)
 
     orchestrator.num_stages = 3
-    orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage]
+    orchestrator.stage_pools = [source_pool, previous_pool, diffusion_pool]
     orchestrator._cfg_tracker = CfgCompanionTracker()
-    orchestrator.stage_vllm_configs = [None, None, None]
-    orchestrator.output_processors = [None, None, None]
 
     params = OmniDiffusionSamplingParams()
     req_state = OrchestratorRequestState(
@@ -175,9 +182,8 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
     )
 
     output = SimpleNamespace(request_id="req-3", finished=True)
-    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state))
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", previous_pool.stage_id, output, req_state))
 
-    assert previous_stage.engine_outputs == [output]
     assert diffusion_stage.calls[0]["kv_sender_info"] == {
         0: {"host": "10.0.0.2", "zmq_port": 50151},
     }
@@ -185,9 +191,10 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
 
 def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
     orchestrator = object.__new__(Orchestrator)
-    sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
     diffusion_stage.custom_process_input_func = lambda *_args, **_kwargs: []
+    sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
+    diffusion_pool = StagePool(1, diffusion_stage)
 
     class _AsyncQueue:
         def __init__(self):
@@ -197,10 +204,8 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
             self.items.append(item)
 
     orchestrator.num_stages = 2
-    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.stage_pools = [sender_pool, diffusion_pool]
     orchestrator._cfg_tracker = CfgCompanionTracker()
-    orchestrator.stage_vllm_configs = [None, None]
-    orchestrator.output_processors = [None, None]
     orchestrator.output_async_queue = _AsyncQueue()
     orchestrator.request_states = {}
     orchestrator._pd_kv_params = {}
@@ -217,7 +222,6 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
     output = SimpleNamespace(request_id="req-empty", finished=True)
     asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-empty", 0, output, req_state))
 
-    assert sender_stage.engine_outputs == [output]
     assert diffusion_stage.calls == []
     assert len(orchestrator.output_async_queue.items) == 1
     terminal_msg = orchestrator.output_async_queue.items[0]
@@ -231,10 +235,11 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
 
 def test_prewarm_diffusion_attaches_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
-    sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+    sender_pool = _build_sender_pool(0, {"host": "10.0.0.3", "zmq_port": 50151})
+    diffusion_pool = StagePool(1, diffusion_stage)
 
-    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.stage_pools = [sender_pool, diffusion_pool]
     orchestrator.num_stages = 2
 
     req_state = OrchestratorRequestState(