test(nvidia_riva): audit #41021 (#48801)

* test(nvidia_riva): audit #41021

Adds an owned Riva ASR gRPC server and provider integration cells for the riva_offline request-body ban, the offline and streaming paths, the per-deployment caller opt-in, and the 413 and 408 error mappings

* test(nvidia_riva): check word timing requests and audio samples in the Riva double

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

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()

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