[Bugfix][TTS] Isolate shared Higgs v3 reference encode from request cancellation (#7076)

Signed-off-by: Allen Wu <allenwu2795@gmail.com>
Co-authored-by: TRAE CLI <traecli@bytedance.com>

tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+import gc
+import threading
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v3 import HiggsAudioV3Adapter
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_ref_code_inflight_survives_request_cancellation():
+    adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace()))
+    encode_started = threading.Event()
+    release_encode = threading.Event()
+    calls = 0
+    expected_codes = torch.arange(16, dtype=torch.long).reshape(2, 8) + 1
+
+    def encode_reference_audio(_wav, _sr):
+        nonlocal calls
+        calls += 1
+        encode_started.set()
+        assert release_encode.wait(timeout=5)
+        return torch.arange(16, dtype=torch.long).reshape(2, 8)
+
+    def apply_delay_pattern(codes):
+        return codes + 1
+
+    async def resolve():
+        return await adapter._resolve_higgs_audio_v3_ref_codes(
+            "ref-a",
+            object(),
+            24000,
+            encode_reference_audio,
+            apply_delay_pattern,
+        )
+
+    async def run():
+        creator = asyncio.create_task(resolve())
+        try:
+            assert await asyncio.wait_for(asyncio.to_thread(encode_started.wait), timeout=5)
+
+            cancelled_waiter = asyncio.create_task(resolve())
+            survivor = asyncio.create_task(resolve())
+            await asyncio.sleep(0)
+
+            cancelled_waiter.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await cancelled_waiter
+
+            creator.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await creator
+
+            release_encode.set()
+            survivor_result = await asyncio.wait_for(survivor, timeout=5)
+            await asyncio.sleep(0)
+            return survivor_result
+        finally:
+            release_encode.set()
+
+    codes, cache_hit, inflight_wait = asyncio.run(run())
+
+    assert calls == 1
+    assert cache_hit is False
+    assert inflight_wait is True
+    assert torch.equal(codes, expected_codes)
+    assert "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight
+    cached = adapter._get_higgs_audio_v3_ref_codes("ref-a")
+    assert cached is not None
+    assert torch.equal(cached, expected_codes)
+
+
+def test_ref_code_inflight_consumes_failure_after_all_callers_cancel():
+    adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace()))
+    encode_started = threading.Event()
+    release_encode = threading.Event()
+    loop_errors = []
+
+    def encode_reference_audio(_wav, _sr):
+        encode_started.set()
+        assert release_encode.wait(timeout=5)
+        raise RuntimeError("encode failed")
+
+    async def resolve():
+        return await adapter._resolve_higgs_audio_v3_ref_codes(
+            "ref-a",
+            object(),
+            24000,
+            encode_reference_audio,
+            lambda codes: codes,
+        )
+
+    async def run():
+        loop = asyncio.get_running_loop()
+        loop.set_exception_handler(lambda _loop, context: loop_errors.append(context))
+        creator = asyncio.create_task(resolve())
+        try:
+            assert await asyncio.wait_for(asyncio.to_thread(encode_started.wait), timeout=5)
+            waiter = asyncio.create_task(resolve())
+            await asyncio.sleep(0)
+
+            creator.cancel()
+            waiter.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await creator
+            with pytest.raises(asyncio.CancelledError):
+                await waiter
+
+            release_encode.set()
+            for _ in range(100):
+                if "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight:
+                    break
+                await asyncio.sleep(0.01)
+            gc.collect()
+            await asyncio.sleep(0)
+        finally:
+            release_encode.set()
+
+    asyncio.run(run())
+
+    assert "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight
+    assert not loop_errors

vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py

@@ -4,7 +4,7 @@
 
 import asyncio
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import torch
@@ -62,7 +62,7 @@ class HiggsAudioV3Adapter(ARTTSAdapter):
             encode_reference_audio,
         )
 
-        wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio)
+        wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio))
         artifact_key = self._get_resolved_ref_audio_artifact_key(cache_key)
         wav = np.asarray(wav_list, dtype=np.float32)
         ref_codes_delayed, cache_hit, inflight_wait = await self._resolve_higgs_audio_v3_ref_codes(
@@ -109,7 +109,7 @@ class HiggsAudioV3Adapter(ARTTSAdapter):
 
         task = self._higgs_audio_v3_ref_code_inflight.get(artifact_key)
         if task is not None:
-            return (await task).clone(), False, True
+            return (await asyncio.shield(task)).clone(), False, True
 
         async def _encode_and_cache() -> torch.Tensor:
             ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr)
@@ -120,12 +120,16 @@ class HiggsAudioV3Adapter(ARTTSAdapter):
 
         task = asyncio.create_task(_encode_and_cache())
         self._higgs_audio_v3_ref_code_inflight[artifact_key] = task
-        try:
-            return (await task).clone(), False, False
-        finally:
-            if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is task:
+
+        def _retire(t: asyncio.Task[torch.Tensor]) -> None:
+            if not t.cancelled():
+                t.exception()
+            if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is t:
                 self._higgs_audio_v3_ref_code_inflight.pop(artifact_key, None)
 
+        task.add_done_callback(_retire)
+        return (await asyncio.shield(task)).clone(), False, False
+
     def _get_higgs_audio_v3_ref_codes(self, artifact_key: str | None) -> torch.Tensor | None:
         if not artifact_key:
             return None