[Core][Diffusion] Add pause_generation(mode="keep") for diffusion stages (#7685)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

tests/diffusion/distributed/test_pause_keep_two_ranks.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Two-rank ``pause_generation(mode="keep")``: every rank runs the pause barrier
+before the ACK, and a barrier failure on a non-zero rank fails the pause
+without aborting or dropping the running batch.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from contextlib import ExitStack
+
+import pytest
+
+from tests.e2e.features.custom_pipeline.test_diffusion_pause_keep import (
+    CUSTOM_PIPELINE_CLASS,
+    MODEL,
+    WORKER_EXTENSION_CLASS,
+    _generate,
+    _pause_round,
+    _wait_for,
+)
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+
+def _all_true(results) -> bool:
+    flat = []
+    stack = list(results)
+    while stack:
+        item = stack.pop()
+        if isinstance(item, list):
+            stack.extend(item)
+        else:
+            flat.append(item)
+    return bool(flat) and all(item is True for item in flat)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_keep_pause_two_ranks_ack_and_rank_failure():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            num_gpus=2,
+            parallel_config={"ulysses_degree": 2, "ulysses_mode": "advanced_uaa"},
+        )
+        after.callback(engine.shutdown)
+        events: dict[str, dict[str, float]] = defaultdict(dict)
+        assert _all_true(await engine.collective_rpc(method="start_counting_synchronize_device"))
+
+        pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(3)]
+        # The result is the all-rank AND, so True means both ranks ran the barrier.
+        assert _all_true(await engine.collective_rpc(method="synchronize_device_seen"))
+
+        assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, True)))
+        a = asyncio.create_task(_generate(engine, "a-fail", events))
+        await _wait_for(lambda: "started" in events["a-fail"])
+        with pytest.raises(RuntimeError, match="pause_scheduler failed"):
+            await engine.pause_generation(mode="keep", clear_cache=False)
+        assert await engine.is_paused()
+        a_output = await asyncio.wait_for(a, 60.0)
+        assert a_output.images, "a failed barrier must not abort the running batch"
+
+        assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, False)))
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+        recovered = await asyncio.wait_for(_generate(engine, "after-recovery", events), 60.0)
+        assert recovered.images
+
+        assert not engine.request_states
+        print(f"[sp2] rounds=3 pause_latency_s min/max={min(pause_latencies):.3f}/{max(pause_latencies):.3f}")