[Refact] Refactor diffusion output streaming (#5166)

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

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -30,7 +30,7 @@ from typing import Any
 import pytest
 
 from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _RpcTask
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode, _RpcTask
 from vllm_omni.diffusion.sched import RequestScheduler
 from vllm_omni.diffusion.sched.interface import RequestBatchSamplingParamsKey
 from vllm_omni.diffusion.worker.utils import RunnerOutput
@@ -91,7 +91,7 @@ class _ConcurrencyTrackingExecutor:
                 # 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"
-                return DiffusionOutput(error=f"result_for_{tag}")
+                return DiffusionOutput(error=f"result_for_{tag}", finished=True)
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
         finally:
@@ -149,12 +149,12 @@ def _make_engine_with_loop(
     engine.scheduler = sched
     engine.step_execution = False
     engine.supports_request_batch = False
+    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
     engine.execute_fn = engine.executor.execute_request
 
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
-    engine._out_queue = {}
-    engine._out_queue_streaming = {}
+    engine._out_streams = {}
     engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
@@ -187,6 +187,15 @@ def _stop_engine(engine: DiffusionEngine) -> None:
     assert not engine.worker_thread.is_alive(), "Busy loop thread did not stop"
 
 
+async def _consume_final_output(generator):
+    final_output = None
+    async for output in generator:
+        final_output = output
+    if final_output is None:
+        raise RuntimeError("Diffusion execution finished without output.")
+    return final_output
+
+
 # ─────────────────────── single-thread invariant ───────────────────────────
 
 
@@ -199,7 +208,7 @@ async def test_executor_only_called_from_busy_loop_thread():
     busy_tid = engine.worker_thread.ident
     try:
         # Per-request path
-        await engine.async_add_req_and_wait_for_response(_make_request("req1"))
+        await _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
         # Raw RPC path (sync from a worker thread)
         result = await asyncio.to_thread(engine.collective_rpc, "ping", args=("a",), unique_reply_rank=0)
         assert result.error == "rpc_result_for_a"
@@ -228,7 +237,7 @@ async def test_executor_calls_never_overlap_under_load():
             return await engine.async_collective_rpc("ping", args=(f"x{i}",), unique_reply_rank=0)
 
         async def _request(i: int):
-            return await engine.async_add_req_and_wait_for_response(_make_request(f"r{i}"))
+            return await _consume_final_output(engine.async_add_req_and_stream_response(_make_request(f"r{i}")))
 
         tasks = [_rpc(i) for i in range(20)] + [_request(i) for i in range(5)]
         results = await asyncio.gather(*tasks)
@@ -510,7 +519,9 @@ async def test_rpc_and_request_results_do_not_swap():
     engine = _make_engine_with_loop(loop, rpc_delay=0.02)
     try:
         rpc_task = asyncio.create_task(engine.async_collective_rpc("ping", args=("rpc1",), unique_reply_rank=0))
-        req_task = asyncio.create_task(engine.async_add_req_and_wait_for_response(_make_request("req1")))
+        req_task = asyncio.create_task(
+            _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
+        )
         rpc_res, req_res = await asyncio.gather(rpc_task, req_task)
     finally:
         _stop_engine(engine)