[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}")

tests/diffusion/test_async_output_worker.py

@@ -497,3 +497,99 @@ class TestWorkerProcShutdown:
         result_mq.shutdown.assert_called_once_with()
         assert proc.mq is None
         assert proc.result_mq is None
+
+
+class TestSynchronizeDeviceBarrier:
+    """The pause barrier must not report a quiet device while outputs still pack."""
+
+    @staticmethod
+    def _rpc(timeout):
+        return {
+            "method": "synchronize_device",
+            "args": (),
+            "kwargs": {"timeout": timeout},
+            "output_rank": 0,
+            "exec_all_ranks": False,
+            "collect_rank_status": False,
+        }
+
+    def test_rejects_undrained_output(self, mocker):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+        mocker.patch.object(proc, "drain_async_outputs", return_value=False)
+
+        with pytest.raises(TimeoutError, match="did not drain"):
+            proc._execute_rpc(self._rpc(0.01))
+
+        proc.worker.execute_method.assert_not_called()
+
+    @pytest.mark.parametrize("timeout", [0.5, None])
+    def test_drains_with_the_call_timeout_then_synchronizes(self, mocker, timeout):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+        call_order: list = []
+        mocker.patch.object(
+            proc,
+            "drain_async_outputs",
+            side_effect=lambda **kw: call_order.append(("drain", kw)) or True,
+        )
+        proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append(("sync", a, kw))
+
+        _, should_reply = proc._execute_rpc(self._rpc(timeout))
+
+        assert call_order == [
+            ("drain", {"timeout": timeout}),
+            ("sync", ("synchronize_device",), {"timeout": timeout}),
+        ]
+        assert should_reply is True
+
+    def test_completes_only_after_output_ready_is_published(self, mocker):
+        proc = _make_worker_proc(step_execution=False)
+        proc.worker = MagicMock()
+
+        mock_platform = mocker.MagicMock()
+        mock_platform.record_device_event.return_value = MagicMock()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
+            mock_platform,
+        )
+        mock_torch = mocker.MagicMock()
+        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
+        release = threading.Event()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=lambda *a, **kw: release.wait(timeout=5.0),
+        )
+
+        proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
+        loop_thread = threading.Thread(target=proc._async_output_loop, daemon=True)
+        loop_thread.start()
+
+        barrier_done = threading.Event()
+
+        def run_barrier():
+            proc._execute_rpc(self._rpc(5.0))
+            barrier_done.set()
+
+        barrier_thread = threading.Thread(target=run_barrier, daemon=True)
+        barrier_thread.start()
+        try:
+            assert not barrier_done.wait(timeout=0.2)
+            proc.worker.execute_method.assert_not_called()
+
+            release.set()
+            assert barrier_done.wait(timeout=5.0)
+
+            kinds = [
+                call.args[0].kind
+                for call in proc.result_mq.enqueue.call_args_list
+                if isinstance(call.args[0], AsyncDiffusionOutput)
+            ]
+            assert kinds == [AsyncOutputKind.COMPUTE_DONE, AsyncOutputKind.OUTPUT_READY]
+            proc.worker.execute_method.assert_called_once_with("synchronize_device", timeout=5.0)
+        finally:
+            release.set()
+            proc._async_output_queue.put(None)
+            loop_thread.join(timeout=2.0)
+            barrier_thread.join(timeout=2.0)

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for the busy-loop-based RPC routing in DiffusionEngine.
 
 These tests guard the invariants established when ``collective_rpc`` was
@@ -60,6 +60,12 @@ class _ConcurrencyTrackingExecutor:
         self.is_failed = False
         self._closed = False
         self.od_config = SimpleNamespace()
+        # Pause tests: a call whose gate is registered (by request id, or by
+        # method name for ``synchronize_device``) blocks inside the executor
+        # until the gate is set; ``version`` tags each call so a test can tell
+        # which "weights" a request ran on.
+        self.gates: dict[str, threading.Event] = {}
+        self.version = "v1"
 
     def collective_rpc(
         self,
@@ -81,16 +87,25 @@ class _ConcurrencyTrackingExecutor:
                     "kwargs": kwargs,
                     "unique_reply_rank": unique_reply_rank,
                     "thread": threading.get_ident(),
+                    "version": self.version,
                 }
             )
         try:
             if self.rpc_delay:
                 time.sleep(self.rpc_delay)
+            if method == "synchronize_device":
+                gate = self.gates.get(method)
+                if gate is not None:
+                    gate.wait(timeout=5.0)
+                return [None]
             # Distinguish per-request execution from raw RPC by method name.
             if method in {"execute_model", "execute_stepwise", "generate"}:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
                 tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
+                gate = self.gates.get(tag)
+                if gate is not None:
+                    gate.wait(timeout=5.0)
                 return DiffusionOutput(error=f"result_for_{tag}", finished=True)
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
@@ -133,21 +148,27 @@ def _make_request(tag: str):
 
 
 def _make_engine_with_loop(
-    loop: asyncio.AbstractEventLoop,
+    loop: asyncio.AbstractEventLoop | None,
     rpc_delay: float = 0.0,
+    *,
+    scheduler: RequestScheduler | None = None,
+    start_loop: bool = True,
 ):
     """Construct a ``DiffusionEngine`` skeleton with a real busy loop.
 
     The engine is wired with a fake executor that asserts no concurrent
-    calls and a real ``RequestScheduler``.
+    calls and a real ``RequestScheduler``. ``start_loop=False`` leaves the
+    engine in its bootstrap state (no busy-loop thread).
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
     engine._closed = False
     engine.od_config = SimpleNamespace(streaming_output=False)
     engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
 
-    sched = RequestScheduler()
-    sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+    sched = scheduler
+    if sched is None:
+        sched = RequestScheduler()
+        sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
     engine.scheduler = sched
     engine.step_execution = False
     engine.supports_request_batch = False
@@ -160,9 +181,14 @@ def _make_engine_with_loop(
     engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
+    engine._scheduling_paused = False
 
     engine.main_loop = loop
     engine.stop_event = threading.Event()
+    if not start_loop:
+        engine._loop_started = False
+        engine.worker_thread = None
+        return engine
     engine._loop_started = True
     engine.worker_thread = threading.Thread(target=engine._busy_loop, daemon=True)
     engine.worker_thread.start()
@@ -170,6 +196,10 @@ def _make_engine_with_loop(
 
 
 def _stop_engine(engine: DiffusionEngine) -> None:
+    if engine.worker_thread is None:
+        return
+    for gate in engine.executor.gates.values():
+        gate.set()
     with engine._cv:
         engine.stop_event.set()
         engine._cv.notify_all()
@@ -565,3 +595,319 @@ async def test_rpc_and_request_results_do_not_swap():
     assert req_res.error == "result_for_req1"
     # Belt and braces: never overlapped on the executor.
     assert engine.executor.max_concurrent == 1
+
+
+# ───────────────────────── scheduler pause (keep) ─────────────────────────
+
+
+class _GatedAdmissionScheduler(RequestScheduler):
+    """Batching wait that only ends when the test opens admission."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.admission_open = threading.Event()
+
+    def get_admission_wait_decision(self, *, now: float, dp_concurrent: bool = False):
+        decision = super().get_admission_wait_decision(now=now, dp_concurrent=dp_concurrent)
+        return decision.__class__(should_wait=True, deadline=None, max_batch=self.max_num_running_reqs)
+
+    def should_end_admission_wait(self, decision, *, now: float, stable_since: float) -> bool:
+        return self.admission_open.is_set()
+
+
+def _pause(engine: DiffusionEngine, timeout: float | None = None):
+    return engine.async_collective_rpc("pause_scheduler", timeout=timeout, kwargs={"mode": "keep"})
+
+
+def _resume(engine: DiffusionEngine):
+    return engine.async_collective_rpc("resume_scheduler")
+
+
+def _methods(engine: DiffusionEngine) -> list[str]:
+    return [call["method"] for call in engine.executor.calls]
+
+
+def _executed(engine: DiffusionEngine) -> list[tuple[str, str]]:
+    return [
+        (call["args"][0].request_id, call["version"])
+        for call in engine.executor.calls
+        if call["method"] == "execute_model"
+    ]
+
+
+def _submit(engine: DiffusionEngine, tag: str) -> asyncio.Task:
+    return asyncio.create_task(_consume_final_output(engine.async_add_req_and_stream_response(_make_request(tag))))
+
+
+async def _wait_until(predicate, timeout: float = 3.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() > deadline:
+            raise AssertionError("condition not reached in time")
+        await asyncio.sleep(0.005)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_retains_waiting_request_and_serves_control_rpcs():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        assert await _pause(engine) is None
+        queued = _submit(engine, "queued")
+        result = await engine.async_collective_rpc("list_loras", unique_reply_rank=0)
+        assert result.error == "rpc_result_for_list_loras"
+        await asyncio.sleep(0.05)
+        assert engine.scheduler.num_waiting_requests() == 1
+        assert _executed(engine) == []
+        assert not queued.done()
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_queued"
+        assert _methods(engine) == ["synchronize_device", "list_loras", "execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_acks_after_running_batch_and_freezes_queue():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    gate = engine.executor.gates["A"] = threading.Event()
+    try:
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+
+        pause = asyncio.create_task(_pause(engine))
+        await asyncio.sleep(0.1)
+        assert not pause.done()
+
+        gate.set()
+        await asyncio.wait_for(pause, 3.0)
+        assert _methods(engine) == ["execute_model", "synchronize_device"]
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert engine.scheduler.num_waiting_requests() == 1
+        assert not b.done()
+
+        engine.executor.version = "v2"
+        await _resume(engine)
+        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
+        assert _executed(engine) == [("A", "v1"), ("B", "v2")]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_ack_does_not_release_later_rpcs_before_the_batch_is_emitted():
+    """A pause ACK must not open the queue to RPCs that touch device memory
+    (sleep) before the batch that just ran has been delivered.
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    barrier_gate = engine.executor.gates["synchronize_device"] = threading.Event()
+    methods_at_emit: list[list[str]] = []
+    emit_outputs = engine._emit_outputs
+
+    def _recording_emit(*args, **kwargs):
+        methods_at_emit.append(_methods(engine))
+        return emit_outputs(*args, **kwargs)
+
+    engine._emit_outputs = _recording_emit
+    try:
+        gate = engine.executor.gates["A"] = threading.Event()
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        pause = asyncio.create_task(_pause(engine))
+        await _wait_until(lambda: engine._scheduling_paused)
+        gate.set()
+        # The barrier holds the busy loop inside the post-execute RPC drain,
+        # so the sleep is queued before the pause ACK is published.
+        await _wait_until(lambda: "synchronize_device" in _methods(engine))
+        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
+        await _wait_until(lambda: not engine._rpc_queue.empty())
+        barrier_gate.set()
+
+        await asyncio.wait_for(pause, 3.0)
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
+    finally:
+        _stop_engine(engine)
+
+    assert _methods(engine) == ["execute_model", "synchronize_device", "handle_sleep_task"]
+    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
+
+
+@pytest.mark.asyncio
+async def test_timed_out_pause_still_ends_the_drain_before_the_batch_is_emitted():
+    """A pause whose caller already gave up is still a drain barrier: its
+    cancelled task must not let a later RPC run before the batch is emitted.
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    methods_at_emit: list[list[str]] = []
+    emit_outputs = engine._emit_outputs
+
+    def _recording_emit(*args, **kwargs):
+        methods_at_emit.append(_methods(engine))
+        return emit_outputs(*args, **kwargs)
+
+    engine._emit_outputs = _recording_emit
+    try:
+        gate = engine.executor.gates["A"] = threading.Event()
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+
+        # The caller times out while A is still executing: the pause future is
+        # cancelled, but its task stays queued ahead of the sleep.
+        with pytest.raises(TimeoutError):
+            await _pause(engine, timeout=0.1)
+        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
+        await _wait_until(lambda: engine._rpc_queue.qsize() == 2)
+        gate.set()
+
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
+    finally:
+        _stop_engine(engine)
+
+    assert _methods(engine) == ["execute_model", "handle_sleep_task"]
+    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]
+
+
+@pytest.mark.asyncio
+async def test_pause_rejects_unsupported_modes_without_closing_the_gate():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        for kwargs in ({"mode": "abort"}, {"mode": "wait"}, {}):
+            with pytest.raises(ValueError):
+                await engine.async_collective_rpc("pause_scheduler", kwargs=kwargs)
+        engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
+        with pytest.raises(NotImplementedError):
+            await _pause(engine)
+        engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
+
+        assert (await asyncio.wait_for(_submit(engine, "still-runs"), 3.0)).error == "result_for_still-runs"
+        assert _methods(engine) == ["execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_keep_pause_during_batching_wait_holds_the_batch():
+    loop = asyncio.get_running_loop()
+    sched = _GatedAdmissionScheduler()
+    sched.initialize(SimpleNamespace(max_num_seqs=4, request_batch_max_wait_ms=1000.0))
+    engine = _make_engine_with_loop(loop, scheduler=sched)
+    try:
+        waiting = _submit(engine, "W")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+        await asyncio.sleep(0.05)
+
+        await asyncio.wait_for(_pause(engine), 3.0)
+        assert _methods(engine) == ["synchronize_device"]
+        assert engine.scheduler.num_waiting_requests() == 1
+
+        await _resume(engine)
+        await asyncio.sleep(0.05)
+        assert _executed(engine) == []
+        sched.admission_open.set()
+        assert (await asyncio.wait_for(waiting, 3.0)).error == "result_for_W"
+    finally:
+        sched.admission_open.set()
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_pause_timeout_leaves_gate_closed_until_explicit_resume():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    gate = engine.executor.gates["A"] = threading.Event()
+    try:
+        a = _submit(engine, "A")
+        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+
+        with pytest.raises(TimeoutError):
+            await _pause(engine, timeout=0.1)
+        gate.set()
+        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
+        await asyncio.sleep(0.1)
+        assert _executed(engine) == [("A", "v1")]
+        assert not b.done()
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
+        assert "synchronize_device" not in _methods(engine)
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_repeated_pause_and_resume_are_idempotent():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        await _resume(engine)
+        await _pause(engine)
+        await _pause(engine)
+        queued = _submit(engine, "Q")
+        await asyncio.sleep(0.05)
+        assert _executed(engine) == []
+        await _resume(engine)
+        await _resume(engine)
+        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_Q"
+        assert _methods(engine) == ["synchronize_device", "synchronize_device", "execute_model"]
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_abort_of_waiting_request_while_paused_delivers_terminal():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        await _pause(engine)
+        b = _submit(engine, "B")
+        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
+        engine.abort("B")
+
+        result = await asyncio.wait_for(b, 3.0)
+        assert result.aborted is True
+        assert engine.scheduler.get_request_state("B") is None
+        assert _executed(engine) == []
+
+        await _resume(engine)
+        assert (await asyncio.wait_for(_submit(engine, "C"), 3.0)).error == "result_for_C"
+        assert _executed(engine) == [("C", "v1")]
+        assert engine.worker_thread.is_alive()
+    finally:
+        _stop_engine(engine)
+
+
+def test_pause_and_resume_before_loop_starts_use_engine_local_path():
+    engine = _make_engine_with_loop(None, start_loop=False)
+
+    assert engine.collective_rpc("pause_scheduler", kwargs={"mode": "keep"}) is None
+    assert _methods(engine) == ["synchronize_device"]
+
+    assert engine.collective_rpc("resume_scheduler") is None
+    assert engine.add_req_and_wait_for_response(_make_request("X")).error == "result_for_X"
+    assert _executed(engine) == [("X", "v1")]
+
+
+@pytest.mark.asyncio
+async def test_requests_without_pause_never_trigger_the_barrier():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        first = await asyncio.wait_for(_submit(engine, "r1"), 3.0)
+        await engine.async_collective_rpc("ping", args=("p",), unique_reply_rank=0)
+        second = await asyncio.wait_for(_submit(engine, "r2"), 3.0)
+    finally:
+        _stop_engine(engine)
+
+    assert (first.error, second.error) == ("result_for_r1", "result_for_r2")
+    assert _methods(engine) == ["execute_model", "ping", "execute_model"]

tests/diffusion/test_diffusion_scheduler.py

@@ -2092,3 +2092,19 @@ class TestStepScheduler:
 
         with pytest.raises(ValueError):
             self.scheduler.add_request(request)
+
+
+class TestPendingFinishedRequestIds:
+    def test_reports_finished_ids_that_still_hold_state_without_clearing_them(self):
+        sched = RequestScheduler()
+        sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
+        sched.add_request(_make_request("a"))
+        sched.add_request(_make_request("b"))
+
+        sched.finish_requests("a", DiffusionRequestStatus.FINISHED_ABORTED)
+        assert sched.pending_finished_request_ids() == {"a"}
+
+        sched.pop_request_state("a")
+        assert sched.pending_finished_request_ids() == set()
+        # The next wave still ships the id to the worker for its own cleanup.
+        assert sched.schedule().finished_req_ids == {"a"}

tests/diffusion/test_diffusion_worker_synchronize_device.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for the device step of the diffusion pause barrier."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_worker(runner: object) -> DiffusionWorker:
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = runner
+    worker.rank = 0
+    return worker
+
+
+def test_synchronize_device_waits_for_prefetch_before_device_sync(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    order: list = []
+    manager = MagicMock()
+    manager.wait_prefetch.side_effect = lambda timeout: order.append(("prefetch", timeout)) or True
+    platform.synchronize.side_effect = lambda: order.append("sync")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    worker.synchronize_device(timeout=2.5)
+
+    assert order == [("prefetch", 2.5), "sync"]
+
+
+def test_synchronize_device_fails_when_prefetch_is_still_outstanding(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    manager = MagicMock()
+    manager.wait_prefetch.return_value = False
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    with pytest.raises(TimeoutError, match="prefetch"):
+        worker.synchronize_device(timeout=0.01)
+
+    platform.synchronize.assert_not_called()
+
+
+def test_synchronize_device_without_transfer_manager_only_syncs(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=None))
+
+    worker.synchronize_device(timeout=None)
+
+    platform.synchronize.assert_called_once_with()