[Bugfix][Qwen3-TTS] Handle codec generations that exhaust their token budget (#6728)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/entrypoints/openai_api/test_serving_speech.py

@@ -45,9 +45,15 @@ from vllm_omni.entrypoints.openai.serving_speech import (
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.base import DEFAULT_TTS_LANGUAGES, PreparedRequest, SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    DEFAULT_TTS_LANGUAGES,
+    PreparedRequest,
+    SpeechServingContext,
+    TTSGenerationError,
+)
 from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers
 from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError
 from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
@@ -2928,6 +2934,40 @@ class TestStreamingResponse:
         assert "text/event-stream" not in response.headers["content-type"]
         assert len(response.content) > 0
 
+    def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monkeypatch):
+        """The raw route must not bypass the shared Qwen3-TTS EOS guard."""
+        speech_server = streaming_app.state.speech_server
+        finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]}
+        captured: dict = {}
+
+        async def prepare(_request, request_id=None):
+            return request_id, object(), finalized_tts_params
+
+        async def generate_chunks(
+            _generator,
+            _request_id,
+            _response_format="pcm",
+            raw_request=None,
+            request_start_s=None,
+            include_sample_rate=False,
+            usage_acc=None,
+            tts_params=None,
+            collect=None,
+        ):
+            captured["tts_params"] = tts_params
+            yield b"\x00\x00"
+
+        monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
+        monkeypatch.setattr(speech_server, "_generate_audio_chunks", generate_chunks)
+
+        response = TestClient(streaming_app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert captured["tts_params"] is finalized_tts_params
+
     def test_sse_streaming(self, streaming_app):
         """stream_format=sse without stream=True returns audio deltas as SSE."""
         client = TestClient(streaming_app)
@@ -4848,6 +4888,57 @@ class TestTTSAsyncOffloading:
         assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
         assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
 
+    @pytest.mark.asyncio
+    async def test_qwen3_tts_nonstream_retries_codec_limit_once(self, qwen3_tts_server, mocker: MockerFixture) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(
+            side_effect=[
+                Qwen3TTSCodecLimitError("codec limit"),
+                (b"RIFF" + b"\x00" * 32, "audio/wav"),
+            ]
+        )
+        request = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
+
+        response = await qwen3_tts_server.create_speech(request)
+
+        assert response.status_code == 200
+        assert qwen3_tts_server._generate_audio_bytes.await_count == 2
+        first_call, retry_call = qwen3_tts_server._generate_audio_bytes.await_args_list
+        assert first_call.args[0] is request
+        assert first_call.kwargs["request_id"] != retry_call.kwargs["request_id"]
+        assert retry_call.args[0].seed is not None
+        assert request.seed is None
+
+    @pytest.mark.asyncio
+    async def test_nonstream_does_not_retry_nonretryable_generation_error(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=TTSGenerationError("invalid generation"))
+
+        response = await qwen3_tts_server.create_speech(OpenAICreateSpeechRequest(input="Hello", task_type="Base"))
+
+        assert response.status_code == 500
+        qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("request_kwargs", [{"seed": 100423}, {"max_new_tokens": 192}])
+    async def test_qwen3_tts_nonstream_does_not_retry_explicit_sampling_controls(
+        self,
+        qwen3_tts_server,
+        mocker: MockerFixture,
+        request_kwargs: dict[str, int],
+    ) -> None:
+        qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+        qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=Qwen3TTSCodecLimitError("codec limit"))
+
+        response = await qwen3_tts_server.create_speech(
+            OpenAICreateSpeechRequest(input="Hello", task_type="Base", **request_kwargs)
+        )
+
+        assert response.status_code == 500
+        qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
         async def failing_generator():
@@ -4862,6 +4953,80 @@ class TestTTSAsyncOffloading:
         assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
         assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
 
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_rejects_qwen3_codec_limit_before_success(self, qwen3_tts_server):
+        async def length_limited_generator():
+            yield SimpleNamespace(
+                multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+                metrics={
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 191,
+                            "finish_reason": "length",
+                        }
+                    }
+                },
+            )
+
+        chunks = qwen3_tts_server._generate_audio_chunks(
+            length_limited_generator(),
+            "req-codec-limit",
+            tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+        )
+        with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
+            async for _ in chunks:
+                pass
+
+    @pytest.mark.asyncio
+    async def test_sse_codec_limit_marks_partial_audio_for_discard(self, qwen3_tts_server):
+        async def length_limited_generator():
+            yield SimpleNamespace(
+                multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+                metrics={
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 191,
+                            "finish_reason": "length",
+                        }
+                    }
+                },
+            )
+
+        events = [
+            event
+            async for event in qwen3_tts_server._generate_audio_sse_events(
+                length_limited_generator(),
+                "req-sse-codec-limit",
+                tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+            )
+        ]
+
+        assert any("event: speech.audio.delta" in event for event in events)
+        assert not any("event: speech.audio.done" in event for event in events)
+        error_event = next(event for event in events if "event: speech.audio.error" in event)
+        payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+        assert payload["error"]["partial_audio"] is True
+        assert payload["error"]["action"] == "discard"
+
+    @pytest.mark.asyncio
+    async def test_sse_error_before_audio_is_not_marked_partial(self, qwen3_tts_server):
+        async def failing_generator():
+            raise RuntimeError("boom")
+            yield  # pragma: no cover
+
+        events = [
+            event
+            async for event in qwen3_tts_server._generate_audio_sse_events(
+                failing_generator(),
+                "req-sse-no-audio",
+            )
+        ]
+
+        error_event = next(event for event in events if "event: speech.audio.error" in event)
+        payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+        assert "partial_audio" not in payload["error"]
+        assert "action" not in payload["error"]
+
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
         async def pcm_generator():

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 import base64
 from types import SimpleNamespace
@@ -44,7 +47,9 @@ def _build_test_app(
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -244,6 +249,7 @@ class TestStreamingSpeechWebSocket:
 
     def test_streaming_multiple_binary_frames(self, mocker: MockerFixture):
         captured_requests = []
+        captured_tts_params = []
 
         speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
@@ -253,11 +259,12 @@ class TestStreamingSpeechWebSocket:
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
-            return "req-stream", object(), {}
+            return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]}
 
         speech_service._prepare_speech_generation = mock_prepare_speech_generation
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
+            captured_tts_params.append(tts_params)
             for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -302,6 +309,7 @@ class TestStreamingSpeechWebSocket:
         assert captured_requests[0].stream is True
         assert captured_requests[0].response_format == "pcm"
         assert captured_requests[0].initial_codec_chunk_frames == 12
+        assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
         assert speech_service._generate_audio_bytes.await_count == 0
 
     def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
@@ -344,7 +352,9 @@ class TestStreamingSpeechWebSocket:
 
         # The forced-aligner stage rides the same generator: its pooling output
         # is surfaced via the ``collect`` channel once the audio has streamed.
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             for chunk in (first_chunk, second_chunk):
                 yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
@@ -427,7 +437,9 @@ class TestStreamingSpeechWebSocket:
         speech_service.forced_aligner_enabled = True
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+        ):
             chunk = b"\x01" * 1000
             yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
@@ -608,7 +620,7 @@ class TestStreamingSpeechWebSocket:
         speech_service.engine_client.abort = mocker.AsyncMock()
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
             yield b"\x01\x02"
             raise RuntimeError("stream boom")
 
@@ -633,6 +645,8 @@ class TestStreamingSpeechWebSocket:
                 assert ws.receive_json() == {
                     "type": "error",
                     "message": "Generation failed for utterance 0, sentence 0: stream boom",
+                    "partial_audio": True,
+                    "action": "discard",
                 }
                 assert ws.receive_json() == {
                     "type": "audio.done",
@@ -702,7 +716,7 @@ class TestStreamingSpeechWebSocket:
         speech_service.engine_client.abort = mocker.AsyncMock()
         speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
             yield b"\x01\x02"
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks

tests/entrypoints/openai_api/test_speech_usage.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for speech token-usage accounting (issue #4646).
 
 These exercise the pure usage logic without a model/server: input-token
@@ -205,6 +206,14 @@ def test_output_tokens_from_stage_metrics():
     assert acc.total() == 77
 
 
+def test_output_tokens_preserves_stage0_finish_reason():
+    acc = SpeechOutputTokenCounter()
+    res = _final_res(stage0_out=128)
+    res.metrics["stage_metrics"]["0"]["finish_reason"] = "length"
+    acc.observe(res)
+    assert acc.stage0_finish_reason == "length"
+
+
 def test_output_tokens_takes_max_across_stages():
     acc = SpeechOutputTokenCounter()
     acc.observe(_final_res(stage0_out=42, stage1_out=0))

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -32,7 +32,11 @@ from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
     MossTTSAdapter,
     MossTTSNanoAdapter,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import (
+    QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY,
+    Qwen3TTSAdapter,
+    Qwen3TTSCodecLimitError,
+)
 from vllm_omni.model_executor.models.indextts2 import prompt_utils
 from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
     INDEXTTS25_TOKENIZER_FILE,
@@ -162,6 +166,56 @@ def test_qwen3_tts_metadata():
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
 
 
+@pytest.mark.parametrize(
+    ("task_type", "text_tokens", "request_cap", "expected_cap"),
+    [
+        ("Base", 0, None, 4096),
+        ("Base", 10, None, 192),
+        ("Base", 23, None, 276),
+        ("Base", 23, 128, 128),
+        ("Base", 23, 512, 512),
+        ("Base", 400, None, 4096),
+        ("CustomVoice", 10, None, 4096),
+        ("CustomVoice", 10, 256, 256),
+    ],
+)
+def test_qwen3_tts_applies_text_scaled_codec_safety_limit(task_type, text_tokens, request_cap, expected_cap, mocker):
+    server = mocker.Mock()
+    server._count_usage_text_tokens.return_value = text_tokens
+    adapter = Qwen3TTSAdapter(SpeechServingContext(server=server))
+    stage_defaults = [SamplingParams(max_tokens=4096, min_tokens=2)]
+    request = OpenAICreateSpeechRequest(
+        input="test text",
+        task_type=task_type,
+        max_new_tokens=request_cap,
+    )
+    prompt: dict[str, Any] = {"additional_information": {}}
+
+    overridden = adapter.apply_sampling_overrides(stage_defaults, request, prompt)
+
+    assert overridden[0].max_tokens == expected_cap
+    assert prompt["additional_information"][QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] == [expected_cap]
+    assert stage_defaults[0].max_tokens == 4096
+
+
+def test_qwen3_tts_rejects_only_length_finished_base_audio(mocker):
+    adapter = Qwen3TTSAdapter(SpeechServingContext(server=mocker.Mock()))
+    params = {
+        "task_type": ["Base"],
+        QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY: [192],
+    }
+
+    # EOS at the exact budget is valid; the token count alone is not a
+    # sufficient failure signal.
+    adapter.validate_generation(params, stage0_finish_reason="stop", output_tokens=192)
+    adapter.validate_generation(params, stage0_finish_reason=None, output_tokens=192)
+
+    # Some frontends expose 191 decoded frames for max_new_tokens=192. The
+    # engine terminal reason still makes this an unambiguous limit failure.
+    with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
+        adapter.validate_generation(params, stage0_finish_reason="length", output_tokens=191)
+
+
 def test_indextts_adapters_are_versioned():
     assert resolve_adapter("indextts2") is IndexTTS2Adapter
     assert resolve_adapter("indextts2_5") is IndexTTS25Adapter