[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.py

@@ -15,7 +15,11 @@ from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _move_tensor_tree_to_cpu
+from vllm_omni.diffusion.diffusion_engine import (
+    DiffusionEngine,
+    DiffusionExecutionMode,
+    _move_tensor_tree_to_cpu,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
     CachedRequestData,
@@ -221,7 +225,7 @@ class TestRequestBatchCapability:
             diffusion_engine_module.supports_request_batch(od_config)
         registry_load.assert_not_called()
 
-    def test_engine_disables_batch_dispatch_for_single_request_pipeline(
+    def test_engine_uses_request_batch_mode_for_single_request_pipeline(
         self,
         monkeypatch: pytest.MonkeyPatch,
         mocker: MockerFixture,
@@ -230,6 +234,7 @@ class TestRequestBatchCapability:
             model_class_name="SinglePipeline",
             custom_pipeline_args=None,
             streaming_output=False,
+            max_num_seqs=1,
         )
         fake_executor = SimpleNamespace(
             execute_request=mocker.Mock(return_value="per-request"),
@@ -255,15 +260,50 @@ class TestRequestBatchCapability:
             "_try_load_model_cls",
             lambda model_class_name: _SingleRequestPipeline,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
         engine = DiffusionEngine(od_config)
-        output = engine.execute_fn(_make_request_mode_sched_output("req-a", "req-b"))
+        output = engine.execute_fn(_make_request_mode_sched_output("req-a"))
 
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
         assert engine.supports_request_batch is False
-        assert output == "per-request"
-        fake_executor.execute_request.assert_called_once()
-        fake_executor.execute_batch.assert_not_called()
+        assert output == "batch"
+        fake_executor.execute_batch.assert_called_once()
+        fake_executor.execute_request.assert_not_called()
+
+    def test_engine_rejects_multi_seq_request_batch_for_single_request_pipeline(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            streaming_output=False,
+            max_num_seqs=2,
+        )
+        fake_executor_cls = mocker.Mock()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+
+        with pytest.raises(ValueError, match="max_num_seqs=1"):
+            DiffusionEngine(od_config)
+        fake_executor_cls.assert_not_called()
 
     @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
     def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
@@ -276,6 +316,7 @@ class TestRequestBatchCapability:
             model_class_name="BatchPipeline",
             custom_pipeline_args=None,
             streaming_output=False,
+            max_num_seqs=2,
         )
         fake_executor = SimpleNamespace(
             execute_request=mocker.Mock(return_value="per-request"),
@@ -301,16 +342,55 @@ class TestRequestBatchCapability:
             "_try_load_model_cls",
             lambda model_class_name: _BatchCapablePipeline,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
         engine = DiffusionEngine(od_config)
         output = engine.execute_fn(_make_request_mode_sched_output(*request_ids))
 
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
         assert engine.supports_request_batch is True
         assert output == "batch"
         fake_executor.execute_batch.assert_called_once()
         fake_executor.execute_request.assert_not_called()
 
+    def test_make_engine_runs_startup_warmup(self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            engine_backend="default",
+            streaming_output=False,
+        )
+        fake_executor = SimpleNamespace(
+            execute_request=mocker.Mock(),
+            execute_batch=mocker.Mock(),
+            execute_step=mocker.Mock(),
+        )
+        fake_executor_cls = mocker.Mock(return_value=fake_executor)
+        warmup = mocker.Mock()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+        monkeypatch.setattr(DiffusionEngine, "run_startup_warmup", warmup)
+
+        engine = DiffusionEngine.make_engine(od_config)
+
+        assert isinstance(engine, DiffusionEngine)
+        warmup.assert_called_once_with()
+
 
 class TestRequestBatchAdmission:
     pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -470,11 +550,20 @@ def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
     assert moved["items"][1][0] == "keep"
 
 
+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
+
+
 @pytest.mark.asyncio
-async def test_async_add_req_and_wait_for_response():
+async def test_async_add_req_and_stream_response():
     engine = object.__new__(DiffusionEngine)
     engine.scheduler = MockScheduler()
-    engine._out_queue = {}
+    engine._out_streams = {}
     engine.abort_queue: queue.Queue[str] = queue.Queue()
     engine._rpc_queue = queue.Queue()
     engine._rpc_lock = threading.RLock()
@@ -485,6 +574,7 @@ async def test_async_add_req_and_wait_for_response():
     engine._loop_started = False
     engine.main_loop = None
     engine.supports_request_batch = False
+    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
 
     engine._finalize_finished_request = lambda rid, out, err: out.result
 
@@ -512,7 +602,7 @@ async def test_async_add_req_and_wait_for_response():
     async def run_task(rid):
         req = SimpleNamespace(request_id=rid)
         start = time.time()
-        res = await engine.async_add_req_and_wait_for_response(req)
+        res = await _consume_final_output(engine.async_add_req_and_stream_response(req))
         return rid, res, time.time() - start
 
     task_ids = [f"req_{i}" for i in range(5)]