diff --git a/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py b/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py new file mode 100644 index 000000000..7992af25f --- /dev/null +++ b/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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py index eb00fe784..1e64f7a63 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py +++ b/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