[large-scale-serving] Integrate OmniCoordinator into stage engine pipeline (#3569)

Signed-off-by: chickeyton <ngton2014@gmail.com>
Signed-off-by: herotai214 <herotai214@gmail.com>
Co-authored-by: herotai214 <herotai214@gmail.com>

tests/entrypoints/test_serve.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import argparse
+from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
@@ -31,143 +33,118 @@ def test_serve_parser_accepts_no_async_chunk() -> None:
     assert args.async_chunk is False
 
 
-def _make_headless_args() -> argparse.Namespace:
-    return argparse.Namespace(
+# ---------------------------------------------------------------------------
+# run_headless validation
+# ---------------------------------------------------------------------------
+
+
+def _make_headless_args(**overrides: Any) -> argparse.Namespace:
+    """Build an argparse.Namespace shaped like the headless CLI passes in.
+
+    Defaults pass every validation gate so individual tests can mutate just
+    the field they're exercising.
+    """
+    defaults = dict(
         model="fake-model",
-        stage_id=3,
+        stage_id=0,
         replica_id=0,
         omni_master_address="127.0.0.1",
         omni_master_port=26000,
+        omni_replica_address=None,
+        omni_dp_size_local=1,
         api_server_count=0,
         worker_backend="multi_process",
         stage_configs_path=None,
+        deploy_config=None,
         log_stats=False,
         disable_log_stats=False,
         stage_init_timeout=600,
+        tokenizer=None,
     )
+    defaults.update(overrides)
+    return argparse.Namespace(**defaults)
 
 
-def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    stage_cfg = mocker.Mock(stage_id=3)
-    stage_cfgs = [stage_cfg]
-    parallel_config = mocker.Mock(
-        data_parallel_size_local=2,
-        data_parallel_rank=4,
-        data_parallel_rank_local=1,
-        node_rank_within_dp=0,
-    )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    vllm_config.needs_dp_coordinator = False
-    executor_class = mocker.Mock()
-    engine_manager = mocker.Mock()
+def test_run_headless_requires_stage_id() -> None:
+    args = _make_headless_args(stage_id=None)
+    with pytest.raises(ValueError, match="--stage-id is required"):
+        run_headless(args)
 
+
+def test_run_headless_requires_master_address() -> None:
+    args = _make_headless_args(omni_master_address=None)
+    with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"):
+        run_headless(args)
+
+
+def test_run_headless_requires_master_port() -> None:
+    args = _make_headless_args(omni_master_port=None)
+    with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"):
+        run_headless(args)
+
+
+def test_run_headless_rejects_multi_api_server_count() -> None:
+    args = _make_headless_args(api_server_count=2)
+    with pytest.raises(ValueError, match="api_server_count can't be set"):
+        run_headless(args)
+
+
+def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
+    args = _make_headless_args(worker_backend="ray")
+    with pytest.raises(ValueError, match="worker_backend=multi_process"):
+        run_headless(args)
+
+
+def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None:
+    """Headless looks up its assigned stage_id in the loaded deploy YAML and
+    fails fast when the launcher's --stage-id doesn't match any entry."""
+    other_stage = SimpleNamespace(stage_id=99)
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
+        return_value=("/fake/stages.yaml", [other_stage]),
     )
-    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
-    mocker.patch(
-        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
-        return_value=(None, None, None),
-    )
-    mock_build_vllm_config = mocker.patch(
-        "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, executor_class),
-    )
-    mock_register = mocker.patch(
-        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
-    )
-    mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
-    mocker.patch("signal.signal")
-    run_headless(args)
 
-    mock_build_vllm_config.assert_called_once_with(
-        stage_cfg,
-        "fake-model",
-        stage_connector_spec={},
-        engine_args_dict={},
-        headless=True,
-    )
-    mock_register.assert_called_once_with(
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_stage_id=3,
-        omni_stage_config=stage_cfg,
-        coordinator=None,
-        replica_id=0,
-    )
-    mock_manager_cls.assert_called_once()
-    manager_kwargs = mock_manager_cls.call_args.kwargs
-    assert manager_kwargs["local_engine_count"] == 2
-    assert manager_kwargs["start_index"] == 4
-    assert manager_kwargs["local_start_index"] == 0
-    assert manager_kwargs["local_client"] is False
-    assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
-    assert manager_kwargs["log_stats"] is False
-    engine_manager.monitor_engine_liveness.assert_called_once_with()
-    engine_manager.shutdown.assert_called_once_with()
+    args = _make_headless_args(stage_id=0)
+    with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
+        run_headless(args)
 
 
-def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.log_stats = True
-    stage_cfg = mocker.Mock(stage_id=3)
-    stage_cfgs = [stage_cfg]
-    parallel_config = mocker.Mock(
-        data_parallel_size_local=2,
-        data_parallel_rank=4,
-        data_parallel_rank_local=1,
-        node_rank_within_dp=0,
-    )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    executor_class = mocker.Mock()
-    engine_manager = mocker.Mock()
-
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
-    )
-    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
-    mocker.patch(
-        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
-        return_value=(None, None, None),
-    )
-    mocker.patch(
-        "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, executor_class),
-    )
-    mocker.patch(
-        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
-    )
-    mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
-    mocker.patch("signal.signal")
-    run_headless(args)
-
-    manager_kwargs = mock_manager_cls.call_args.kwargs
-    assert manager_kwargs["log_stats"] is True
+# ---------------------------------------------------------------------------
+# run_headless happy paths
+# ---------------------------------------------------------------------------
 
 
-def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.replica_id = 2
-    stage_cfg = mocker.Mock(stage_id=3)
-    parallel_config = mocker.Mock(
+def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace:
+    """Build a stage config that satisfies every attribute run_headless reads.
+
+    Notably ``engine_args`` is a real dict (not a Mock) so
+    ``get_stage_devices_per_replica`` can call ``.get("tensor_parallel_size")``
+    and feed the result through ``int()`` without TypeError.
+    """
+    return SimpleNamespace(
+        stage_id=stage_id,
+        stage_type=stage_type,
+        # No "devices" key -> split_devices_for_replicas skipped, each replica
+        # inherits the launcher's CUDA_VISIBLE_DEVICES.
+        runtime=None,
+        engine_args={},
+    )
+
+
+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."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(0, stage_type="llm")
+    parallel_config = SimpleNamespace(
         data_parallel_size_local=1,
         data_parallel_rank=0,
         data_parallel_rank_local=0,
         node_rank_within_dp=0,
     )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    vllm_config.needs_dp_coordinator = False
+    vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False)
     engine_manager = mocker.Mock()
 
     mocker.patch(
@@ -175,82 +152,247 @@ def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None:
         return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
     mocker.patch(
         "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
         return_value=(None, None, None),
     )
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
     mocker.patch(
         "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, mocker.Mock()),
+        return_value=(vllm_config, object),
     )
     mock_register = mocker.patch(
         "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address="tcp://127.0.0.1:26100",
+        ),
+    )
+    mock_manager_cls = mocker.patch(
+        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        return_value=engine_manager,
     )
-    mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
     mocker.patch("signal.signal")
 
-    run_headless(args)
+    run_headless(_make_headless_args(stage_id=0))
 
-    assert mock_register.call_args.kwargs["replica_id"] == 2
+    # 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).
+    assert mock_register.call_count == 1
+    kwargs = mock_register.call_args.kwargs
+    assert kwargs["omni_master_address"] == "127.0.0.1"
+    assert kwargs["omni_master_port"] == 26000
+    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 mock_manager_cls.call_count == 1
+    mgr_kwargs = mock_manager_cls.call_args.kwargs
+    assert mgr_kwargs["local_engine_count"] == 1
+    assert mgr_kwargs["local_client"] is False
+    assert mgr_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
+    assert mgr_kwargs["omni_stage_id"] == 0
+    assert mgr_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100"
+    assert mgr_kwargs["omni_replica_base_id"] == 0
+
+    engine_manager.monitor_engine_liveness.assert_called_once_with()
+    engine_manager.shutdown.assert_called_once_with()
 
 
-def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.replica_id = 1
-    stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion")
-    stage_cfg.engine_args = mocker.Mock()
-    stage_cfg.engine_input_source = []
-    stage_cfgs = [stage_cfg]
-    metadata = mocker.Mock(stage_id=3, stage_type="diffusion")
+def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: MockerFixture) -> None:
+    """``--omni-dp-size-local=N`` must spawn N managers, each with its own
+    master-assigned replica_id, and join all of them before returning."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(0, stage_type="llm")
+    parallel_config = SimpleNamespace(
+        data_parallel_size_local=1,
+        data_parallel_rank=0,
+        data_parallel_rank_local=0,
+        node_rank_within_dp=0,
+    )
+    vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False)
+    manager_a = mocker.Mock()
+    manager_b = mocker.Mock()
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        return_value=("/fake/stages.yaml", [stage_cfg]),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.build_vllm_config",
+        return_value=(vllm_config, object),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        side_effect=[
+            StageRegistrationResponse(
+                handshake_address=f"tcp://127.0.0.1:2700{idx}",
+                input_address=f"tcp://127.0.0.1:2710{idx}",
+                output_address=f"tcp://127.0.0.1:2720{idx}",
+                replica_id=idx,
+                coordinator_router_address=None,
+            )
+            for idx in (0, 1)
+        ],
+    )
+    mock_manager_cls = mocker.patch(
+        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        side_effect=[manager_a, manager_b],
+    )
+    mocker.patch("signal.signal")
+
+    run_headless(_make_headless_args(stage_id=0, omni_dp_size_local=2))
+
+    assert mock_manager_cls.call_count == 2
+    assigned_ids = [call.kwargs["omni_replica_base_id"] for call in mock_manager_cls.call_args_list]
+    assert assigned_ids == [0, 1]
+
+    # Multi-replica path joins the monitor threads instead of calling
+    # ``monitor_engine_liveness`` synchronously on the main thread, but every
+    # manager must still be shut down in the finally block.
+    manager_a.shutdown.assert_called_once_with()
+    manager_b.shutdown.assert_called_once_with()
+
+
+def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) -> None:
+    """Diffusion headless: registers as auto-assign, spawns a single
+    ``StageDiffusionProc`` per local replica, and waits for it via
+    ``multiprocessing.connection.wait``."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(1, stage_type="diffusion")
     od_config = mocker.Mock()
-    proc = mocker.Mock()
-    proc.exitcode = 0
+    proc = mocker.Mock(sentinel=object(), exitcode=0)
     proc.is_alive.return_value = False
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
-    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
     mocker.patch(
         "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
         return_value=(None, None, None),
     )
-    mocker.patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata)
-    mock_inject_stage_info = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
+    )
+    mock_inject = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
     mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config)
     mock_register = mocker.patch(
         "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            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, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        return_value=(proc, None, None, None),
     )
     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(args)
 
-    mock_inject_stage_info.assert_called_once_with(stage_cfg, 3)
-    mock_register.assert_called_once_with(
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_stage_id=3,
-        omni_stage_config=stage_cfg,
-        return_addresses=True,
-        replica_id=1,
-    )
-    mock_spawn.assert_called_once_with(
-        "fake-model",
-        od_config,
-        handshake_address="tcp://127.0.0.1:26001",
-        request_address="tcp://127.0.0.1:26002",
-        response_address="tcp://127.0.0.1:26003",
-    )
+    run_headless(_make_headless_args(stage_id=1))
+
+    mock_inject.assert_called_once_with(stage_cfg, 1)
+
+    reg_kwargs = mock_register.call_args.kwargs
+    assert reg_kwargs["omni_master_address"] == "127.0.0.1"
+    assert reg_kwargs["omni_master_port"] == 26000
+    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
+
+    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)
-    proc.join.assert_called_once_with()
+
+
+def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None:
+    """A diffusion replica that exits with a non-zero code must surface as a
+    RuntimeError from ``run_headless`` (the head needs the signal to roll
+    back its own stage init)."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(1, stage_type="diffusion")
+    proc = mocker.Mock(sentinel=object(), exitcode=137, name="proc-stage1-rep0")
+    proc.is_alive.return_value = False
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        return_value=("/fake/stages.yaml", [stage_cfg]),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=mocker.Mock())
+    mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address=None,
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
+        return_value=(proc, None, None, None),
+    )
+    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"):
+        run_headless(_make_headless_args(stage_id=1))