[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