[Bugfix] Restore user config/runtime stage init timeout (#2519)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,5 +1,6 @@
 import importlib
 import os
+import threading
 import types
 
 import pytest
@@ -86,6 +87,144 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
             os.environ[env_var] = old_env
 
 
+def test_initialize_stages_passes_stage_init_timeout_to_diffusion_handshake(monkeypatch):
+    """Regression test for stage_init_timeout passing to complete_diffusion_handshake
+    in the diffusion stage path.
+    """
+    import vllm_omni.diffusion.data as diffusion_data_mod
+    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+    import vllm_omni.engine.async_omni_engine as engine_mod
+    from vllm_omni.platforms import current_omni_platform
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.num_stages = 1
+    engine.async_chunk = False
+    engine.diffusion_batch_size = 1
+    engine.single_stage_mode = False
+    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion", engine_args={})]
+
+    metadata = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        runtime_cfg={"devices": "0"},
+        prompt_expand_func=None,
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params=None,
+        custom_process_input_func=None,
+        engine_input_source=None,
+        cfg_kv_collect_func=None,
+    )
+
+    captured_timeout = None
+    device_env_var = current_omni_platform.device_control_env_var
+    prev_device_env = os.environ.get(device_env_var)
+    os.environ[device_env_var] = "0"
+
+    monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
+    monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    monkeypatch.setattr(
+        engine_mod,
+        "finalize_initialized_stages",
+        lambda stage_clients, _input_processor: (
+            stage_clients,
+            [types.SimpleNamespace()],
+            [{"final_output_type": "image"}],
+        ),
+    )
+    monkeypatch.setattr(
+        diffusion_data_mod.OmniDiffusionConfig,
+        "from_kwargs",
+        classmethod(lambda cls, **kwargs: types.SimpleNamespace(parallel_config=types.SimpleNamespace(world_size=1))),
+    )
+    monkeypatch.setattr(
+        client_mod,
+        "spawn_diffusion_proc",
+        lambda model, od_cfg: (object(), "ipc://handshake", "ipc://request", "ipc://response"),
+    )
+
+    def _capture_handshake_timeout(proc, handshake_address, handshake_timeout):
+        nonlocal captured_timeout
+        captured_timeout = handshake_timeout
+
+    monkeypatch.setattr(client_mod, "complete_diffusion_handshake", _capture_handshake_timeout)
+    monkeypatch.setattr(
+        client_mod.zmq,
+        "Context",
+        lambda: types.SimpleNamespace(socket=lambda _: types.SimpleNamespace(connect=lambda _: None)),
+    )
+
+    try:
+        engine._initialize_stages(stage_init_timeout=302)
+    finally:
+        if prev_device_env is None:
+            os.environ.pop(device_env_var, None)
+        else:
+            os.environ[device_env_var] = prev_device_env
+
+    assert captured_timeout == 302
+
+
+def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
+    """Regression test for stage_init_timeout reaching complete_stage_handshake
+    in the LLM stage path.
+    """
+    import vllm_omni.engine.async_omni_engine as engine_mod
+    from vllm_omni.platforms import current_omni_platform
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.single_stage_mode = False
+    engine._omni_master_server = None
+
+    metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"})
+    fake_vllm_config = types.SimpleNamespace()
+    fake_addresses = types.SimpleNamespace()
+    fake_proc = types.SimpleNamespace()
+
+    captured_timeout = None
+
+    device_env_var = current_omni_platform.device_control_env_var
+    prev_device_env = os.environ.get(device_env_var)
+    os.environ[device_env_var] = "0"
+
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
+    monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object))
+    monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
+    monkeypatch.setattr(
+        engine_mod,
+        "spawn_stage_core",
+        lambda **_: (fake_addresses, fake_proc, "ipc://handshake"),
+    )
+
+    def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout):
+        nonlocal captured_timeout
+        captured_timeout = handshake_timeout
+
+    monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout)
+
+    try:
+        engine._launch_llm_stage(
+            stage_cfg=types.SimpleNamespace(engine_args={}),
+            metadata=metadata,
+            stage_connector_spec={},
+            stage_init_timeout=302,
+            llm_stage_launch_lock=threading.Lock(),
+        )
+    finally:
+        if prev_device_env is None:
+            os.environ.pop(device_env_var, None)
+        else:
+            os.environ[device_env_var] = prev_device_env
+
+    assert captured_timeout == 302
+
+
 def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
     """Regression test for GLM-Image t2i preprocessing path.
 

tests/engine/test_single_stage_mode.py

@@ -1459,6 +1459,7 @@ class TestLaunchLlmStageSingleStageMode:
 
         fake_proc = Mock()
         fake_handshake_address = "ipc:///tmp/fake-handshake"
+        stage_init_timeout = 60
 
         with (
             patch("vllm_omni.engine.async_omni_engine.setup_stage_devices"),
@@ -1484,7 +1485,7 @@ class TestLaunchLlmStageSingleStageMode:
                 stage_cfg=_make_stage_cfg(0),
                 metadata=metadata,
                 stage_connector_spec={},
-                stage_init_timeout=60,
+                stage_init_timeout=stage_init_timeout,
                 llm_stage_launch_lock=threading.Lock(),
             )
 
@@ -1498,6 +1499,7 @@ class TestLaunchLlmStageSingleStageMode:
             fake_handshake_address,
             fake_addresses,
             fake_vllm_config,
+            stage_init_timeout,
         )
         mock_omni.assert_not_called()
         assert isinstance(result, StartedLlmStage)