vllm-omni/tests/engine/test_output_metadata_snapshots.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Snapshot metadata must not accumulate as if it were generated content."""

import pytest
import torch
from vllm.outputs import PoolingRequestOutput
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine import FinishReason

from vllm_omni.outputs.mm_outputs import MultimodalPayload
from vllm_omni.outputs.output_modality import OutputModality
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
RATE_KEYS = ("sr", "sample_rate", "audio_sample_rate")


def _state(kind: RequestOutputKind) -> OmniRequestState:
    return OmniRequestState(
        request_id="audio",
        external_req_id="audio",
        parent_req=None,
        request_index=0,
        lora_request=None,
        prompt=None,
        prompt_token_ids=[0],
        prompt_embeds=None,
        logprobs_processor=None,
        detokenizer=None,
        max_tokens_param=None,
        arrival_time=0.0,
        queue=None,
        log_stats=False,
        stream_interval=1,
        output_kind=kind,
    )


@pytest.mark.parametrize("key", RATE_KEYS)
@pytest.mark.parametrize("representation", ["integer", "scalar_tensor", "vector_tensor", "metadata_tensor"])
def test_rate_is_latest_snapshot_before_consolidation(key, representation):
    payload = MultimodalPayload()
    for index, rate in enumerate((16000, 24000, 48000)):
        value = (
            rate if representation == "integer" else torch.tensor([rate] if representation == "vector_tensor" else rate)
        )
        incoming = (
            MultimodalPayload(tensors={"audio": torch.full((2,), index)}, metadata={key: value})
            if representation == "metadata_tensor"
            else MultimodalPayload.from_dict({"audio": torch.full((2,), index), key: value})
        )
        assert incoming is not None
        payload = payload.merged_with(incoming)
        snapshot = payload[key]
        assert not isinstance(snapshot, list)
        assert int(snapshot) == rate
        assert payload.get(key) is payload.to_dict()[key]
    payload = payload.merged_with(MultimodalPayload(tensors={"audio": torch.full((2,), 3)}))
    assert int(payload[key]) == 48000  # Missing metadata means retain, not clear.
    payload.consolidate_tensors(OutputModality.AUDIO)
    payload.consolidate_metadata()
    torch.testing.assert_close(payload["audio"], torch.arange(4).repeat_interleave(2))


@pytest.mark.parametrize("key", RATE_KEYS)
@pytest.mark.parametrize("tensor_first", [False, True])
def test_rate_representation_change_removes_stale_partition(key, tensor_first):
    first = torch.tensor(16000) if tensor_first else 16000
    last = 24000 if tensor_first else torch.tensor(24000)
    payload = MultimodalPayload.from_dict({key: first})
    incoming = MultimodalPayload.from_dict({key: last})
    assert payload is not None and incoming is not None
    incoming_keys = (set(incoming.tensors), set(incoming.metadata))
    payload = payload.merged_with(incoming)
    assert int(payload[key]) == 24000
    assert payload[key] is payload.to_dict()[key]
    assert (key in payload.tensors) != (key in payload.metadata)
    assert incoming[key] is last
    assert incoming_keys == (set(incoming.tensors), set(incoming.metadata))


@pytest.mark.parametrize("key", RATE_KEYS)
def test_self_merge_retains_snapshot_and_content(key):
    payload = MultimodalPayload.from_dict({key: torch.tensor(24000), "audio": torch.tensor([1.0, 2.0])})
    assert payload is not None
    payload = payload.merged_with(payload)
    assert int(payload[key]) == 24000
    payload.consolidate_tensors(OutputModality.AUDIO)
    torch.testing.assert_close(payload["audio"], torch.tensor([1.0, 2.0, 1.0, 2.0]))


@pytest.mark.parametrize("key", RATE_KEYS)
@pytest.mark.parametrize("kind", [RequestOutputKind.DELTA, RequestOutputKind.CUMULATIVE])
def test_long_stream_keeps_rate_bounded_without_losing_audio(key, kind):
    state = _state(kind)
    for index in range(1000):
        state.add_multimodal_tensor(
            {"model_outputs": torch.full((3,), index / 1000), key: torch.tensor(24000)}, "audio"
        )
        output = state.make_request_output([], None, None, None)
        assert output is not None and not isinstance(output, PoolingRequestOutput)
        payload = output.outputs[0].multimodal_output
        assert isinstance(payload[key], torch.Tensor) and payload[key].numel() == 1
        audio = payload["audio"]
        assert audio.numel() == (3 if kind == RequestOutputKind.DELTA else 3 * (index + 1))
        torch.testing.assert_close(audio[-3:], torch.full((3,), index / 1000))
        assert ("audio" in state.mm_accumulated) is (kind == RequestOutputKind.CUMULATIVE)
    final = state.make_request_output([], None, FinishReason.STOP, None)
    assert final is not None and final.finished
    if kind == RequestOutputKind.DELTA:
        assert "audio" not in final.outputs[0].multimodal_output


@pytest.mark.parametrize("pending", [0, 3])
def test_abort_flushes_only_pending_content_with_latest_rate(pending):
    state = _state(RequestOutputKind.DELTA)
    processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False, output_modality=OutputModality.AUDIO)
    processor.request_states[state.request_id] = state
    processor.external_req_ids[state.external_req_id] = [state.request_id]
    for rate in (16000, 24000):
        state.add_multimodal_tensor({"model_outputs": torch.ones(3), "sr": torch.tensor(rate)}, "audio")
        output = state.make_request_output([], None, None, None)
        assert output is not None
    if pending:
        state.add_multimodal_tensor({"model_outputs": torch.ones(pending)}, "audio")
    aborted, outputs = processor.abort_requests_collecting_outputs([state.request_id], internal=True)
    assert aborted == [state.request_id] and len(outputs) == 1 and outputs[0].finished
    payload = outputs[0].outputs[0].multimodal_output
    assert int(payload["sr"]) == 24000
    assert (payload["audio"].numel() if "audio" in payload else 0) == pending
    assert state.request_id not in processor.request_states


def test_non_snapshot_tensors_still_accumulate_and_empty_merge_keeps_identity():
    first = MultimodalPayload.from_dict({"latent": torch.tensor([1.0]), "audio": torch.tensor([2.0])})
    second = MultimodalPayload.from_dict({"latent": torch.tensor([3.0]), "audio": torch.tensor([4.0])})
    assert first is not None and second is not None
    assert MultimodalPayload().merged_with(first) is first
    merged = first.merged_with(second)
    merged.consolidate_tensors(OutputModality.AUDIO)
    torch.testing.assert_close(merged["latent"], torch.tensor([1.0, 3.0]))
    torch.testing.assert_close(merged["audio"], torch.tensor([2.0, 4.0]))