wehub-resource-sync
534 lines
19 KiB
Python
534 lines
19 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
|
|
import torch
|
|
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.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.sched.interface import (
|
|
CachedRequestData,
|
|
NewRequestData,
|
|
)
|
|
from vllm_omni.diffusion.sched.interface import (
|
|
DiffusionSchedulerOutput as RealDiffusionSchedulerOutput,
|
|
)
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
|
|
@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_request_ids(self):
|
|
ids = [req.request_id for req in self.scheduled_new_reqs]
|
|
if self.scheduled_cached_reqs:
|
|
ids.extend(self.scheduled_cached_reqs.request_ids)
|
|
return ids
|
|
|
|
@property
|
|
def is_empty(self):
|
|
return len(self.scheduled_request_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]
|
|
|
|
|
|
class _BatchCapablePipeline:
|
|
supports_request_batch = True
|
|
|
|
|
|
class _SingleRequestPipeline:
|
|
pass
|
|
|
|
|
|
class _SingleRequestOverridePipeline(_BatchCapablePipeline):
|
|
def forward(self, req, prompt_ids=None):
|
|
return DiffusionOutput(output=None)
|
|
|
|
|
|
def _make_request_mode_sched_output(*request_ids: str) -> RealDiffusionSchedulerOutput:
|
|
new_reqs = [
|
|
NewRequestData(
|
|
request_id=request_id,
|
|
req=OmniDiffusionRequest(
|
|
prompt=f"prompt_{request_id}",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
|
request_id=request_id,
|
|
),
|
|
)
|
|
for request_id in request_ids
|
|
]
|
|
return RealDiffusionSchedulerOutput(
|
|
step_id=0,
|
|
scheduled_new_reqs=new_reqs,
|
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
|
finished_req_ids=set(),
|
|
num_running_reqs=len(new_reqs),
|
|
num_waiting_reqs=0,
|
|
)
|
|
|
|
|
|
class TestRequestBatchCapability:
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
def test_supports_request_batch_uses_registered_model_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
od_config = SimpleNamespace(model_class_name="BatchPipeline", custom_pipeline_args=None)
|
|
|
|
monkeypatch.setattr(
|
|
diffusion_engine_module.DiffusionModelRegistry,
|
|
"_try_load_model_cls",
|
|
lambda model_class_name: _BatchCapablePipeline if model_class_name == "BatchPipeline" else None,
|
|
)
|
|
|
|
assert diffusion_engine_module.supports_request_batch(od_config) is True
|
|
|
|
def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
od_config = SimpleNamespace(
|
|
model_class_name="SinglePipeline",
|
|
custom_pipeline_args={"pipeline_class": _BatchCapablePipeline},
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
diffusion_engine_module.DiffusionModelRegistry,
|
|
"_try_load_model_cls",
|
|
lambda model_class_name: _SingleRequestPipeline,
|
|
)
|
|
|
|
assert diffusion_engine_module.supports_request_batch(od_config) is True
|
|
|
|
def test_supports_request_batch_uses_custom_pipeline_class_name(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
od_config = SimpleNamespace(
|
|
model_class_name="SinglePipeline",
|
|
custom_pipeline_args={"pipeline_class": "test.module.BatchPipeline"},
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
diffusion_engine_module,
|
|
"resolve_obj_by_qualname",
|
|
lambda qualname: _BatchCapablePipeline if qualname == "test.module.BatchPipeline" else None,
|
|
)
|
|
monkeypatch.setattr(
|
|
diffusion_engine_module.DiffusionModelRegistry,
|
|
"_try_load_model_cls",
|
|
lambda model_class_name: _SingleRequestPipeline,
|
|
)
|
|
|
|
assert diffusion_engine_module.supports_request_batch(od_config) is True
|
|
|
|
def test_supports_request_batch_uses_only_explicit_pipeline_attribute_for_custom_override(
|
|
self, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
od_config = SimpleNamespace(
|
|
model_class_name="BatchPipeline",
|
|
custom_pipeline_args={"pipeline_class": _SingleRequestOverridePipeline},
|
|
)
|
|
monkeypatch.setattr(
|
|
diffusion_engine_module.DiffusionModelRegistry,
|
|
"_try_load_model_cls",
|
|
lambda model_class_name: None,
|
|
)
|
|
|
|
assert diffusion_engine_module.supports_request_batch(od_config) is True
|
|
|
|
def test_supports_request_batch_honors_explicit_false_on_custom_override(
|
|
self, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
class _ExplicitlyUnsupportedOverride(_BatchCapablePipeline):
|
|
supports_request_batch = False
|
|
|
|
def forward(self, req, prompt_ids=None):
|
|
return DiffusionOutput(output=None)
|
|
|
|
od_config = SimpleNamespace(
|
|
model_class_name="BatchPipeline",
|
|
custom_pipeline_args={"pipeline_class": _ExplicitlyUnsupportedOverride},
|
|
)
|
|
monkeypatch.setattr(
|
|
diffusion_engine_module.DiffusionModelRegistry,
|
|
"_try_load_model_cls",
|
|
lambda model_class_name: None,
|
|
)
|
|
|
|
assert diffusion_engine_module.supports_request_batch(od_config) is False
|
|
|
|
def test_supports_request_batch_rejects_invalid_custom_pipeline_class_name(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
od_config = SimpleNamespace(
|
|
model_class_name="BatchPipeline",
|
|
custom_pipeline_args={"pipeline_class": "test.module.MissingPipeline"},
|
|
)
|
|
|
|
def fail_resolve(qualname):
|
|
raise ImportError(qualname)
|
|
|
|
monkeypatch.setattr(diffusion_engine_module, "resolve_obj_by_qualname", fail_resolve)
|
|
registry_load = mocker.Mock(return_value=_BatchCapablePipeline)
|
|
monkeypatch.setattr(
|
|
diffusion_engine_module.DiffusionModelRegistry,
|
|
"_try_load_model_cls",
|
|
registry_load,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Failed to resolve custom diffusion pipeline class"):
|
|
diffusion_engine_module.supports_request_batch(od_config)
|
|
registry_load.assert_not_called()
|
|
|
|
def test_engine_disables_batch_dispatch_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,
|
|
)
|
|
fake_executor = SimpleNamespace(
|
|
execute_request=mocker.Mock(return_value="per-request"),
|
|
execute_batch=mocker.Mock(return_value="batch"),
|
|
execute_step=mocker.Mock(return_value="step"),
|
|
)
|
|
fake_executor_cls = mocker.Mock(return_value=fake_executor)
|
|
|
|
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, "_dummy_run", lambda self: None)
|
|
|
|
engine = DiffusionEngine(od_config)
|
|
output = engine.execute_fn(_make_request_mode_sched_output("req-a", "req-b"))
|
|
|
|
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()
|
|
|
|
@pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
|
|
def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
|
|
self,
|
|
request_ids: tuple[str, ...],
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
od_config = SimpleNamespace(
|
|
model_class_name="BatchPipeline",
|
|
custom_pipeline_args=None,
|
|
streaming_output=False,
|
|
)
|
|
fake_executor = SimpleNamespace(
|
|
execute_request=mocker.Mock(return_value="per-request"),
|
|
execute_batch=mocker.Mock(return_value="batch"),
|
|
execute_step=mocker.Mock(return_value="step"),
|
|
)
|
|
fake_executor_cls = mocker.Mock(return_value=fake_executor)
|
|
|
|
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: _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.supports_request_batch is True
|
|
assert output == "batch"
|
|
fake_executor.execute_batch.assert_called_once()
|
|
fake_executor.execute_request.assert_not_called()
|
|
|
|
|
|
class TestRequestBatchAdmission:
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
def test_config_rejects_negative_request_batch_max_wait_ms(self) -> None:
|
|
with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
|
|
OmniDiffusionConfig(model="test", request_batch_max_wait_ms=-1.0)
|
|
|
|
def test_config_normalizes_request_batch_max_wait_ms_to_float(self) -> None:
|
|
config = OmniDiffusionConfig(model="test", request_batch_max_wait_ms=5)
|
|
|
|
assert config.request_batch_max_wait_ms == 5.0
|
|
assert isinstance(config.request_batch_max_wait_ms, float)
|
|
|
|
def test_scheduler_exposes_waiting_and_running_counts(self) -> None:
|
|
from vllm_omni.diffusion.sched import RequestScheduler
|
|
|
|
od_config = SimpleNamespace(max_num_seqs=4)
|
|
scheduler = RequestScheduler()
|
|
scheduler.initialize(od_config)
|
|
|
|
assert scheduler.num_waiting_requests() == 0
|
|
assert scheduler.num_running_requests() == 0
|
|
|
|
scheduler.add_request(
|
|
OmniDiffusionRequest(
|
|
prompt="prompt_a",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
|
request_id="req-a",
|
|
)
|
|
)
|
|
scheduler.add_request(
|
|
OmniDiffusionRequest(
|
|
prompt="prompt_b",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
|
request_id="req-b",
|
|
)
|
|
)
|
|
assert scheduler.num_waiting_requests() == 2
|
|
assert scheduler.num_running_requests() == 0
|
|
|
|
scheduler.schedule()
|
|
assert scheduler.num_waiting_requests() == 0
|
|
assert scheduler.num_running_requests() == 2
|
|
|
|
def test_request_batch_admission_exits_early_when_waiting_queue_stable(self) -> None:
|
|
from vllm_omni.diffusion.sched import RequestScheduler
|
|
|
|
od_config = SimpleNamespace(
|
|
max_num_seqs=32,
|
|
request_batch_max_wait_ms=1000.0,
|
|
step_execution=False,
|
|
)
|
|
scheduler = RequestScheduler()
|
|
scheduler.initialize(od_config)
|
|
for idx in range(2):
|
|
scheduler.add_request(
|
|
OmniDiffusionRequest(
|
|
prompt=f"prompt_{idx}",
|
|
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
|
request_id=f"req-{idx}",
|
|
)
|
|
)
|
|
|
|
engine = object.__new__(DiffusionEngine)
|
|
engine.od_config = od_config
|
|
engine.scheduler = scheduler
|
|
engine.step_execution = False
|
|
engine.supports_request_batch = True
|
|
engine.stop_event = threading.Event()
|
|
engine._rpc_lock = threading.RLock()
|
|
engine._cv = threading.Condition(engine._rpc_lock)
|
|
|
|
start = time.monotonic()
|
|
with engine._cv:
|
|
engine._wait_for_request_batch_admission_locked()
|
|
waited_s = time.monotonic() - start
|
|
|
|
# Stable-window exit (~50ms), not the full 1000ms deadline.
|
|
assert waited_s < 0.5
|
|
assert waited_s >= 0.04
|
|
assert scheduler.num_waiting_requests() == 2
|
|
assert scheduler.num_running_requests() == 0
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.cpu
|
|
def test_move_tensor_tree_keeps_cpu_tensor_identity() -> None:
|
|
tensor = torch.arange(8, dtype=torch.float32)
|
|
|
|
moved = _move_tensor_tree_to_cpu(tensor)
|
|
|
|
assert moved is tensor
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.cpu
|
|
def test_move_tensor_tree_preserves_nested_structure_without_mutating_input() -> None:
|
|
tensor = torch.arange(4, dtype=torch.float32)
|
|
nested_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
|
|
sentinel = object()
|
|
payload = {
|
|
"tensor": tensor,
|
|
"list": [nested_tensor, sentinel],
|
|
"tuple": ({"inner": tensor}, "metadata"),
|
|
"scalar": 3,
|
|
}
|
|
|
|
moved = _move_tensor_tree_to_cpu(payload)
|
|
|
|
assert moved is not payload
|
|
assert set(moved) == {"tensor", "list", "tuple", "scalar"}
|
|
assert moved["list"] is not payload["list"]
|
|
assert moved["tuple"] is not payload["tuple"]
|
|
assert moved["tuple"][0] is not payload["tuple"][0]
|
|
assert moved["tensor"] is tensor
|
|
assert moved["list"][0] is nested_tensor
|
|
assert moved["list"][1] is sentinel
|
|
assert moved["tuple"][0]["inner"] is tensor
|
|
assert moved["tuple"][1] == "metadata"
|
|
assert moved["scalar"] == 3
|
|
assert payload["list"][0] is nested_tensor
|
|
assert payload["list"][1] is sentinel
|
|
assert payload["tuple"][0]["inner"] is tensor
|
|
assert payload["tuple"][1] == "metadata"
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.cpu
|
|
def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None:
|
|
value = object()
|
|
|
|
moved = _move_tensor_tree_to_cpu(value)
|
|
|
|
assert moved is value
|
|
|
|
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.cuda
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
|
def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
|
|
tensor = torch.arange(8, dtype=torch.float32, device="cuda")
|
|
other = torch.arange(4, dtype=torch.int64, device="cuda")
|
|
payload = {"tensor": tensor, "items": [other, ("keep", tensor)]}
|
|
|
|
moved = _move_tensor_tree_to_cpu(payload)
|
|
|
|
assert moved["tensor"].device.type == "cpu"
|
|
assert moved["items"][0].device.type == "cpu"
|
|
assert moved["items"][1][1].device.type == "cpu"
|
|
torch.testing.assert_close(moved["tensor"], tensor.cpu())
|
|
torch.testing.assert_close(moved["items"][0], other.cpu())
|
|
torch.testing.assert_close(moved["items"][1][1], tensor.cpu())
|
|
assert moved["items"][1][0] == "keep"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_add_req_and_wait_for_response():
|
|
engine = object.__new__(DiffusionEngine)
|
|
engine.scheduler = MockScheduler()
|
|
engine._out_queue = {}
|
|
engine.abort_queue: queue.Queue[str] = queue.Queue()
|
|
engine._rpc_queue = queue.Queue()
|
|
engine._rpc_lock = threading.RLock()
|
|
engine._cv = threading.Condition(engine._rpc_lock)
|
|
engine._init_lock = asyncio.Lock()
|
|
engine._closed = False
|
|
engine.od_config = SimpleNamespace(streaming_output=False)
|
|
engine._loop_started = False
|
|
engine.main_loop = None
|
|
engine.supports_request_batch = False
|
|
|
|
engine._finalize_finished_request = lambda rid, out, err: out.result
|
|
|
|
def mock_execute_batch(sched_output):
|
|
request_ids = sched_output.scheduled_request_ids
|
|
|
|
time.sleep(1)
|
|
|
|
class MockRunnerOutput:
|
|
def __init__(self, ids):
|
|
self.request_id = ids
|
|
self.step_index = 0
|
|
self.finished = True
|
|
self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
|
|
|
|
def get_request_output(self, rid):
|
|
return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
|
|
|
|
return MockRunnerOutput(request_ids)
|
|
|
|
engine.execute_fn = mock_execute_batch
|
|
|
|
await engine._check_and_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]
|
|
try:
|
|
results = await asyncio.gather(*tasks)
|
|
finally:
|
|
with engine._cv:
|
|
engine.stop_event.set()
|
|
engine._cv.notify_all()
|
|
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
|