[Bugfix][Qwen3-TTS] Mode-aware ref_audio artifact readiness (fixes #5049) (#5157)

tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E regression for Qwen3-TTS ref_audio artifact mode collision (#5049).
+
+Reproduces the issue's exact sequence on a live server: an
+``x_vector_only_mode=true`` request followed by an ICL
+(``x_vector_only_mode=false`` + ``ref_text``) request using the **same**
+``ref_audio``.
+
+Before the fix, the ref_audio artifact readiness was mode-agnostic: the
+x-vector-only request (which caches a speaker embedding but no ``ref_code``)
+marked the artifact "ready", so the ICL request took the artifact-only serving
+path (``ref_audio`` stripped before dispatch), found no ``ref_code`` and no
+audio to recompute from, and raised in ``build_prompt_embeds`` ->
+``EngineDeadError`` — taking the whole server down rather than failing the one
+request.
+
+After the fix (readiness keyed by ``(artifact_key, x_vector_only)``), the ICL
+request no longer matches the x-vector artifact, so it sends ``ref_audio`` and
+the worker computes ``ref_code``. Both requests succeed and the server stays up.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0
+
+# Vendored under tests/assets/qwen3_tts/clone_2.wav (see test_qwen3_tts_base.py).
+REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+INPUT_TEXT = "The weather is nice today, perfect for a walk in the park."
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
+            server_args=["--trust-remote-code"],
+        ),
+        id="async_chunk",
+    )
+]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_xvector_then_icl_same_ref_audio_keeps_engine_alive(omni_server, openai_client) -> None:
+    """Regression for #5049: x-vector-only then ICL on the same ref_audio.
+
+    Both requests must return audio and the engine must stay alive (pre-fix the
+    second request killed EngineCore).
+    """
+    base_request = {
+        "model": omni_server.model,
+        "input": INPUT_TEXT,
+        "stream": False,
+        "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+        "response_format": "wav",
+        "task_type": "Base",
+        "ref_audio": REF_AUDIO_URL,
+        # Assert non-empty audio came back (not just HTTP 200).
+        "min_audio_bytes": 1,
+    }
+
+    # 1) x-vector-only request: caches a speaker embedding (no ref_code) and, on
+    #    completion, marks the ref_audio artifact "ready".
+    openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True})
+
+    # 2) ICL request with the SAME ref_audio. Pre-fix this reused the x-vector
+    #    artifact via the artifact-only path, hit the missing ref_code, and
+    #    killed EngineCore (#5049). It must now succeed with audio.
+    openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": False, "ref_text": REF_TEXT})
+
+    # 3) The server must still be serving after the ICL request (pre-fix this
+    #    would fail with EngineDeadError / connection refused).
+    openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True})

tests/entrypoints/openai_api/test_serving_speech.py

@@ -3682,7 +3682,7 @@ class TestTTSAsyncOffloading:
             24000,
             artifact_key,
         )
-        qwen3_tts_server._ref_audio_model_artifact_ready.add(artifact_key)
+        qwen3_tts_server._ref_audio_model_artifact_ready.add((artifact_key, False))
         qwen3_tts_server._codec_frame_rate = 25.0
         qwen3_tts_server._tts_tokenizer = lambda _text, padding=False: {"input_ids": list(range(10))}
         qwen3_tts_server.engine_client.model_config.hf_config.talker_config = SimpleNamespace(
@@ -3712,10 +3712,10 @@ class TestTTSAsyncOffloading:
         qwen3_tts_server._ref_audio_resolve_cache_max_bytes = 1_000_000
 
         qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a")
-        qwen3_tts_server._ref_audio_model_artifact_ready.add("artifact-a")
+        qwen3_tts_server._ref_audio_model_artifact_ready.add(("artifact-a", False))
         qwen3_tts_server._put_resolved_ref_audio("ref-b", [0.0] * 8, 24000, "artifact-b")
 
-        assert "artifact-a" not in qwen3_tts_server._ref_audio_model_artifact_ready
+        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
@@ -3724,13 +3724,13 @@ class TestTTSAsyncOffloading:
             raise ValueError("boom")
             yield  # pragma: no cover
 
-        qwen3_tts_server._request_ref_audio_artifact_keys["req-fail"] = "artifact-fail"
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-fail"] = ("artifact-fail", False)
 
         with pytest.raises(ValueError, match="boom"):
             await anext(qwen3_tts_server._generate_audio_chunks(failing_generator(), "req-fail"))
 
         assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
-        assert "artifact-fail" not in qwen3_tts_server._ref_audio_model_artifact_ready
+        assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
 
     @pytest.mark.asyncio
     async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
@@ -3743,22 +3743,48 @@ class TestTTSAsyncOffloading:
             )
             await asyncio.sleep(0)
 
-        qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = "artifact-close"
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = ("artifact-close", False)
 
         stream = qwen3_tts_server._generate_audio_chunks(pcm_generator(), "req-close")
         assert await anext(stream)
         await stream.aclose()
 
         assert "req-close" not in qwen3_tts_server._request_ref_audio_artifact_keys
-        assert "artifact-close" not in qwen3_tts_server._ref_audio_model_artifact_ready
+        assert ("artifact-close", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
 
     def test_qwen3_ref_audio_artifact_ready_requires_live_resolve_cache_entry(self, qwen3_tts_server):
-        qwen3_tts_server._request_ref_audio_artifact_keys["req-evicted"] = "artifact-evicted"
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-evicted"] = ("artifact-evicted", False)
 
         qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-evicted")
 
         assert "req-evicted" not in qwen3_tts_server._request_ref_audio_artifact_keys
-        assert "artifact-evicted" not in qwen3_tts_server._ref_audio_model_artifact_ready
+        assert ("artifact-evicted", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
+
+    def test_qwen3_xvector_ready_artifact_does_not_enable_icl_artifact_only(self, qwen3_tts_server):
+        # An x-vector-only artifact (speaker embedding, no ref_code) must not enable
+        # the artifact-only path for a later ICL request with the same ref_audio (#5049).
+        qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a")
+        qwen3_tts_server._track_ref_audio_artifact_warmup("req-xvec", "artifact-a", x_vector_only=True)
+        qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-xvec")
+
+        icl_params = {"task_type": ["Base"], "x_vector_only_mode": [False]}
+        assert qwen3_tts_server._qwen3_tts_can_use_ref_audio_artifact_only(icl_params, "artifact-a") is False
+
+    def test_qwen3_xvector_ready_artifact_still_reusable_by_xvector_request(self, qwen3_tts_server):
+        qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a")
+        qwen3_tts_server._track_ref_audio_artifact_warmup("req-xvec", "artifact-a", x_vector_only=True)
+        qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-xvec")
+
+        xvec_params = {"task_type": ["Base"], "x_vector_only_mode": [True]}
+        assert qwen3_tts_server._qwen3_tts_can_use_ref_audio_artifact_only(xvec_params, "artifact-a") is True
+
+    def test_qwen3_icl_ready_artifact_enables_icl_artifact_only(self, qwen3_tts_server):
+        qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a")
+        qwen3_tts_server._track_ref_audio_artifact_warmup("req-icl", "artifact-a", x_vector_only=False)
+        qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-icl")
+
+        icl_params = {"task_type": ["Base"], "x_vector_only_mode": [False]}
+        assert qwen3_tts_server._qwen3_tts_can_use_ref_audio_artifact_only(icl_params, "artifact-a") is True
 
     def test_shutdown_is_idempotent(self, mocker: MockerFixture):
         """Calling shutdown() twice should not raise."""

tests/helpers/runtime.py

@@ -1795,6 +1795,7 @@ class OpenAIClientHandler:
             "instructions",
             "speed",
             "stream_format",
+            "x_vector_only_mode",
         ):
             if key in request_config:
                 extra_body[key] = request_config[key]

vllm_omni/entrypoints/openai/serving_speech.py

@@ -322,8 +322,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         self._ref_audio_resolve_cache_bytes = 0
         self._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES
         self._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES
-        self._ref_audio_model_artifact_ready: set[str] = set()
-        self._request_ref_audio_artifact_keys: dict[str, str] = {}
+        # Readiness is keyed by (artifact_key, x_vector_only). An x-vector-only
+        # request caches a speaker embedding but no ref_code, so its artifact
+        # must not satisfy a later ICL request that needs ref_code (#5049).
+        self._ref_audio_model_artifact_ready: set[tuple[str, bool]] = set()
+        self._request_ref_audio_artifact_keys: dict[str, tuple[str, bool]] = {}
         self._higgs_audio_v3_ref_code_cache: OrderedDict[str, tuple[torch.Tensor, int]] = OrderedDict()
         self._higgs_audio_v3_ref_code_cache_bytes = 0
         self._higgs_audio_v3_ref_code_inflight: dict[str, asyncio.Task[torch.Tensor]] = {}
@@ -2069,23 +2072,35 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _discard_ref_audio_artifact_ready_if_unreferenced(self, artifact_key: str) -> None:
         if artifact_key and all(entry[3] != artifact_key for entry in self._ref_audio_resolve_cache.values()):
-            self._ref_audio_model_artifact_ready.discard(artifact_key)
+            self._ref_audio_model_artifact_ready = {
+                (key, mode) for (key, mode) in self._ref_audio_model_artifact_ready if key != artifact_key
+            }
+
+    @staticmethod
+    def _tts_x_vector_only(tts_params: dict[str, Any]) -> bool:
+        return bool((tts_params.get("x_vector_only_mode") or [False])[0])
 
     def _qwen3_tts_can_use_ref_audio_artifact_only(self, tts_params: dict[str, Any], artifact_key: str | None) -> bool:
         if self._tts_model_type != "qwen3_tts":
             return False
-        if not artifact_key or artifact_key not in self._ref_audio_model_artifact_ready:
+        x_vector_only = self._tts_x_vector_only(tts_params)
+        if not artifact_key or (artifact_key, x_vector_only) not in self._ref_audio_model_artifact_ready:
             return False
         return (tts_params.get("task_type") or ["CustomVoice"])[0] == "Base"
 
-    def _track_ref_audio_artifact_warmup(self, request_id: str, artifact_key: str | None) -> None:
+    def _track_ref_audio_artifact_warmup(
+        self, request_id: str, artifact_key: str | None, x_vector_only: bool = False
+    ) -> None:
         if artifact_key:
-            self._request_ref_audio_artifact_keys[request_id] = artifact_key
+            self._request_ref_audio_artifact_keys[request_id] = (artifact_key, bool(x_vector_only))
 
     def _mark_ref_audio_artifact_ready_for_request(self, request_id: str) -> None:
-        artifact_key = self._request_ref_audio_artifact_keys.pop(request_id, None)
+        tracked = self._request_ref_audio_artifact_keys.pop(request_id, None)
+        if tracked is None:
+            return
+        artifact_key, x_vector_only = tracked
         if artifact_key and any(entry[3] == artifact_key for entry in self._ref_audio_resolve_cache.values()):
-            self._ref_audio_model_artifact_ready.add(artifact_key)
+            self._ref_audio_model_artifact_ready.add((artifact_key, x_vector_only))
 
     def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None:
         self._request_ref_audio_artifact_keys.pop(request_id, None)
@@ -3234,7 +3249,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             sampling_params_list=sampling_params_list,
             output_modalities=["audio"],
         )
-        self._track_ref_audio_artifact_warmup(request_id, qwen3_ref_audio_warmup_artifact_key)
+        self._track_ref_audio_artifact_warmup(
+            request_id,
+            qwen3_ref_audio_warmup_artifact_key,
+            x_vector_only=self._tts_x_vector_only(tts_params),
+        )
         return request_id, generator, tts_params
 
     async def _generate_pcm_chunks(self, generator, request_id: str, *, include_sample_rate: bool = False):