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

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

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)