118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
|
# 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
|