diff --git a/tests/integration/_support/riva.py b/tests/integration/_support/riva.py new file mode 100644 index 00000000000..8552bc2197d --- /dev/null +++ b/tests/integration/_support/riva.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import io +import threading +import wave +from array import array +from collections.abc import Iterable, Iterator +from concurrent import futures +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from queue import SimpleQueue +from typing import Final + +import grpc +from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc + +SAMPLE_RATE_HZ: Final = 16000 +OFFLINE_ONLY_REFUSAL: Final = ( + "Error: Unavailable model requested given these parameters: language_code=en-US; sample_rate=16000; type=online;" +) +_STALL_LIMIT_SECONDS: Final = 120 + + +class Rpc(Enum): + RECOGNIZE = "Recognize" + STREAMING = "StreamingRecognize" + + +@dataclass(frozen=True, slots=True) +class Word: + word: str + start_ms: int + end_ms: int + + +@dataclass(frozen=True, slots=True) +class RivaCall: + rpc: Rpc + audio: bytes + cancelled: bool + + +@dataclass(frozen=True, slots=True) +class RivaServer: + address: str + calls: SimpleQueue[RivaCall] + + def drain(self) -> tuple[RivaCall, ...]: + return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize())) + + +def _clip_samples(seconds: float) -> array[int]: + return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * SAMPLE_RATE_HZ)))) + + +def pcm_wav(seconds: float) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(SAMPLE_RATE_HZ) + writer.writeframes(_clip_samples(seconds).tobytes()) + return buffer.getvalue() + + +def pcm_bytes(seconds: float) -> int: + return int(seconds * SAMPLE_RATE_HZ) * 2 + + +def carries_clip(audio: bytes, seconds: float) -> bool: + sent: Final = _clip_samples(seconds) + if len(audio) != len(sent) * 2: + return False + received: Final = array("h", audio) + return all(abs(left - right) <= 1 for left, right in zip(sent, received, strict=True)) + + +class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): + def __init__(self, words: tuple[Word, ...], offline_only: bool, stall: bool, calls: SimpleQueue[RivaCall]) -> None: + self._words: Final = words + self._offline_only: Final = offline_only + self._stall: Final = stall + self._calls: Final = calls + + def _alternative(self, config: riva_asr_pb2.RecognitionConfig) -> riva_asr_pb2.SpeechRecognitionAlternative: + timed: Final = self._words if config.enable_word_time_offsets else () + return riva_asr_pb2.SpeechRecognitionAlternative( + transcript=" ".join(word.word for word in self._words), + confidence=0.9, + words=[ + riva_asr_pb2.WordInfo(word=word.word, start_time=word.start_ms, end_time=word.end_ms) for word in timed + ], + ) + + def Recognize( + self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext + ) -> riva_asr_pb2.RecognizeResponse: + if self._stall: + released: Final = threading.Event() + context.add_callback(released.set) + released.wait(timeout=_STALL_LIMIT_SECONDS) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active())) + return riva_asr_pb2.RecognizeResponse() + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False)) + return riva_asr_pb2.RecognizeResponse( + results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])] + ) + + def StreamingRecognize( + self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext + ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: + if self._offline_only: + self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False)) + context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL) + requests: Final = tuple(request_iterator) + config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig() + self._calls.put( + RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False) + ) + yield riva_asr_pb2.StreamingRecognizeResponse( + results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)] + ) + + +@contextmanager +def riva_server( + words: tuple[Word, ...], + *, + offline_only: bool = False, + stall: bool = False, + max_receive_bytes: int | None = None, +) -> Iterator[RivaServer]: + calls: Final[SimpleQueue[RivaCall]] = SimpleQueue() + options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),) + server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=8), options=options) + riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(words, offline_only, stall, calls), server) + port: Final = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + yield RivaServer(f"127.0.0.1:{port}", calls) + finally: + server.stop(grace=None).wait() diff --git a/tests/integration/providers/test_nvidia_riva_wire.py b/tests/integration/providers/test_nvidia_riva_wire.py new file mode 100644 index 00000000000..ef36c449761 --- /dev/null +++ b/tests/integration/providers/test_nvidia_riva_wire.py @@ -0,0 +1,137 @@ +from collections.abc import Mapping +from typing import Final + +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.openai_wire import chat_reply +from integration._support.riva import RivaServer, Rpc, Word, carries_clip, pcm_bytes, pcm_wav, riva_server +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL: Final = "nvidia_riva/nvidia/parakeet-0.6b-tdt" +_WORDS: Final = (Word("riva", 0, 400), Word("offline", 400, 900), Word("works", 900, 1400)) +_TRANSCRIPT: Final = "riva offline works" +_CLIP_SECONDS: Final = 0.5 +_REFUSED_FIELD: Final = "riva_offline is not allowed in request body" + + +def _riva_model(scenario: Scenario, server: RivaServer, **parameters: JsonValue) -> str: + return scenario.model( + model=_MODEL, + api_base=server.address, + model_info={"mode": "audio_transcription"}, + **parameters, + ) + + +def _transcribe( + gateway: Gateway, key: str, model: str, seconds: float = _CLIP_SECONDS, **fields: str +) -> httpx.Response: + return gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, **fields}, + {"file": ("clip.wav", pcm_wav(seconds), "audio/wav")}, + key=key, + ) + + +def _sdk(gateway: Gateway, key: str) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=key, max_retries=0) + + +def test_a_form_riva_offline_on_transcriptions_is_refused_before_riva_is_called(gateway: Gateway) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server) + refused: Final = _transcribe(gateway, scenario.key(), model, riva_offline="true") + assert refused.status_code == 401, refused.text + assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text + assert server.drain() == () + + +def test_a_json_riva_offline_on_chat_is_refused_before_the_provider_is_called(gateway: Gateway) -> None: + def respond(_: Request) -> Reply: + return chat_reply("chatcmpl-riva-audit", "gpt-4o-mini", "served", stream=False) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + refused: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "riva flag on chat"}], "riva_offline": True}, + key=scenario.key(), + ) + assert refused.status_code == 401, refused.text + assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text + assert wire.drain() == () + + +def test_an_offline_deployment_sends_the_whole_clip_in_one_recognize_call(gateway: Gateway) -> None: + with riva_server(_WORDS, offline_only=True) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, riva_offline=True) + transcription: Final = _sdk(gateway, scenario.key()).audio.transcriptions.create( + model=model, + file=("clip.wav", pcm_wav(_CLIP_SECONDS), "audio/wav"), + response_format="verbose_json", + timestamp_granularities=["word"], + ) + assert transcription.text == _TRANSCRIPT + assert [(word.word, word.start, word.end) for word in transcription.words or ()] == [ + (word.word, word.start_ms / 1000, word.end_ms / 1000) for word in _WORDS + ] + assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [ + (Rpc.RECOGNIZE, True) + ] + + +def test_a_deployment_without_riva_offline_still_streams_the_clip(gateway: Gateway) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, timeout=30) + served: Final = _transcribe(gateway, scenario.key(), model) + assert served.status_code == 200, served.text + assert served.json()["text"] == _TRANSCRIPT, served.text + assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [ + (Rpc.STREAMING, True) + ] + + +@pytest.mark.parametrize( + ("sent", "deployment", "status", "rpcs"), + ( + pytest.param("true", {}, 200, (Rpc.RECOGNIZE,), id="true-on-a-streaming-deployment"), + pytest.param("false", {"riva_offline": True}, 200, (Rpc.STREAMING,), id="false-on-an-offline-deployment"), + pytest.param("yes", {}, 400, (), id="malformed"), + pytest.param("yes", {"drop_params": True}, 200, (Rpc.STREAMING,), id="malformed-under-drop-params"), + ), +) +def test_an_opted_in_caller_riva_offline_picks_the_rpc( + gateway: Gateway, sent: str, deployment: Mapping[str, JsonValue], status: int, rpcs: tuple[Rpc, ...] +) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, configurable_clientside_auth_params=["riva_offline"], **deployment) + answered: Final = _transcribe(gateway, scenario.key(), model, riva_offline=sent) + assert answered.status_code == status, answered.text + assert tuple(call.rpc for call in server.drain()) == rpcs + + +def test_offline_audio_over_the_servers_message_limit_answers_413(gateway: Gateway) -> None: + seconds: Final = 2.0 + limit: Final = pcm_bytes(seconds) // 2 + with ( + riva_server(_WORDS, offline_only=True, max_receive_bytes=limit) as server, + gateway.scenario() as scenario, + ): + model: Final = _riva_model(scenario, server, riva_offline=True) + refused: Final = _transcribe(gateway, scenario.key(), model, seconds) + assert refused.status_code == 413, refused.text + assert server.drain() == () + + +def test_a_stalled_offline_recognize_answers_408_at_the_deployment_timeout(gateway: Gateway) -> None: + with riva_server(_WORDS, offline_only=True, stall=True) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, riva_offline=True, timeout=2) + timed_out: Final = _transcribe(gateway, scenario.key(), model) + assert timed_out.status_code == 408, timed_out.text + calls: Final = eventually(server.drain, lambda observed: len(observed) > 0, seconds=30) + assert all(call.rpc is Rpc.RECOGNIZE and call.cancelled for call in calls), calls