397 lines
16 KiB
Python
397 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.sampling_params import RequestOutputKind
|
|
from vllm.v1.engine import FinishReason
|
|
|
|
from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
|
|
from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
|
|
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
|
|
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
|
|
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _CodecPool:
|
|
final_output = True
|
|
|
|
def __init__(self, output_kind, registered, replica_id):
|
|
self.output_processor = MultimodalOutputProcessor(
|
|
tokenizer=None, log_stats=False, engine_core_output_type="audio"
|
|
)
|
|
self.output_kind = output_kind
|
|
self.replica_id = replica_id
|
|
self.bound_replica = None
|
|
if registered:
|
|
self.register()
|
|
|
|
def register(self):
|
|
self.bound_replica = self.replica_id
|
|
self.output_processor.request_states["r"] = OmniRequestState(
|
|
request_id="r",
|
|
external_req_id="r",
|
|
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=self.output_kind,
|
|
)
|
|
self.output_processor.external_req_ids["r"].append("r")
|
|
|
|
def get_bound_replica_id(self, request_id):
|
|
return self.bound_replica
|
|
|
|
async def process_llm_raw_outputs(self, replica_id, raw, **kwargs):
|
|
assert replica_id == self.replica_id
|
|
return self.output_processor.process_outputs(raw.outputs, raw.timestamp).request_outputs
|
|
|
|
|
|
def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, replica_id=2, final_stage_id=1):
|
|
obj = Orchestrator.__new__(Orchestrator)
|
|
obj.stage_pools = [SimpleNamespace(final_output=False), _CodecPool(output_kind, registered, replica_id)]
|
|
obj.request_states = {"r": OrchestratorRequestState(request_id="r", final_stage_id=final_stage_id)}
|
|
obj.output_async_queue = asyncio.Queue()
|
|
obj._cfg_tracker = CfgCompanionTracker()
|
|
obj._pd_kv_params = {}
|
|
obj._running_counter = None
|
|
obj._abort_request_ids = AsyncMock(return_value=[])
|
|
obj._release_request_bindings = Mock()
|
|
obj._finish_raw_terminal_requests = AsyncMock()
|
|
|
|
async def process(stage, replica, raw, terminals):
|
|
await obj._route_upstream_first_audio(stage, replica, raw)
|
|
return await obj.stage_pools[stage].process_llm_raw_outputs(replica, raw)
|
|
|
|
async def route(stage, replica, outputs):
|
|
for output in outputs:
|
|
await obj.output_async_queue.put(
|
|
OutputMessage(
|
|
request_id=output.request_id,
|
|
stage_id=stage,
|
|
replica_id=replica,
|
|
engine_outputs=output,
|
|
metrics=None,
|
|
finished=output.finished,
|
|
)
|
|
)
|
|
|
|
obj._process_llm_stage_outputs = process
|
|
obj._handle_processed_outputs = route
|
|
return obj
|
|
|
|
|
|
def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)):
|
|
return OmniEngineCoreOutputs(
|
|
outputs=[
|
|
OmniEngineCoreOutput(
|
|
request_id="r",
|
|
new_token_ids=[],
|
|
finish_reason=FinishReason.STOP if terminal else None,
|
|
multimodal_output={
|
|
"model_outputs": torch.tensor(samples, dtype=torch.float32),
|
|
"sr": torch.tensor(24000),
|
|
"_omni_first_audio": torch.tensor(first),
|
|
"_omni_first_audio_required": torch.tensor(required),
|
|
},
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
async def _codec_output(obj, raw, stage=1):
|
|
replica = obj.stage_pools[stage].replica_id
|
|
processed = await obj._process_llm_stage_outputs(stage, replica, raw, set())
|
|
await obj._handle_processed_outputs(stage, replica, processed)
|
|
|
|
|
|
def _messages(obj):
|
|
messages = []
|
|
while not obj.output_async_queue.empty():
|
|
messages.append(obj.output_async_queue.get_nowait())
|
|
return messages
|
|
|
|
|
|
def _audio(output):
|
|
audio = output.outputs[0].multimodal_output["audio"]
|
|
return torch.cat(audio) if isinstance(audio, list) else audio
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("kind", list(RequestOutputKind))
|
|
@pytest.mark.parametrize("suffix_overtook", [False, True])
|
|
async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook):
|
|
obj = _orchestrator(kind, final_stage_id=0)
|
|
obj.stage_pools = [obj.stage_pools[1]]
|
|
suffix = _raw(required=True, samples=(20, 21))
|
|
terminal = _raw(required=True, terminal=True, samples=(30, 31))
|
|
if suffix_overtook:
|
|
await _codec_output(obj, suffix, stage=0)
|
|
await _codec_output(obj, terminal, stage=0)
|
|
assert obj.output_async_queue.empty()
|
|
await obj._route_upstream_first_audio(0, 2, _raw(first=True, samples=(10, 11)))
|
|
if not suffix_overtook:
|
|
await _codec_output(obj, suffix, stage=0)
|
|
await _codec_output(obj, terminal, stage=0)
|
|
messages = _messages(obj)
|
|
expected = {
|
|
RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]],
|
|
RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]],
|
|
RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]],
|
|
}
|
|
assert [_audio(message.engine_outputs).tolist() for message in messages] == expected[kind]
|
|
assert all((message.stage_id, message.replica_id) == (0, 2) for message in messages)
|
|
assert messages[-1].finished
|
|
assert not obj.request_states["r"].pending_first_audio_outputs
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio():
|
|
obj = _orchestrator(final_stage_id=0)
|
|
obj.stage_pools = [obj.stage_pools[1]]
|
|
await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
|
|
assert obj.request_states["r"].pending_first_audio_outputs
|
|
await obj._cleanup_request_ids(["r"], abort=True)
|
|
await obj._route_upstream_first_audio(0, 2, _raw(first=True))
|
|
assert not obj.request_states and obj.output_async_queue.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal():
|
|
obj = _orchestrator(final_stage_id=0)
|
|
obj.stage_pools = [obj.stage_pools[1]]
|
|
await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
|
|
obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
|
|
await obj._route_upstream_first_audio(0, 2, _raw(first=True))
|
|
error = obj.output_async_queue.get_nowait()
|
|
assert isinstance(error, ErrorMessage) and error.stage_id == 0
|
|
assert not obj.request_states
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("kind", list(RequestOutputKind))
|
|
@pytest.mark.parametrize("codec_overtook", [False, True])
|
|
async def test_first_audio_obeys_codec_output_kind(kind, codec_overtook):
|
|
obj = _orchestrator(kind)
|
|
suffix = _raw(required=True, samples=(20, 21))
|
|
terminal = _raw(terminal=True, samples=(30, 31))
|
|
if codec_overtook:
|
|
await _codec_output(obj, suffix)
|
|
await _codec_output(obj, terminal)
|
|
assert obj.output_async_queue.empty()
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
|
|
if kind == RequestOutputKind.FINAL_ONLY and not codec_overtook:
|
|
assert obj.output_async_queue.empty()
|
|
if not codec_overtook:
|
|
await _codec_output(obj, suffix)
|
|
await _codec_output(obj, terminal)
|
|
messages = _messages(obj)
|
|
# The source replica is zero; output ownership belongs to codec replica two.
|
|
assert all((message.stage_id, message.replica_id) == (1, 2) for message in messages)
|
|
actual = [_audio(message.engine_outputs).tolist() for message in messages]
|
|
expected = {
|
|
RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]],
|
|
RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]],
|
|
RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]],
|
|
}
|
|
assert actual == expected[kind]
|
|
assert messages[-1].finished
|
|
assert obj.request_states["r"].pending_first_audio_outputs == []
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True))
|
|
assert obj.output_async_queue.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_first_audio_waits_for_codec_registration_and_ignores_duplicates():
|
|
obj = _orchestrator(registered=False)
|
|
req_state = obj.request_states["r"]
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(90, 91)))
|
|
assert obj.output_async_queue.empty()
|
|
assert not req_state.upstream_first_audio
|
|
assert req_state.pending_upstream_first_audio is not None
|
|
obj.stage_pools[1].register()
|
|
await obj._flush_upstream_first_audio(req_state)
|
|
messages = _messages(obj)
|
|
assert len(messages) == 1
|
|
assert _audio(messages[0].engine_outputs).tolist() == [10, 11]
|
|
assert req_state.upstream_first_audio
|
|
assert req_state.pending_upstream_first_audio is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("processing_failed", [False, True])
|
|
async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, processing_failed):
|
|
obj = _orchestrator(registered=False, final_stage_id=2 if processing_failed else 1)
|
|
req_state = obj.request_states["r"]
|
|
req_state.sampling_params_list = [None] * (req_state.final_stage_id + 1)
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
|
|
pool = obj.stage_pools[1]
|
|
pool.stage_type = "llm"
|
|
pool.stage_vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None))
|
|
obj.stage_pools[0].get_bound_replica_id = lambda request_id: 0
|
|
obj._stage_receives_async_chunks = lambda stage: True
|
|
obj._build_payload_sender_info = Mock(return_value={})
|
|
obj._on_stage_submitted = Mock()
|
|
obj._emit_tx_edge = Mock()
|
|
mocker.patch(
|
|
"vllm_omni.engine.orchestrator.build_engine_core_request_from_tokens",
|
|
return_value=SimpleNamespace(request_id="r"),
|
|
)
|
|
|
|
async def submit(*args, **kwargs):
|
|
pool.register()
|
|
|
|
pool.submit_initial = submit
|
|
if processing_failed:
|
|
pool.process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
|
|
later_pool = SimpleNamespace(submit_initial=AsyncMock())
|
|
obj.stage_pools.append(later_pool)
|
|
submitted = await obj._prewarm_async_chunk_stages("r", SimpleNamespace(prompt_token_ids=[0]), req_state)
|
|
assert submitted is not processing_failed
|
|
messages = _messages(obj)
|
|
assert len(messages) == 1
|
|
if processing_failed:
|
|
assert isinstance(messages[0], ErrorMessage)
|
|
later_pool.submit_initial.assert_not_called()
|
|
assert "r" not in obj.request_states
|
|
else:
|
|
assert _audio(messages[0].engine_outputs).tolist() == [10, 11]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pending_first_audio_is_discarded_on_cancellation():
|
|
obj = _orchestrator(registered=False)
|
|
req_state = obj.request_states["r"]
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True))
|
|
await obj._cleanup_request_ids(["r"], abort=True)
|
|
obj.stage_pools[1].register()
|
|
await obj._flush_upstream_first_audio(req_state)
|
|
assert obj.output_async_queue.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancellation_while_routing_first_audio_does_not_release_codec_outputs():
|
|
obj = _orchestrator()
|
|
await obj._route_upstream_first_audio(1, 2, _raw(required=True, terminal=True))
|
|
|
|
async def cancel(stage, replica, outputs):
|
|
await obj._cleanup_request_ids(["r"], abort=True)
|
|
|
|
obj._handle_processed_outputs = cancel
|
|
obj._process_llm_stage_outputs = AsyncMock()
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True))
|
|
obj._process_llm_stage_outputs.assert_not_called()
|
|
assert obj.output_async_queue.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_ending_at_source_does_not_emit_first_audio():
|
|
obj = _orchestrator(final_stage_id=0)
|
|
raw = _raw(first=True)
|
|
await obj._route_upstream_first_audio(0, 0, raw)
|
|
assert not raw.outputs and obj.output_async_queue.empty()
|
|
assert obj.request_states["r"].pending_upstream_first_audio is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_later_aligner_receives_full_codec_waveform():
|
|
from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner
|
|
|
|
obj = _orchestrator(RequestOutputKind.CUMULATIVE, final_stage_id=2)
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
|
|
await _codec_output(obj, _raw(required=True, terminal=True, samples=(20, 21)))
|
|
messages = _messages(obj)
|
|
assert all(message.stage_id == 1 for message in messages)
|
|
inputs = code2wav2aligner(
|
|
[messages[-1].engine_outputs], prompt={"additional_information": {"text": ["Hello world"]}}
|
|
)
|
|
assert inputs[0]["multi_modal_data"]["audio"][0].tolist() == [10, 11, 20, 21]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_first_audio_processing_failure_fails_only_its_request():
|
|
obj = _orchestrator()
|
|
obj.request_states["healthy"] = OrchestratorRequestState(request_id="healthy", final_stage_id=1)
|
|
obj.stage_pools[1].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True))
|
|
error = obj.output_async_queue.get_nowait()
|
|
assert isinstance(error, ErrorMessage)
|
|
assert error.request_id == "r" and error.stage_id == 1
|
|
assert "processing failed" in error.error
|
|
assert set(obj.request_states) == {"healthy"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unmarked_upstream_audio_and_regular_codec_outputs_are_untouched():
|
|
obj = _orchestrator()
|
|
for stage in (0, 1):
|
|
raw = _raw()
|
|
await obj._route_upstream_first_audio(stage, 0, raw)
|
|
assert len(raw.outputs) == 1
|
|
assert torch.equal(raw.outputs[0].multimodal_output["model_outputs"], torch.ones(4))
|
|
assert obj.output_async_queue.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_request_does_not_receive_late_first_audio():
|
|
obj = _orchestrator()
|
|
obj.request_states.clear()
|
|
raw = _raw(first=True)
|
|
await obj._route_upstream_first_audio(0, 0, raw)
|
|
assert not raw.outputs and obj.output_async_queue.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_first_audio_supports_both_speech_and_chat_consumers():
|
|
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
|
|
|
obj = _orchestrator()
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
|
|
result = obj.output_async_queue.get_nowait().engine_outputs
|
|
assert not result.finished
|
|
assert _audio(result).tolist() == [10, 11]
|
|
payload, key = OmniOpenAIServingSpeech._extract_audio_output(result)
|
|
assert payload is not None and key is not None
|
|
audio = payload[key]
|
|
audio = torch.cat(audio) if isinstance(audio, list) else audio
|
|
assert audio.tolist() == [10, 11]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("audio", [None, [], torch.empty(0)])
|
|
@pytest.mark.parametrize("codec_overtook", [False, True])
|
|
async def test_invalid_first_audio_reports_error_and_releases_request(audio, codec_overtook):
|
|
obj = _orchestrator()
|
|
obj.request_states["healthy"] = OrchestratorRequestState(request_id="healthy", final_stage_id=1)
|
|
if codec_overtook:
|
|
await obj._route_upstream_first_audio(1, 2, _raw(required=True, terminal=True))
|
|
assert obj.request_states["r"].pending_first_audio_outputs
|
|
first = _raw(first=True)
|
|
first.outputs[0].multimodal_output["model_outputs"] = audio
|
|
await obj._route_upstream_first_audio(0, 0, first)
|
|
|
|
error = obj.output_async_queue.get_nowait()
|
|
assert isinstance(error, ErrorMessage)
|
|
assert error.request_id == "r" and "first-audio" in error.error
|
|
assert not first.outputs and set(obj.request_states) == {"healthy"}
|
|
obj._abort_request_ids.assert_awaited_once_with(["r"])
|
|
obj._release_request_bindings.assert_called_once_with(["r"])
|
|
await obj._route_upstream_first_audio(0, 0, _raw(first=True))
|
|
assert obj.output_async_queue.empty()
|