[Bugfix] Make the async-output wait bound configurable and default it higher (#6255)

Signed-off-by: ivanusto <ivanusto@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/design/feature/async_diffusion_output.md

@@ -124,8 +124,8 @@ When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_
 1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
    When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
 
-2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
-   The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (30.0s). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter.
+2. **Timeout Diagnostics (`_async_output_timeout()`)**:
+   The engine waits for `wait_output_ready()` with a timeout of `_async_output_timeout()` — 600.0s by default, overridable per run with the `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` environment variable (a non-numeric or non-positive value is ignored with a warning and the default applies, since this resolves on the request path). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter. The packing itself takes milliseconds, but the wait is queued behind the step's GPU work, so its wall-clock duration tracks step time; the bound is generous because a hung engine is surfaced by the worker monitor and `check_health()`, not by this timeout.
 
 3. **Queue Serialization & Atomic Future Resolution**:
    - `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).

recipes/MiniMaxAI/MiniMax-H3-4090.md

@@ -110,9 +110,10 @@ outputs of each pair agreed to within 115 bytes.
 Four-GPU T2VA repeats: 273.8 s and 269.5 s, peaking at 15,560 MiB and
 15,612 MiB. Four-GPU Ref2VA is a single successful end-to-end run (545.2 s,
 16,448 MiB). Follow-up Ref2VA requests on this build completed diffusion but
-then hit the engine's hardcoded 30 s async output wait
-(`_ASYNC_OUTPUT_TIMEOUT` in `diffusion_engine.py`) during post-compute D2H,
-so they are not reported as repeats.
+then hit the engine's async output wait during post-compute D2H, so they are
+not reported as repeats. That wait was a hardcoded 30 s when these numbers were
+taken; it is now `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` (default 600 s), so a rerun on
+current main should not lose these repeats to the bound.
 
 Relative to the two-GPU baseline on the same shape, four-GPU TP2×USP2 was about
 1.6× faster for both tasks, while per-GPU peak HBM stayed in the mid-15 GiB

tests/diffusion/test_async_output_timeout.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the async-output wait bound (_async_output_timeout)."""
+
+import logging
+
+import pytest
+
+from vllm_omni.diffusion.diffusion_engine import (
+    _ASYNC_OUTPUT_TIMEOUT_DEFAULT,
+    _ASYNC_OUTPUT_TIMEOUT_ENV,
+    _async_output_timeout,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestAsyncOutputTimeoutIsConfigurable:
+    def test_default_is_generous(self, monkeypatch):
+        """The bound covers a copy queued behind a full denoise step, so the
+        default must not be tight enough to abort a slow but healthy render.
+        """
+        monkeypatch.delenv(_ASYNC_OUTPUT_TIMEOUT_ENV, raising=False)
+        assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT == 600.0
+
+    def test_env_var_overrides_the_default(self, monkeypatch):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "1800")
+        assert _async_output_timeout() == 1800.0
+
+    def test_env_var_accepts_a_float(self, monkeypatch):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "45.5")
+        assert _async_output_timeout() == 45.5
+
+    def test_value_is_not_frozen_at_import(self, monkeypatch):
+        """Resolved per call rather than captured in a module constant, so the
+        value never goes stale relative to os.environ and no module reload is
+        needed to exercise it.
+        """
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "60")
+        assert _async_output_timeout() == 60.0
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "120")
+        assert _async_output_timeout() == 120.0
+
+
+class TestAsyncOutputTimeoutRejectsBadValues:
+    """This runs on the request path, so a typo in the environment must degrade
+    to the default rather than fail the generation.
+    """
+
+    @pytest.mark.parametrize("value", ["", "abc", "30s", "None"])
+    def test_non_numeric_falls_back_to_the_default(self, monkeypatch, caplog, value):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value)
+        with caplog.at_level(logging.WARNING):
+            assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+
+    @pytest.mark.parametrize("value", ["0", "-1", "-0.5"])
+    def test_non_positive_falls_back_to_the_default(self, monkeypatch, caplog, value):
+        monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value)
+        with caplog.at_level(logging.WARNING):
+            assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT

vllm_omni/diffusion/diffusion_engine.py

@@ -7,6 +7,7 @@ import asyncio
 import concurrent.futures
 import copy
 import inspect
+import os
 import queue
 import threading
 import time
@@ -61,7 +62,50 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-_ASYNC_OUTPUT_TIMEOUT = 30.0  # seconds
+_ASYNC_OUTPUT_TIMEOUT_ENV = "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT"
+_ASYNC_OUTPUT_TIMEOUT_DEFAULT = 600.0  # seconds
+
+
+def _async_output_timeout() -> float:
+    """Seconds to wait for one step's background D2H/SHM copy.
+
+    The copy itself finishes in milliseconds, but it is queued behind the GPU
+    work for that step, so the wall-clock wait tracks step time — a single-GPU
+    box legitimately runs tens of seconds per step on large shapes. A tight
+    bound therefore does not catch a hung engine (worker death and a dead
+    result pump are surfaced by the worker monitor and ``check_health``); it
+    only aborts renders that are still making progress, throwing away the
+    denoise that already completed. The default matches
+    ``_DLO_DP_WAVE_TIMEOUT_S`` in the same subsystem.
+
+    Resolved here rather than at import so a malformed value degrades to the
+    default instead of raising on the request path: this runs inside
+    ``step_streaming``/``add_req_and_wait_for_response``, where a typo in the
+    environment must not start failing generations.
+    """
+    raw = os.environ.get(_ASYNC_OUTPUT_TIMEOUT_ENV)
+    if raw is None:
+        return _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+    try:
+        timeout = float(raw)
+    except ValueError:
+        logger.warning_once(
+            "Ignoring %s=%r: not a number. Using the default %.1fs.",
+            _ASYNC_OUTPUT_TIMEOUT_ENV,
+            raw,
+            _ASYNC_OUTPUT_TIMEOUT_DEFAULT,
+        )
+        return _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+    if timeout <= 0:
+        logger.warning_once(
+            "Ignoring %s=%r: must be positive. Using the default %.1fs.",
+            _ASYNC_OUTPUT_TIMEOUT_ENV,
+            raw,
+            _ASYNC_OUTPUT_TIMEOUT_DEFAULT,
+        )
+        return _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+    return timeout
+
 
 __all__ = [
     "DiffusionEngine",
@@ -386,13 +430,16 @@ class DiffusionEngine:
             # Async mode: wait for background D2H/SHM to complete.
             if output.async_output_id:
                 fut = self.executor.wait_output_ready(output.async_output_id)
+                timeout = _async_output_timeout()
                 try:
-                    output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=_ASYNC_OUTPUT_TIMEOUT)
+                    output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout)
                 except (TimeoutError, asyncio.TimeoutError):
                     describe = getattr(self.executor, "describe_pending_state", None)
                     logger.error(
-                        "Timed out after %.0fs waiting for async output; executor state: %s",
-                        _ASYNC_OUTPUT_TIMEOUT,
+                        "Timed out after %.1fs waiting for async output; set %s to a larger value "
+                        "to allow slower steps. Executor state: %s",
+                        timeout,
+                        _ASYNC_OUTPUT_TIMEOUT_ENV,
                         describe(output.async_output_id) if describe else "unavailable",
                     )
                     raise
@@ -906,7 +953,7 @@ class DiffusionEngine:
                     )
                     if output.async_output_id:
                         fut = self.executor.wait_output_ready(output.async_output_id)
-                        output = fut.result(timeout=_ASYNC_OUTPUT_TIMEOUT)
+                        output = fut.result(timeout=_async_output_timeout())
                     return output
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: