[Refactor] Omni Stage Runtime and Distributed Replica Control Plane (#3855)

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Chenguang ZHENG <645327136@qq.com>

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -625,11 +625,11 @@ class TestStageDiffusionClientErrorPropagation:
         client._tasks = {}
         client._shutting_down = False
         client._engine_dead = engine_dead
-        client._owns_process = True
-        client._proc = MagicMock(
+        proc = MagicMock(
             is_alive=MagicMock(return_value=proc_alive),
             exitcode=1,
         )
+        client._proc_manager = SimpleNamespace(proc=proc)
         client._request_socket = MagicMock()
         client._response_socket = MagicMock()
         client._encoder = MagicMock()
@@ -672,7 +672,7 @@ class TestStageDiffusionClientErrorPropagation:
         assert client.get_diffusion_output_nowait() is None
 
     def test_check_health_raises_when_proc_dead(self):
-        """``check_health`` detects a dead subprocess via ``_proc.is_alive()``
+        """``check_health`` detects a dead subprocess via the manager's proc
         and raises ``EngineDeadError``, setting ``_engine_dead`` as a
         side effect."""
         client = self._make_client(proc_alive=False)
@@ -699,7 +699,7 @@ class TestStageDiffusionClientErrorPropagation:
         ``get_diffusion_output_nowait`` returns ``None`` and sets
         ``_shutting_down`` instead of raising."""
         client = self._make_client(proc_alive=False)
-        client._proc.exitcode = 137  # SIGKILL (128 + 9)
+        client._proc_manager.proc.exitcode = 137  # SIGKILL (128 + 9)
         client._response_socket.recv.side_effect = zmq.Again
 
         result = client.get_diffusion_output_nowait()
@@ -727,7 +727,6 @@ class TestStageDiffusionClientErrorPropagation:
                 metadata,
                 "tcp://req",
                 "tcp://resp",
-                proc=None,
                 batch_size=1,
             )
 
@@ -836,7 +835,7 @@ class TestStageDiffusionClientProcMonitor:
         client._engine_dead = False
 
         proc = _make_short_lived_process()
-        client._proc = proc
+        client._proc_manager = SimpleNamespace(proc=proc)
 
         client._start_proc_monitor()
         proc.join(5)

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
+import threading
 import time
 
 import pytest
@@ -48,6 +49,33 @@ def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None:
         _recv_replica_list(sub, timeout_ms=50)
 
 
+def test_omni_coordinator_wait_for_shutdown_unblocks_on_close():
+    router_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    pub_addr = get_engine_client_zmq_addr(
+        local_only=False,
+        host="127.0.0.1",
+        port=0,
+    )
+    coordinator = OmniCoordinator(
+        router_zmq_addr=router_addr,
+        pub_zmq_addr=pub_addr,
+        heartbeat_timeout=1000.0,
+    )
+
+    waiter = threading.Thread(target=coordinator.wait_for_shutdown)
+    waiter.start()
+    time.sleep(0.05)
+    assert waiter.is_alive()
+
+    coordinator.close()
+    waiter.join(timeout=1.0)
+    assert not waiter.is_alive()
+
+
 def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
     """Rapid updates should be coalesced into fewer PUB messages."""
     router_addr = get_engine_client_zmq_addr(

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,21 +1,20 @@
-import concurrent.futures
+import contextlib
 import importlib
 import os
-import threading
 import time
 import types
-from contextlib import contextmanager
 
 import pytest
 
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
+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,
 )
+from vllm_omni.engine.stage_runtime import StageRuntime
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -63,7 +62,7 @@ def _make_diffusion_metadata(stage_id: int, *, replica_id: int = 0, final_output
 def _make_llm_plan(
     stage_idx: int,
     *,
-    configured_stage_id: int,
+    stage_id: int,
     vllm_config: object,
     num_replicas: int = 1,
     final_output: bool = False,
@@ -73,7 +72,7 @@ def _make_llm_plan(
     replicas: list[ReplicaInitPlan] = []
     for replica_id in range(num_replicas):
         stage_cfg = types.SimpleNamespace(
-            stage_id=configured_stage_id,
+            stage_id=stage_id,
             stage_type="llm",
             runtime=types.SimpleNamespace(devices=str(replica_id)),
             engine_args={},
@@ -85,7 +84,7 @@ def _make_llm_plan(
                 launch_mode="local",
                 stage_cfg=stage_cfg,
                 metadata=_make_llm_metadata(
-                    configured_stage_id,
+                    stage_id,
                     replica_id=replica_id,
                     final_output=final_output,
                     final_output_type=final_output_type,
@@ -99,7 +98,7 @@ def _make_llm_plan(
         )
     return LogicalStageInitPlan(
         stage_idx=stage_idx,
-        configured_stage_id=configured_stage_id,
+        stage_id=stage_id,
         replicas=replicas,
     )
 
@@ -107,13 +106,13 @@ def _make_llm_plan(
 def _make_diffusion_plan(
     stage_idx: int,
     *,
-    configured_stage_id: int,
+    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_id=stage_id,
             stage_type="diffusion",
             runtime=types.SimpleNamespace(devices=str(replica_id)),
             engine_args={},
@@ -124,14 +123,14 @@ def _make_diffusion_plan(
                 num_replicas=num_replicas,
                 launch_mode="local",
                 stage_cfg=stage_cfg,
-                metadata=_make_diffusion_metadata(configured_stage_id, replica_id=replica_id),
+                metadata=_make_diffusion_metadata(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,
+        stage_id=stage_id,
         replicas=replicas,
     )
 
@@ -165,7 +164,7 @@ 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(
+    cleanup_clients = StageRuntime._collect_initialized_clients_for_cleanup(
         stage_pools=[types.SimpleNamespace(clients=[shared, None])],
         initialized_clients_by_stage={0: [shared], 1: [extra]},
     )
@@ -173,35 +172,36 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
     assert cleanup_clients == [shared, extra]
 
 
-def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
+def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+    from vllm_omni.engine.stage_engine_startup import StageReplicaResources
     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 = []
+    runtime = StageRuntime(
+        stage_configs=[],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
 
-    plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
+    plan = _make_diffusion_plan(0, 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"
+    runtime._init_visible_devices_baseline = "0,1"
 
-    @contextmanager
-    def _fake_stage_runtime_setup(_stage_id, _runtime_cfg):
-        current_omni_platform.set_device_control_env_var("1")
-        yield
-
-    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _fake_stage_runtime_setup)
-    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace())
+    monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
+    monkeypatch.setattr(
+        runtime_mod,
+        "launch_diffusion_stage_replica",
+        lambda **_: (types.SimpleNamespace(), StageReplicaResources()),
+    )
 
     try:
-        engine._initialize_diffusion_replica(plan, stage_init_timeout=1, stage_launch_lock=threading.Lock())
+        runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1)
         assert os.environ.get(env_var) == "0,1"
     finally:
         if old_env is None:
@@ -210,68 +210,63 @@ def test_initialize_diffusion_replica_restores_device_visibility_after_local_ini
             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
+def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+    from vllm_omni.engine.stage_engine_startup import StageReplicaResources
 
-    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 = []
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=4,
+        async_chunk=False,
+    )
 
-    plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
+    plan = _make_diffusion_plan(0, stage_id=0).replicas[0]
 
     captured: dict[str, object] = {}
 
-    @contextmanager
-    def _noop_stage_runtime_setup(*_):
-        yield
+    monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
 
-    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
-    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
+    def _capture_launch_diffusion_stage_replica(**kwargs):
+        captured["stage_id"] = kwargs["metadata"].stage_id
+        captured["stage_init_timeout"] = kwargs["stage_init_timeout"]
+        captured["batch_size"] = kwargs["batch_size"]
+        captured["use_inline"] = kwargs["use_inline"]
+        captured["omni_master_server"] = kwargs["omni_master_server"]
+        return types.SimpleNamespace(), StageReplicaResources()
 
-    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(runtime_mod, "launch_diffusion_stage_replica", _capture_launch_diffusion_stage_replica)
 
-    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())
+    runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=302)
 
     assert captured == {
         "stage_id": 0,
         "stage_init_timeout": 302,
         "batch_size": 4,
         "use_inline": True,
+        "omni_master_server": None,
     }
 
 
-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
+def test_stage_runtime_initializes_stage_pools(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_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()]
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace(), types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
 
     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),
+        _make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True),
+        _make_llm_plan(1, stage_id=1, vllm_config=cfg1, final_output=True),
     ]
 
     stage0_client_r0 = types.SimpleNamespace(
@@ -302,51 +297,43 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu
 
     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(runtime, "_prepare_stage_plans", lambda: stage_plans)
+    monkeypatch.setattr(runtime, "_initialize_stage_replicas", lambda *_: initialized_clients)
     monkeypatch.setattr(
-        engine_mod,
+        runtime_mod,
         "build_llm_stage_output_processor",
         lambda plan, _cfg, **_kw: 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)
+    runtime.initialize()
 
-    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 == [
-        StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm"),
-        StageRuntimeInfo(final_output=True, final_output_type=None, stage_type="llm"),
-    ]
+    assert len(runtime.stage_pools) == 2
+    assert runtime.stage_pools[0].stage_client is stage0_client_r0
+    assert runtime.stage_pools[1].stage_client is stage1_client_r0
+    assert runtime.stage_pools[0].stage_vllm_config is cfg0
+    assert runtime.stage_pools[1].stage_vllm_config is cfg1
+    assert runtime.stage_pools[0].output_processor is stage0_output_processor
+    assert runtime.stage_pools[1].output_processor is stage1_output_processor
 
 
 def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
+    import vllm_omni.engine.stage_runtime as runtime_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")
-        ),
-    ]
+    runtime = StageRuntime(
+        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")
+            ),
+        ],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
 
     metadata_by_stage = {
         0: _make_llm_metadata(0),
@@ -354,78 +341,46 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
     }
 
     monkeypatch.setattr(
-        engine_mod,
+        runtime_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(runtime_mod, "get_stage_connector_spec", lambda **_: {})
+    monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {})
     monkeypatch.setattr(
-        engine_mod,
+        runtime_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(
+    stage_plans = runtime._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 [plan.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_reuses_long_lived_executor(monkeypatch):
-    engine = object.__new__(AsyncOmniEngine)
-    engine._stage_init_executor = None
-
-    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
-    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)]
-
-    created_executors = []
-
-    class FakeExecutor:
-        def __init__(self, *args, **kwargs):
-            self.shutdown_calls = []
-            created_executors.append(self)
-
-        def submit(self, fn, *args, **kwargs):
-            future = concurrent.futures.Future()
-            future.set_result(fn(*args, **kwargs))
-            return future
-
-        def shutdown(self, wait=False, cancel_futures=False):
-            self.shutdown_calls.append((wait, cancel_futures))
-
-    monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor)
-    monkeypatch.setattr(
-        engine,
-        "_initialize_replica",
-        lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}",
-    )
-
-    first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
-    second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
-
-    assert first == {0: ["client-0"]}
-    assert second == {0: ["client-0"]}
-    assert len(created_executors) == 1
-    assert engine._stage_init_executor is created_executors[0]
-
-
 def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
-    engine = object.__new__(AsyncOmniEngine)
+    runtime = StageRuntime(
+        stage_configs=[],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=123,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
 
     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),
+        _make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2),
+        _make_llm_plan(1, stage_id=1, vllm_config=cfg1, num_replicas=2),
     ]
 
     clients = {
@@ -435,13 +390,13 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
         (1, 1): types.SimpleNamespace(name="stage1-replica1"),
     }
 
-    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+    def _initialize_replica(plan, _stage_init_timeout):
         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)
+    monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica)
 
-    initialized_clients = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=123)
+    initialized_clients = runtime._initialize_stage_replicas(stage_plans, stage_init_timeout=123)
 
     assert initialized_clients == {
         0: [clients[(0, 0)], clients[(0, 1)]],
@@ -449,148 +404,120 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
     }
 
 
-def test_shutdown_closes_long_lived_stage_init_executor():
-    engine = object.__new__(AsyncOmniEngine)
-    engine._shutdown_called = False
-    engine._weak_finalizer = None
-    engine._omni_master_server = None
-    engine._coordinator_runtime = None
-    engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False)
+def test_remote_replicas_use_distinct_init_group_keys():
+    runtime = StageRuntime(
+        stage_configs=[],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=123,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
+    plan = _make_llm_plan(
+        1,
+        stage_id=1,
+        vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)),
+        num_replicas=2,
+    )
 
-    class FakeQueue:
-        def __init__(self):
-            self.closed = False
-            self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None)
+    for replica in plan.replicas:
+        replica.launch_mode = "remote"
+        replica.metadata.runtime_cfg = None
 
-        def close(self):
-            self.closed = True
-
-    class FakeExecutor:
-        def __init__(self):
-            self.shutdown_calls = []
-
-        def shutdown(self, wait=False, cancel_futures=False):
-            self.shutdown_calls.append((wait, cancel_futures))
-
-    engine.request_queue = FakeQueue()
-    engine.output_queue = FakeQueue()
-    engine.rpc_output_queue = FakeQueue()
-    engine._stage_init_executor = FakeExecutor()
-
-    engine.shutdown()
-
-    assert engine._stage_init_executor is None
-    assert engine.request_queue.closed is True
-    assert engine.output_queue.closed is True
-    assert engine.rpc_output_queue.closed is True
+    assert [runtime._replica_init_group_key(replica) for replica in plan.replicas] == [
+        "remote:1:0",
+        "remote: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"
-    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._coordinator_runtime = None
-    engine.stage_configs = [types.SimpleNamespace()]
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
 
     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)]
+    stage_plans = [_make_llm_plan(0, 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, "compute_replica_layout", lambda _cfgs: ([2], {}))
-    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
 
-    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+    def _initialize_replica(plan, _stage_init_timeout):
         if plan.replica_id == 0:
             return initialized_client
         time.sleep(0.05)
         raise RuntimeError("replica launch failed")
 
-    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
+    monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica)
 
     captured_cleanup: list[list[object]] = []
 
     def _capture_shutdown(clients):
         captured_cleanup.append(list(clients))
 
-    monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
+    monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown)
 
     with pytest.raises(RuntimeError, match="replica launch failed"):
-        engine._initialize_stages(stage_init_timeout=1)
+        runtime.initialize()
 
     assert captured_cleanup == [[initialized_client]]
 
 
 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
-
-    engine = object.__new__(AsyncOmniEngine)
-    engine.model = "dummy-model"
-    engine.config_path = "dummy-config"
-    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._coordinator_runtime = None
-    engine.stage_configs = [types.SimpleNamespace()]
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
 
     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)]
+    stage_plans = [_make_llm_plan(0, 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, "compute_replica_layout", lambda _cfgs: ([2], {}))
-    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
 
-    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
+    def _initialize_stage_replicas(_stage_plans, _stage_init_timeout):
+        exc = RuntimeError("replica launch failed")
+        exc._initialized_clients_by_stage = {0: [None, initialized_client]}
+        raise exc
 
-    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
+    monkeypatch.setattr(runtime, "_initialize_stage_replicas", _initialize_stage_replicas)
 
     captured_cleanup: list[list[object]] = []
 
     def _capture_shutdown(clients):
         captured_cleanup.append(list(clients))
 
-    monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
+    monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown)
 
     with pytest.raises(RuntimeError, match="replica launch failed"):
-        engine._initialize_stages(stage_init_timeout=1)
+        runtime.initialize()
 
     assert captured_cleanup == [[initialized_client]]
 
 
-def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
-    import vllm_omni.engine.async_omni_engine as engine_mod
+def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
     from vllm_omni.platforms import current_omni_platform
 
-    engine = object.__new__(AsyncOmniEngine)
-    engine.model = "dummy-model"
-    engine.single_stage_mode = False
-    engine._omni_master_server = None
-    engine.stage_configs = []
+    runtime = StageRuntime(
+        stage_configs=[],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=302,
+        diffusion_batch_size=1,
+        async_chunk=False,
+    )
 
     fake_vllm_config = types.SimpleNamespace()
     fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None)
-    fake_proc = types.SimpleNamespace()
     captured_timeout: int | None = None
 
     plan = ReplicaInitPlan(
@@ -603,34 +530,38 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
         omni_kv_connector=(None, None, None),
         stage_vllm_config=fake_vllm_config,
         executor_class=object,
+        engine_args_dict={},
     )
 
     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"
 
-    @contextmanager
-    def _noop_stage_runtime_setup(*_):
-        yield
-
-    monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
-    monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
-    monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
-    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):
+    def _capture_acquire_device_locks(*_args):
         nonlocal captured_timeout
-        captured_timeout = handshake_timeout
+        captured_timeout = _args[2]
+        return []
 
-    monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout)
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", _capture_acquire_device_locks)
+
+    from vllm_omni.engine.stage_engine_startup import StageReplicaResources
+
+    @contextlib.contextmanager
+    def _fake_launch_stage_replica(**_kwargs):
+        yield StageReplicaResources(
+            manager=types.SimpleNamespace(shutdown=lambda: None),
+            addresses=fake_addresses,
+        )
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
     monkeypatch.setattr(
-        engine_mod.StageEngineCoreClientBase,
+        runtime_mod.StageEngineCoreClientBase,
         "make_async_mp_client",
         staticmethod(lambda **_: types.SimpleNamespace(shutdown=lambda: None)),
     )
 
     try:
-        engine._initialize_llm_replica(plan, 302, threading.Lock())
+        runtime._initialize_local_llm_replica(plan, 302)
     finally:
         if prev_device_env is None:
             os.environ.pop(device_env_var, None)

tests/engine/test_membership_controller.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus
+from vllm_omni.engine.membership_controller import MembershipController
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeHub:
+    def __init__(self, snapshots=None):
+        self.snapshots = list(snapshots or [])
+        self.closed = False
+
+    def get_replica_list(self):
+        if self.snapshots:
+            return self.snapshots.pop(0)
+        return SimpleNamespace(replicas=[])
+
+    def close(self):
+        self.closed = True
+
+
+class FakePool:
+    def __init__(self, stage_id: int):
+        self.stage_id = stage_id
+        self.clients = []
+        self.added = []
+        self.removed = []
+        self.invalidated = []
+        self.hub = None
+        self.lb = None
+
+    def attach_hub(self, hub):
+        self.hub = hub
+
+    def attach_load_balancer(self, lb):
+        self.lb = lb
+
+    def add_client(self, input_addr, client):
+        self.added.append((input_addr, client))
+        self.clients.append(client)
+        return len(self.clients) - 1
+
+    def invalidate_addr(self, input_addr):
+        self.invalidated.append(input_addr)
+        return ["req-1", "req-2"]
+
+    def remove_client(self, input_addr):
+        self.removed.append(input_addr)
+        return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown"))
+
+
+def _snapshot(*replicas):
+    return SimpleNamespace(replicas=list(replicas))
+
+
+def _replica(stage_id: int, input_addr: str, status=ReplicaStatus.UP):
+    return SimpleNamespace(stage_id=stage_id, input_addr=input_addr, status=status)
+
+
+def _controller(monkeypatch, pool, hub, remote_replica_factory=None):
+    import vllm_omni.engine.membership_controller as membership_mod
+
+    monkeypatch.setattr(membership_mod, "OmniCoordClientForHub", lambda _addr: hub)
+    if remote_replica_factory is None:
+
+        def remote_replica_factory(stage_id, replica_id):
+            return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"})
+
+    return MembershipController(
+        stage_pools=[pool],
+        coordinator_pub_address="tcp://127.0.0.1:12345",
+        load_balancer_factory=lambda: object(),
+        remote_replica_factory=remote_replica_factory,
+    )
+
+
+@pytest.mark.asyncio
+async def test_watch_replica_list_unregisters_disappeared_replicas(monkeypatch):
+    pool = FakePool(stage_id=0)
+    hub = FakeHub(
+        snapshots=[
+            _snapshot(_replica(0, "tcp://gone")),
+            _snapshot(),
+        ]
+    )
+    controller = _controller(monkeypatch, pool, hub)
+    controller.WATCH_INTERVAL_S = 0
+    unregistered = []
+
+    async def _handle_unregister(stage_id, input_addr, *args, **kwargs):
+        unregistered.append((stage_id, input_addr))
+        controller._shutdown_event.set()
+
+    controller.handle_unregister = _handle_unregister  # type: ignore[method-assign]
+
+    await asyncio.wait_for(controller._watch_replica_list(), timeout=1)
+    await controller.drain_tasks(timeout=1)
+
+    assert unregistered == [(0, "tcp://gone")]
+
+
+@pytest.mark.asyncio
+async def test_shutdown_closes_hub_then_cancels_watcher(monkeypatch):
+    pool = FakePool(stage_id=0)
+    hub = FakeHub()
+    controller = _controller(monkeypatch, pool, hub)
+    watcher = asyncio.create_task(asyncio.sleep(10))
+    controller._watcher_task = watcher
+
+    controller.shutdown()
+    await asyncio.sleep(0)
+
+    assert hub.closed is True
+    assert watcher.cancelled()
+
+
+@pytest.mark.asyncio
+async def test_drain_tasks_waits_for_membership_tasks(monkeypatch):
+    pool = FakePool(stage_id=0)
+    controller = _controller(monkeypatch, pool, FakeHub())
+    completed = []
+
+    async def _task():
+        await asyncio.sleep(0)
+        completed.append(True)
+
+    controller._spawn_task(_task(), label="unit")
+    await controller.drain_tasks(timeout=1)
+
+    assert completed == [True]
+
+
+@pytest.mark.asyncio
+async def test_do_register_offloads_remote_factory(monkeypatch):
+    pool = FakePool(stage_id=0)
+    factory_thread_ids = []
+
+    def _factory(stage_id, replica_id):
+        factory_thread_ids.append(threading.get_ident())
+        return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"})
+
+    controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
+
+    await controller._do_register(0, 1)
+
+    assert factory_thread_ids
+    assert factory_thread_ids[0] != threading.get_ident()
+    assert pool.added[0][0] == "tcp://stage-0-replica-1"

tests/engine/test_orchestrator.py

@@ -11,11 +11,9 @@ from types import SimpleNamespace
 from typing import Any
 
 import janus
-import psutil
 import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
-from vllm.v1.engine.core_client import AsyncMPClient
 
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
@@ -27,7 +25,6 @@ from vllm_omni.engine.messages import (
     StageSubmissionMessage,
 )
 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
@@ -180,42 +177,6 @@ class FakeOutputProcessor:
         return None
 
 
-class _FakeProc:
-    pid = 1234
-
-    def __init__(self):
-        self.terminated = False
-        self.killed = False
-        self.join_calls = []
-
-    def is_alive(self):
-        return not self.terminated and not self.killed
-
-    def terminate(self):
-        self.terminated = True
-
-    def kill(self):
-        self.killed = True
-
-    def join(self, timeout=None):
-        self.join_calls.append(timeout)
-
-
-class _FakeChildProc:
-    def __init__(self):
-        self.terminated = False
-        self.killed = False
-
-    def is_running(self):
-        return not self.terminated and not self.killed
-
-    def terminate(self):
-        self.terminated = True
-
-    def kill(self):
-        self.killed = True
-
-
 def _sampling_params(max_tokens: int = 4) -> SamplingParams:
     return SamplingParams(max_tokens=max_tokens)
 
@@ -437,61 +398,6 @@ async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, requ
     orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids))
 
 
-def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch):
-    fake_proc = _FakeProc()
-    fake_child = _FakeChildProc()
-
-    class FakePsutilProcess:
-        def __init__(self, pid):
-            assert pid == fake_proc.pid
-
-        def children(self, recursive=True):
-            assert recursive
-            return [fake_child]
-
-    def fail_base_shutdown(self, **kwargs):
-        raise RuntimeError("base shutdown failed")
-
-    monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
-    monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: (list(procs), []))
-    monkeypatch.setattr(AsyncMPClient, "shutdown", fail_base_shutdown)
-
-    client = object.__new__(StageEngineCoreClient)
-    client._proc = fake_proc
-
-    with pytest.raises(RuntimeError, match="base shutdown failed"):
-        client.shutdown()
-
-    assert fake_proc.terminated
-    assert fake_proc.join_calls == [5]
-    assert fake_child.terminated
-
-
-def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch):
-    fake_proc = _FakeProc()
-    fake_child = _FakeChildProc()
-
-    class FakePsutilProcess:
-        def __init__(self, pid):
-            assert pid == fake_proc.pid
-
-        def children(self, recursive=True):
-            assert recursive
-            return [fake_child]
-
-    monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
-    monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs)))
-    monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None)
-
-    client = object.__new__(StageEngineCoreClient)
-    client._proc = fake_proc
-
-    client.shutdown()
-
-    assert fake_child.terminated
-    assert fake_child.killed
-
-
 @pytest.fixture
 def orchestrator_factory():
     fixtures: list[OrchestratorFixture] = []

tests/engine/test_stage_engine_core_client.py

@@ -1,15 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for StageEngineCoreClient.check_health().
-
-Uses object.__new__ to construct a minimal client — check_health only
-touches self.resources, self.stage_id, and self._proc.
-"""
+"""Tests for StageEngineCoreClient.check_health()."""
 
 from __future__ import annotations
 
 from types import SimpleNamespace
-from unittest.mock import MagicMock
 
 import pytest
 from vllm.v1.engine.exceptions import EngineDeadError
@@ -19,28 +14,19 @@ from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_client(*, engine_dead=False, proc_alive=True):
+def _make_client(*, engine_dead=False):
     client = object.__new__(StageEngineCoreClient)
     client.stage_id = 0
     client.resources = SimpleNamespace(engine_dead=engine_dead)
-    client._proc = MagicMock(is_alive=MagicMock(return_value=proc_alive), exitcode=1)
     return client
 
 
 def test_check_health_passes_when_alive():
-    client = _make_client(engine_dead=False, proc_alive=True)
+    client = _make_client(engine_dead=False)
     client.check_health()  # no exception
 
 
 def test_check_health_raises_when_resources_engine_dead():
-    client = _make_client(engine_dead=True, proc_alive=True)
+    client = _make_client(engine_dead=True)
     with pytest.raises(EngineDeadError, match="engine core is dead"):
         client.check_health()
-
-
-def test_check_health_raises_when_proc_not_alive():
-    client = _make_client(engine_dead=False, proc_alive=False)
-    with pytest.raises(EngineDeadError, match="not alive"):
-        client.check_health()
-    # Verify it set resources.engine_dead as a side effect
-    assert client.resources.engine_dead is True

tests/entrypoints/test_serve.py

@@ -119,7 +119,7 @@ def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace:
 def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: MockerFixture) -> None:
     """LLM headless: each loop iteration registers with auto-assigned
     replica_id (master picks a free slot) and spawns one
-    ``OmniCoreEngineProcManager`` per local replica."""
+    ``StageEngineCoreProcManager`` per local replica."""
     from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
 
     stage_cfg = _make_stage_cfg(0, stage_type="llm")
@@ -159,7 +159,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
         ),
     )
     mock_manager_cls = mocker.patch(
-        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        "vllm_omni.engine.stage_engine_core_proc_manager.StageEngineCoreProcManager",
         return_value=engine_manager,
     )
     mocker.patch("signal.signal")
@@ -168,8 +168,8 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
 
     # The launcher must request auto-assignment (replica_id=None) and the
     # full response so it can wire the master-allocated coordinator into the
-    # spawned subprocess. ``replica_binds_sockets=False`` is required for LLM
-    # because the head binds all three sockets (handshake, input, output).
+    # spawned subprocess. LLM uses head-owned sockets: the head binds all
+    # three sockets (handshake, input, output) and the worker connects.
     assert mock_register.call_count == 1
     kwargs = mock_register.call_args.kwargs
     assert kwargs["omni_master_address"] == "127.0.0.1"
@@ -177,8 +177,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
     assert kwargs["omni_stage_id"] == 0
     assert kwargs["omni_stage_config"] is stage_cfg
     assert kwargs["replica_id"] is None
-    assert kwargs["return_full_response"] is True
-    assert kwargs["replica_binds_sockets"] is False
+    assert "socket_ownership" not in kwargs
 
     assert mock_manager_cls.call_count == 1
     mgr_kwargs = mock_manager_cls.call_args.kwargs
@@ -239,7 +238,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo
         ],
     )
     mock_manager_cls = mocker.patch(
-        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        "vllm_omni.engine.stage_engine_core_proc_manager.StageEngineCoreProcManager",
         side_effect=[manager_a, manager_b],
     )
     mocker.patch("signal.signal")
@@ -294,18 +293,24 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
             coordinator_router_address="tcp://127.0.0.1:26100",
         ),
     )
-    mock_spawn = mocker.patch(
-        "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
-        return_value=(proc, None, None, None),
+    fake_manager = SimpleNamespace(
+        proc=proc,
+        addresses=SimpleNamespace(
+            inputs=["tcp://127.0.0.1:26002"],
+            outputs=["tcp://127.0.0.1:26003"],
+        ),
+        shutdown=mocker.Mock(),
+    )
+    mock_manager = mocker.patch(
+        "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager.launch_headless",
+        return_value=fake_manager,
     )
-    mock_handshake = mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake")
     # Replace the blocking wait with one that returns the only proc's sentinel
     # immediately so the test does not hang.
     mocker.patch(
         "multiprocessing.connection.wait",
         side_effect=lambda sentinels: [sentinels[0]],
     )
-    mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc")
     mocker.patch("signal.signal")
 
     run_headless(_make_headless_args(stage_id=1))
@@ -321,17 +326,15 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
     assert reg_kwargs["omni_stage_id"] == 1
     assert reg_kwargs["omni_stage_config"] is stage_cfg
     assert reg_kwargs["replica_id"] is None
-    assert reg_kwargs["return_full_response"] is True
+    assert "socket_ownership" not in reg_kwargs
 
-    spawn_kwargs = mock_spawn.call_args.kwargs
-    assert spawn_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
-    assert spawn_kwargs["request_address"] == "tcp://127.0.0.1:26002"
-    assert spawn_kwargs["response_address"] == "tcp://127.0.0.1:26003"
-    assert spawn_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100"
-    assert spawn_kwargs["omni_stage_id"] == 1
-    assert spawn_kwargs["omni_replica_id"] == 0
-
-    mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600)
+    manager_kwargs = mock_manager.call_args.kwargs
+    assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
+    assert manager_kwargs["addresses"].inputs == ["tcp://127.0.0.1:26002"]
+    assert manager_kwargs["addresses"].outputs == ["tcp://127.0.0.1:26003"]
+    assert manager_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100"
+    assert manager_kwargs["omni_stage_id"] == 1
+    assert manager_kwargs["omni_replica_id"] == 0
 
 
 def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None:
@@ -371,15 +374,13 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
         ),
     )
     mocker.patch(
-        "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
-        return_value=(proc, None, None, None),
+        "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager.launch_headless",
+        return_value=SimpleNamespace(proc=proc, shutdown=mocker.Mock()),
     )
-    mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake")
     mocker.patch(
         "multiprocessing.connection.wait",
         side_effect=lambda sentinels: [sentinels[0]],
     )
-    mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc")
     mocker.patch("signal.signal")
 
     with pytest.raises(RuntimeError, match=r"exited with code 137"):

tests/entrypoints/test_stage_utils.py

@@ -4,7 +4,7 @@ import sys
 import pytest
 from pytest_mock import MockerFixture
 
-from vllm_omni.entrypoints.stage_utils import _map_device_list, set_stage_devices
+from vllm_omni.entrypoints.stage_utils import _map_device_list, resolve_stage_physical_devices, set_stage_devices
 
 
 def _make_dummy_torch(call_log):
@@ -207,6 +207,13 @@ def test_map_device_list_multi_replica_offset_cvd():
     assert _map_device_list(1, ["2", "3"], visible) == ["4", "5"]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3")
+    assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4"
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_raises_on_non_numeric():

vllm_omni/diffusion/stage_diffusion_client.py

@@ -1,8 +1,7 @@
 """Stage Diffusion Client for vLLM-Omni multi-stage runtime.
 
-Spawns StageDiffusionProc in a subprocess and communicates via ZMQ
-(PUSH/PULL) to expose the same interface the Orchestrator expects
-from any stage client.
+Owns the frontend-side ZMQ sockets for StageDiffusionProc and exposes the
+interface the Orchestrator expects from a stage client.
 """
 
 from __future__ import annotations
@@ -23,15 +22,14 @@ from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.stage_diffusion_proc import (
     StageDiffusionProc,
-    complete_diffusion_handshake,
-    spawn_diffusion_proc,
+    StageDiffusionProcManager,
 )
 from vllm_omni.distributed.omni_connectors.utils.serialization import (
     OmniMsgpackDecoder,
     OmniMsgpackEncoder,
 )
 from vllm_omni.engine.stage_client import StageClientBase
-from vllm_omni.engine.stage_init_utils import StageMetadata, terminate_alive_proc
+from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
@@ -54,8 +52,17 @@ def create_diffusion_client(
         from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
 
         return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
-    return StageDiffusionClient(
-        model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size
+    proc_manager = StageDiffusionProcManager(
+        model=model,
+        od_config=od_config,
+        stage_init_timeout=stage_init_timeout,
+    )
+    return StageDiffusionClient.from_addresses(
+        metadata,
+        request_address=proc_manager.addresses.inputs[0],
+        response_address=proc_manager.addresses.outputs[0],
+        proc_manager=proc_manager,
+        batch_size=batch_size,
     )
 
 
@@ -74,16 +81,20 @@ class StageDiffusionClient(StageClientBase):
 
     def __init__(
         self,
-        model: str,
-        od_config: OmniDiffusionConfig,
         metadata: StageMetadata,
-        stage_init_timeout: int,
+        request_address: str,
+        response_address: str,
+        *,
+        proc_manager: StageDiffusionProcManager | None = None,
         batch_size: int = 1,
     ) -> None:
-        # Spawn StageDiffusionProc subprocess and wait for READY.
-        proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config)
-        complete_diffusion_handshake(proc, handshake_address, stage_init_timeout)
-        self._initialize_client(metadata, request_address, response_address, proc=proc, batch_size=batch_size)
+        self._initialize_client(
+            metadata,
+            request_address,
+            response_address,
+            proc_manager=proc_manager,
+            batch_size=batch_size,
+        )
 
     @classmethod
     def from_addresses(
@@ -92,19 +103,17 @@ class StageDiffusionClient(StageClientBase):
         request_address: str,
         response_address: str,
         *,
-        proc: Any = None,
+        proc_manager: StageDiffusionProcManager | None = None,
         batch_size: int = 1,
     ) -> StageDiffusionClient:
         """Create a client for an already-running diffusion subprocess."""
-        client = cls.__new__(cls)
-        client._initialize_client(
+        return cls(
             metadata,
             request_address,
             response_address,
-            proc=proc,
+            proc_manager=proc_manager,
             batch_size=batch_size,
         )
-        return client
 
     def _initialize_client(
         self,
@@ -112,37 +121,12 @@ class StageDiffusionClient(StageClientBase):
         request_address: str,
         response_address: str,
         *,
-        proc: Any,
+        proc_manager: StageDiffusionProcManager | None = None,
         batch_size: int,
     ) -> None:
-        self.stage_id = metadata.stage_id
-        self.replica_id = metadata.replica_id
-        self.final_output = metadata.final_output
-        self.final_output_type = metadata.final_output_type
-        self.model_stage = getattr(metadata, "model_stage", None)
-        self.default_sampling_params = metadata.default_sampling_params
-        self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False)
-        self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None)
-        self.engine_input_source = getattr(metadata, "engine_input_source", [])
-        self._proc = proc
-        self._owns_process = proc is not None
-        # Expose the ZMQ addresses on the instance so callers (e.g.
-        # ``StagePool._client_input_addr``) can identify the diffusion
-        # replica by its bound address.
-        self.request_address = request_address
-        self.response_address = response_address
-
-        self._zmq_ctx = zmq.Context()
-        self._request_socket = self._zmq_ctx.socket(zmq.PUSH)
-        self._request_socket.connect(request_address)
-        self._response_socket = self._zmq_ctx.socket(zmq.PULL)
-        self._response_socket.connect(response_address)
-
-        self._response_poller = zmq.asyncio.Poller()
-        self._response_poller.register(self._response_socket, zmq.POLLIN)
-
-        self._encoder = OmniMsgpackEncoder()
-        self._decoder = OmniMsgpackDecoder()
+        self._set_stage_metadata(metadata)
+        self._proc_manager = proc_manager
+        self._connect_transport(request_address, response_address)
 
         self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
         self._rpc_results: dict[str, Any] = {}
@@ -151,19 +135,48 @@ class StageDiffusionClient(StageClientBase):
         self._shutting_down = False
         self._engine_dead: bool = False
 
-        # Background thread to detect silent process death (SIGKILL, segfault)
-        # where the subprocess cannot send the ZMQ death sentinel.
-        # Mirrors MPClient.start_engine_core_monitor() in vLLM.
-        self._start_proc_monitor()
+        if self._proc_manager is not None:
+            self._start_proc_monitor()
 
         logger.info(
             "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)",
             self.stage_id,
             self.replica_id,
-            self._owns_process,
+            self._proc_manager is not None,
             batch_size,
         )
 
+    def _set_stage_metadata(self, metadata: StageMetadata) -> None:
+        self.stage_id = metadata.stage_id
+        self.replica_id = metadata.replica_id
+        self.final_output = metadata.final_output
+        self.final_output_type = metadata.final_output_type
+        self.model_stage = metadata.model_stage
+        self.default_sampling_params = metadata.default_sampling_params
+        self.prompt_expand_func = metadata.prompt_expand_func
+        self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False)
+        self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None)
+        self.engine_input_source = getattr(metadata, "engine_input_source", [])
+
+    def _connect_transport(self, request_address: str, response_address: str) -> None:
+        # Expose the ZMQ addresses on the instance so callers (e.g.
+        # ``StagePool._client_input_addr``) can identify the diffusion
+        # replica by its bound address.
+        self.request_address = request_address
+        self.response_address = response_address
+
+        self._zmq_ctx = zmq.Context()
+        self._request_socket = self._zmq_ctx.socket(zmq.PUSH)
+        self._request_socket.bind(request_address)
+        self._response_socket = self._zmq_ctx.socket(zmq.PULL)
+        self._response_socket.bind(response_address)
+
+        self._response_poller = zmq.asyncio.Poller()
+        self._response_poller.register(self._response_socket, zmq.POLLIN)
+
+        self._encoder = OmniMsgpackEncoder()
+        self._decoder = OmniMsgpackDecoder()
+
     # ------------------------------------------------------------------
     # Process monitor (mirrors vLLM's MPClient.start_engine_core_monitor)
     # ------------------------------------------------------------------
@@ -175,7 +188,10 @@ class StageDiffusionClient(StageClientBase):
         (e.g. SIGKILL, segfault), this thread sets ``_engine_dead`` so
         subsequent calls raise ``EngineDeadError``.
         """
-        proc = self._proc
+        # Background thread to detect silent process death (SIGKILL, segfault)
+        # where the subprocess cannot send the ZMQ death sentinel.
+        # Mirrors MPClient.start_engine_core_monitor() in vLLM.
+        proc = self._proc_manager.proc
         self_ref = weakref.ref(self)
 
         def _monitor() -> None:
@@ -413,9 +429,12 @@ class StageDiffusionClient(StageClientBase):
                 if self._shutting_down:
                     return None
                 raise EngineDeadError()
-            if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive():
+            if self._proc_manager is None:
+                return None
+            proc = self._proc_manager.proc
+            if not self._shutting_down and not proc.is_alive():
                 self._engine_dead = True
-                exitcode = self._proc.exitcode
+                exitcode = proc.exitcode
                 # One final drain – the last ZMQ frame may have arrived
                 # between the first drain and the is_alive() check.
                 self._drain_responses()
@@ -490,11 +509,12 @@ class StageDiffusionClient(StageClientBase):
                 result = self._rpc_results.pop(rpc_id, None)
                 if result is not None:
                     return result
-                if self._engine_dead or (self._owns_process and self._proc is not None and not self._proc.is_alive()):
+                proc = self._proc_manager.proc
+                if self._engine_dead or not proc.is_alive():
                     self._engine_dead = True
                     raise EngineDeadError(
                         f"StageDiffusionProc died while waiting for "
-                        f"collective_rpc '{method}' (exit code {self._proc.exitcode})"
+                        f"collective_rpc '{method}' (exit code {proc.exitcode})"
                     )
                 if deadline is not None and time.monotonic() > deadline:
                     raise TimeoutError(f"collective_rpc_async '{method}' timed out after {timeout}s")
@@ -518,10 +538,13 @@ class StageDiffusionClient(StageClientBase):
         """
         if self._engine_dead:
             raise EngineDeadError(f"Stage-{self.stage_id} diffusion subprocess is dead")
-        if self._proc is not None and not self._proc.is_alive():
+        if self._proc_manager is None:
+            return
+        proc = self._proc_manager.proc
+        if not proc.is_alive():
             self._engine_dead = True
             raise EngineDeadError(
-                f"Stage-{self.stage_id} diffusion subprocess is not alive (exit code: {self._proc.exitcode})."
+                f"Stage-{self.stage_id} diffusion subprocess is not alive (exit code: {proc.exitcode})."
             )
 
     def shutdown(self) -> None:
@@ -531,9 +554,8 @@ class StageDiffusionClient(StageClientBase):
         except Exception:
             pass
 
-        if self._owns_process and self._proc is not None and self._proc.is_alive():
-            self._proc.join(timeout=10)
-            terminate_alive_proc(self._proc)
+        if self._proc_manager is not None and self._proc_manager.proc.is_alive():
+            self._proc_manager.shutdown(timeout=10)
 
         self._request_socket.close(linger=0)
         self._response_socket.close(linger=0)

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -8,10 +8,10 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import multiprocessing.connection
 import signal
-import time
 from concurrent.futures import ThreadPoolExecutor
-from multiprocessing.process import BaseProcess
+from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any
 
 import msgspec
@@ -22,6 +22,8 @@ from PIL import Image
 from vllm.logger import init_logger
 from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
 from vllm.utils.system_utils import get_mp_context
+from vllm.v1.engine.core import EngineCoreProc
+from vllm.v1.engine.utils import CoreEngine, EngineZmqAddresses, wait_for_engine_startup
 from vllm.v1.utils import shutdown
 
 from vllm_omni.diffusion.data import DiffusionRequestAbortedError
@@ -360,10 +362,10 @@ class StageDiffusionProc:
         ctx = zmq.asyncio.Context()
 
         request_socket = ctx.socket(zmq.PULL)
-        request_socket.bind(request_address)
+        request_socket.connect(request_address)
 
         response_socket = ctx.socket(zmq.PUSH)
-        response_socket.bind(response_address)
+        response_socket.connect(response_address)
 
         encoder = OmniMsgpackEncoder()
         decoder = OmniMsgpackDecoder()
@@ -612,15 +614,51 @@ class StageDiffusionProc:
     # Subprocess entry point
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _open_startup_handshake(
+        handshake_address: str,
+        *,
+        local_client: bool,
+        headless: bool,
+    ) -> tuple[zmq.Context, zmq.Socket, EngineZmqAddresses]:
+        ctx = zmq.Context()
+        socket = ctx.socket(zmq.DEALER)
+        socket.setsockopt(zmq.IDENTITY, (0).to_bytes(2, "little"))
+        socket.connect(handshake_address)
+        addresses = EngineCoreProc.startup_handshake(
+            socket,
+            local_client=local_client,
+            headless=headless,
+            parallel_config=None,
+        )
+        return ctx, socket, addresses
+
+    @staticmethod
+    def _send_startup_ready(
+        handshake_socket: zmq.Socket,
+        *,
+        local_client: bool,
+        headless: bool,
+    ) -> None:
+        handshake_socket.send(
+            msgspec.msgpack.encode(
+                {
+                    "status": "READY",
+                    "local": local_client,
+                    "headless": headless,
+                }
+            )
+        )
+
     @classmethod
     def run_diffusion_proc(
         cls,
         model: str,
         od_config: OmniDiffusionConfig,
         handshake_address: str,
-        request_address: str,
-        response_address: str,
         *,
+        local_client: bool,
+        headless: bool,
         omni_coordinator_address: str | None = None,
         omni_stage_id: int | None = None,
         omni_replica_id: int = 0,
@@ -651,21 +689,32 @@ class StageDiffusionProc:
 
         proc = cls(model, od_config)
         coord_client: OmniCoordClientForStage | None = None
+        handshake_ctx: zmq.Context | None = None
+        handshake_socket: zmq.Socket | None = None
         try:
+            handshake_ctx, handshake_socket, addresses = cls._open_startup_handshake(
+                handshake_address,
+                local_client=local_client,
+                headless=headless,
+            )
+            request_address = addresses.inputs[0]
+            response_address = addresses.outputs[0]
+
             proc.initialize()
 
-            # Send READY via handshake socket
-            handshake_ctx = zmq.Context()
-            handshake_socket = handshake_ctx.socket(zmq.DEALER)
-            handshake_socket.connect(handshake_address)
-            handshake_socket.send(msgspec.msgpack.encode({"status": "READY"}))
+            cls._send_startup_ready(
+                handshake_socket,
+                local_client=local_client,
+                headless=headless,
+            )
             handshake_socket.close()
             handshake_ctx.term()
+            handshake_socket = None
+            handshake_ctx = None
 
-            # Wire OmniCoordClientForStage *after* READY so that the head
-            # has bound its head-side request/response sockets — the
-            # address pair we report is the same pair this proc binds to
-            # (request/response addresses passed in).
+            # Wire OmniCoordClientForStage *after* READY. The address pair is
+            # owned by the frontend client; this proc connects to it as the
+            # backend runtime.
             if omni_coordinator_address is not None:
                 if omni_stage_id is None:
                     raise ValueError("omni_stage_id must be provided when omni_coordinator_address is set")
@@ -688,7 +737,6 @@ class StageDiffusionProc:
                     omni_coordinator_address,
                 )
 
-            # Run async event loop
             asyncio.run(proc.run_loop(request_address, response_address))
 
         except SystemExit:
@@ -698,105 +746,138 @@ class StageDiffusionProc:
             logger.exception("StageDiffusionProc encountered a fatal error.")
             raise
         finally:
+            if handshake_socket is not None:
+                handshake_socket.close(linger=0)
+            if handshake_ctx is not None:
+                handshake_ctx.term()
             if coord_client is not None:
                 with contextlib.suppress(RuntimeError):
                     coord_client.close()
             proc.close()
 
 
-# -- Free functions for backward compatibility with StageDiffusionClient ------
+class StageDiffusionProcManager:
+    """Owns a StageDiffusionProc subprocess.
 
-
-def spawn_diffusion_proc(
-    model: str,
-    od_config: OmniDiffusionConfig,
-    handshake_address: str | None = None,
-    request_address: str | None = None,
-    response_address: str | None = None,
-    *,
-    omni_coordinator_address: str | None = None,
-    omni_stage_id: int | None = None,
-    omni_replica_id: int = 0,
-) -> tuple[BaseProcess, str, str, str]:
-    """Spawn a StageDiffusionProc subprocess.
-
-    Returns ``(proc, handshake_address, request_address, response_address)``.
-
-    Pass ``omni_coordinator_address`` / ``omni_stage_id`` / ``omni_replica_id``
-    to have the subprocess publish heartbeats to an OmniCoordinator.
+    Mirrors the small process-lifecycle surface used by vLLM's
+    CoreEngineProcManager while keeping diffusion's custom wire protocol.
     """
-    handshake_address = handshake_address or get_open_zmq_ipc_path()
-    request_address = request_address or get_open_zmq_ipc_path()
-    response_address = response_address or get_open_zmq_ipc_path()
 
-    ctx = get_mp_context()
-    proc = ctx.Process(
-        target=StageDiffusionProc.run_diffusion_proc,
-        name="StageDiffusionProc",
-        kwargs={
-            "model": model,
-            "od_config": od_config,
-            "handshake_address": handshake_address,
-            "request_address": request_address,
-            "response_address": response_address,
-            "omni_coordinator_address": omni_coordinator_address,
-            "omni_stage_id": omni_stage_id,
-            "omni_replica_id": omni_replica_id,
-        },
-    )
-    proc.start()
-    # Wait for the process to become alive before returning.
-    deadline = time.monotonic() + 10
-    while not proc.is_alive():
-        if proc.exitcode is not None:
-            raise RuntimeError(f"StageDiffusionProc failed to start (exit code {proc.exitcode})")
-        if time.monotonic() > deadline:
-            raise TimeoutError("StageDiffusionProc did not become alive within 10s")
-        time.sleep(0.01)
-    return proc, handshake_address, request_address, response_address
+    def __init__(
+        self,
+        *,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        stage_init_timeout: int,
+        handshake_address: str | None = None,
+        addresses: EngineZmqAddresses | None = None,
+        omni_coordinator_address: str | None = None,
+        omni_stage_id: int | None = None,
+        omni_replica_id: int = 0,
+    ) -> None:
+        handshake_address = handshake_address or get_open_zmq_ipc_path()
+        addresses = addresses or EngineZmqAddresses(
+            inputs=[get_open_zmq_ipc_path()],
+            outputs=[get_open_zmq_ipc_path()],
+        )
 
+        ctx = get_mp_context()
+        proc = ctx.Process(
+            target=StageDiffusionProc.run_diffusion_proc,
+            name="StageDiffusionProc",
+            kwargs={
+                "model": model,
+                "od_config": od_config,
+                "handshake_address": handshake_address,
+                "local_client": True,
+                "headless": False,
+                "omni_coordinator_address": omni_coordinator_address,
+                "omni_stage_id": omni_stage_id,
+                "omni_replica_id": omni_replica_id,
+            },
+        )
+        proc.start()
+        self.proc = proc
+        self.addresses = addresses
+        self.manager_stopped = False
+        self.failed_proc_name: str | None = None
 
-def complete_diffusion_handshake(
-    proc: BaseProcess,
-    handshake_address: str,
-    handshake_timeout: int,
-) -> None:
-    """Wait for the diffusion subprocess to signal READY.
+        self._wait_until_started(handshake_address, stage_init_timeout)
 
-    On failure the process is terminated before re-raising.
-    """
-    try:
-        _perform_diffusion_handshake(proc, handshake_address, handshake_timeout)
-    except Exception:
-        shutdown([proc])
-        raise
+    @classmethod
+    def launch_headless(
+        cls,
+        *,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        handshake_address: str,
+        addresses: EngineZmqAddresses,
+        omni_coordinator_address: str | None,
+        omni_stage_id: int,
+        omni_replica_id: int,
+    ) -> StageDiffusionProcManager:
+        """Launch a headless diffusion backend that connects to head-owned sockets."""
+        self = cls.__new__(cls)
+        ctx = get_mp_context()
+        proc = ctx.Process(
+            target=StageDiffusionProc.run_diffusion_proc,
+            name="StageDiffusionProc",
+            kwargs={
+                "model": model,
+                "od_config": od_config,
+                "handshake_address": handshake_address,
+                "local_client": False,
+                "headless": True,
+                "omni_coordinator_address": omni_coordinator_address,
+                "omni_stage_id": omni_stage_id,
+                "omni_replica_id": omni_replica_id,
+            },
+        )
+        proc.start()
+        self.proc = proc
+        self.addresses = addresses
+        self.manager_stopped = False
+        self.failed_proc_name = None
+        return self
 
-
-def _perform_diffusion_handshake(
-    proc: BaseProcess,
-    handshake_address: str,
-    handshake_timeout: int,
-) -> None:
-    """Run the handshake with the diffusion subprocess."""
-    with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket:
-        poller = zmq.Poller()
-        poller.register(handshake_socket, zmq.POLLIN)
-        poller.register(proc.sentinel, zmq.POLLIN)
-
-        timeout_ms = handshake_timeout * 1000
-        while True:
-            events = dict(poller.poll(timeout=timeout_ms))
-            if not events:
-                raise TimeoutError(
-                    f"Timed out waiting for READY from StageDiffusionProc after {handshake_timeout}s. "
-                    f"This typically indicates model loading or warmup is taking too long. "
-                    f"Consider increasing `stage_init_timeout` for large models."
+    def _wait_until_started(self, handshake_address: str, stage_init_timeout: int) -> None:
+        try:
+            with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket:
+                wait_for_engine_startup(
+                    handshake_socket,
+                    self.addresses,
+                    [CoreEngine(index=0, local=True)],
+                    SimpleNamespace(
+                        data_parallel_size_local=1,
+                        data_parallel_hybrid_lb=False,
+                        data_parallel_external_lb=False,
+                    ),
+                    False,
+                    None,
+                    self,
+                    None,
                 )
-            if handshake_socket in events:
-                identity, raw = handshake_socket.recv_multipart()
-                msg = msgspec.msgpack.decode(raw)
-                if msg.get("status") == "READY":
-                    return
-                raise RuntimeError(f"Expected READY, got: {msg}")
-            if proc.exitcode is not None:
-                raise RuntimeError(f"StageDiffusionProc died during handshake (exit code {proc.exitcode})")
+        except Exception:
+            shutdown([self.proc])
+            raise
+
+    def shutdown(self, timeout: float | None = None) -> None:
+        self.manager_stopped = True
+        shutdown([self.proc], timeout=timeout)
+
+    def sentinels(self) -> list[int]:
+        return [self.proc.sentinel]
+
+    def finished_procs(self) -> dict[str, int]:
+        if self.proc.exitcode is None:
+            return {}
+        return {self.proc.name: self.proc.exitcode}
+
+    def monitor_engine_liveness(self) -> None:
+        try:
+            multiprocessing.connection.wait([self.proc.sentinel])
+        except Exception:
+            return
+        if self.proc.exitcode not in (None, 0) and not self.manager_stopped:
+            self.failed_proc_name = self.proc.name
+        self.shutdown()

vllm_omni/distributed/omni_coordinator/__init__.py

@@ -8,11 +8,13 @@ from .load_balancer import (
     RandomBalancer,
     RoundRobinBalancer,
     Task,
-    build_load_balancer_factory,
 )
 from .messages import ReplicaEvent, ReplicaInfo, ReplicaList, ReplicaStatus
 from .omni_coord_client_for_hub import OmniCoordClientForHub
-from .omni_coord_client_for_stage import OmniCoordClientForStage
+from .omni_coord_client_for_stage import (
+    OmniCoordClientForStage,
+    create_stage_coord_client,
+)
 from .omni_coordinator import OmniCoordinator
 from .runtime import OmniCoordinatorRuntime
 
@@ -24,6 +26,7 @@ __all__ = [
     "ReplicaInfo",
     "ReplicaList",
     "OmniCoordClientForStage",
+    "create_stage_coord_client",
     "OmniCoordClientForHub",
     "Task",
     "LoadBalancer",
@@ -31,5 +34,4 @@ __all__ = [
     "RandomBalancer",
     "RoundRobinBalancer",
     "LeastQueueLengthBalancer",
-    "build_load_balancer_factory",
 ]

vllm_omni/distributed/omni_coordinator/load_balancer.py

@@ -6,7 +6,6 @@ from __future__ import annotations
 import random
 import threading
 from abc import ABC, abstractmethod
-from collections.abc import Callable
 from enum import Enum
 from typing import Any, TypedDict
 
@@ -122,22 +121,6 @@ class LeastQueueLengthBalancer(LoadBalancer):
         return random.choice(candidates)
 
 
-def build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]:
-    """Translate ``--omni-lb-policy`` (string) into a per-pool LB factory."""
-    try:
-        normalized = LoadBalancingPolicy(policy)
-    except ValueError as exc:
-        valid = ", ".join(p.value for p in LoadBalancingPolicy)
-        raise ValueError(f"unknown --omni-lb-policy {policy!r} (valid: {valid})") from exc
-    if normalized is LoadBalancingPolicy.RANDOM:
-        return RandomBalancer
-    if normalized is LoadBalancingPolicy.ROUND_ROBIN:
-        return RoundRobinBalancer
-    if normalized is LoadBalancingPolicy.LEAST_QUEUE_LENGTH:
-        return LeastQueueLengthBalancer
-    raise ValueError(f"unhandled load balancing policy {normalized!r}")
-
-
 __all__ = [
     "Task",
     "LoadBalancingPolicy",
@@ -145,5 +128,4 @@ __all__ = [
     "RandomBalancer",
     "RoundRobinBalancer",
     "LeastQueueLengthBalancer",
-    "build_load_balancer_factory",
 ]

vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py

@@ -230,3 +230,30 @@ class OmniCoordClientForStage:
             except zmq.ZMQError:
                 pass
             self._closed = True
+
+
+def create_stage_coord_client(
+    *,
+    coord_zmq_addr: str,
+    input_addr: str,
+    output_addr: str,
+    stage_id: int,
+    queue_length_getter: Callable[[], int] | None = None,
+) -> OmniCoordClientForStage:
+    """Create a stage coordinator client with an optional heartbeat hook."""
+    client = OmniCoordClientForStage(
+        coord_zmq_addr=coord_zmq_addr,
+        input_addr=input_addr,
+        output_addr=output_addr,
+        stage_id=stage_id,
+    )
+    if queue_length_getter is not None:
+
+        def _refresh_queue_length() -> None:
+            try:
+                client._queue_length = max(int(queue_length_getter()), 0)
+            except Exception:
+                pass
+
+        client._on_heartbeat = _refresh_queue_length
+    return client

vllm_omni/distributed/omni_coordinator/omni_coordinator.py

@@ -190,6 +190,15 @@ class OmniCoordinator:
         except zmq.ZMQError:
             pass
 
+    def wait_for_shutdown(self) -> None:
+        """Block until the coordinator is asked to stop.
+
+        ``OmniCoordinatorRuntime`` runs the coordinator in a child process.
+        The child needs a small, explicit blocking lifecycle API so it does
+        not return immediately after construction.
+        """
+        self._stop_event.wait()
+
     def _parse_replica_event(self, data: dict[str, Any]) -> ReplicaEvent | None:
         """Parse wire payload dict into ReplicaEvent. Returns None if invalid."""
         try:

vllm_omni/distributed/omni_coordinator/runtime.py

@@ -3,35 +3,91 @@
 
 """Lifecycle wrapper around :class:`OmniCoordinator`.
 
-``OmniCoordinatorRuntime`` is the single-purpose owner of the head-side
-coordinator process artifacts: it picks two free TCP ports, constructs an
-:class:`OmniCoordinator` bound to them, exposes the resulting addresses, and
-provides a single ``close()`` method to tear everything down.
+``OmniCoordinatorRuntime`` spawns the coordinator as an independent process
+(matching vLLM's DPCoordinator pattern). Physical isolation prevents GIL
+contention and makes direct-object coupling impossible.
 
 The ROUTER address is later handed to :class:`OmniMasterServer` so it can be
 published to registering replicas; the PUB address is handed to the
-``Orchestrator``, which constructs its :class:`OmniCoordClientForHub` against
-it.
+``MembershipController``, which constructs its :class:`OmniCoordClientForHub`
+against it.
 """
 
 from __future__ import annotations
 
 import logging
+import multiprocessing
+import multiprocessing.connection
+import os
+import signal
+import weakref
+from typing import Any
 
 from vllm.utils.network_utils import get_open_ports_list
 
-from .omni_coordinator import OmniCoordinator
+from vllm_omni.distributed.omni_coordinator.omni_coordinator import OmniCoordinator
 
 logger = logging.getLogger(__name__)
 
 
-class OmniCoordinatorRuntime:
-    """Own one :class:`OmniCoordinator` and the two ports it binds.
+def run_omni_coordinator_proc(
+    router_zmq_addr: str,
+    pub_zmq_addr: str,
+    heartbeat_timeout: float,
+    ready_pipe: Any,
+) -> None:
+    """Main loop running inside the coordinator child process."""
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
 
-    Constructor binds; :meth:`close` tears down. The class deliberately does
-    not expose the coordinator instance — callers should consume the
-    coordinator only via its wire protocol through
-    :class:`OmniCoordClientForStage` and :class:`OmniCoordClientForHub`.
+    coordinator = OmniCoordinator(
+        router_zmq_addr=router_zmq_addr,
+        pub_zmq_addr=pub_zmq_addr,
+        heartbeat_timeout=heartbeat_timeout,
+    )
+
+    ready_pipe.send("ready")
+    ready_pipe.close()
+
+    coordinator.wait_for_shutdown()
+
+
+def _get_coordinator_mp_context() -> multiprocessing.context.BaseContext:
+    """Return the multiprocessing context used for OmniCoordinator.
+
+    For the current vllm-omni startup path, ``spawn`` is too expensive: the
+    child re-imports a heavy CLI / model stack before it can acknowledge the
+    ready pipe, which can exceed the coordinator startup timeout. Prefer
+    ``fork`` on platforms that support it.
+
+    ``fork`` must happen before the parent initializes CUDA or owns long-lived
+    ZMQ sockets; otherwise the child inherits unsafe state. If coordinator
+    startup moves later, switch this to ``spawn``/``forkserver``.
+
+    TODO: make the coordinator child entry cheap and spawn-safe, then revisit
+    whether ``spawn`` is still needed here.
+    """
+    if os.name != "nt" and "fork" in multiprocessing.get_all_start_methods():
+        return multiprocessing.get_context("fork")
+    return multiprocessing.get_context("spawn")
+
+
+def _shutdown_proc(proc: multiprocessing.Process) -> None:
+    """Best-effort process termination for weakref finalizer."""
+    if proc.is_alive():
+        proc.terminate()
+        proc.join(timeout=5)
+        if proc.is_alive():
+            proc.kill()
+            proc.join(timeout=2)
+
+
+class OmniCoordinatorRuntime:
+    """Own one :class:`OmniCoordinator` running in a child process.
+
+    Constructor spawns the process; :meth:`close` tears it down.
+    The class deliberately does not expose the coordinator instance —
+    callers consume it only via ZMQ through :class:`OmniCoordClientForStage`
+    and :class:`OmniCoordClientForHub`.
     """
 
     def __init__(
@@ -50,25 +106,54 @@ class OmniCoordinatorRuntime:
         self.pub_address: str = f"tcp://{host}:{pub_port}"
 
         self._closed = False
-        self._coordinator = OmniCoordinator(
-            router_zmq_addr=self.router_address,
-            pub_zmq_addr=self.pub_address,
-            heartbeat_timeout=heartbeat_timeout,
+
+        ctx = _get_coordinator_mp_context()
+        parent_conn, child_conn = ctx.Pipe(duplex=False)
+
+        self._proc: multiprocessing.Process = ctx.Process(
+            target=run_omni_coordinator_proc,
+            kwargs={
+                "router_zmq_addr": self.router_address,
+                "pub_zmq_addr": self.pub_address,
+                "heartbeat_timeout": heartbeat_timeout,
+                "ready_pipe": child_conn,
+            },
+            daemon=True,
+            name="OmniCoordinator",
         )
+        self._proc.start()
+        child_conn.close()
+
+        ready = multiprocessing.connection.wait([parent_conn, self._proc.sentinel], timeout=30)
+        if not ready:
+            self._proc.terminate()
+            self._proc.join(timeout=5)
+            raise RuntimeError("OmniCoordinator process failed to start within 30s")
+
+        try:
+            status = parent_conn.recv()
+        except EOFError:
+            raise RuntimeError("OmniCoordinator process died during startup") from None
+        finally:
+            parent_conn.close()
+
+        if status != "ready":
+            raise RuntimeError(f"OmniCoordinator unexpected status: {status}")
+
+        self._finalizer = weakref.finalize(self, _shutdown_proc, self._proc)
 
         logger.info(
-            "[OmniCoordinatorRuntime] Started (router=%s pub=%s heartbeat_timeout=%.1fs)",
+            "[OmniCoordinatorRuntime] Started (pid=%d router=%s pub=%s heartbeat_timeout=%.1fs)",
+            self._proc.pid,
             self.router_address,
             self.pub_address,
             heartbeat_timeout,
         )
 
     def close(self) -> None:
-        """Tear down the underlying coordinator. Idempotent."""
+        """Tear down the coordinator process. Idempotent."""
         if self._closed:
             return
         self._closed = True
-        try:
-            self._coordinator.close()
-        except Exception:
-            logger.exception("[OmniCoordinatorRuntime] coordinator close failed")
+        _shutdown_proc(self._proc)
+        self._finalizer.detach()

vllm_omni/engine/membership_controller.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""MembershipController: distributed replica lifecycle management.
+
+Extracted from Orchestrator to keep request-flow code free of distributed
+concerns. Owns the OmniCoordClientForHub, watches for replica
+disappearances, and handles register/unregister by building head-side
+clients via an injected factory and mutating StagePool membership.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.distributed.omni_coordinator import (
+    LoadBalancer,
+)
+from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus
+from vllm_omni.distributed.omni_coordinator.omni_coord_client_for_hub import OmniCoordClientForHub
+from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage
+from vllm_omni.engine.stage_pool import StagePool
+
+logger = init_logger(__name__)
+
+RemoteReplicaFactory = Callable[[int, int], Any]
+
+
+class MembershipController:
+    """Manages dynamic replica attach/detach for distributed mode.
+
+    Constructed by DistStageRuntime and passed to Orchestrator.
+    The Orchestrator delegates register/unregister messages here.
+    """
+
+    WATCH_INTERVAL_S: float = 0.5
+
+    def __init__(
+        self,
+        stage_pools: list[StagePool],
+        coordinator_pub_address: str,
+        load_balancer_factory: Callable[[], LoadBalancer],
+        remote_replica_factory: RemoteReplicaFactory,
+    ) -> None:
+        self._stage_pools = stage_pools
+        self._remote_replica_factory = remote_replica_factory
+        self._membership_tasks: set[asyncio.Task[None]] = set()
+        self._shutdown_event = asyncio.Event()
+        self._watcher_task: asyncio.Task[None] | None = None
+        self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None
+        self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None
+
+        self._hub = OmniCoordClientForHub(coordinator_pub_address)
+        factory = load_balancer_factory
+        for pool in self._stage_pools:
+            pool.attach_hub(self._hub)
+            pool.attach_load_balancer(factory())
+
+    def start(self) -> asyncio.Task[None]:
+        """Start the replica watcher as a background task. Returns the task."""
+        self._watcher_task = asyncio.create_task(self._watch_replica_list(), name="membership-watcher")
+        return self._watcher_task
+
+    async def handle_register(self, stage_id: int, replica_id: int) -> None:
+        """Handle a register_remote_replica message (fire-and-forget)."""
+        self._spawn_task(
+            self._do_register(stage_id, replica_id),
+            label=f"register-s{stage_id}-r{replica_id}",
+        )
+
+    async def handle_unregister(
+        self,
+        stage_id: int,
+        input_addr: str,
+        output_queue: asyncio.Queue[EngineQueueMessage] | None = None,
+        cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None,
+    ) -> None:
+        """Handle an unregister_remote_replica message."""
+        pool = self._pool_for_stage_id(stage_id)
+        if pool is None:
+            return
+        effective_output_queue = output_queue if output_queue is not None else self._output_queue
+        effective_cleanup_callback = cleanup_callback if cleanup_callback is not None else self._cleanup_callback
+        affected = pool.invalidate_addr(input_addr)
+        self._detach_replica(stage_id, input_addr)
+        if affected and effective_cleanup_callback is not None:
+            await effective_cleanup_callback(affected)
+        if affected and effective_output_queue is not None:
+            for req_id in affected:
+                await effective_output_queue.put(
+                    ErrorMessage(error="stage replica disappeared", request_id=req_id, stage_id=stage_id)
+                )
+
+    def shutdown(self) -> None:
+        """Signal stop and close the hub."""
+        self._shutdown_event.set()
+        if self._hub is not None:
+            self._hub.close()
+            self._hub = None
+        if self._watcher_task is not None and not self._watcher_task.done():
+            self._watcher_task.cancel()
+
+    async def drain_tasks(self, timeout: float = 10.0) -> None:
+        """Wait for in-flight membership tasks to complete."""
+        if self._membership_tasks:
+            await asyncio.wait(self._membership_tasks, timeout=timeout)
+
+    # ---- Internal ----
+
+    def _spawn_task(self, coro: Awaitable[None], *, label: str) -> None:
+        task = asyncio.create_task(coro, name=f"membership-{label}")
+        self._membership_tasks.add(task)
+
+        def _on_done(t: asyncio.Task[None]) -> None:
+            self._membership_tasks.discard(t)
+            if t.cancelled():
+                return
+            exc = t.exception()
+            if exc is not None:
+                logger.error("[MembershipController] %s task crashed", label, exc_info=exc)
+
+        task.add_done_callback(_on_done)
+
+    async def _watch_replica_list(self) -> None:
+        """Convert hub replica disappearances into unregister actions."""
+        last_up: set[tuple[int, str]] = set()
+        while not self._shutdown_event.is_set():
+            try:
+                snap = self._hub.get_replica_list()
+                current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP}
+                for stage_id, addr in last_up - current:
+                    self._spawn_task(
+                        self.handle_unregister(stage_id, addr),
+                        label=f"unregister-s{stage_id}",
+                    )
+                last_up = current
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                logger.exception("[MembershipController] _watch_replica_list iteration failed")
+
+            try:
+                await asyncio.sleep(self.WATCH_INTERVAL_S)
+            except asyncio.CancelledError:
+                raise
+
+    async def _do_register(self, stage_id: int, replica_id: int) -> None:
+        pool = self._pool_for_stage_id(stage_id)
+        if pool is None:
+            logger.warning("[MembershipController] register: stage_id %d out of range", stage_id)
+            return
+        client = await asyncio.to_thread(self._remote_replica_factory, stage_id, replica_id)
+        input_addr = StagePool._client_input_addr(client)
+        if input_addr is None:
+            raise RuntimeError(f"remote replica factory for stage {stage_id} produced a client without input address")
+        pool.add_client(input_addr, client)
+        logger.info(
+            "[MembershipController] attached remote replica stage=%d replica=%d addr=%s",
+            stage_id,
+            replica_id,
+            input_addr,
+        )
+
+    def _detach_replica(self, stage_id: int, input_addr: str) -> None:
+        pool = self._pool_for_stage_id(stage_id)
+        if pool is None:
+            return
+        client = pool.remove_client(input_addr)
+        if client is None:
+            return
+        try:
+            client.shutdown()
+        except Exception:
+            logger.exception("[MembershipController] failed to shutdown client stage=%d addr=%s", stage_id, input_addr)
+        logger.info("[MembershipController] detached replica stage=%d addr=%s", stage_id, input_addr)
+
+    def install_unregister_handlers(
+        self,
+        *,
+        output_queue: asyncio.Queue[EngineQueueMessage],
+        cleanup_callback: Callable[[list[str]], Awaitable[None]],
+    ) -> None:
+        """Install shared cleanup sinks for watcher-driven unregister events."""
+        self._output_queue = output_queue
+        self._cleanup_callback = cleanup_callback
+
+    def _pool_for_stage_id(self, stage_id: int) -> StagePool | None:
+        if not (0 <= stage_id < len(self._stage_pools)):
+            return None
+        return self._stage_pools[stage_id]

vllm_omni/engine/orchestrator.py

@@ -5,19 +5,14 @@ Runs inside a background thread with its own asyncio event loop.
 Owns logical request progression across stage pools and handles
 stage-to-stage transfer logic.
 
-In distributed mode (``coordinator_pub_address`` provided), it also
-owns the single :class:`OmniCoordClientForHub`, runs a
-:meth:`_watch_replica_list` task that converts replica disappearances
-into ``unregister_remote_replica`` control messages, and handles the
-``register_remote_replica`` / ``unregister_remote_replica`` flow that
-attaches / detaches head-side stage clients for headless replicas.
+Distributed membership (replica attach/detach, hub monitoring) is
+handled by :class:`MembershipController`, which is injected optionally.
 """
 
 from __future__ import annotations
 
 import asyncio
 import time as _time
-from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -31,14 +26,9 @@ from vllm.v1.engine import EngineCoreOutputs
 from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.metrics.stats import IterationStats
 
-from vllm_omni.distributed.omni_coordinator import (
-    LoadBalancer,
-    OmniCoordClientForHub,
-    RandomBalancer,
-    ReplicaStatus,
-)
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+from vllm_omni.engine.membership_controller import MembershipController
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
     AddCompanionRequestMessage,
@@ -59,16 +49,6 @@ from vllm_omni.metrics.prometheus import OmniRequestCounter
 from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
 from vllm_omni.outputs import OmniRequestOutput
 
-# Factory signature for building a head-side stage client for a
-# *dynamically attached* (auto-assigned) remote replica.
-#
-# Receives ``(stage_id, replica_id)`` and returns an awaitable yielding the
-# constructed client (any type — it must satisfy the shape expected by the
-# matching :class:`StagePool`, i.e. expose ``client_addresses["input_address"]``
-# or ``request_address``, plus the usual ``add_request_async`` /
-# ``get_output_async`` / ``shutdown`` surface).
-RemoteReplicaFactory = Callable[[int, int], Awaitable[Any]]
-
 logger = init_logger(__name__)
 
 
@@ -159,10 +139,6 @@ class StreamingInputState:
 class Orchestrator:
     """Runs inside a background thread's asyncio event loop."""
 
-    # Cadence at which the replica-list watcher polls for disappearances.
-    _WATCH_REPLICA_INTERVAL_S: float = 0.5
-    _WATCH_REPLICA_IDLE_INTERVAL_S: float = 1.0
-
     # Class-level defaults so tests that bypass __init__ via object.__new__
     # don't AttributeError when transfer / counter emit paths access them.
     _running_counter: OmniRequestCounter | None = None
@@ -178,10 +154,8 @@ class Orchestrator:
         *,
         async_chunk: bool = False,
         pd_config: dict[str, Any] | None = None,
+        membership_controller: MembershipController | None = None,
         running_counter: OmniRequestCounter | None = None,
-        coordinator_pub_address: str | None = None,
-        load_balancer_factory: Callable[[], LoadBalancer] | None = None,
-        remote_replica_factory: RemoteReplicaFactory | None = None,
         transfer_emitter: Any = None,
         log_stats: bool = False,
     ) -> None:
@@ -212,27 +186,8 @@ class Orchestrator:
         self._fatal_error: str | None = None
         self._fatal_error_stage_id: int | None = None
 
-        # Background tasks for fire-and-forget message handlers (currently
-        # only ``register_remote_replica`` and ``unregister_remote_replica``).
-        # Held as a set so each task's reference survives the loop and the
-        # task can self-deregister on completion.
-        self._membership_tasks: set[asyncio.Task[None]] = set()
-
-        # Distributed-mode wiring. The hub is constructed on the
-        # orchestrator's asyncio loop because it spawns a SUB background
-        # thread; building it from another thread would race the
-        # ``_init_done`` event.
-        self._hub: OmniCoordClientForHub | None = (
-            OmniCoordClientForHub(coordinator_pub_address) if coordinator_pub_address is not None else None
-        )
-        self._remote_replica_factory = remote_replica_factory
-        # Inject hub + per-pool LB into each StagePool so they can run
-        # distributed dispatch via ``StagePool.pick``.
-        if self._hub is not None:
-            factory = load_balancer_factory or RandomBalancer
-            for pool in self.stage_pools:
-                pool.attach_hub(self._hub)
-                pool.attach_load_balancer(factory())
+        # Distributed membership (optional, injected by DistStageRuntime)
+        self._membership = membership_controller
 
     def _init_metrics_state(
         self,
@@ -308,16 +263,22 @@ class Orchestrator:
             self._orchestration_output_handler(),
             name="orchestrator-stage-output-handler",
         )
-        # The replica watcher only runs in distributed mode. It's still
-        # created in both cases so ``run()`` has a uniform task graph;
-        # ``_watch_replica_list`` is a no-op poll when ``self._hub`` is None.
-        watch_task = asyncio.create_task(
-            self._watch_replica_list(),
-            name="orchestrator-replica-watcher",
-        )
+
+        # Start membership watcher if distributed mode is active.
+        membership_watcher: asyncio.Task[None] | None = None
+        if self._membership is not None:
+            self._membership.install_unregister_handlers(
+                output_queue=self.output_async_queue,
+                cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True),
+            )
+            membership_watcher = self._membership.start()
+
+        tasks = [request_task, output_task]
+        if membership_watcher is not None:
+            tasks.append(membership_watcher)
 
         try:
-            await asyncio.gather(request_task, output_task, watch_task)
+            await asyncio.gather(*tasks)
         except asyncio.CancelledError:
             raise
         except EngineDeadError as e:
@@ -332,46 +293,23 @@ class Orchestrator:
             raise
         finally:
             self._shutdown_event.set()
-            for task in (request_task, output_task, watch_task):
+            for task in tasks:
                 if not task.done():
                     task.cancel()
             try:
-                await asyncio.gather(request_task, output_task, watch_task, return_exceptions=True)
+                await asyncio.gather(*tasks, return_exceptions=True)
             except Exception:
                 pass
 
-            # If a fatal error caused the shutdown, drain any pending
-            # add_request messages that were never processed and broadcast
-            # fatal error responses so callers are not left hanging.
             if self._fatal_error is not None:
                 await self._drain_pending_requests_on_fatal()
 
-            # Wait briefly for any in-flight membership handlers (register /
-            # unregister remote replica) to finish so they don't leave the
-            # head-side pool in a half-attached state. Cancel anything that
-            # hasn't completed in time; the generic pending-task sweep below
-            # will collect the cancellations.
-            if self._membership_tasks:
-                try:
-                    await asyncio.wait_for(
-                        asyncio.gather(*self._membership_tasks, return_exceptions=True),
-                        timeout=10.0,
-                    )
-                except (asyncio.TimeoutError, Exception):
-                    for t in self._membership_tasks:
-                        if not t.done():
-                            t.cancel()
+            if self._membership is not None:
+                await self._membership.drain_tasks(timeout=10.0)
+                self._membership.shutdown()
 
             self._shutdown_stages()
 
-            # Close the hub last so any in-flight dispatch still has access.
-            if self._hub is not None:
-                try:
-                    self._hub.close()
-                except RuntimeError:
-                    pass
-                self._hub = None
-
             loop = asyncio.get_running_loop()
             pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()]
             for task in pending:
@@ -379,27 +317,6 @@ class Orchestrator:
             if pending:
                 await asyncio.gather(*pending, return_exceptions=True)
 
-    # ---- Background task helpers ----
-
-    def _spawn_membership_task(self, coro: Awaitable[None], *, label: str) -> None:
-        """Run a fire-and-forget membership-change coroutine.
-
-        Holds a strong reference until completion (asyncio would otherwise
-        garbage-collect a bare task), and logs any uncaught exception.
-        """
-        task = asyncio.create_task(coro, name=f"orchestrator-{label}")
-        self._membership_tasks.add(task)
-
-        def _on_done(t: asyncio.Task[None]) -> None:
-            self._membership_tasks.discard(t)
-            if t.cancelled():
-                return
-            exc = t.exception()
-            if exc is not None:
-                logger.error("[Orchestrator] %s task crashed", label, exc_info=exc)
-
-        task.add_done_callback(_on_done)
-
     # ---- Request handling ----
 
     async def _request_handler(self) -> None:
@@ -419,18 +336,11 @@ class Orchestrator:
             elif msg_type == "collective_rpc":
                 await self._handle_collective_rpc(msg)
             elif isinstance(msg, RegisterRemoteReplicaMessage):
-                # Dynamic-attach involves a ~5s blocking handshake (run in a
-                # thread by ``_build_remote_replica``); ``await`` here would
-                # block the queue and stall the next ``add_request`` until
-                # the attach finishes. Dispatch as a background task so the
-                # main message loop keeps draining.
-                self._spawn_membership_task(self._handle_register_remote_replica(msg), label="register_remote_replica")
+                if self._membership is not None:
+                    await self._membership.handle_register(msg.stage_id, msg.replica_id)
             elif isinstance(msg, UnregisterRemoteReplicaMessage):
-                # Symmetric with register: keep the main queue flowing.
-                self._spawn_membership_task(
-                    self._handle_unregister_remote_replica(msg),
-                    label="unregister_remote_replica",
-                )
+                if self._membership is not None:
+                    await self._membership.handle_unregister(msg.stage_id, msg.input_addr)
             elif isinstance(msg, ShutdownRequestMessage):
                 logger.info("[Orchestrator] Received shutdown signal")
                 self._shutdown_event.set()
@@ -441,7 +351,8 @@ class Orchestrator:
                     for client in pool.clients:
                         if hasattr(client, "_shutting_down"):
                             client._shutting_down = True
-                self._shutdown_stages()
+                # Stage teardown runs once in run()'s finally after the
+                # orchestration loop observes _shutdown_event and exits.
                 break
             else:
                 logger.warning("[Orchestrator] Unknown message type: %s", msg_type)
@@ -1300,7 +1211,7 @@ class Orchestrator:
         for next_input in next_inputs:
             # Only AR thinker stages consume encoder mm_features; downstream
             # (talker/code2wav/…) must not see them (avoids encoder-cache misses).
-            model_stage = getattr(next_client, "model_stage", None)
+            model_stage = getattr(getattr(next_pool.stage_vllm_config, "model_config", None), "model_stage", None)
             mm_features = req_state.mm_features if model_stage == "thinker" else None
             request = build_engine_core_request_from_tokens(
                 request_id=req_id,
@@ -1495,129 +1406,6 @@ class Orchestrator:
                 )
             self.request_states.pop(req_id, None)
 
-    # ---- Distributed-mode replica attach / detach ----
-
-    async def _watch_replica_list(self) -> None:
-        """Convert hub replica disappearances into unregister control messages."""
-        last_up: set[tuple[int, str]] = set()
-        while not self._shutdown_event.is_set():
-            if self._hub is None:
-                # No coordinator wired up; sleep coarsely and re-check shutdown.
-                try:
-                    await asyncio.sleep(self._WATCH_REPLICA_IDLE_INTERVAL_S)
-                except asyncio.CancelledError:
-                    raise
-                continue
-
-            try:
-                snap = self._hub.get_replica_list()
-                current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP}
-                for stage_id, addr in last_up - current:
-                    await self.request_async_queue.put(
-                        UnregisterRemoteReplicaMessage(
-                            stage_id=stage_id,
-                            input_addr=addr,
-                        )
-                    )
-                last_up = current
-            except asyncio.CancelledError:
-                raise
-            except Exception:
-                logger.exception("[Orchestrator] _watch_replica_list iteration failed")
-
-            try:
-                await asyncio.sleep(self._WATCH_REPLICA_INTERVAL_S)
-            except asyncio.CancelledError:
-                raise
-
-    async def _handle_register_remote_replica(self, msg: RegisterRemoteReplicaMessage) -> None:
-        """Bind a head-side client for a newly registered remote replica."""
-        stage_id = int(msg.stage_id)
-        replica_id = int(msg.replica_id)
-        if not (0 <= stage_id < self.num_stages):
-            logger.warning(
-                "[Orchestrator] register_remote_replica: stage_id %d out of range (num_stages=%d)",
-                stage_id,
-                self.num_stages,
-            )
-            return
-        if self._remote_replica_factory is None:
-            logger.warning(
-                "[Orchestrator] register_remote_replica received for stage=%d replica=%d but no factory installed",
-                stage_id,
-                replica_id,
-            )
-            return
-
-        try:
-            await self._attach_remote_replica(stage_id, replica_id)
-        except Exception:
-            logger.exception(
-                "[Orchestrator] failed to attach remote replica stage=%d replica=%d",
-                stage_id,
-                replica_id,
-            )
-
-    async def _handle_unregister_remote_replica(self, msg: UnregisterRemoteReplicaMessage) -> None:
-        """Tear down the head-side client for a vanished remote replica."""
-        stage_id = int(msg.stage_id)
-        input_addr = str(msg.input_addr)
-        if not (0 <= stage_id < self.num_stages):
-            return
-        pool = self.stage_pools[stage_id]
-        affected = pool.invalidate_addr(input_addr)
-        self._detach_remote_replica(stage_id, input_addr)
-        if affected:
-            await self._cleanup_request_ids(affected, abort=True)
-            for req_id in affected:
-                await self.output_async_queue.put(
-                    ErrorMessage(
-                        error="stage replica disappeared",
-                        request_id=req_id,
-                        stage_id=stage_id,
-                    )
-                )
-
-    async def _attach_remote_replica(self, stage_id: int, replica_id: int) -> None:
-        """Build a head-side stage client via the injected factory and register it."""
-        factory = self._remote_replica_factory
-        if factory is None:
-            return
-        pool = self.stage_pools[stage_id]
-        client = await factory(stage_id, replica_id)
-        input_addr = StagePool._client_input_addr(client)
-        if input_addr is None:
-            raise RuntimeError(
-                f"remote replica factory for stage {stage_id} produced a client without a discoverable input address"
-            )
-        pool.add_client(input_addr, client)
-        logger.info(
-            "[Orchestrator] attached remote replica stage=%d replica=%d addr=%s",
-            stage_id,
-            replica_id,
-            input_addr,
-        )
-
-    def _detach_remote_replica(self, stage_id: int, input_addr: str) -> None:
-        """Shut down + remove the head-side client at ``input_addr``."""
-        pool = self.stage_pools[stage_id]
-        client = pool.remove_client(input_addr)
-        if client is None:
-            return
-        try:
-            client.shutdown()
-        except Exception:
-            logger.exception(
-                "[Orchestrator] failed to shutdown client for stage=%d addr=%s",
-                stage_id,
-                input_addr,
-            )
-        logger.info(
-            "[Orchestrator] detached remote replica stage=%d addr=%s",
-            stage_id,
-            input_addr,
-        )
-
     def _shutdown_stages(self) -> None:
         """Shutdown all stage pools."""
         if self._stages_shutdown:

vllm_omni/engine/stage_client.py

@@ -26,9 +26,11 @@ class StageClient(Protocol):
     stage_id: int
     replica_id: int
     stage_type: str
+    model_stage: str | None
     final_output: bool
     final_output_type: FinalOutputModalityType | None
     default_sampling_params: OmniSamplingParams
+    prompt_expand_func: Any | None
     requires_multimodal_data: bool
     custom_process_input_func: Any | None
     engine_input_source: Sequence[int]
@@ -66,8 +68,6 @@ class StagePoolClient(StageClient, Protocol):
 class StagePoolLLMClient(StagePoolClient, Protocol):
     """Pool-facing API for LLM-style stages."""
 
-    model_stage: str | None
-
     async def add_request_async(self, request: EngineCoreRequest) -> None: ...
 
     async def get_output_async(self) -> EngineCoreOutputs: ...

vllm_omni/engine/stage_engine_core_client.py

@@ -6,18 +6,12 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture.
 
 from __future__ import annotations
 
-import contextlib
 import inspect
-import multiprocessing.connection
 import os
 import socket
-import threading
-import time
-import weakref
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
-import psutil
 import vllm.v1.engine as _vllm_engine_module
 import vllm.v1.engine.core_client as _vllm_core_client_module
 from vllm.logger import init_logger
@@ -43,8 +37,6 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-SHUTDOWN_TIMEOUT_S = 5
-
 
 def _default_process_engine_inputs(
     source_outputs: list[Any],
@@ -78,10 +70,8 @@ class StageEngineCoreClientBase(StageClientBase):
     - outputs_queue, output_queue_task
     - All utility methods (get_output_async, abort_requests_async, etc.)
 
-    The subprocess is spawned externally via ``spawn_stage_core`` /
-    ``complete_stage_handshake`` from *stage_engine_core_proc.py*.
-    In single-stage CLI mode, the client may instead attach to an
-    ``engine_manager`` / ``coordinator`` pair created elsewhere.
+    The stage engine subprocesses are owned by vLLM-style engine managers
+    (`resources.engine_manager`), matching `MPClient` ownership.
     """
 
     replica_id: int = 0
@@ -93,7 +83,6 @@ class StageEngineCoreClientBase(StageClientBase):
         log_stats: bool = False,
         metadata: StageMetadata | None = None,
         client_addresses: dict[str, str] | None = None,
-        proc: Any = None,
         engine_manager: Any = None,
         coordinator: Any = None,
         client_count: int = 1,
@@ -107,7 +96,6 @@ class StageEngineCoreClientBase(StageClientBase):
             log_stats=log_stats,
             metadata=metadata,
             client_addresses=client_addresses,
-            proc=proc,
             engine_manager=engine_manager,
             coordinator=coordinator,
             client_count=client_count,
@@ -125,7 +113,6 @@ class StageEngineCoreClientBase(StageClientBase):
         executor_class: type,
         log_stats: bool = False,
         client_addresses: dict[str, str] | None = None,
-        proc: Any = None,
         client_count: int = 1,
         client_index: int = 0,
         *,
@@ -140,10 +127,8 @@ class StageEngineCoreClientBase(StageClientBase):
         via helpers in stage_init_utils.py. This constructor just stores metadata
         and calls super().__init__().
 
-        The subprocess is spawned externally via ``spawn_stage_core`` /
-        ``complete_stage_handshake`` (see *stage_engine_core_proc.py*).
-        The resulting ``proc`` handle is passed in so this client can
-        manage the process lifecycle on shutdown.
+        The subprocess lifecycle is owned by the engine manager attached to
+        vLLM's background resources, not by this client directly.
         """
         # -------- Stage metadata (public fields used at runtime) --------
         self.replica_id = 0
@@ -151,19 +136,17 @@ class StageEngineCoreClientBase(StageClientBase):
             self.stage_id = metadata.stage_id
             self.replica_id = getattr(metadata, "replica_id", 0)
             self.stage_type = metadata.stage_type
-            self.engine_output_type = metadata.engine_output_type
+            self.model_stage = metadata.model_stage
             self.is_comprehension = metadata.is_comprehension
             self.requires_multimodal_data = metadata.requires_multimodal_data
             self.engine_input_source = metadata.engine_input_source
             self.final_output = metadata.final_output
             self.final_output_type = metadata.final_output_type
             self.default_sampling_params = metadata.default_sampling_params
+            self.prompt_expand_func = metadata.prompt_expand_func
             self.custom_process_input_func = metadata.custom_process_input_func
-            self.model_stage = metadata.model_stage
 
         self.engine_outputs: Any = None
-        self._proc = proc
-        self._shutting_down = False
         self.client_addresses = dict(client_addresses or {})
         self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None)
         self._kv_sender_host = self._resolve_contact_host()
@@ -203,6 +186,7 @@ class StageEngineCoreClientBase(StageClientBase):
             )
             if engine_manager is not None:
                 self.resources.engine_manager = engine_manager
+                self.start_engine_core_monitor()
             if coordinator is not None:
                 self.resources.coordinator = coordinator
         except Exception:
@@ -226,9 +210,6 @@ class StageEngineCoreClientBase(StageClientBase):
 
         self._initialize_kv_sender_endpoint()
 
-        if self._proc is not None:
-            self._start_proc_monitor()
-
         logger.info(
             "[%s] stage-%s [rep-%s] EngineCore running",
             client_name,
@@ -236,64 +217,6 @@ class StageEngineCoreClientBase(StageClientBase):
             self.replica_id,
         )
 
-    def _start_proc_monitor(self) -> None:
-        """Start a daemon thread that watches the subprocess sentinel.
-
-        When the subprocess dies without sending the ZMQ ``ENGINE_CORE_DEAD``
-        sentinel (e.g. SIGKILL, segfault, OOM-killer), this thread sets
-        ``resources.engine_dead`` so subsequent calls raise
-        ``EngineDeadError``.
-        """
-        proc = self._proc
-        resources_ref = weakref.ref(self.resources)
-        stage_id = self.stage_id
-        replica_id = self.replica_id
-
-        def _monitor() -> None:
-            try:
-                multiprocessing.connection.wait([proc.sentinel])
-            except Exception:
-                return
-
-            # Give multiprocessing a brief chance to publish a stable exitcode
-            # before emitting the monitor log. Without this, the sentinel can
-            # fire first and the log may misleadingly print ``None``.
-            with contextlib.suppress(Exception):
-                proc.join(timeout=0)
-            exitcode = proc.exitcode
-            if exitcode is None:
-                deadline = time.monotonic() + 0.2
-                while exitcode is None and time.monotonic() < deadline:
-                    time.sleep(0.01)
-                    with contextlib.suppress(Exception):
-                        proc.join(timeout=0)
-                    exitcode = proc.exitcode
-
-            resources = resources_ref()
-            if resources is None or resources.engine_dead or self._shutting_down:
-                return
-            if exitcode == 0:
-                logger.info(
-                    "[StageEngineCoreClient] stage-%s [rep-%s] subprocess exited cleanly.",
-                    stage_id,
-                    replica_id,
-                )
-                return
-            resources.engine_dead = True
-            logger.error(
-                "[StageEngineCoreClient] stage-%s [rep-%s] subprocess died unexpectedly (exit code %s).",
-                stage_id,
-                replica_id,
-                exitcode,
-            )
-
-        t = threading.Thread(
-            target=_monitor,
-            daemon=True,
-            name=f"StageCoreProcMonitor-{stage_id}",
-        )
-        t.start()
-
     def check_health(self) -> None:
         """Raise ``EngineDeadError`` if the stage subprocess is dead.
 
@@ -302,9 +225,6 @@ class StageEngineCoreClientBase(StageClientBase):
         """
         if self.resources.engine_dead:
             raise EngineDeadError(f"Stage-{self.stage_id} engine core is dead")
-        if self._proc is not None and not self._proc.is_alive():
-            self.resources.engine_dead = True
-            raise EngineDeadError(f"Stage-{self.stage_id} subprocess is not alive (exit code {self._proc.exitcode})")
 
     # ==================== Overrides ====================
 
@@ -335,7 +255,16 @@ class StageEngineCoreClientBase(StageClientBase):
                 return None
 
     def _resolve_contact_host(self) -> str | None:
-        """Resolve a routable host for this stage from its client addresses."""
+        """Resolve a routable host for this stage from its client addresses.
+
+        For remote LLM replicas the ZMQ sockets are bound on the head node,
+        so the ZMQ addresses contain the head's IP. The ``replica_host`` key
+        (injected by DistStageRuntime) carries the actual replica IP where
+        the KV connector binds — prefer it when available.
+        """
+        replica_host = self.client_addresses.get("replica_host")
+        if replica_host:
+            return replica_host
         for key in ("input_address", "output_address", "stats_update_address"):
             address = self.client_addresses.get(key)
             if not address:
@@ -503,41 +432,6 @@ class StageEngineCoreClientBase(StageClientBase):
             kwargs=kwargs,
         )
 
-    def shutdown(self, timeout: float | None = None) -> None:
-        """Shutdown managed resources and any externally spawned subprocess."""
-        self._shutting_down = True
-        child_procs: list[psutil.Process] = []
-        if self._proc is not None and self._proc.pid is not None:
-            try:
-                child_procs = psutil.Process(self._proc.pid).children(recursive=True)
-            except psutil.Error:
-                child_procs = []
-
-        try:
-            super().shutdown(timeout=timeout)
-        finally:
-            if self._proc is not None and self._proc.is_alive():
-                self._proc.terminate()
-                self._proc.join(timeout=SHUTDOWN_TIMEOUT_S)
-                if self._proc.is_alive():
-                    self._proc.kill()
-                    self._proc.join(timeout=SHUTDOWN_TIMEOUT_S)
-
-            alive_children = [proc for proc in child_procs if proc.is_running()]
-            for proc in alive_children:
-                try:
-                    proc.terminate()
-                except psutil.Error:
-                    pass
-            _, still_alive = psutil.wait_procs(alive_children, timeout=SHUTDOWN_TIMEOUT_S)
-            for proc in still_alive:
-                try:
-                    proc.kill()
-                except psutil.Error:
-                    pass
-            # The process handle is no longer reliable after best-effort cleanup.
-            self._proc = None
-
 
 class StageEngineCoreClient(StageEngineCoreClientBase, AsyncMPClient):
     """Stage async client backed by vLLM's ``AsyncMPClient``."""

vllm_omni/engine/stage_engine_core_proc.py

@@ -10,40 +10,28 @@ from __future__ import annotations
 import contextlib
 import os
 import signal
-from multiprocessing.process import BaseProcess
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
-import msgspec
 import vllm.v1.engine.core as _vllm_engine_core_module
-import zmq
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import (
     maybe_register_config_serialize_by_value,
 )
-from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
 from vllm.utils.system_utils import (
     decorate_logs,
-    get_mp_context,
     set_process_title,
 )
 from vllm.v1.engine import EngineCoreRequestType
 from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState
 from vllm.v1.engine.utils import (
-    EngineHandshakeMetadata,
     EngineZmqAddresses,
     SignalCallback,
-    get_engine_zmq_addresses,
 )
-from vllm.v1.utils import shutdown
 
-from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
+from vllm_omni.distributed.omni_coordinator import create_stage_coord_client
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.stage_init_utils import set_death_signal
 
-if TYPE_CHECKING:
-    from vllm.config import VllmConfig
-    from vllm.v1.executor import Executor
-
 logger = init_logger(__name__)
 
 
@@ -106,7 +94,7 @@ class StageEngineCoreProc(EngineCoreProc):
             )
 
         engine_core: StageEngineCoreProc | None = None
-        coord_client: OmniCoordClientForStage | None = None
+        coord_client = None
         try:
             # NOTE: previous revisions hardcoded data_parallel_size=1 here
             # (TODO referencing issue #984). The hardcoding has been removed
@@ -154,29 +142,17 @@ class StageEngineCoreProc(EngineCoreProc):
                         "EngineCore handshake did not populate input/output addresses; "
                         "cannot start OmniCoordClientForStage"
                     )
-                coord_client = OmniCoordClientForStage(
+                scheduler = getattr(engine_core, "scheduler", None)
+                if scheduler is None:
+                    raise RuntimeError("EngineCore scheduler is not initialized")
+                coord_client = create_stage_coord_client(
                     coord_zmq_addr=omni_coordinator_address,
                     input_addr=addresses.inputs[0],
                     output_addr=addresses.outputs[0],
                     stage_id=int(omni_stage_id),
+                    queue_length_getter=scheduler.get_num_unfinished_requests,
                 )
 
-                def _refresh_queue_length() -> None:
-                    """Pre-heartbeat hook: refresh queue_length from scheduler."""
-                    scheduler = getattr(engine_core, "scheduler", None)
-                    if scheduler is None:
-                        return
-                    try:
-                        coord_client._queue_length = int(  # type: ignore[union-attr]
-                            scheduler.get_num_unfinished_requests()
-                        )
-                    except Exception:
-                        # Live scheduler stats are best-effort — heartbeats
-                        # must not fail because of a stats lookup error.
-                        pass
-
-                coord_client._on_heartbeat = _refresh_queue_length
-
             def wakeup_engine() -> None:
                 engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None))
 
@@ -212,113 +188,3 @@ class StageEngineCoreProc(EngineCoreProc):
                     coord_client.close()
             if engine_core is not None:
                 engine_core.shutdown()
-
-
-def spawn_stage_core(
-    vllm_config: VllmConfig,
-    executor_class: type[Executor],
-    log_stats: bool = False,
-    omni_stage_id: int | None = None,
-    omni_replica_id: int = 0,
-) -> tuple[EngineZmqAddresses, BaseProcess, str]:
-    """Spawn a *StageEngineCoreProc* subprocess without performing the handshake.
-
-    Must be called while the correct device env vars are set (e.g. under
-    the stage-launch lock).  Call ``complete_stage_handshake`` afterwards.
-
-    Returns ``(addresses, process, handshake_address)``.
-    """
-    addresses = get_engine_zmq_addresses(vllm_config)
-    handshake_address = get_open_zmq_ipc_path()
-
-    ctx = get_mp_context()
-    proc = ctx.Process(
-        target=StageEngineCoreProc.run_stage_core,
-        name="StageEngineCoreProc",
-        kwargs={
-            "vllm_config": vllm_config,
-            "local_client": True,
-            "handshake_address": handshake_address,
-            "executor_class": executor_class,
-            "log_stats": log_stats,
-            "dp_rank": 0,
-            "local_dp_rank": 0,
-            "omni_stage_id": omni_stage_id,
-            "omni_replica_id": omni_replica_id,
-        },
-    )
-    proc.start()
-    return addresses, proc, handshake_address
-
-
-def complete_stage_handshake(
-    proc: BaseProcess,
-    handshake_address: str,
-    addresses: EngineZmqAddresses,
-    vllm_config: VllmConfig,
-    handshake_timeout: int,
-) -> None:
-    """Perform the HELLO/INIT/READY handshake with an already-spawned proc.
-
-    On failure the process is terminated before re-raising.
-    """
-    try:
-        _perform_handshake(proc, handshake_address, addresses, vllm_config, handshake_timeout)
-    except Exception:
-        shutdown([proc])
-        raise
-
-
-def _perform_handshake(
-    proc: BaseProcess,
-    handshake_address: str,
-    addresses: EngineZmqAddresses,
-    vllm_config: VllmConfig,
-    handshake_timeout: int,
-) -> None:
-    """Run the HELLO / INIT / READY handshake with the subprocess."""
-    with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket:
-        poller = zmq.Poller()
-        poller.register(handshake_socket, zmq.POLLIN)
-        poller.register(proc.sentinel, zmq.POLLIN)
-
-        identity, msg = _recv(poller, handshake_socket, proc, "HELLO", handshake_timeout)
-        if msg.get("status") != "HELLO":
-            raise RuntimeError(f"Expected HELLO, got: {msg}")
-
-        init_payload = EngineHandshakeMetadata(
-            addresses=addresses,
-            parallel_config={},
-        )
-        handshake_socket.send_multipart([identity, msgspec.msgpack.encode(init_payload)])
-
-        identity, msg = _recv(poller, handshake_socket, proc, "READY", handshake_timeout)
-        if msg.get("status") != "READY":
-            raise RuntimeError(f"Expected READY, got: {msg}")
-        num_gpu_blocks = msg.get("num_gpu_blocks")
-        if num_gpu_blocks is not None:
-            vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
-
-
-def _recv(
-    poller: zmq.Poller,
-    handshake_socket: zmq.Socket,
-    proc: BaseProcess,
-    expected: str,
-    timeout_s: int = 600,
-) -> tuple[bytes, dict]:
-    """Wait for one handshake message; raise if the process dies first."""
-    timeout_ms = timeout_s * 1000
-    while True:
-        events = dict(poller.poll(timeout=timeout_ms))
-        if not events:
-            raise TimeoutError(
-                f"Timed out waiting for {expected} from StageEngineCoreProc after {timeout_s}s. "
-                f"This typically indicates model loading or initialization is taking too long. "
-                f"Consider increasing `stage_init_timeout` for large models."
-            )
-        if handshake_socket in events:
-            identity, raw = handshake_socket.recv_multipart()
-            return identity, msgspec.msgpack.decode(raw)
-        if proc.exitcode is not None:
-            raise RuntimeError(f"StageEngineCoreProc died during {expected} (exit code {proc.exitcode})")

vllm_omni/engine/stage_engine_core_proc_manager.py

@@ -47,7 +47,7 @@ except ImportError:  # pragma: no cover - depends on vLLM build
     set_device_control_env_var = None  # type: ignore[assignment]
 
 
-class OmniCoreEngineProcManager(CoreEngineProcManager):
+class StageEngineCoreProcManager(CoreEngineProcManager):
     """Spawn :class:`StageEngineCoreProc` subprocesses with omni kwargs.
 
     The body mirrors :class:`CoreEngineProcManager.__init__` because the
@@ -102,7 +102,7 @@ class OmniCoreEngineProcManager(CoreEngineProcManager):
         # Intra-replica vLLM DP mesh (i.e. ``data_parallel_size`` ranks sharing
         # one engine, one DPCoordinator, one set of weights). Distinct from
         # the omni-level notion of multiple independent replicas of a stage —
-        # those each spawn their own OmniCoreEngineProcManager and never join
+        # those each spawn their own StageEngineCoreProcManager and never join
         # a vLLM DP group across replicas.
         has_intra_replica_dp = vllm_config.parallel_config.data_parallel_size > 1
 

vllm_omni/engine/stage_init_utils.py

@@ -51,6 +51,7 @@ class ReplicaInitPlan:
     omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None]
     stage_vllm_config: Any | None = None
     executor_class: type | None = None
+    engine_args_dict: dict[str, Any] | None = None
 
 
 @dataclass
@@ -58,65 +59,8 @@ class LogicalStageInitPlan:
     """Startup plan for one logical stage."""
 
     stage_idx: int
-    configured_stage_id: int
-    replicas: list[ReplicaInitPlan]
-
-
-@dataclass
-class StageRemoteFactoryContext:
-    """Per-stage context cached by AsyncOmniEngine for dynamic replica attach.
-
-    Populated once during ``_bootstrap_orchestrator`` from the per-stage
-    init plans. ``_build_remote_replica`` consumes it to construct the
-    right head-side stage client when a headless replica registers.
-    """
-
     stage_id: int
-    stage_type: str
-    stage_cfg: Any
-    base_metadata: Any
-    # LLM-only fields:
-    vllm_config: Any | None = None
-    executor_class: type | None = None
-    # Diffusion-only fields:
-    diffusion_batch_size: int = 1
-
-
-def capture_stage_factory_contexts(
-    stage_plans: Sequence[LogicalStageInitPlan],
-    diffusion_batch_size: int,
-) -> dict[int, StageRemoteFactoryContext]:
-    """Snapshot per-stage construction context for dynamic replica attach.
-
-    Called once after ``_initialize_stages`` finishes. The captured
-    context holds everything ``_build_remote_replica`` needs to build a
-    fresh head-side client when a new headless replica registers
-    (vllm_config / executor_class for LLM, batch_size for diffusion,
-    plus the base stage metadata).
-
-    Per-replica fields like ``replica_id`` are filled in at build time,
-    not at capture time.
-    """
-    contexts: dict[int, StageRemoteFactoryContext] = {}
-    for plan in stage_plans:
-        if not plan.replicas:
-            # Stage was declared but has zero replicas locally; we still
-            # want to be able to attach incoming headless ones, so use
-            # the stage_cfg-derived context if any replica plan exists.
-            continue
-        template = plan.replicas[0]
-        stage_id = int(plan.configured_stage_id)
-        stage_type = template.metadata.stage_type or "llm"
-        contexts[stage_id] = StageRemoteFactoryContext(
-            stage_id=stage_id,
-            stage_type=stage_type,
-            stage_cfg=template.stage_cfg,
-            base_metadata=template.metadata,
-            vllm_config=template.stage_vllm_config,
-            executor_class=template.executor_class,
-            diffusion_batch_size=diffusion_batch_size,
-        )
-    return contexts
+    replicas: list[ReplicaInitPlan]
 
 
 def _resolve_model_to_local_path(model: str) -> str:
@@ -359,6 +303,26 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[
         logger.debug("Failed to inject stage info into omni_kv_config: %s", e)
 
 
+def inject_omni_kv_connector_config(
+    engine_args_dict: dict[str, Any],
+    omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None],
+    stage_id: int,
+) -> None:
+    """Inject resolved connector config into a stage engine-args dict."""
+    omni_conn_cfg, omni_from, omni_to = omni_kv_connector
+    if not omni_conn_cfg:
+        return
+
+    omni_kv = engine_args_dict.get("omni_kv_config") or {}
+    if not isinstance(omni_kv, dict):
+        omni_kv = dict(omni_kv)
+    omni_kv["connector_config"] = omni_conn_cfg
+    omni_kv["omni_from_stage"] = omni_from
+    omni_kv["omni_to_stage"] = omni_to
+    omni_kv.setdefault("stage_id", stage_id)
+    engine_args_dict["omni_kv_config"] = omni_kv
+
+
 @dataclass
 class StageMetadata:
     """Lightweight stage attributes extracted from stage_config."""
@@ -1061,32 +1025,6 @@ def release_device_locks(lock_fds: list[int]) -> None:
             pass
 
 
-def acquire_diffusion_device_locks(
-    stage_id: int,
-    od_config: Any,
-    stage_init_timeout: int,
-) -> list[int]:
-    """Acquire init locks for the GPU set used by a diffusion stage.
-
-    Diffusion stages express their device count via ``OmniDiffusionConfig``'s
-    ``parallel_config.world_size`` rather than the LLM-style
-    ``tensor_parallel_size`` knob, so adapt to the shape that
-    ``acquire_device_locks`` understands.
-    """
-    parallel_config = getattr(od_config, "parallel_config", None)
-    world_size = getattr(parallel_config, "world_size", 1)
-    try:
-        world_size = max(1, int(world_size))
-    except (TypeError, ValueError):
-        world_size = 1
-
-    return acquire_device_locks(
-        stage_id,
-        {"tensor_parallel_size": world_size},
-        stage_init_timeout,
-    )
-
-
 def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any:
     """Load omni transfer config from an explicit path or resolved model config.
 

vllm_omni/entrypoints/cli/serve.py

@@ -9,10 +9,7 @@ import argparse
 import json
 import os
 import signal
-import threading
-from multiprocessing import connection
 from types import FrameType
-from typing import Any
 
 import uvloop
 from vllm.entrypoints.cli.types import CLISubcommand
@@ -697,33 +694,24 @@ def run_headless(args: TrackingNamespace) -> None:
     headless invocations can coexist) and reports heartbeats to the head's
     OmniCoordinator.
     """
-    from vllm.v1.engine.coordinator import DPCoordinator
     from vllm.v1.executor.multiproc_executor import MultiprocExecutor
     from vllm.version import __version__ as VLLM_VERSION
 
-    from vllm_omni.diffusion.stage_diffusion_proc import (
-        complete_diffusion_handshake,
-        spawn_diffusion_proc,
-    )
     from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage
-    from vllm_omni.engine.omni_core_engine_proc_manager import OmniCoreEngineProcManager
-    from vllm_omni.engine.stage_engine_startup import register_stage_with_omni_master
+    from vllm_omni.engine.stage_engine_startup import (
+        get_headless_replica_devices,
+        launch_headless_diffusion_replicas,
+        launch_headless_llm_replicas,
+    )
     from vllm_omni.engine.stage_init_utils import (
-        build_diffusion_config,
         build_engine_args_dict,
         build_vllm_config,
-        extract_stage_metadata,
         get_stage_connector_spec,
-        get_stage_devices_per_replica,
-        inject_kv_stage_info,
+        inject_omni_kv_connector_config,
         load_omni_transfer_config_for_model,
         prepare_engine_environment,
-        setup_stage_devices,
-        split_devices_for_replicas,
-        terminate_alive_proc,
     )
-    from vllm_omni.entrypoints.utils import inject_omni_kv_config, load_and_resolve_stage_configs
-    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
 
     model = args.model
     stage_id: int | None = args.stage_id
@@ -778,155 +766,25 @@ def run_headless(args: TrackingNamespace) -> None:
         )
 
     prepare_engine_environment()
-    omni_transfer_config = load_omni_transfer_config_for_model(model, config_path)
-    omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
-
-    # When ``--omni-dp-size-local > 1``, slice the YAML's ``devices:`` field
-    # into per-replica subsets so each subprocess we spawn below sees a
-    # narrowed ``CUDA_VISIBLE_DEVICES`` and doesn't stack on cuda:0. Mirrors
-    # the head-side per-replica device application at
-    # ``async_omni_engine.py`` (setup_stage_devices around each launch).
-    runtime_cfg = getattr(stage_cfg, "runtime", None)
-    devices_str: str | None = None
-    if runtime_cfg is not None:
-        devices_str = (
-            runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
-        )
-    devices_per_replica = get_stage_devices_per_replica(stage_cfg)
-    if devices_str:
-        # Always remap YAML's logical devices through setup_stage_devices,
-        # even for omni_dp_size_local==1. The launcher's CUDA_VISIBLE_DEVICES
-        # is dropped from the engine-subprocess env between vllm-serve and
-        # OmniCoreEngineProcManager.Process, so the worker would otherwise
-        # default cuda:0 to physical GPU 0 and collide with a co-located
-        # head on the same host (see hyi3_multi_host_1 reproducer).
-        per_replica_devices: list[str | None] = split_devices_for_replicas(
-            devices_str, omni_dp_size_local, devices_per_replica, stage_id
-        )
-        logger.info(
-            "[Headless] Stage %d: %d local replicas, devices_per_replica=%d, per-replica devices: %s",
-            stage_id,
-            omni_dp_size_local,
-            devices_per_replica,
-            per_replica_devices,
-        )
-    else:
-        per_replica_devices = [None] * omni_dp_size_local
-    device_control_env = current_omni_platform.device_control_env_var
+    per_replica_devices = get_headless_replica_devices(stage_cfg, stage_id, omni_dp_size_local)
 
     if stage_cfg.stage_type == "diffusion":
-        metadata = extract_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
-        # from the loaded deploy config so heterogeneous KV routing keys match
-        # the head process (e.g. from_tp=2, to_tp=1).
-        inject_kv_stage_info(stage_cfg, stage_id, stage_configs)
-        od_config = build_diffusion_config(model, stage_cfg, metadata)
-
-        logger.info(
-            "[Headless] Launching %d diffusion replica(s) for stage %d via OmniMasterServer at %s:%d",
-            omni_dp_size_local,
-            stage_id,
-            omni_master_address,
-            omni_master_port,
+        launch_headless_diffusion_replicas(
+            model=model,
+            stage_cfg=stage_cfg,
+            stage_configs=stage_configs,
+            stage_id=stage_id,
+            omni_master_address=omni_master_address,
+            omni_master_port=omni_master_port,
+            omni_dp_size_local=omni_dp_size_local,
+            per_replica_devices=per_replica_devices,
+            config_path=config_path,
+            replica_bind_address=omni_replica_address,
         )
+        return
 
-        procs: list[Any] = []
-        try:
-            for _rep_idx in range(omni_dp_size_local):
-                # Always auto-assign: headless processes carry no knowledge
-                # of their per-replica id and the master server is the sole
-                # authority on the per-stage id namespace.
-                response = register_stage_with_omni_master(
-                    omni_master_address=omni_master_address,
-                    omni_master_port=omni_master_port,
-                    omni_stage_id=stage_id,
-                    omni_stage_config=stage_cfg,
-                    replica_id=None,
-                    return_full_response=True,
-                    replica_bind_address=omni_replica_address,
-                )
-                # Apply this replica's CUDA_VISIBLE_DEVICES (only when
-                # ``--omni-dp-size-local > 1`` and the YAML's stage devices
-                # field is set). The spawned subprocess inherits the env at
-                # spawn time; we restore the parent env afterwards so the
-                # next replica's setup sees the same baseline.
-                previous_visible_devices = os.environ.get(device_control_env)
-                try:
-                    if per_replica_devices[_rep_idx] is not None:
-                        setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]})
-                    # Each StageDiffusionProc starts its own
-                    # torch.distributed group bound to
-                    # ``od_config.master_port``. Without an explicit
-                    # per-replica override all spawned subprocesses
-                    # share the value ``OmniDiffusionConfig.__post_init__``
-                    # picked once (and the second binder hits EADDRINUSE
-                    # on ``init_process_group``). We can't use
-                    # kernel-ephemeral allocation either, because the
-                    # master server's pre-allocated ZMQ ports (returned
-                    # by ``register_stage_with_omni_master``) also live
-                    # in the ephemeral range and are not actually bound
-                    # until the headless ``_perform_diffusion_handshake``
-                    # runs — so picking an ephemeral port here can steal
-                    # a port the master server already promised to a
-                    # sibling headless. Use ``settle_port`` from a base
-                    # above the Linux default ephemeral range
-                    # (32768-60999) so torch.distributed master ports
-                    # never overlap with ZMQ allocations.
-                    if omni_dp_size_local > 1:
-                        od_config.master_port = od_config.settle_port(
-                            61000 + _rep_idx * 100,
-                            port_inc=37,
-                        )
-                    proc, _, _, _ = spawn_diffusion_proc(
-                        model,
-                        od_config,
-                        handshake_address=response.handshake_address,
-                        request_address=response.input_address,
-                        response_address=response.output_address,
-                        omni_coordinator_address=response.coordinator_router_address,
-                        omni_stage_id=stage_id,
-                        omni_replica_id=response.replica_id,
-                    )
-                finally:
-                    if previous_visible_devices is None:
-                        current_omni_platform.unset_device_control_env_var()
-                    else:
-                        current_omni_platform.set_device_control_env_var(previous_visible_devices)
-                complete_diffusion_handshake(proc, response.handshake_address, args.stage_init_timeout)
-                procs.append(proc)
-                logger.info(
-                    "[Headless] Diffusion replica id=%d for stage %d is up (coord=%s)",
-                    response.replica_id,
-                    stage_id,
-                    response.coordinator_router_address,
-                )
-
-            # Block on the sentinel set so any replica crash is detected
-            # immediately (the previous per-proc join loop only noticed
-            # crashes in registration order). Any exit triggers fleet
-            # shutdown via the finally block; non-zero exits propagate.
-            sentinel_to_proc = {p.sentinel: p for p in procs}
-            died = connection.wait(list(sentinel_to_proc.keys()))
-            first = sentinel_to_proc[died[0]]
-            logger.info(
-                "[Headless] Diffusion replica %s exited (code=%s); shutting down stage %d.",
-                first.name,
-                first.exitcode,
-                stage_id,
-            )
-            if first.exitcode not in (None, 0):
-                raise RuntimeError(
-                    f"Diffusion stage {stage_id} replica {first.name!r} exited with code {first.exitcode}"
-                )
-            return
-        finally:
-            logger.info("[Headless] Shutting down %d diffusion replica(s) for stage %d.", len(procs), stage_id)
-            for p in procs:
-                if p.is_alive():
-                    terminate_alive_proc(p)
-
+    omni_transfer_config = load_omni_transfer_config_for_model(model, config_path)
+    omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
     stage_connector_spec = get_stage_connector_spec(
         omni_transfer_config=omni_transfer_config,
         stage_id=stage_id,
@@ -944,17 +802,7 @@ def run_headless(args: TrackingNamespace) -> None:
         cli_tokenizer=getattr(args, "tokenizer", None),
     )
 
-    # Inject omni KV connector config so the engine runner can initialize the
-    # correct connector (sender/receiver role, type, addresses, etc.).
-    if omni_conn_cfg:
-        omni_kv = engine_args_dict.get("omni_kv_config") or {}
-        if not isinstance(omni_kv, dict):
-            omni_kv = dict(omni_kv)
-        omni_kv["connector_config"] = omni_conn_cfg
-        omni_kv["omni_from_stage"] = omni_from
-        omni_kv["omni_to_stage"] = omni_to
-        omni_kv.setdefault("stage_id", stage_id)
-        engine_args_dict["omni_kv_config"] = omni_kv
+    inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id)
 
     vllm_config, executor_class = build_vllm_config(
         stage_cfg,
@@ -964,10 +812,6 @@ def run_headless(args: TrackingNamespace) -> None:
         headless=True,
     )
     parallel_config = vllm_config.parallel_config
-    local_engine_count = parallel_config.data_parallel_size_local
-
-    if local_engine_count <= 0:
-        raise ValueError("data_parallel_size_local must be > 0 in headless mode")
 
     shutdown_requested = False
 
@@ -994,126 +838,22 @@ def run_headless(args: TrackingNamespace) -> None:
         executor.start_worker_monitor(inline=True)
         return
 
-    dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0
-    coordinator = None
-    if vllm_config.needs_dp_coordinator and dp_rank == 0:
-        coordinator = DPCoordinator(
-            parallel_config,
-            enable_wave_coordination=vllm_config.model_config.is_moe,
-        )
-        logger.info(
-            "[Headless] Started DP Coordinator process for stage %d (PID: %d)",
-            stage_id,
-            coordinator.proc.pid,
-        )
-
-    logger.info(
-        "[Headless] Launching %d omni replica(s) (vLLM dp_size_local=%d each) for stage %d "
-        "via OmniMasterServer at %s:%d",
-        omni_dp_size_local,
-        local_engine_count,
-        stage_id,
-        omni_master_address,
-        omni_master_port,
-    )
-
-    # One OmniMasterServer registration per omni replica; each registration
-    # yields its own (handshake, input, output) allocation and the head's
-    # OmniCoordinator ROUTER address. We then spawn one
-    # OmniCoreEngineProcManager per replica so its subprocess gets the
-    # right replica id wired into its OmniCoordClientForStage.
     log_stats = bool(args.log_stats)
     if args.disable_log_stats:
         log_stats = False
 
-    engine_managers: list[Any] = []
-    monitor_threads: list[threading.Thread] = []
-
-    def _monitor_target(mgr: Any) -> None:
-        try:
-            mgr.monitor_engine_liveness()
-        except Exception:
-            logger.exception("[Headless] monitor_engine_liveness raised")
-
-    try:
-        for _rep_idx in range(omni_dp_size_local):
-            # Always auto-assign: see the diffusion branch comment above
-            # for the rationale (headless owns no replica-id namespace).
-            response = register_stage_with_omni_master(
-                omni_master_address=omni_master_address,
-                omni_master_port=omni_master_port,
-                omni_stage_id=stage_id,
-                omni_stage_config=stage_cfg,
-                coordinator=coordinator,
-                replica_id=None,
-                return_full_response=True,
-                replica_bind_address=omni_replica_address,
-                # LLM headless: the head binds *all* three sockets —
-                # handshake ROUTER (``connect_remote_engine_cores``),
-                # input ROUTER and output PULL (``CoreClient`` —
-                # ``make_zmq_socket`` defaults bind=True for PULL).
-                # The remote LLM worker is purely a connector: it
-                # opens 3 outbound TCP connections to the master's
-                # host. So the master must keep every address on
-                # its own host; rewriting any of them to this
-                # replica's NIC makes the head's ``bind`` go
-                # EADDRNOTAVAIL on a cross-host launch.
-                replica_binds_sockets=False,
-            )
-            # Per-replica CUDA_VISIBLE_DEVICES, same pattern as the diffusion
-            # branch above. OmniCoreEngineProcManager.__init__ spawns its
-            # subprocesses via context.Process inside the constructor, so we
-            # must set the env *before* instantiation and restore after.
-            previous_visible_devices = os.environ.get(device_control_env)
-            try:
-                if per_replica_devices[_rep_idx] is not None:
-                    setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]})
-                mgr = OmniCoreEngineProcManager(
-                    local_engine_count=local_engine_count,
-                    start_index=dp_rank,
-                    local_start_index=0,
-                    vllm_config=vllm_config,
-                    local_client=False,
-                    handshake_address=response.handshake_address,
-                    executor_class=executor_class,
-                    log_stats=log_stats,
-                    omni_stage_id=stage_id,
-                    omni_coordinator_address=response.coordinator_router_address,
-                    omni_replica_base_id=response.replica_id,
-                )
-            finally:
-                if previous_visible_devices is None:
-                    current_omni_platform.unset_device_control_env_var()
-                else:
-                    current_omni_platform.set_device_control_env_var(previous_visible_devices)
-            engine_managers.append(mgr)
-            logger.info(
-                "[Headless] Stage %d replica id=%d up (coord=%s)",
-                stage_id,
-                response.replica_id,
-                response.coordinator_router_address,
-            )
-
-        # Run all managers' liveness monitors in parallel. Each blocks
-        # until its own subprocesses exit (or fail).
-        if len(engine_managers) == 1:
-            engine_managers[0].monitor_engine_liveness()
-        else:
-            for mgr in engine_managers:
-                t = threading.Thread(target=_monitor_target, args=(mgr,), name=f"omni-replica-monitor-{id(mgr):x}")
-                t.start()
-                monitor_threads.append(t)
-            for t in monitor_threads:
-                t.join()
-    finally:
-        logger.info("[Headless] Shutting down stage %d (%d managers).", stage_id, len(engine_managers))
-        for mgr in engine_managers:
-            try:
-                mgr.shutdown()
-            except Exception:
-                logger.exception("[Headless] engine manager shutdown failed")
-        if coordinator is not None:
-            coordinator.shutdown()
+    launch_headless_llm_replicas(
+        vllm_config=vllm_config,
+        executor_class=executor_class,
+        log_stats=log_stats,
+        omni_master_address=omni_master_address,
+        omni_master_port=omni_master_port,
+        stage_id=stage_id,
+        stage_config=stage_cfg,
+        omni_dp_size_local=omni_dp_size_local,
+        per_replica_devices=per_replica_devices,
+        replica_bind_address=omni_replica_address,
+    )
 
 
 def cmd_init() -> list[CLISubcommand]:

vllm_omni/entrypoints/stage_utils.py

@@ -11,9 +11,39 @@ from vllm_omni.platforms import current_omni_platform
 logger = logging.getLogger(__name__)
 
 
+def resolve_stage_physical_devices(
+    stage_id: int,
+    devices: str | int | None,
+    *,
+    visible_baseline: str | None = None,
+) -> str | None:
+    """Map logical stage devices to physical IDs without mutating process env.
+
+    When ``visible_baseline`` is provided it is used for logical-to-physical
+    mapping instead of the current device-control env var. This keeps parallel
+    stage initialization correct even if another thread has temporarily
+    narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage.
+    """
+    if devices in (None, "cpu"):
+        return None
+
+    if not isinstance(devices, (int, str)):
+        raise TypeError(f"Expected str or int device IDs for stage initialization, got type {type(devices)}")
+
+    device_list = _parse_device_list(devices)
+    env_var = current_omni_platform.device_control_env_var
+    visible_devices = visible_baseline if visible_baseline is not None else os.environ.get(env_var)
+    if visible_devices is not None:
+        visible_device_list = _parse_device_list(visible_devices)
+        device_list = _map_device_list(stage_id, device_list, visible_device_list)
+    return ",".join(device_list)
+
+
 def set_stage_devices(
     stage_id: int,
     devices: str | int | None,
+    *,
+    visible_baseline: str | None = None,
 ) -> str | None:
     """Configure per-stage device visibility and current device (CUDA or NPU).
 
@@ -42,19 +72,18 @@ def set_stage_devices(
         The list of physical devices that were set for the given stage
         or None if we have no passed devices / are using cpu.
     """
-    env_var = current_omni_platform.device_control_env_var
-    vis = os.environ.get(env_var)
-
     if devices in (None, "cpu"):
         logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
         return None
 
     elif isinstance(devices, (int, str)):
-        device_list = _parse_device_list(devices)
-        if vis is not None:
-            visible_device_list = _parse_device_list(vis)
-            device_list = _map_device_list(stage_id, device_list, visible_device_list)
-        device_str = ",".join(device_list)
+        device_str = resolve_stage_physical_devices(
+            stage_id,
+            devices,
+            visible_baseline=visible_baseline,
+        )
+        if device_str is None:
+            return None
         current_omni_platform.set_device_control_env_var(device_str)
         return device_str