[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>

examples/offline_inference/text_to_image/text_to_image.py

@@ -242,6 +242,18 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable logging of diffusion pipeline stats.",
     )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=600,
+        help="Timeout for initializing a single stage in seconds (default: 600s)",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Timeout for initializing a single stage in seconds (default: 600s)",
+    )
     parser.add_argument(
         "--use-system-prompt",
         type=str,
@@ -346,6 +358,8 @@ def main():
         "mode": "text-to-image",
         "log_stats": args.log_stats,
         "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
+        "init_timeout": args.init_timeout,
+        "stage_init_timeout": args.stage_init_timeout,
         **lora_args,
         **quant_kwargs,
     }

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)

vllm_omni/diffusion/stage_diffusion_client.py

@@ -50,11 +50,12 @@ class StageDiffusionClient:
         model: str,
         od_config: OmniDiffusionConfig,
         metadata: StageMetadata,
+        stage_init_timeout: int,
         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)
+        complete_diffusion_handshake(proc, handshake_address, stage_init_timeout)
         self._initialize_client(metadata, request_address, response_address, proc=proc, batch_size=batch_size)
 
     @classmethod

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -39,8 +39,6 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-_HANDSHAKE_POLL_TIMEOUT_S = 600
-
 
 class StageDiffusionProc:
     """Subprocess entry point for diffusion inference.
@@ -619,13 +617,14 @@ def spawn_diffusion_proc(
 def complete_diffusion_handshake(
     proc: BaseProcess,
     handshake_address: str,
+    handshake_timeout: int,
 ) -> None:
     """Wait for the diffusion subprocess to signal READY.
 
     On failure the process is terminated before re-raising.
     """
     try:
-        _perform_diffusion_handshake(proc, handshake_address)
+        _perform_diffusion_handshake(proc, handshake_address, handshake_timeout)
     except Exception:
         shutdown([proc])
         raise
@@ -634,6 +633,7 @@ def complete_diffusion_handshake(
 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:
@@ -641,11 +641,15 @@ def _perform_diffusion_handshake(
         poller.register(handshake_socket, zmq.POLLIN)
         poller.register(proc.sentinel, zmq.POLLIN)
 
-        timeout_ms = _HANDSHAKE_POLL_TIMEOUT_S * 1000
+        timeout_ms = handshake_timeout * 1000
         while True:
             events = dict(poller.poll(timeout=timeout_ms))
             if not events:
-                raise TimeoutError("Timed out waiting for READY from StageDiffusionProc")
+                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."
+                )
             if handshake_socket in events:
                 identity, raw = handshake_socket.recv_multipart()
                 msg = msgspec.msgpack.decode(raw)

vllm_omni/engine/async_omni_engine.py

@@ -430,7 +430,9 @@ class AsyncOmniEngine:
                         else:
                             assert proc is not None
                             assert handshake_address is not None
-                            complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
+                            complete_stage_handshake(
+                                proc, handshake_address, addresses, vllm_config, stage_init_timeout
+                            )
                         logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
                     finally:
                         if previous_visible_devices is None:
@@ -759,6 +761,7 @@ class AsyncOmniEngine:
                                         self.model,
                                         stage_cfg,
                                         metadata,
+                                        stage_init_timeout=stage_init_timeout,
                                         batch_size=self.diffusion_batch_size,
                                     )
                                 logger.info(

vllm_omni/engine/stage_engine_core_proc.py

@@ -37,8 +37,6 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-_HANDSHAKE_POLL_TIMEOUT_S = 600
-
 
 class StageEngineCoreProc(EngineCoreProc):
     """Stage-specific engine core process for vLLM-Omni.
@@ -145,13 +143,14 @@ def complete_stage_handshake(
     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)
+        _perform_handshake(proc, handshake_address, addresses, vllm_config, handshake_timeout)
     except Exception:
         shutdown([proc])
         raise
@@ -162,6 +161,7 @@ def _perform_handshake(
     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:
@@ -169,7 +169,7 @@ def _perform_handshake(
         poller.register(handshake_socket, zmq.POLLIN)
         poller.register(proc.sentinel, zmq.POLLIN)
 
-        identity, msg = _recv(poller, handshake_socket, proc, "HELLO")
+        identity, msg = _recv(poller, handshake_socket, proc, "HELLO", handshake_timeout)
         if msg.get("status") != "HELLO":
             raise RuntimeError(f"Expected HELLO, got: {msg}")
 
@@ -179,7 +179,7 @@ def _perform_handshake(
         )
         handshake_socket.send_multipart([identity, msgspec.msgpack.encode(init_payload)])
 
-        identity, msg = _recv(poller, handshake_socket, proc, "READY")
+        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")
@@ -192,13 +192,18 @@ def _recv(
     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 = _HANDSHAKE_POLL_TIMEOUT_S * 1000
+    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")
+            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)

vllm_omni/engine/stage_init_utils.py

@@ -335,7 +335,7 @@ def build_vllm_config(
 def acquire_device_locks(
     stage_id: int,
     engine_args_dict: dict[str, Any],
-    stage_init_timeout: int = 300,
+    stage_init_timeout: int,
 ) -> list[int]:
     """Acquire exclusive file locks on devices needed by this stage.
 
@@ -528,6 +528,7 @@ def initialize_diffusion_stage(
     model: str,
     stage_cfg: Any,
     metadata: StageMetadata,
+    stage_init_timeout: int,
     batch_size: int = 1,
 ) -> Any:
     """Build a diffusion stage client.
@@ -536,6 +537,7 @@ def initialize_diffusion_stage(
         model: Model name or path.
         stage_cfg: Stage configuration.
         metadata: Extracted stage metadata.
+        stage_init_timeout: Timeout in seconds for stage initialization handshake
         batch_size: Maximum number of requests to batch together in the
             diffusion engine.  Passed through to ``StageDiffusionClient``
             and ultimately to ``AsyncOmni``.
@@ -543,7 +545,9 @@ def initialize_diffusion_stage(
     from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 
     od_config = build_diffusion_config(model, stage_cfg, metadata)
-    return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+    return StageDiffusionClient(
+        model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size
+    )
 
 
 def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None: