[Refactor][OutputProcessor 2/8]: OmniRequestOutput should inherit RequestOutput and no nested wrap-up (#5146)
Signed-off-by: Boao Shi <aoibosh@connect.hku.hk> Signed-off-by: bowieshi <aoibosh@connect.hku.hk> Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: R2-Y <ruiruyang2@gmail.com> Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com> Signed-off-by: Shi Boao Bowie <aoibosh@connect.hku.hk> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: ruirui(rein) yang <73573651+R2-Y@users.noreply.github.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
89 files changed
+944
-639
No files matched your search
@@ -50,7 +50,7 @@ async def generate(
|
||||
output_modalities=["text"],
|
||||
):
|
||||
stage_id = omni_output.stage_id
|
||||
out = omni_output.request_output
|
||||
out = omni_output
|
||||
if stage_id == 0:
|
||||
num_tokens = sum(len(output.token_ids) for output in out.outputs)
|
||||
count = num_tokens
|
||||
|
||||
@@ -4,171 +4,269 @@
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestOmniRequestOutput:
|
||||
"""Tests for OmniRequestOutput class."""
|
||||
# ── helpers ────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_from_diffusion(self):
|
||||
"""Test creating output from diffusion model."""
|
||||
images = [Image.new("RGB", (64, 64), color="red")]
|
||||
output = OmniRequestOutput.from_diffusion(
|
||||
request_id="test-123",
|
||||
images=images,
|
||||
prompt="a cat",
|
||||
|
||||
def _make_text_request_output(
|
||||
request_id: str = "req-1",
|
||||
text: str = "hello world",
|
||||
token_ids: list[int] | None = None,
|
||||
finished: bool = True,
|
||||
) -> RequestOutput:
|
||||
"""Build a minimal vLLM RequestOutput for testing."""
|
||||
return RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="test prompt",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
prompt_logprobs=None,
|
||||
outputs=[
|
||||
CompletionOutput(
|
||||
index=0,
|
||||
text=text,
|
||||
token_ids=token_ids or [10, 11, 12],
|
||||
cumulative_logprob=-1.0,
|
||||
logprobs=None,
|
||||
finish_reason="stop" if finished else None,
|
||||
stop_reason=None,
|
||||
)
|
||||
],
|
||||
finished=finished,
|
||||
metrics=None,
|
||||
lora_request=None,
|
||||
)
|
||||
|
||||
|
||||
# ── from_stage_output ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestFromStageOutput:
|
||||
"""Tests for OmniRequestOutput.from_stage_output()."""
|
||||
|
||||
def test_copies_request_output_content(self):
|
||||
"""Content from a vLLM RequestOutput is flattened onto the new object."""
|
||||
source = _make_text_request_output(text="copied text")
|
||||
out = OmniRequestOutput.from_stage_output(source, request_id="req-1", stage_id=0, final_output_type="text")
|
||||
|
||||
assert out.prompt == "test prompt"
|
||||
assert out.prompt_token_ids == [1, 2, 3]
|
||||
assert len(out.outputs) == 1
|
||||
assert out.outputs[0].text == "copied text"
|
||||
assert out.finished is True
|
||||
assert out.stage_id == 0
|
||||
assert out.final_output_type == "text"
|
||||
|
||||
def test_copies_omni_content(self):
|
||||
"""Diffusion fields are copied when source is another OmniRequestOutput."""
|
||||
image = Image.new("RGB", (64, 64), color="blue")
|
||||
inner = OmniRequestOutput.from_diffusion(request_id="req-diff", images=[image], prompt="a cat")
|
||||
outer = OmniRequestOutput.from_stage_output(inner, request_id="req-diff", stage_id=0, final_output_type="image")
|
||||
|
||||
assert len(outer.images) == 1
|
||||
assert outer.images[0] is image
|
||||
assert outer.prompt == "a cat"
|
||||
assert outer.is_diffusion_output
|
||||
|
||||
def test_with_duck_typed_source(self):
|
||||
"""A duck-typed source (SimpleNamespace-style) copies correctly."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
source = SimpleNamespace(
|
||||
outputs=[],
|
||||
prompt="duck prompt",
|
||||
prompt_token_ids=[7, 8, 9],
|
||||
finished=False,
|
||||
)
|
||||
out = OmniRequestOutput.from_stage_output(source, request_id="duck-1", stage_id=2)
|
||||
|
||||
assert out.prompt == "duck prompt"
|
||||
assert out.prompt_token_ids == [7, 8, 9]
|
||||
assert out.finished is False
|
||||
assert out.stage_id == 2
|
||||
|
||||
def test_request_id_always_copied_from_source(self):
|
||||
"""request_id is always copied from source, even when passed in kwargs."""
|
||||
source = _make_text_request_output(request_id="source-id")
|
||||
out = OmniRequestOutput.from_stage_output(source, request_id="kwargs-id", stage_id=1)
|
||||
# Source wins — request_id is content, not a control field.
|
||||
assert out.request_id == "source-id"
|
||||
|
||||
def test_passes_control_fields(self):
|
||||
"""Control fields (metrics, stage_durations, peak_memory_mb) are set."""
|
||||
source = _make_text_request_output()
|
||||
out = OmniRequestOutput.from_stage_output(
|
||||
source,
|
||||
request_id="r1",
|
||||
stage_id=3,
|
||||
metrics={"ttft": 0.5},
|
||||
stage_durations={"prefill": 0.12},
|
||||
peak_memory_mb=1024.0,
|
||||
)
|
||||
|
||||
assert out.metrics == {"ttft": 0.5}
|
||||
assert out.stage_durations == {"prefill": 0.12}
|
||||
assert out.peak_memory_mb == 1024.0
|
||||
|
||||
def test_finished_field_copied(self):
|
||||
"""The *finished* flag is copied from the source when present."""
|
||||
source = _make_text_request_output(finished=False)
|
||||
out = OmniRequestOutput.from_stage_output(source, request_id="r1")
|
||||
assert out.finished is False
|
||||
|
||||
source2 = _make_text_request_output(finished=True)
|
||||
out2 = OmniRequestOutput.from_stage_output(source2, request_id="r2")
|
||||
assert out2.finished is True
|
||||
|
||||
def test_outputs_field_copied(self):
|
||||
"""The *outputs* list is copied from the source."""
|
||||
source = _make_text_request_output(text="hello", token_ids=[1, 2, 3])
|
||||
out = OmniRequestOutput.from_stage_output(source, request_id="r1")
|
||||
|
||||
assert len(out.outputs) == 1
|
||||
assert out.outputs[0].text == "hello"
|
||||
assert out.outputs[0].token_ids == [1, 2, 3]
|
||||
|
||||
|
||||
# ── dataclass safety ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestDataclassSafety:
|
||||
def test_from_stage_output_is_classmethod(self):
|
||||
"""from_stage_output must be a classmethod."""
|
||||
assert callable(OmniRequestOutput.from_stage_output)
|
||||
|
||||
|
||||
# ── msgpack round-trip ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestMsgpackRoundTrip:
|
||||
"""Tests that OmniRequestOutput survives msgpack encode → decode without
|
||||
infinite recursion or data loss."""
|
||||
|
||||
def test_round_trip_pipeline_output(self):
|
||||
"""A pipeline output round-trips without data loss."""
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import (
|
||||
OmniMsgpackDecoder,
|
||||
OmniMsgpackEncoder,
|
||||
)
|
||||
|
||||
source = _make_text_request_output(text="round trip")
|
||||
out = OmniRequestOutput.from_stage_output(
|
||||
source,
|
||||
request_id="rt-1",
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
)
|
||||
|
||||
encoder = OmniMsgpackEncoder()
|
||||
decoder = OmniMsgpackDecoder()
|
||||
|
||||
encoded = encoder.encode(out)
|
||||
decoded = decoder.decode(encoded)
|
||||
|
||||
assert isinstance(decoded, OmniRequestOutput)
|
||||
assert decoded.request_id == "req-1"
|
||||
assert decoded.stage_id == 0
|
||||
assert decoded.final_output_type == "text"
|
||||
assert decoded.finished is True
|
||||
assert decoded.outputs[0].text == "round trip"
|
||||
assert decoded.prompt == "test prompt"
|
||||
assert decoded.prompt_token_ids == [1, 2, 3]
|
||||
|
||||
def test_round_trip_diffusion_output(self):
|
||||
"""A diffusion output round-trips without data loss."""
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import (
|
||||
OmniMsgpackDecoder,
|
||||
OmniMsgpackEncoder,
|
||||
)
|
||||
|
||||
image = Image.new("RGB", (64, 64), color="red")
|
||||
out = OmniRequestOutput.from_diffusion(
|
||||
request_id="diff-1",
|
||||
images=[image],
|
||||
prompt="a red image",
|
||||
metrics={"steps": 50},
|
||||
)
|
||||
assert output.request_id == "test-123"
|
||||
assert output.images == images
|
||||
assert output.prompt == "a cat"
|
||||
assert output.metrics == {"steps": 50}
|
||||
assert output.is_diffusion_output
|
||||
assert output.num_images == 1
|
||||
|
||||
def test_from_pipeline(self, mocker: MockerFixture):
|
||||
"""Test creating output from pipeline stage."""
|
||||
mock_request_output = mocker.Mock()
|
||||
mock_request_output.request_id = "pipeline-123"
|
||||
mock_request_output.prompt_token_ids = [1, 2, 3]
|
||||
mock_request_output.outputs = [mocker.Mock()]
|
||||
mock_request_output.encoder_prompt_token_ids = None
|
||||
mock_request_output.prompt_logprobs = None
|
||||
mock_request_output.num_cached_tokens = 10
|
||||
mock_request_output.kv_transfer_params = None
|
||||
mock_request_output.multimodal_output = {"image": mocker.Mock()}
|
||||
|
||||
output = OmniRequestOutput.from_pipeline(
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
request_output=mock_request_output,
|
||||
multimodal_output={"audio": None},
|
||||
)
|
||||
|
||||
assert output.request_id == "pipeline-123"
|
||||
assert output.stage_id == 0
|
||||
assert output.final_output_type == "text"
|
||||
assert output.is_pipeline_output
|
||||
encoder = OmniMsgpackEncoder()
|
||||
decoder = OmniMsgpackDecoder()
|
||||
|
||||
def test_prompt_token_ids_property(self, mocker: MockerFixture):
|
||||
"""Test prompt_token_ids property for streaming compatibility."""
|
||||
mock_request_output = mocker.Mock()
|
||||
mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
|
||||
encoded = encoder.encode(out)
|
||||
decoded = decoder.decode(encoded)
|
||||
|
||||
output = OmniRequestOutput.from_pipeline(
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
request_output=mock_request_output,
|
||||
assert isinstance(decoded, OmniRequestOutput)
|
||||
assert decoded.request_id == "diff-1"
|
||||
assert decoded.prompt == "a red image"
|
||||
assert decoded.num_images == 1
|
||||
assert decoded.images[0].size == (64, 64)
|
||||
assert decoded.metrics == {"steps": 50}
|
||||
|
||||
def test_round_trip_does_not_recurse(self):
|
||||
"""Encoding a pipeline output must not trigger infinite recursion."""
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import (
|
||||
OmniMsgpackEncoder,
|
||||
)
|
||||
|
||||
assert output.prompt_token_ids == [1, 2, 3, 4, 5]
|
||||
source = _make_text_request_output(text="no recursion")
|
||||
out = OmniRequestOutput.from_stage_output(source, request_id="rec-1", stage_id=0)
|
||||
|
||||
def test_prompt_token_ids_none_when_no_request_output(self):
|
||||
"""Test prompt_token_ids returns None when no request_output."""
|
||||
output = OmniRequestOutput.from_diffusion(
|
||||
request_id="test-123",
|
||||
images=[],
|
||||
prompt="a cat",
|
||||
)
|
||||
assert output.prompt_token_ids is None
|
||||
encoder = OmniMsgpackEncoder()
|
||||
encoded = encoder.encode(out)
|
||||
assert isinstance(encoded, bytes)
|
||||
assert len(encoded) > 0
|
||||
|
||||
def test_outputs_property(self, mocker: MockerFixture):
|
||||
"""Test outputs property for chat completion compatibility."""
|
||||
mock_output = mocker.Mock()
|
||||
mock_request_output = mocker.Mock()
|
||||
mock_request_output.outputs = [mock_output]
|
||||
|
||||
output = OmniRequestOutput.from_pipeline(
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
request_output=mock_request_output,
|
||||
def test_legacy_wire_format_deserialized(self):
|
||||
"""When the wire format contains a nested ``request_output`` key
|
||||
(pre-inheritance legacy), it is merged via _copy_content_from."""
|
||||
from vllm_omni.distributed.omni_connectors.utils.serialization import (
|
||||
OmniMsgpackDecoder,
|
||||
)
|
||||
|
||||
assert output.outputs == [mock_output]
|
||||
# Build a legacy-style dict: OmniRequestOutput fields + a nested
|
||||
# request_output dict containing RequestOutput-like content.
|
||||
legacy_dict = {
|
||||
"request_id": "legacy-1",
|
||||
"finished": True,
|
||||
"stage_id": 1,
|
||||
"final_output_type": "text",
|
||||
"images": [],
|
||||
"_multimodal_output": {},
|
||||
"_custom_output": {},
|
||||
"stage_durations": {},
|
||||
"peak_memory_mb": 0.0,
|
||||
# --- nested request_output key (legacy) ---
|
||||
"request_output": {
|
||||
"request_id": "legacy-1",
|
||||
"prompt": "legacy prompt",
|
||||
"prompt_token_ids": [1, 2],
|
||||
"prompt_logprobs": None,
|
||||
"outputs": [],
|
||||
"finished": True,
|
||||
"metrics": {},
|
||||
"lora_request": None,
|
||||
"encoder_prompt": None,
|
||||
"encoder_prompt_token_ids": None,
|
||||
"num_cached_tokens": None,
|
||||
"kv_transfer_params": None,
|
||||
},
|
||||
}
|
||||
|
||||
def test_outputs_empty_when_no_request_output(self):
|
||||
"""Test outputs returns empty list when no request_output."""
|
||||
output = OmniRequestOutput.from_diffusion(
|
||||
request_id="test-123",
|
||||
images=[],
|
||||
prompt="a cat",
|
||||
)
|
||||
assert output.outputs == []
|
||||
# Manually run the decoder's post_process to test legacy path.
|
||||
# The decoder's _decode_omni_request_output pops "request_output"
|
||||
# and merges it via _copy_content_from.
|
||||
decoder = OmniMsgpackDecoder()
|
||||
result = decoder._post_process(legacy_dict)
|
||||
|
||||
def test_encoder_prompt_token_ids_property(self, mocker: MockerFixture):
|
||||
"""Test encoder_prompt_token_ids property."""
|
||||
mock_request_output = mocker.Mock()
|
||||
mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
|
||||
|
||||
output = OmniRequestOutput.from_pipeline(
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
request_output=mock_request_output,
|
||||
)
|
||||
|
||||
assert output.encoder_prompt_token_ids == [10, 20, 30]
|
||||
|
||||
def test_num_cached_tokens_property(self, mocker: MockerFixture):
|
||||
"""Test num_cached_tokens property."""
|
||||
mock_request_output = mocker.Mock()
|
||||
mock_request_output.num_cached_tokens = 42
|
||||
|
||||
output = OmniRequestOutput.from_pipeline(
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
request_output=mock_request_output,
|
||||
)
|
||||
|
||||
assert output.num_cached_tokens == 42
|
||||
|
||||
def test_multimodal_output_property(self, mocker: MockerFixture):
|
||||
"""Test multimodal_output property."""
|
||||
mock_request_output = mocker.Mock()
|
||||
mock_audio = mocker.Mock()
|
||||
expected_output = {"audio": mock_audio}
|
||||
mock_request_output.outputs = []
|
||||
mock_request_output.multimodal_output = expected_output
|
||||
|
||||
output = OmniRequestOutput.from_pipeline(
|
||||
stage_id=0,
|
||||
final_output_type="audio",
|
||||
request_output=mock_request_output,
|
||||
)
|
||||
|
||||
assert output.multimodal_output is expected_output
|
||||
|
||||
def test_to_dict_diffusion(self):
|
||||
"""Test to_dict for diffusion output."""
|
||||
output = OmniRequestOutput.from_diffusion(
|
||||
request_id="test-123",
|
||||
images=[Image.new("RGB", (64, 64), color="red")],
|
||||
prompt="a cat",
|
||||
metrics={"steps": 50},
|
||||
)
|
||||
result = output.to_dict()
|
||||
|
||||
assert result["request_id"] == "test-123"
|
||||
assert result["finished"] is True
|
||||
assert result["final_output_type"] == "image"
|
||||
assert result["num_images"] == 1
|
||||
assert result["prompt"] == "a cat"
|
||||
|
||||
def test_to_dict_pipeline(self, mocker: MockerFixture):
|
||||
"""Test to_dict for pipeline output."""
|
||||
mock_request_output = mocker.Mock()
|
||||
mock_request_output.request_id = "pipeline-123"
|
||||
|
||||
output = OmniRequestOutput.from_pipeline(
|
||||
stage_id=0,
|
||||
final_output_type="text",
|
||||
request_output=mock_request_output,
|
||||
)
|
||||
result = output.to_dict()
|
||||
|
||||
assert result["request_id"] == "pipeline-123"
|
||||
assert result["finished"] is True
|
||||
assert result["final_output_type"] == "text"
|
||||
assert result["stage_id"] == 0
|
||||
assert isinstance(result, OmniRequestOutput)
|
||||
assert result.request_id == "legacy-1"
|
||||
assert result.finished is True
|
||||
# Content merged from legacy nested key
|
||||
assert result.prompt == "legacy prompt"
|
||||
assert result.prompt_token_ids == [1, 2]
|
||||
Reference in new issue
Block a user