vllm-omni/tests/worker_v2/test_first_audio_sender.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import queue
from types import SimpleNamespace

import pytest
import torch

from vllm_omni.engine import OmniEngineCoreOutputs
from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink

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


def _scheduler(**client_by_request):
    return SimpleNamespace(
        requests={request_id: SimpleNamespace(client_index=client) for request_id, client in client_by_request.items()}
    )


def test_sink_groups_outputs_by_client_and_carries_audio():
    output_queue: queue.Queue = queue.Queue()
    sink = engine_output_queue_sink(output_queue, _scheduler(a=0, b=1, c=0))
    sample_rate = torch.tensor(24000, dtype=torch.int32)
    rows = [torch.full((4,), float(i)) for i in range(3)]

    sink.prepare(["a", "b", "c"])(["a", "b", "c"], rows, sample_rate)

    delivered = {}
    while not output_queue.empty():
        client_index, outputs = output_queue.get_nowait()
        assert isinstance(outputs, OmniEngineCoreOutputs)
        delivered[client_index] = outputs.outputs
    assert sorted(delivered) == [0, 1]
    assert [o.request_id for o in delivered[0]] == ["a", "c"]
    assert [o.request_id for o in delivered[1]] == ["b"]
    for output in delivered[0] + delivered[1]:
        # First audio carries no tokens and does not finish the request.
        assert output.new_token_ids == []
        assert output.finish_reason is None
        assert output.multimodal_output["sr"] is sample_rate
    assert torch.equal(delivered[0][1].multimodal_output["model_outputs"], rows[2])


def test_sink_skips_requests_that_already_left_the_scheduler():
    output_queue: queue.Queue = queue.Queue()
    sink = engine_output_queue_sink(output_queue, _scheduler(kept=0))

    sink.prepare(["gone", "kept"])(["gone", "kept"], [torch.zeros(2), torch.ones(2)], torch.tensor(24000))
    sink.prepare(["gone"])(["gone"], [torch.zeros(2)], torch.tensor(24000))

    client_index, outputs = output_queue.get_nowait()
    assert client_index == 0
    assert [o.request_id for o in outputs.outputs] == ["kept"]
    assert output_queue.empty()


class _DoneEvent:
    def record(self):
        pass

    def synchronize(self):
        pass


def _run_sender(items):
    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender

    output_queue: queue.Queue = queue.Queue()
    sink = engine_output_queue_sink(output_queue, _scheduler(**{rid: 0 for item in items for rid in item[2]}))
    sender = FirstAudioSender(sink)
    for item in items:
        sender._queue.put((*item, sink.prepare(item[2])))
    sender.close()
    delivered = []
    while not output_queue.empty():
        _, outputs = output_queue.get_nowait()
        delivered.append(
            ([o.request_id for o in outputs.outputs], [o.multimodal_output["model_outputs"] for o in outputs.outputs])
        )
    return delivered


def test_sender_drops_rows_whose_frame_is_not_audio():
    pcm = torch.arange(6, dtype=torch.float32).reshape(3, 2)
    sample_rate = torch.tensor(24000, dtype=torch.int32)
    valid = torch.tensor([True, False, True])

    delivered = _run_sender(
        [
            (_DoneEvent(), pcm, ["a", "b", "c"], sample_rate, valid),
            # Every row invalid: nothing reaches the sink.
            (_DoneEvent(), pcm[:1], ["d"], sample_rate, torch.tensor([False])),
            # No validity: every row is delivered (previous behaviour).
            (_DoneEvent(), pcm[:1], ["e"], sample_rate, None),
        ]
    )

    assert [ids for ids, _rows in delivered] == [["a", "c"], ["e"]]
    assert torch.equal(delivered[0][1][1], pcm[2])


def test_prepared_delivery_survives_scheduler_completion():
    output_queue: queue.Queue = queue.Queue()
    scheduler = _scheduler(r=3)
    sink = engine_output_queue_sink(output_queue, scheduler)
    deliver = sink.prepare(["r"])
    scheduler.requests.clear()
    deliver(["r"], [torch.ones(2)], torch.tensor(24000))
    client, outputs = output_queue.get_nowait()
    assert client == 3
    assert outputs.outputs[0].multimodal_output["_omni_first_audio"]


def test_delivery_failure_emits_request_error():
    from vllm.v1.engine import FinishReason

    output_queue: queue.Queue = queue.Queue()
    delivery = engine_output_queue_sink(output_queue, _scheduler(r=0)).prepare(["r"])
    delivery.fail(["r"])
    _, outputs = output_queue.get_nowait()
    assert outputs.outputs[0].finish_reason == FinishReason.ERROR


@pytest.mark.parametrize("client_index", [0, 3])
def test_submit_reports_only_requests_with_a_prepared_route(monkeypatch, client_index):
    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender

    # Exercise the routing contract without allocating CUDA or pinned memory.
    original_empty = torch.empty
    monkeypatch.setattr(torch, "empty", lambda *args, pin_memory=False, **kwargs: original_empty(*args, **kwargs))
    monkeypatch.setattr(torch.cuda, "Event", _DoneEvent)
    output_queue: queue.Queue = queue.Queue()
    sender = FirstAudioSender(engine_output_queue_sink(output_queue, _scheduler(kept=client_index)))
    pcm = torch.arange(6, dtype=torch.float32).reshape(3, 2)
    try:
        accepted = sender.submit(["gone", "kept", "also-gone"], pcm, torch.tensor(24000))
    finally:
        sender.close()

    assert accepted == ["kept"]
    client, outputs = output_queue.get_nowait()
    assert client == client_index
    assert [output.request_id for output in outputs.outputs] == ["kept"]
    assert torch.equal(outputs.outputs[0].multimodal_output["model_outputs"], pcm[1])
    assert output_queue.empty()


def test_submit_without_any_route_leaves_audio_to_the_codec(monkeypatch):
    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender

    def unexpected_copy(*args, **kwargs):
        pytest.fail("No delivery is possible; do not start a device copy")

    monkeypatch.setattr(torch, "empty", unexpected_copy)
    monkeypatch.setattr(torch.cuda, "Event", unexpected_copy)
    output_queue: queue.Queue = queue.Queue()
    sender = FirstAudioSender(engine_output_queue_sink(output_queue, _scheduler()))
    try:
        accepted = sender.submit(["gone"], torch.ones(1, 2), torch.tensor(24000))
    finally:
        sender.close()

    assert accepted == []
    assert output_queue.empty()


@pytest.mark.parametrize("sink", [lambda ids, rows, sr: None, SimpleNamespace(prepare=None)])
def test_sender_rejects_sink_without_route_preparation(sink):
    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender

    with pytest.raises(TypeError, match="prepare"):
        FirstAudioSender(sink)


def test_sender_failure_only_aborts_valid_undelivered_rows():
    from vllm.v1.engine import FinishReason

    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender

    class FailSecondClientOnce(queue.Queue):
        def put_nowait(self, item):
            client, outputs = item
            if client == 1 and outputs.outputs[0].finish_reason is None:
                raise RuntimeError("client enqueue failed")
            super().put_nowait(item)

    output_queue = FailSecondClientOnce()
    sink = engine_output_queue_sink(output_queue, _scheduler(eos=1, sent=0, failed=1, pending=2, next=3))
    sender = FirstAudioSender(sink)
    ids = ["eos", "sent", "failed", "pending", "gone"]
    sender._queue.put(
        (
            _DoneEvent(),
            torch.ones(5, 2),
            ids,
            torch.tensor(24000),
            torch.tensor([False, True, True, True, True]),
            sink.prepare(ids),
        )
    )
    sender._queue.put((_DoneEvent(), torch.ones(1, 2), ["next"], torch.tensor(24000), None, sink.prepare(["next"])))
    sender.close()

    actual: list[tuple[int, str, FinishReason | None]] = []
    while not output_queue.empty():
        client, outputs = output_queue.get_nowait()
        actual.extend((client, o.request_id, o.finish_reason) for o in outputs.outputs)
    assert actual == [
        (0, "sent", None),
        (1, "failed", FinishReason.ERROR),
        (2, "pending", FinishReason.ERROR),
        (3, "next", None),
    ]


def test_sender_copy_failure_reports_all_prepared_routes():
    from vllm.v1.engine import FinishReason

    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender

    class FailedEvent:
        def synchronize(self):
            raise RuntimeError("copy failed before validity became readable")

    output_queue: queue.Queue = queue.Queue()
    sink = engine_output_queue_sink(output_queue, _scheduler(r=3))
    sender = FirstAudioSender(sink)
    sender._queue.put(
        (FailedEvent(), torch.ones(2, 2), ["r", "gone"], torch.tensor(24000), None, sink.prepare(["r", "gone"]))
    )
    sender.close()
    client, outputs = output_queue.get_nowait()
    assert client == 3
    assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)]
    assert output_queue.empty()


def test_final_codec_delivery_omits_upstream_marker_and_freezes_route():
    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY

    outputs: queue.Queue = queue.Queue()
    scheduler = _scheduler(r=3)
    delivery = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False).prepare(["r"])
    scheduler.requests.clear()
    pcm = torch.ones(4)
    delivery(["r"], [pcm], torch.tensor(24000))
    client, batch = outputs.get_nowait()
    assert client == 3
    assert FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
    assert batch.outputs[0].multimodal_output["model_outputs"] is pcm
    assert batch.outputs[0].finish_reason is None