[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/conftest.py

@@ -75,6 +75,7 @@ class OmniServerParams(NamedTuple):
     use_omni: bool = True
     use_stage_cli: bool = False
     init_timeout: int | None = None
+    stage_init_timeout: int | None = None  # None defers to the server's own default (300 s)
 
 
 def assert_image_diffusion_response(
@@ -1768,8 +1769,8 @@ 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 params.use_omni and params.stage_init_timeout is not None:
+            server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
         if params.init_timeout is not None:
             server_args = [*server_args, "--init-timeout", str(params.init_timeout)]
         if params.use_stage_cli:
@@ -3257,7 +3258,7 @@ def omni_runner(request, model_prefix):
     with _omni_server_lock:
         model, stage_config_path = request.param
         model = model_prefix + model
-        with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
+        with OmniRunner(model, seed=42, stage_configs_path=stage_config_path) as runner:
             print("OmniRunner started successfully")
             yield runner
             print("OmniRunner stopping...")

tests/dfx/perf/scripts/run_benchmark.py

@@ -23,6 +23,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 CONFIG_FILE_PATH = str(Path(__file__).parent.parent / "tests" / "test.json")
 BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
+STAGE_INIT_TIMEOUT = 600
 
 
 STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs"
@@ -43,7 +44,7 @@ def omni_server(request):
 
         print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        server_args = ["--stage-init-timeout", "120", "--init-timeout", "900"]
+        server_args = ["--stage-init-timeout", str(STAGE_INIT_TIMEOUT), "--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/accuracy/conftest.py

@@ -114,8 +114,8 @@ class AccuracyServerConfig:
         params = self.generate_params
         model = self.model_prefix + params.model
         server_args = params.server_args or []
-        if params.use_omni:
-            server_args = ["--stage-init-timeout", "120", *server_args]
+        if params.use_omni and params.stage_init_timeout is not None:
+            server_args = ["--stage-init-timeout", str(params.stage_init_timeout), *server_args]
         with OmniServer(
             model,
             server_args,
@@ -226,6 +226,7 @@ def _build_accuracy_server_config(
             server_args=generate_server_args,
             env_dict={"CUDA_VISIBLE_DEVICES": shared_gpu},
             use_omni=True,
+            stage_init_timeout=300,
         ),
         judge_params=OmniServerParams(
             model=judge_model,

tests/e2e/online_serving/test_bagel_online.py

@@ -47,7 +47,7 @@ test_params = [
     OmniServerParams(
         model=MODEL,
         stage_config_path=STAGE_CONFIGS_PATH,
-        server_args=["--stage-init-timeout", "300"],
+        stage_init_timeout=300,
     ),
 ]
 

tests/e2e/online_serving/test_dynin_omni_expansion.py

@@ -30,7 +30,7 @@ T2I_PROMPT = "A high quality detailed living room interior photo."
 T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving."
 I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors."
 
-TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG)]
+TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, stage_init_timeout=600)]
 _STAGE_COUNT = 3
 _I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False}
 

tests/e2e/online_serving/test_qwen3_tts_batch.py

@@ -30,6 +30,7 @@ from tests.conftest import (
 from tests.utils import hardware_test
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
+STAGE_INIT_TIMEOUT_S = 120
 
 
 def get_stage_config(name: str = "qwen3_tts.yaml"):
@@ -47,7 +48,7 @@ def omni_server():
             "--stage-configs-path",
             stage_config_path,
             "--stage-init-timeout",
-            "120",
+            str(STAGE_INIT_TIMEOUT_S),
             "--trust-remote-code",
             "--enforce-eager",
             "--disable-log-stats",
@@ -337,7 +338,7 @@ def omni_server_batch2():
             "--stage-configs-path",
             config_path,
             "--stage-init-timeout",
-            "120",
+            str(STAGE_INIT_TIMEOUT_S),
             "--trust-remote-code",
             "--enforce-eager",
             "--disable-log-stats",

tests/e2e/online_serving/test_qwen3_tts_speaker_embedding.py

@@ -23,6 +23,7 @@ from tests.utils import hardware_test
 
 MODEL_BASE = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 MODEL_BASE_1_7B = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+STAGE_INIT_TIMEOUT_S = 120
 
 # A synthetic 1024-dim speaker embedding (all 0.1 — not a real voice, but
 # exercises the full code path through the talker's _build_prompt_embeds).
@@ -47,7 +48,7 @@ def _server_args():
         "--stage-configs-path",
         get_stage_config(),
         "--stage-init-timeout",
-        "120",
+        str(STAGE_INIT_TIMEOUT_S),
         "--trust-remote-code",
         "--enforce-eager",
         "--disable-log-stats",

tests/e2e/online_serving/test_qwen3_tts_websocket.py

@@ -19,6 +19,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
+STAGE_INIT_TIMEOUT_S = 120
 
 
 def get_stage_config() -> str:
@@ -37,7 +38,7 @@ def omni_server():
             "--stage-configs-path",
             stage_config_path,
             "--stage-init-timeout",
-            "120",
+            str(STAGE_INIT_TIMEOUT_S),
             "--trust-remote-code",
             "--enforce-eager",
             "--disable-log-stats",

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)