vllm-omni/tests/engine/test_omni_request_output.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for OmniRequestOutput class."""

import pytest
from PIL import Image
from vllm.outputs import CompletionOutput, RequestOutput

from vllm_omni.outputs import OmniRequestOutput

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


# ── helpers ────────────────────────────────────────────────────────────────


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},
            multimodal_output={"audio": None},
        )

        encoder = OmniMsgpackEncoder()
        decoder = OmniMsgpackDecoder()

        encoded = encoder.encode(out)
        decoded = decoder.decode(encoded)

        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,
        )

        source = _make_text_request_output(text="no recursion")
        out = OmniRequestOutput.from_stage_output(source, request_id="rec-1", stage_id=0)

        encoder = OmniMsgpackEncoder()
        encoded = encoder.encode(out)
        assert isinstance(encoded, bytes)
        assert len(encoded) > 0

    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,
        )

        # 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,
            },
        }

        # 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)

        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]