[Feat] DiffusionEngine Support async batch infer (#2729)

Signed-off-by: Semmer <semmer@live.cn>
Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: jader <yjader@foxmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/diffusion/test_diffusion_engine.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+
+@dataclass
+class DiffusionSchedulerOutput:
+    step_id: int
+    scheduled_new_reqs: list = field(default_factory=list)
+    scheduled_cached_reqs: Any = None
+    finished_req_ids: set = field(default_factory=set)
+    num_running_reqs: int = 0
+    num_waiting_reqs: int = 0
+
+    @property
+    def scheduled_req_ids(self):
+        ids = [req.request_id for req in self.scheduled_new_reqs]
+        if self.scheduled_cached_reqs:
+            ids.extend(self.scheduled_cached_reqs.sched_req_ids)
+        return ids
+
+    @property
+    def is_empty(self):
+        return len(self.scheduled_req_ids) == 0
+
+
+class MockScheduler:
+    def __init__(self):
+        self._waiting_queue = []
+        self._step_id = 0
+
+    def add_request(self, request):
+        self._waiting_queue.append(request)
+        return request.request_id
+
+    def has_requests(self):
+        return len(self._waiting_queue) > 0
+
+    def schedule(self) -> DiffusionSchedulerOutput:
+        if not self._waiting_queue:
+            return DiffusionSchedulerOutput(step_id=self._step_id)
+
+        batch = []
+        while self._waiting_queue:
+            req = self._waiting_queue.pop(0)
+            batch.append(SimpleNamespace(request_id=req.request_id))
+
+        output = DiffusionSchedulerOutput(step_id=self._step_id, scheduled_new_reqs=batch)
+        self._step_id += 1
+        return output
+
+    def update_from_output(self, sched_output, runner_output):
+        # assume all new req finished
+        return [req.request_id for req in sched_output.scheduled_new_reqs]
+
+
+@pytest.mark.asyncio
+async def test_async_add_req_and_wait_for_response():
+    from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+    engine = object.__new__(DiffusionEngine)
+    engine.scheduler = MockScheduler()
+    engine._out_queue = {}
+    engine.abort_queue: queue.Queue[str] = queue.Queue()
+    engine._rpc_lock = threading.Lock()
+
+    engine._finalize_finished_request = lambda rid, out, err: out.result
+
+    def mock_execute_batch(sched_output):
+        req_ids = sched_output.scheduled_req_ids
+
+        time.sleep(1)
+
+        class MockRunnerOutput:
+            def __init__(self, ids):
+                self.req_id = ids
+                self.step_index = 0
+                self.finished = True
+                self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
+
+            def get_req_output(self, rid):
+                return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
+
+        return MockRunnerOutput(req_ids)
+
+    engine.execute_fn = mock_execute_batch
+
+    engine.stop_event = threading.Event()
+    engine.start_background_loop()
+
+    async def run_task(rid):
+        req = SimpleNamespace(request_id=rid)
+        start = time.time()
+        res = await engine.async_add_req_and_wait_for_response(req)
+        return rid, res, time.time() - start
+
+    task_ids = [f"req_{i}" for i in range(5)]
+    tasks = [run_task(rid) for rid in task_ids]
+    results = await asyncio.gather(*tasks)
+
+    engine.stop_event.set()
+    engine.worker_thread.join(timeout=5)
+    assert len(results) == 5
+    for rid, res, elapsed in results:
+        assert rid in res.result_data
+
+    eps = 0.5
+    latencies = [r[2] for r in results]
+    assert max(latencies) - min(latencies) < eps