Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/engine/test_cfg_companion_lifecycle.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Deterministic tests for the CFG-companion output lifecycle.
 
 Regression suite for the companion-output race behind the nightly Bagel
@@ -94,6 +94,9 @@ class _FakePool:
     def release_bindings(self, request_ids):
         pass
 
+    async def release_request_resources(self, request_ids):
+        pass
+
 
 def _make_orchestrator(num_stages: int = 2) -> Orchestrator:
     orch = Orchestrator.__new__(Orchestrator)

tests/engine/test_orchestrator_first_audio.py

@@ -114,10 +114,10 @@ def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)):
     )
 
 
-async def _codec_output(obj, raw):
-    replica = obj.stage_pools[1].replica_id
-    processed = await obj._process_llm_stage_outputs(1, replica, raw, set())
-    await obj._handle_processed_outputs(1, replica, processed)
+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):
@@ -132,6 +132,57 @@ def _audio(output):
     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])

tests/engine/test_stage_engine_core_proc.py

@@ -1,9 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import patch
 
+import pytest
 from vllm.v1.engine.core import EngineCoreProc
+from vllm.v1.executor.uniproc_executor import UniProcExecutor
 
-from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc
+from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc, _bind_first_audio_sink
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "first_decoder,stream_decoder,stream_first_audio,expected",
+    [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)],
+)
+def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
+    mocker,
+    first_decoder,
+    stream_decoder,
+    stream_first_audio,
+    expected,
+):
+    executor = UniProcExecutor.__new__(UniProcExecutor)
+    executor.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)
+    )
+    executor.driver_worker = mocker.Mock()
+    runner = executor.driver_worker.worker.model_runner
+    runner.model.first_frame_decoder = object() if first_decoder else None
+    runner.model.stream_decoder = object() if stream_decoder else None
+    runner.model.stream_first_audio = stream_first_audio
+    runner._omni_data_plane = None
+    assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
+    assert runner.model_state.set_first_audio_sink.call_count == int(expected)
 
 
 def test_preprocess_add_request_preserves_omni_fields():
@@ -26,3 +59,64 @@ def test_preprocess_add_request_preserves_omni_fields():
     assert current_wave == 3
     assert result.external_req_id == "external"
     assert result.additional_information == {"conditioning": "payload"}
+
+
+def test_codec_uses_generic_first_audio_binding(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+
+    def hook(*args):
+        return False
+
+    def bind(sink):
+        calls["sink"] = sink
+        return hook
+
+    plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value))
+    model = SimpleNamespace(bind_first_chunk_fast_path=bind)
+    runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane)
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    assert calls["hook"] is hook
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    client, batch = outputs.get_nowait()
+    assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
+
+
+def test_first_audio_binding_preserves_talker_marker(monkeypatch):
+    import queue
+
+    import torch
+
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+    from vllm_omni.engine import stage_engine_core_proc as module
+
+    calls: dict[str, Any] = {}
+    runner = SimpleNamespace(
+        model=SimpleNamespace(first_frame_decoder=object()),
+        model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)),
+    )
+    executor = SimpleNamespace(
+        vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
+        driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
+    )
+    monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
+    outputs: queue.Queue = queue.Queue()
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)})
+    assert module._bind_first_audio_sink(executor, outputs, scheduler)
+    calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
+    assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY]
+    executor.vllm_config.parallel_config.tensor_parallel_size = 2
+    assert not module._bind_first_audio_sink(executor, outputs, scheduler)

tests/engine/test_stage_pool_collective_rpc.py

@@ -175,6 +175,46 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_release_request_resources_skips_without_async_chunk():
+    async def run() -> None:
+        call = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=call)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)),
+        )
+
+        await pool.release_request_resources(["req-1"])
+        call.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch):
+    async def run() -> None:
+        async def hang(*_args):
+            await asyncio.Event().wait()
+
+        live = AsyncMock()
+        pool = StagePool(
+            0,
+            [SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)],  # type: ignore[list-item]
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)),
+        )
+        monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2)
+
+        release = asyncio.create_task(pool.release_request_resources(["req-1"]))
+        await asyncio.sleep(0.05)
+        live.assert_awaited_once_with("omni_release_request_resources", ["req-1"])
+        assert not release.done()
+
+        await asyncio.wait_for(release, timeout=1.0)
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
 @pytest.mark.parametrize("failure", ["error", "timeout", "missing"])