[Entrypoint][Refactor]Stage CLI Refactor (#2020)

Signed-off-by: wuhang <wuhang6@huawei.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

tests/conftest.py

@@ -73,6 +73,8 @@ class OmniServerParams(NamedTuple):
     server_args: list[str] | None = None
     env_dict: dict[str, str] | None = None
     use_omni: bool = True
+    use_stage_cli: bool = False
+    init_timeout: int | None = None
 
 
 def assert_image_diffusion_response(
@@ -1546,6 +1548,183 @@ class OmniServer:
         cleanup_dist_env_and_memory()
 
 
+class OmniServerStageCli(OmniServer):
+    """Omni server harness that exercises the stage CLI flow."""
+
+    def __init__(
+        self,
+        model: str,
+        stage_config_path: str,
+        serve_args: list[str] | None = None,
+        *,
+        stage_ids: list[int] | None = None,
+        port: int | None = None,
+        env_dict: dict[str, str] | None = None,
+    ) -> None:
+        super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True)
+        self.stage_config_path = stage_config_path
+        self.master_port = get_open_port()
+        self.visible_device_list = self._load_visible_device_list(env_dict)
+        self.stage_runtime_devices = self._load_stage_runtime_devices(stage_config_path)
+        self.stage_ids = stage_ids or self._load_stage_ids(stage_config_path)
+        if 0 not in self.stage_ids:
+            raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}")
+        self.stage_procs: dict[int, subprocess.Popen] = {}
+        self.proc = None
+
+    @staticmethod
+    def _load_stage_ids(stage_config_path: str) -> list[int]:
+        with open(stage_config_path, encoding="utf-8") as f:
+            cfg = yaml.safe_load(f) or {}
+
+        stage_ids = [stage["stage_id"] for stage in cfg.get("stage_args", []) if "stage_id" in stage]
+        if not stage_ids:
+            raise ValueError(f"No stage IDs found in config: {stage_config_path}")
+        return stage_ids
+
+    @staticmethod
+    def _load_stage_runtime_devices(stage_config_path: str) -> dict[int, str]:
+        with open(stage_config_path, encoding="utf-8") as f:
+            cfg = yaml.safe_load(f) or {}
+
+        runtime_devices: dict[int, str] = {}
+        for stage in cfg.get("stage_args", []):
+            stage_id = stage.get("stage_id")
+            devices = stage.get("runtime", {}).get("devices")
+            if stage_id is not None and devices:
+                runtime_devices[int(stage_id)] = str(devices)
+        return runtime_devices
+
+    @classmethod
+    def _parse_device_list(cls, devices: str | int) -> list[str]:
+        if isinstance(devices, int):
+            if devices < 0:
+                raise ValueError("Device IDs must be non-negative integers")
+            return [str(devices)]
+        return [token.strip() for token in str(devices).split(",") if token.strip()]
+
+    @classmethod
+    def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None:
+        env = os.environ.copy()
+        if env_dict is not None:
+            env.update(env_dict)
+
+        env_var = getattr(current_omni_platform, "device_control_env_var", None)
+        if env_var and env_var in env:
+            return [token.strip() for token in env[env_var].split(",") if token.strip()]
+        return None
+
+    @classmethod
+    def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str:
+        device_list = cls._parse_device_list(devices)
+
+        if visible_device_list is None:
+            return ",".join(device_list)
+
+        if not all(device.isdigit() for device in device_list):
+            raise ValueError("Logical devices must be non-negative integers")
+
+        logical_ids = [int(device) for device in device_list]
+        if logical_ids and max(logical_ids) >= len(visible_device_list):
+            raise ValueError(
+                f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices"
+            )
+
+        return ",".join(visible_device_list[idx] for idx in logical_ids)
+
+    def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None:
+        mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices)
+        env_var = getattr(current_omni_platform, "device_control_env_var", None)
+        if env_var:
+            env[env_var] = mapped_devices
+
+    def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]:
+        cmd = [
+            sys.executable,
+            "-m",
+            "vllm_omni.entrypoints.cli.main",
+            "serve",
+            self.model,
+            "--omni",
+            "--stage-configs-path",
+            self.stage_config_path,
+            "--stage-id",
+            str(stage_id),
+            "--omni-master-address",
+            self.host,
+            "--omni-master-port",
+            str(self.master_port),
+        ]
+
+        if headless:
+            cmd.append("--headless")
+        else:
+            cmd += ["--host", self.host, "--port", str(self.port)]
+
+        cmd += self.serve_args
+        return cmd
+
+    def _launch_stage(self, stage_id: int, *, headless: bool) -> None:
+        env = os.environ.copy()
+        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        if self.env_dict is not None:
+            env.update(self.env_dict)
+
+        devices = self.stage_runtime_devices.get(stage_id)
+        if devices:
+            self._set_stage_device_env(stage_id, env, devices)
+
+        cmd = self._build_stage_cmd(stage_id, headless=headless)
+        print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}")
+        proc = subprocess.Popen(
+            cmd,
+            env=env,
+            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        )
+        self.stage_procs[stage_id] = proc
+        if stage_id == 0:
+            self.proc = proc
+
+    def _ensure_stage_processes_alive(self) -> None:
+        for stage_id, proc in self.stage_procs.items():
+            ret = proc.poll()
+            if ret is not None:
+                raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.")
+
+    def _start_server(self) -> None:
+        ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]]
+
+        self._launch_stage(0, headless=False)
+        time.sleep(2)
+        self._ensure_stage_processes_alive()
+
+        for stage_id in ordered_stage_ids[1:]:
+            self._launch_stage(stage_id, headless=True)
+
+        max_wait = 1200
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            self._ensure_stage_processes_alive()
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                sock.settimeout(1)
+                result = sock.connect_ex((self.host, self.port))
+                if result == 0:
+                    print(f"OmniServerStageCli ready on {self.host}:{self.port}")
+                    return
+            time.sleep(2)
+
+        raise RuntimeError(f"OmniServerStageCli failed to start within {max_wait} seconds")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for stage_id in sorted(self.stage_procs, reverse=True):
+            proc = self.stage_procs[stage_id]
+            if proc.poll() is None:
+                self._kill_process_tree(proc.pid)
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
+        cleanup_dist_env_and_memory()
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--run-level",
@@ -1568,9 +1747,11 @@ _omni_server_lock = threading.Lock()
 
 @pytest.fixture(scope="module")
 def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]:
-    """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses session scope so the server starts only once for the entire test session.
-    Multi-stage initialization can take 10-20+ minutes.
+    """Start vLLM-Omni through the standard or stage-CLI launcher.
+
+    The fixture stays module-scoped because multi-stage initialization is costly.
+    The ``use_stage_cli`` flag on ``OmniServerParams`` routes the setup through the
+    stage-CLI harness while still reusing the same fixture grouping semantics.
     """
     with _omni_server_lock:
         params: OmniServerParams = request.param
@@ -1589,28 +1770,47 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
         server_args = params.server_args or []
         if params.use_omni:
             server_args = ["--stage-init-timeout", "120", *server_args]
-        if stage_config_path is not None:
-            server_args += ["--stage-configs-path", stage_config_path]
+        if params.init_timeout is not None:
+            server_args = [*server_args, "--init-timeout", str(params.init_timeout)]
+        if params.use_stage_cli:
+            if not params.use_omni:
+                raise ValueError("omni_server with use_stage_cli=True requires use_omni=True")
+            if stage_config_path is None:
+                raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path")
 
-        with (
-            OmniServer(
+            with OmniServerStageCli(
                 model,
+                stage_config_path,
                 server_args,
                 port=port,
                 env_dict=params.env_dict,
-                use_omni=params.use_omni,
-            )
-            if port
-            else OmniServer(
-                model,
-                server_args,
-                env_dict=params.env_dict,
-                use_omni=params.use_omni,
-            )
-        ) as server:
-            print("OmniServer started successfully")
-            yield server
-            print("OmniServer stopping...")
+            ) as server:
+                print("OmniServer started successfully")
+                yield server
+                print("OmniServer stopping...")
+        else:
+            if stage_config_path is not None:
+                server_args += ["--stage-configs-path", stage_config_path]
+
+            with (
+                OmniServer(
+                    model,
+                    server_args,
+                    port=port,
+                    env_dict=params.env_dict,
+                    use_omni=params.use_omni,
+                )
+                if port
+                else OmniServer(
+                    model,
+                    server_args,
+                    env_dict=params.env_dict,
+                    use_omni=params.use_omni,
+                )
+            ) as server:
+                print("OmniServer started successfully")
+                yield server
+                print("OmniServer stopping...")
 
         print("OmniServer stopped")
 
@@ -2653,10 +2853,11 @@ class OpenAIClientHandler:
 
 
 @pytest.fixture
-def openai_client(omni_server: OmniServer, run_level: str):
+def openai_client(request: pytest.FixtureRequest, run_level: str):
     """Create OpenAIClientHandler fixture to facilitate communication with OmniServer
     with encapsulated request sending, concurrent requests, response handling, and validation."""
-    return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level)
+    server = request.getfixturevalue("omni_server")
+    return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level)
 
 
 class OmniRunner:

tests/dfx/perf/scripts/run_benchmark.py

@@ -43,7 +43,7 @@ def omni_server(request):
 
         print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        server_args = ["--stage-init-timeout", "120"]
+        server_args = ["--stage-init-timeout", "120", "--init-timeout", "900"]
         if stage_config_path:
             server_args = ["--stage-configs-path", stage_config_path] + server_args
         with OmniServer(model, server_args) as server:

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -2,13 +2,11 @@
 End-to-end diffusion coverage for FLUX.2-dev in online serving mode.
 
 Coverage:
-- Cache-DiT cache acceleration backend
 - CPU offload
 
-This test verifies that FLUX.2-dev can be launched with the Cache-DiT backend
-and CPU offload enabled, accepts text-to-image requests through the
-OpenAI-compatible API, and returns valid generated images with the requested
-resolution.
+This test verifies that FLUX.2-dev can be launched with CPU offload enabled,
+accepts text-to-image requests through the OpenAI-compatible API, and returns
+valid generated images with the requested resolution.
 
 assert_diffusion_response validates successful generation and the expected
 image resolution.
@@ -33,19 +31,17 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
 
 
 def _get_flux_2_dev_feature_cases(model: str):
-    """Return FLUX.2-dev diffusion feature cases for Cache-DiT + CPU offload."""
+    """Return FLUX.2-dev diffusion feature cases for CPU offload."""
 
     return [
         pytest.param(
             OmniServerParams(
                 model=model,
                 server_args=[
-                    "--cache-backend",
-                    "cache_dit",
                     "--enable-cpu-offload",
                 ],
             ),
-            id="cache_dit_cpu_offload",
+            id="cpu_offload",
             marks=SINGLE_CARD_FEATURE_MARKS,
         ),
         pytest.param(
@@ -76,7 +72,7 @@ def test_flux_2_dev(
     omni_server: OmniServer,
     openai_client: OpenAIClientHandler,
 ):
-    """Validate FLUX.2-dev online serving with Cache-DiT and CPU offload."""
+    """Validate FLUX.2-dev online serving with CPU offload."""
 
     messages = dummy_messages_from_mix_data(content_text=PROMPT)
 

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -67,13 +67,17 @@ if current_omni_platform.is_xpu():
 
 # Create parameter combinations for model and stage config
 test_params = [
-    pytest.param(OmniServerParams(model=model, stage_config_path=default_path), id="default"),
-    pytest.param(OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path)), id="async_chunk"),
+    pytest.param(OmniServerParams(model=model, stage_config_path=default_path, use_stage_cli=True), id="default"),
+    pytest.param(
+        OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path), use_stage_cli=True),
+        id="async_chunk",
+    ),
 ]
 
 test_token_params = [
     pytest.param(
-        OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path)), id="batch_token_64"
+        OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path), use_stage_cli=True),
+        id="batch_token_64",
     )
 ]
 

tests/engine/test_arg_utils.py

@@ -24,6 +24,18 @@ def test_sync_config_is_omni():
     assert isinstance(cfg, OmniModelConfig)
 
 
+def test_default_stage_id_is_concrete_int():
+    """Ensure `stage_id` stays safe for downstream arithmetic/indexing."""
+    engine_args = OmniEngineArgs()
+
+    assert engine_args.stage_id == 0
+    assert isinstance(engine_args.stage_id, int)
+    assert engine_args.log_stats is False
+
+    cfg = engine_args.create_model_config()
+    assert cfg.stage_id == 0
+
+
 def test_multimodal_kwarg_overrides():
     """Ensure that overrides in the multimodal config are preserved."""
     # Get a different value than the default for a multimodal field

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,3 +1,4 @@
+import importlib
 import os
 import types
 
@@ -8,6 +9,17 @@ from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
+    """Regression test for forward references in make_async_mp_client."""
+    import vllm_omni.engine.stage_engine_core_client as client_mod
+
+    importlib.reload(client_mod)
+
+    assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
+        "StageEngineCoreClient | DPLBStageEngineCoreClient"
+    )
+
+
 def test_initialize_stages_restores_device_visibility_after_diffusion_init(monkeypatch):
     """Regression test for stage device env leakage across stage init.
 
@@ -23,6 +35,9 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
     engine.num_stages = 1
     engine.async_chunk = False
     engine.diffusion_batch_size = 1
+    engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
+    engine._omni_master_server = None
     engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")]
 
     env_var = current_omni_platform.device_control_env_var
@@ -49,7 +64,7 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
         current_omni_platform.set_device_control_env_var("1")
 
     monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
-    monkeypatch.setattr(engine_mod, "_inject_kv_stage_info", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
     monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: diffusion_client)
     monkeypatch.setattr(
         engine_mod,
@@ -101,7 +116,11 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
             self.vllm_config = vllm_config
             self.renderer = renderer
 
-    monkeypatch.setattr(engine_mod, "StageEngineCoreClient", DummyStageEngineCoreClient)
+    monkeypatch.setattr(
+        engine_mod.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        staticmethod(lambda **kwargs: DummyStageEngineCoreClient(**kwargs)),
+    )
     monkeypatch.setattr(engine_mod, "MultimodalOutputProcessor", DummyOutputProcessor)
     monkeypatch.setattr(engine_mod, "InputProcessor", DummyInputProcessor)
     monkeypatch.setattr(engine_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)

tests/entrypoints/test_serve.py

@@ -0,0 +1,195 @@
+"""Unit tests for the Omni serve CLI helpers."""
+
+from __future__ import annotations
+
+import argparse
+from unittest.mock import Mock, patch
+
+import pytest
+
+from vllm_omni.entrypoints.cli.serve import run_headless
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_headless_args() -> argparse.Namespace:
+    return argparse.Namespace(
+        model="fake-model",
+        stage_id=3,
+        omni_master_address="127.0.0.1",
+        omni_master_port=26000,
+        api_server_count=0,
+        worker_backend="multi_process",
+        stage_configs_path=None,
+        log_stats=False,
+        disable_log_stats=False,
+    )
+
+
+def test_run_headless_registers_stage_once_and_launches_all_local_engines() -> None:
+    args = _make_headless_args()
+    stage_cfg = Mock(stage_id=3)
+    stage_cfgs = [stage_cfg]
+    parallel_config = Mock(
+        data_parallel_size_local=2,
+        data_parallel_rank=4,
+        data_parallel_rank_local=1,
+        node_rank_within_dp=0,
+    )
+    vllm_config = Mock(parallel_config=parallel_config)
+    executor_class = Mock()
+    engine_manager = Mock()
+
+    with (
+        patch(
+            "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+            return_value=("/fake/stages.yaml", stage_cfgs),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
+        patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
+        patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}),
+        patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}),
+        patch(
+            "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+            return_value=(None, None, None),
+        ),
+        patch(
+            "vllm_omni.engine.stage_init_utils.build_vllm_config",
+            return_value=(vllm_config, executor_class),
+        ) as mock_build_vllm_config,
+        patch(
+            "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+            return_value="tcp://127.0.0.1:26001",
+        ) as mock_register,
+        patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls,
+        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,
+    )
+    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.join_first.assert_called_once_with()
+    engine_manager.shutdown.assert_called_once_with()
+
+
+def test_run_headless_honors_explicit_log_stats_flag() -> None:
+    args = _make_headless_args()
+    args.log_stats = True
+    stage_cfg = Mock(stage_id=3)
+    stage_cfgs = [stage_cfg]
+    parallel_config = Mock(
+        data_parallel_size_local=2,
+        data_parallel_rank=4,
+        data_parallel_rank_local=1,
+        node_rank_within_dp=0,
+    )
+    vllm_config = Mock(parallel_config=parallel_config)
+    executor_class = Mock()
+    engine_manager = Mock()
+
+    with (
+        patch(
+            "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+            return_value=("/fake/stages.yaml", stage_cfgs),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
+        patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
+        patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={}),
+        patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={}),
+        patch(
+            "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+            return_value=(None, None, None),
+        ),
+        patch(
+            "vllm_omni.engine.stage_init_utils.build_vllm_config",
+            return_value=(vllm_config, executor_class),
+        ),
+        patch(
+            "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+            return_value="tcp://127.0.0.1:26001",
+        ),
+        patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager) as mock_manager_cls,
+        patch("signal.signal"),
+    ):
+        run_headless(args)
+
+    manager_kwargs = mock_manager_cls.call_args.kwargs
+    assert manager_kwargs["log_stats"] is True
+
+
+def test_run_headless_launches_diffusion_stage_via_omni_master() -> None:
+    args = _make_headless_args()
+    stage_cfg = Mock(stage_id=3, stage_type="diffusion")
+    stage_cfg.engine_args = Mock()
+    stage_cfg.engine_input_source = []
+    stage_cfgs = [stage_cfg]
+    metadata = Mock(stage_id=3)
+    od_config = Mock()
+    proc = Mock()
+    proc.exitcode = 0
+    proc.is_alive.return_value = False
+
+    with (
+        patch(
+            "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+            return_value=("/fake/stages.yaml", stage_cfgs),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment"),
+        patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=Mock()),
+        patch(
+            "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+            return_value=(None, None, None),
+        ),
+        patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata),
+        patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info") as mock_inject_stage_info,
+        patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config),
+        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"),
+        ) as mock_register,
+        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"),
+        ) as mock_spawn,
+        patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") as mock_handshake,
+        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,
+    )
+    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",
+    )
+    mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001")
+    proc.join.assert_called_once_with()

tests/entrypoints/test_utils.py

@@ -8,6 +8,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
+from vllm_omni.config.yaml_util import create_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -16,6 +17,7 @@ from vllm_omni.entrypoints.utils import (
     _filter_dict_like_object,
     filter_dataclass_kwargs,
     load_and_resolve_stage_configs,
+    load_stage_configs_from_yaml,
     resolve_model_config_path,
 )
 
@@ -322,3 +324,70 @@ class TestLoadAndResolveStageConfigs:
         assert config_path is None
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
+
+
+class TestLoadStageConfigsFromYaml:
+    """Regression tests for stage-config loading and merging."""
+
+    def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
+        yaml_config = create_config(
+            {
+                "async_chunk": True,
+                "stage_args": [
+                    {
+                        "stage_id": 0,
+                        "runtime": {"device": 0},
+                        "engine_args": {
+                            "parallel_config": {"tensor_parallel_size": 4},
+                        },
+                    }
+                ],
+            }
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.utils.load_yaml_config",
+            return_value=yaml_config,
+        )
+
+        stages = load_stage_configs_from_yaml(
+            "fake.yaml",
+            base_engine_args={
+                "parallel_config": {
+                    "tensor_parallel_size": 1,
+                    "pipeline_parallel_size": 2,
+                },
+                "model": "base-model",
+            },
+        )
+
+        merged_engine_args = stages[0]["engine_args"]
+        assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
+        assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
+        assert merged_engine_args["model"] == "base-model"
+        assert merged_engine_args["async_chunk"] is True
+
+    def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
+        yaml_config = create_config(
+            {
+                "stage_args": [
+                    {
+                        "stage_id": 0,
+                        "engine_args": {
+                            "nested": {"override": 2},
+                        },
+                    }
+                ],
+            }
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.utils.load_yaml_config",
+            return_value=yaml_config,
+        )
+
+        stages = load_stage_configs_from_yaml(
+            "fake.yaml",
+            base_engine_args={"nested": {"base": 1}},
+        )
+
+        assert stages[0]["engine_args"]["nested"]["base"] == 1
+        assert stages[0]["engine_args"]["nested"]["override"] == 2