690 lines
23 KiB
Python
690 lines
23 KiB
Python
from __future__ import annotations
|
|
|
|
import queue
|
|
from collections.abc import Callable
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from vllm.entrypoints.openai.models.protocol import BaseModelPath
|
|
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
|
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
|
|
|
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
THREE_STAGE_META = [
|
|
{"stage_type": "llm", "final_output": True, "final_output_type": "text"},
|
|
{"stage_type": "llm", "final_output": False, "final_output_type": None},
|
|
{"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
|
|
]
|
|
|
|
DIFFUSION_ONLY_META = [
|
|
{"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
|
|
]
|
|
|
|
LLM_DIFFUSION_META = [
|
|
{"stage_type": "llm", "final_output": True, "final_output_type": "text"},
|
|
{"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
|
|
]
|
|
|
|
|
|
class FakeEngineOutput:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
payload: str,
|
|
finished: bool,
|
|
images: list[str] | None = None,
|
|
stage_durations: dict[str, float] | None = None,
|
|
) -> None:
|
|
self.payload = payload
|
|
self.finished = finished
|
|
self.images = images or []
|
|
self.stage_durations = stage_durations or {}
|
|
|
|
|
|
def make_output_msg(
|
|
request_id: str,
|
|
stage_id: int,
|
|
*,
|
|
payload: str,
|
|
output_finished: bool,
|
|
finished: bool | None = None,
|
|
images: list[str] | None = None,
|
|
metrics: Any = None,
|
|
) -> dict[str, Any]:
|
|
if finished is None:
|
|
finished = output_finished
|
|
return {
|
|
"type": "output",
|
|
"request_id": request_id,
|
|
"stage_id": stage_id,
|
|
"engine_outputs": FakeEngineOutput(
|
|
payload=payload,
|
|
finished=output_finished,
|
|
images=images,
|
|
),
|
|
"finished": finished,
|
|
"metrics": metrics,
|
|
}
|
|
|
|
|
|
class FakeAsyncOmniEngine:
|
|
def __init__(
|
|
self,
|
|
model: str = "dummy-model",
|
|
*,
|
|
stage_metadata: list[dict[str, Any]] | None = None,
|
|
default_sampling_params_list: list[Any] | None = None,
|
|
on_add_request: Callable[[FakeAsyncOmniEngine, dict[str, Any]], None] | None = None,
|
|
rpc_results: list[Any] | None = None,
|
|
**_: Any,
|
|
) -> None:
|
|
self.model = model
|
|
self.config_path = None
|
|
self.stage_configs: list[Any] = []
|
|
self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]]
|
|
self.num_stages = len(self.stage_metadata)
|
|
self.default_sampling_params_list = default_sampling_params_list or [
|
|
SamplingParams(max_tokens=8) for _ in range(self.num_stages)
|
|
]
|
|
self.supported_tasks = ("generate",)
|
|
self.stage_clients = [SimpleNamespace(is_comprehension=False) for _ in range(self.num_stages)]
|
|
self.stage_vllm_configs = [None for _ in range(self.num_stages)]
|
|
self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
|
|
self.input_processor = None
|
|
|
|
self.output_q: queue.Queue[dict[str, Any]] = queue.Queue()
|
|
self.submitted: list[dict[str, Any]] = []
|
|
self.aborted: list[list[str]] = []
|
|
self.rpc_results = rpc_results or []
|
|
self.on_add_request = on_add_request
|
|
self.shutdown_called = False
|
|
self._alive = True
|
|
|
|
def add_request(
|
|
self,
|
|
request_id: str,
|
|
prompt: Any,
|
|
sampling_params_list: list[Any] | None = None,
|
|
final_stage_id: int = 0,
|
|
arrival_time: float | None = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
msg = {
|
|
"request_id": request_id,
|
|
"prompt": prompt,
|
|
"sampling_params_list": sampling_params_list,
|
|
"final_stage_id": final_stage_id,
|
|
"arrival_time": arrival_time,
|
|
}
|
|
self.submitted.append(msg)
|
|
if self.on_add_request is not None:
|
|
self.on_add_request(self, msg)
|
|
|
|
async def add_request_async(self, *args, **kwargs) -> None:
|
|
self.add_request(*args, **kwargs)
|
|
|
|
def try_get_output(self, timeout: float = 0.001) -> dict[str, Any] | None:
|
|
try:
|
|
return self.output_q.get_nowait()
|
|
except queue.Empty:
|
|
return None
|
|
|
|
async def try_get_output_async(self) -> dict[str, Any] | None:
|
|
return self.try_get_output()
|
|
|
|
def get_stage_metadata(self, stage_id: int) -> dict[str, Any]:
|
|
return self.stage_metadata[stage_id]
|
|
|
|
def abort(self, request_ids: list[str]) -> None:
|
|
self.aborted.append(list(request_ids))
|
|
|
|
async def abort_async(self, request_ids: list[str]) -> None:
|
|
self.abort(request_ids)
|
|
|
|
async def collective_rpc_async(self, **_: Any) -> list[Any]:
|
|
return list(self.rpc_results)
|
|
|
|
def is_alive(self) -> bool:
|
|
return self._alive
|
|
|
|
def shutdown(self) -> None:
|
|
self.shutdown_called = True
|
|
self._alive = False
|
|
|
|
|
|
def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) -> None:
|
|
monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", lambda *args, **kwargs: engine)
|
|
monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
|
|
|
|
|
|
def _stage_spec(
|
|
stage_id: int,
|
|
*,
|
|
payloads: list[str],
|
|
finished: bool = False,
|
|
image_payloads: list[str] | None = None,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"stage_id": stage_id,
|
|
"payloads": payloads,
|
|
"finished": finished,
|
|
"image_payloads": image_payloads or [],
|
|
}
|
|
|
|
|
|
def _enqueue_outputs(
|
|
engine: FakeAsyncOmniEngine,
|
|
msg: dict[str, Any],
|
|
*,
|
|
stage_specs: list[dict[str, Any]],
|
|
) -> None:
|
|
request_id = msg["request_id"]
|
|
for spec in stage_specs:
|
|
payloads = spec["payloads"]
|
|
image_payloads = spec.get("image_payloads", [])
|
|
last_idx = len(payloads) - 1
|
|
|
|
for idx, payload_tmpl in enumerate(payloads):
|
|
images = []
|
|
if idx < len(image_payloads):
|
|
images = [image_payloads[idx].format(request_id=request_id, idx=idx)]
|
|
|
|
engine.output_q.put_nowait(
|
|
make_output_msg(
|
|
request_id,
|
|
spec["stage_id"],
|
|
payload=payload_tmpl.format(request_id=request_id, idx=idx),
|
|
output_finished=(idx == last_idx),
|
|
finished=bool(spec.get("finished")) and idx == last_idx,
|
|
images=images,
|
|
)
|
|
)
|
|
|
|
|
|
def _enqueue_omni_final_only_outputs(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
|
sampling_params_list = msg["sampling_params_list"]
|
|
llm_streaming = any(params.output_kind != RequestOutputKind.FINAL_ONLY for params in sampling_params_list[:2])
|
|
stage0_count = 3 if llm_streaming else 1
|
|
stage1_count = 3 if llm_streaming else 1
|
|
|
|
_enqueue_outputs(
|
|
engine,
|
|
msg,
|
|
stage_specs=[
|
|
_stage_spec(0, payloads=[f"{{request_id}}-stage0-{idx}" for idx in range(stage0_count)]),
|
|
_stage_spec(1, payloads=[f"{{request_id}}-stage1-{idx}" for idx in range(stage1_count)]),
|
|
_stage_spec(
|
|
2,
|
|
payloads=["{request_id}-stage2-final"],
|
|
finished=True,
|
|
image_payloads=["{request_id}-img-final"],
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
def _enqueue_omni_llm_diffusion_outputs(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
|
sampling_params_list = msg["sampling_params_list"]
|
|
llm_streaming = sampling_params_list[0].output_kind != RequestOutputKind.FINAL_ONLY
|
|
stage0_count = 3 if llm_streaming else 1
|
|
|
|
_enqueue_outputs(
|
|
engine,
|
|
msg,
|
|
stage_specs=[
|
|
_stage_spec(0, payloads=[f"{{request_id}}-text-{idx}" for idx in range(stage0_count)]),
|
|
_stage_spec(
|
|
1,
|
|
payloads=["{request_id}-image-final"],
|
|
finished=True,
|
|
image_payloads=["{request_id}-image"],
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
def _enqueue_async_three_stage_outputs(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
|
_enqueue_outputs(
|
|
engine,
|
|
msg,
|
|
stage_specs=[
|
|
_stage_spec(0, payloads=[f"{{request_id}}-stage0-{idx}" for idx in range(3)]),
|
|
_stage_spec(1, payloads=[f"{{request_id}}-stage1-{idx}" for idx in range(3)]),
|
|
_stage_spec(
|
|
2,
|
|
payloads=[f"{{request_id}}-stage2-{idx}" for idx in range(3)],
|
|
finished=True,
|
|
image_payloads=[f"{{request_id}}-img-{idx}" for idx in range(3)],
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
def _enqueue_async_finish_outputs(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
|
_enqueue_outputs(
|
|
engine,
|
|
msg,
|
|
stage_specs=[
|
|
_stage_spec(0, payloads=["{request_id}-stage0"]),
|
|
_stage_spec(
|
|
2,
|
|
payloads=["{request_id}-stage2-final"],
|
|
finished=True,
|
|
image_payloads=["{request_id}-img-final"],
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
def _enqueue_async_diffusion_only_output(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
|
_enqueue_outputs(
|
|
engine,
|
|
msg,
|
|
stage_specs=[
|
|
_stage_spec(
|
|
0,
|
|
payloads=["{request_id}-diffusion-final"],
|
|
finished=True,
|
|
image_payloads=["{request_id}-image"],
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
def _enqueue_async_llm_diffusion_outputs(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
|
_enqueue_outputs(
|
|
engine,
|
|
msg,
|
|
stage_specs=[
|
|
_stage_spec(0, payloads=[f"{{request_id}}-text-{idx}" for idx in range(3)]),
|
|
_stage_spec(
|
|
1,
|
|
payloads=["{request_id}-image-final"],
|
|
finished=True,
|
|
image_payloads=["{request_id}-image"],
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
def _enqueue_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
|
engine.output_q.put_nowait(
|
|
{
|
|
"type": "error",
|
|
"request_id": msg["request_id"],
|
|
"stage_id": 0,
|
|
"error": "engine boom",
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_supported_tasks_returns_engine_supported_tasks():
|
|
omni = object.__new__(AsyncOmni)
|
|
omni.engine = SimpleNamespace(supported_tasks=("generate", "speech"))
|
|
|
|
supported_tasks = await omni.get_supported_tasks()
|
|
|
|
assert supported_tasks == ("generate", "speech")
|
|
|
|
|
|
def test_model_config_and_vllm_config_forward_from_comprehension_stage():
|
|
model_config = SimpleNamespace(model="Qwen/Qwen3-TTS")
|
|
vllm_config = SimpleNamespace(model_config=model_config)
|
|
renderer = SimpleNamespace(name="renderer")
|
|
input_processor = SimpleNamespace(renderer=renderer)
|
|
io_processor = SimpleNamespace(name="io-processor")
|
|
omni = object.__new__(AsyncOmni)
|
|
omni.engine = SimpleNamespace(
|
|
stage_clients=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)],
|
|
stage_vllm_configs=[None, vllm_config],
|
|
)
|
|
omni.input_processor = input_processor
|
|
omni.io_processor = io_processor
|
|
|
|
assert omni.vllm_config is vllm_config
|
|
assert omni.model_config is model_config
|
|
assert omni.renderer is renderer
|
|
assert omni.input_processor is input_processor
|
|
assert omni.io_processor is io_processor
|
|
|
|
|
|
def test_openai_serving_models_can_consume_async_omni_compat_attrs():
|
|
model_config = SimpleNamespace(model="Qwen/Qwen3-TTS", max_model_len=32768)
|
|
vllm_config = SimpleNamespace(model_config=model_config)
|
|
renderer = SimpleNamespace(name="renderer")
|
|
input_processor = SimpleNamespace(renderer=renderer)
|
|
io_processor = SimpleNamespace(name="io-processor")
|
|
omni = object.__new__(AsyncOmni)
|
|
omni.engine = SimpleNamespace(
|
|
stage_clients=[SimpleNamespace(is_comprehension=True)],
|
|
stage_vllm_configs=[vllm_config],
|
|
)
|
|
omni.input_processor = input_processor
|
|
omni.io_processor = io_processor
|
|
|
|
serving_models = OpenAIServingModels(
|
|
engine_client=omni,
|
|
base_model_paths=[BaseModelPath(name="tts-model", model_path="Qwen/Qwen3-TTS")],
|
|
)
|
|
|
|
assert serving_models.model_config is model_config
|
|
assert serving_models.renderer is renderer
|
|
assert serving_models.io_processor is io_processor
|
|
assert serving_models.input_processor is input_processor
|
|
|
|
|
|
def test_get_diffusion_od_config_returns_diffusion_stage_config():
|
|
diffusion_od_config = object()
|
|
omni = object.__new__(AsyncOmni)
|
|
omni.engine = SimpleNamespace(
|
|
stage_clients=[
|
|
SimpleNamespace(stage_type="llm"),
|
|
SimpleNamespace(stage_type="diffusion", _engine=SimpleNamespace(od_config=diffusion_od_config)),
|
|
]
|
|
)
|
|
|
|
assert omni.get_diffusion_od_config() is diffusion_od_config
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.MonkeyPatch):
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=THREE_STAGE_META,
|
|
on_add_request=lambda eng, msg: _enqueue_outputs(
|
|
eng,
|
|
msg,
|
|
stage_specs=[
|
|
_stage_spec(1, payloads=["non-final"]),
|
|
_stage_spec(2, payloads=["final"], finished=True, image_payloads=["final-img"]),
|
|
],
|
|
),
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = AsyncOmni("dummy-model")
|
|
try:
|
|
outputs = []
|
|
async for output in app.generate(prompt="hello", request_id="req-1"):
|
|
outputs.append(output)
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert [output.stage_id for output in outputs] == [2]
|
|
assert [output.request_output.payload for output in outputs] == ["final"]
|
|
assert "req-1" not in app.request_states
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_accepts_multiple_final_stage_streams(monkeypatch: pytest.MonkeyPatch):
|
|
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_async_three_stage_outputs)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = AsyncOmni("dummy-model")
|
|
try:
|
|
outputs = []
|
|
async for output in app.generate(prompt="hello", request_id="req-1"):
|
|
outputs.append(output)
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert [output.stage_id for output in outputs] == [0, 0, 0, 2, 2, 2]
|
|
assert [output.request_output.payload for output in outputs] == [
|
|
"req-1-stage0-0",
|
|
"req-1-stage0-1",
|
|
"req-1-stage0-2",
|
|
"req-1-stage2-0",
|
|
"req-1-stage2-1",
|
|
"req-1-stage2-2",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_stops_on_final_stage_finished(monkeypatch: pytest.MonkeyPatch):
|
|
# Intentionally jump from stage 0 to stage 2: stage 1 is a non-final stage
|
|
# and should be filtered out from the client-visible output stream.
|
|
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_async_finish_outputs)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = AsyncOmni("dummy-model")
|
|
try:
|
|
outputs = []
|
|
async for output in app.generate(prompt="hello", request_id="req-1"):
|
|
outputs.append(output)
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert [output.request_output.payload for output in outputs] == [
|
|
"req-1-stage0",
|
|
"req-1-stage2-final",
|
|
]
|
|
assert "req-1" not in app.request_states
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_diffusion_only_yields_single_image_output(monkeypatch: pytest.MonkeyPatch):
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=DIFFUSION_ONLY_META,
|
|
on_add_request=_enqueue_async_diffusion_only_output,
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = AsyncOmni("dummy-model")
|
|
try:
|
|
outputs = []
|
|
async for output in app.generate(prompt="hello", request_id="req-1"):
|
|
outputs.append(output)
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert len(outputs) == 1
|
|
assert outputs[0].stage_id == 0
|
|
assert outputs[0].final_output_type == "image"
|
|
assert outputs[0].images == ["req-1-image"]
|
|
assert outputs[0].request_output.payload == "req-1-diffusion-final"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatch: pytest.MonkeyPatch):
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=LLM_DIFFUSION_META,
|
|
on_add_request=_enqueue_async_llm_diffusion_outputs,
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = AsyncOmni("dummy-model")
|
|
try:
|
|
outputs = []
|
|
async for output in app.generate(prompt="hello", request_id="req-1"):
|
|
outputs.append(output)
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert [output.stage_id for output in outputs] == [0, 0, 0, 1]
|
|
assert [output.final_output_type for output in outputs] == ["text", "text", "text", "image"]
|
|
assert [output.request_output.payload for output in outputs] == [
|
|
"req-1-text-0",
|
|
"req-1-text-1",
|
|
"req-1-text-2",
|
|
"req-1-image-final",
|
|
]
|
|
assert outputs[-1].images == ["req-1-image"]
|
|
assert "req-1" not in app.request_states
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
|
|
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = AsyncOmni("dummy-model")
|
|
try:
|
|
app.request_states["req-1"] = object()
|
|
await app.abort("req-1")
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert engine.aborted == [["req-1"]]
|
|
assert "req-1" not in app.request_states
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_propagates_engine_error(monkeypatch: pytest.MonkeyPatch):
|
|
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_error_message)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = AsyncOmni("dummy-model")
|
|
try:
|
|
with pytest.raises(RuntimeError, match="engine boom"):
|
|
async for _ in app.generate(prompt="hello", request_id="req-1"):
|
|
pass
|
|
finally:
|
|
app.shutdown()
|
|
|
|
|
|
def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch):
|
|
sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=THREE_STAGE_META,
|
|
default_sampling_params_list=sampling_params,
|
|
on_add_request=_enqueue_omni_final_only_outputs,
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = Omni("dummy-model")
|
|
outputs = list(app.generate(["p1", "p2"], py_generator=True, use_tqdm=False))
|
|
|
|
assert len(outputs) == 4
|
|
assert [output.stage_id for output in outputs] == [0, 2, 0, 2]
|
|
assert [output.request_output.payload for output in outputs] == [
|
|
f"{engine.submitted[0]['request_id']}-stage0-0",
|
|
f"{engine.submitted[0]['request_id']}-stage2-final",
|
|
f"{engine.submitted[1]['request_id']}-stage0-0",
|
|
f"{engine.submitted[1]['request_id']}-stage2-final",
|
|
]
|
|
assert engine.shutdown_called is True
|
|
|
|
|
|
def test_omni_generate_returns_list_when_not_using_generator(monkeypatch: pytest.MonkeyPatch):
|
|
sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=THREE_STAGE_META,
|
|
default_sampling_params_list=sampling_params,
|
|
on_add_request=_enqueue_omni_final_only_outputs,
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = Omni("dummy-model")
|
|
try:
|
|
outputs = app.generate(["p1", "p2"], py_generator=False, use_tqdm=False)
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert isinstance(outputs, list)
|
|
assert len(outputs) == 4
|
|
assert [output.stage_id for output in outputs] == [0, 2, 0, 2]
|
|
|
|
|
|
def test_omni_generate_diffusion_only_yields_single_image_per_request(monkeypatch: pytest.MonkeyPatch):
|
|
sampling_params = [SamplingParams(max_tokens=8)]
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=DIFFUSION_ONLY_META,
|
|
default_sampling_params_list=sampling_params,
|
|
on_add_request=_enqueue_async_diffusion_only_output,
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = Omni("dummy-model")
|
|
outputs = list(app.generate(["p1", "p2"], py_generator=True, use_tqdm=False))
|
|
|
|
assert len(outputs) == 2
|
|
assert [output.stage_id for output in outputs] == [0, 0]
|
|
assert [output.final_output_type for output in outputs] == ["image", "image"]
|
|
assert [output.request_output.payload for output in outputs] == [
|
|
f"{engine.submitted[0]['request_id']}-diffusion-final",
|
|
f"{engine.submitted[1]['request_id']}-diffusion-final",
|
|
]
|
|
assert [output.images for output in outputs] == [
|
|
[f"{engine.submitted[0]['request_id']}-image"],
|
|
[f"{engine.submitted[1]['request_id']}-image"],
|
|
]
|
|
assert engine.shutdown_called is True
|
|
|
|
|
|
def test_omni_generate_llm_diffusion_yields_final_text_then_image_per_request(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
sampling_params = [SamplingParams(max_tokens=8) for _ in range(2)]
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=LLM_DIFFUSION_META,
|
|
default_sampling_params_list=sampling_params,
|
|
on_add_request=_enqueue_omni_llm_diffusion_outputs,
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = Omni("dummy-model")
|
|
outputs = list(app.generate(["p1", "p2"], py_generator=True, use_tqdm=False))
|
|
|
|
assert len(outputs) == 4
|
|
assert [output.stage_id for output in outputs] == [0, 1, 0, 1]
|
|
assert [output.final_output_type for output in outputs] == ["text", "image", "text", "image"]
|
|
assert [output.request_output.payload for output in outputs] == [
|
|
f"{engine.submitted[0]['request_id']}-text-0",
|
|
f"{engine.submitted[0]['request_id']}-image-final",
|
|
f"{engine.submitted[1]['request_id']}-text-0",
|
|
f"{engine.submitted[1]['request_id']}-image-final",
|
|
]
|
|
assert [output.images for output in outputs] == [
|
|
[],
|
|
[f"{engine.submitted[0]['request_id']}-image"],
|
|
[],
|
|
[f"{engine.submitted[1]['request_id']}-image"],
|
|
]
|
|
assert engine.submitted[0]["sampling_params_list"][0].output_kind == RequestOutputKind.FINAL_ONLY
|
|
assert engine.shutdown_called is True
|
|
|
|
|
|
def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
|
|
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = Omni("dummy-model")
|
|
try:
|
|
app.request_states["req-1"] = object()
|
|
app.abort("req-1")
|
|
finally:
|
|
app.shutdown()
|
|
|
|
assert engine.aborted == [["req-1"]]
|
|
assert "req-1" not in app.request_states
|
|
|
|
|
|
def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch):
|
|
sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
|
|
original_diffusion_output_kind = sampling_params[2].output_kind
|
|
engine = FakeAsyncOmniEngine(
|
|
stage_metadata=THREE_STAGE_META,
|
|
default_sampling_params_list=sampling_params,
|
|
on_add_request=_enqueue_omni_final_only_outputs,
|
|
)
|
|
_patch_engine(monkeypatch, engine)
|
|
|
|
app = Omni("dummy-model")
|
|
try:
|
|
outputs = list(app.generate(["p1"], py_generator=True, use_tqdm=False))
|
|
finally:
|
|
if not engine.shutdown_called:
|
|
app.shutdown()
|
|
|
|
submitted_params = engine.submitted[0]["sampling_params_list"]
|
|
assert submitted_params[0].output_kind == RequestOutputKind.FINAL_ONLY
|
|
assert submitted_params[1].output_kind == RequestOutputKind.FINAL_ONLY
|
|
assert submitted_params[2].output_kind == original_diffusion_output_kind
|
|
assert len(outputs) == 2
|