feat: add automatic recovery for streaming ASR/TTS

riva/client/__init__.py

@@ -4,6 +4,7 @@
 from riva.client.asr import (
     AudioChunkFileIterator,
     ASRService,
+    ResilientStreamingASR,
     add_audio_file_specs_to_config,
     add_word_boosting_to_config,
     add_speaker_diarization_to_config,
@@ -39,5 +40,10 @@
 from riva.client.proto.riva_audio_pb2 import AudioEncoding
 from riva.client.proto.riva_nlp_pb2 import AnalyzeIntentOptions
 from riva.client.proto.riva_nmt_pb2 import StreamingTranslateSpeechToSpeechConfig, TranslationConfig, SynthesizeSpeechConfig, StreamingTranslateSpeechToTextConfig
-from riva.client.tts import SpeechSynthesisService
+from riva.client.retry import (
+    RETRYABLE_GRPC_CODES,
+    exponential_backoff,
+    is_retryable_grpc_error,
+)
+from riva.client.tts import SpeechSynthesisService, ResilientStreamingTTS
 from riva.client.nmt import NeuralMachineTranslationClient

riva/client/asr.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: MIT
 
 import io
+import logging
 import os
 import sys
 import time
@@ -11,15 +12,21 @@
 import wave
 from itertools import groupby
 from pathlib import Path
-from typing import Callable, Dict, Generator, Iterable, List, Optional, TextIO, Union
+from collections import deque
+from typing import Callable, Deque, Dict, Generator, Iterable, List, Optional, TextIO, Union
 
+import grpc
 from google.protobuf.json_format import MessageToJson
 from grpc._channel import _MultiThreadedRendezvous
 
 import riva.client
 import riva.client.proto.riva_asr_pb2 as rasr
 import riva.client.proto.riva_asr_pb2_grpc as rasr_srv
 from riva.client.auth import Auth
+from riva.client.retry import exponential_backoff, is_retryable_grpc_error
+
+
+LOGGER = logging.getLogger(__name__)
 
 
 def get_wav_file_parameters(input_file: Union[str, os.PathLike]) -> Dict[str, Union[int, float]]:
@@ -483,3 +490,154 @@ def offline_recognize(
         request = rasr.RecognizeRequest(config=config, audio=audio_bytes)
         func = self.stub.Recognize.future if future else self.stub.Recognize
         return func(request, metadata=self.auth.get_auth_metadata())
+
+
+
+class ResilientStreamingASR:
+    """A resilient wrapper around :class:`ASRService` for streaming recognition.
+
+    This class buffers recent audio and automatically reconnects on transient
+    gRPC failures, replaying buffered audio so that recognition can continue
+    with a bounded audio lookback. Recovery is best effort: audio older than
+    the configured lookback window may not be replayed, and applications must
+    tolerate repeated transcripts after a reconnect.
+
+    Example:
+        >>> auth = Auth(uri="localhost:50051")
+        >>> asr = ASRService(auth)
+        >>> config = StreamingRecognitionConfig(
+        ...     config=RecognitionConfig(enable_automatic_punctuation=True),
+        ...     interim_results=True,
+        ... )
+        >>> resilient_asr = ResilientStreamingASR(asr, config)
+        >>> for response in resilient_asr.stream(audio_chunks):
+        ...     print(response)
+
+    Args:
+        asr_service: The underlying :class:`ASRService` instance.
+        streaming_config: Configuration for streaming recognition.
+        max_retries: Maximum number of reconnection attempts per failure.
+        lookback_seconds: Duration of audio to replay after reconnecting.
+            A larger value improves recovery at the cost of higher latency.
+        sample_rate_hz: Sample rate of linear PCM audio.
+        audio_channel_count: Number of PCM audio channels.
+        sample_width_bytes: Bytes per PCM sample.
+        base_delay: Initial backoff delay in seconds.
+        max_delay: Maximum backoff delay in seconds.
+    """
+
+    def __init__(
+        self,
+        asr_service: ASRService,
+        streaming_config: rasr.StreamingRecognitionConfig,
+        max_retries: int = 3,
+        lookback_seconds: float = 2.0,
+        sample_rate_hz: int = 16000,
+        audio_channel_count: int = 1,
+        sample_width_bytes: int = 2,
+        base_delay: float = 1.0,
+        max_delay: float = 60.0,
+    ) -> None:
+        self.asr_service = asr_service
+        self.streaming_config = streaming_config
+        self.max_retries = max_retries
+        self.base_delay = base_delay
+        self.max_delay = max_delay
+
+        if lookback_seconds <= 0:
+            raise ValueError("lookback_seconds must be greater than zero")
+        if min(sample_rate_hz, audio_channel_count, sample_width_bytes) <= 0:
+            raise ValueError("PCM format values must be greater than zero")
+        self._lookback_max_bytes = int(
+            lookback_seconds * sample_rate_hz * audio_channel_count * sample_width_bytes
+        )
+        self._audio_buffer: Deque[bytes] = deque()
+        self._buffered_bytes = 0
+        self._retry_count = 0
+
+    def _buffered_request_generator(
+        self,
+        audio_source: Iterable[bytes],
+    ) -> Generator[rasr.StreamingRecognizeRequest, None, None]:
+        """Yield the config message, buffered audio, then new audio.
+
+        Each chunk from *audio_source* is appended to the lookback buffer
+        before being yielded so that it is available for the next reconnect.
+        """
+        yield rasr.StreamingRecognizeRequest(streaming_config=self.streaming_config)
+
+        # Replay buffered audio from previous (partial) stream
+        for chunk in self._audio_buffer:
+            yield rasr.StreamingRecognizeRequest(audio_content=chunk)
+
+        for chunk in audio_source:
+            self._append_audio(chunk)
+            yield rasr.StreamingRecognizeRequest(audio_content=chunk)
+
+    def _append_audio(self, chunk: bytes) -> None:
+        self._audio_buffer.append(chunk)
+        self._buffered_bytes += len(chunk)
+        while self._audio_buffer and self._buffered_bytes > self._lookback_max_bytes:
+            self._buffered_bytes -= len(self._audio_buffer.popleft())
+
+    def stream(
+        self,
+        audio_source: Iterable[bytes],
+    ) -> Generator[rasr.StreamingRecognizeResponse, None, None]:
+        """Stream audio for recognition with automatic recovery.
+
+        Args:
+            audio_source: An iterable of raw audio chunks.
+
+        Yields:
+            :obj:`StreamingRecognizeResponse` objects. A reconnect replays
+            the configured audio lookback, so callers should deduplicate
+            transcripts if their application requires exactly-once output.
+
+        Raises:
+            :obj:`grpc.RpcError`: If a non-retryable error occurs or the
+            maximum number of retries is exceeded.
+        """
+        audio_iterator = iter(audio_source)
+        attempt = 0
+        while True:
+            try:
+                generator = self._buffered_request_generator(audio_iterator)
+                for response in self.asr_service.stub.StreamingRecognize(
+                    generator, metadata=self.asr_service.auth.get_auth_metadata()
+                ):
+                    yield response
+                # Stream completed normally.
+                if self._retry_count > 0:
+                    LOGGER.info("Streaming ASR recovered after %d retry(s).", self._retry_count)
+                return
+
+            except grpc.RpcError as exc:
+                if not is_retryable_grpc_error(exc):
+                    LOGGER.warning(
+                        "Non-retryable gRPC error in streaming ASR: %s – %s",
+                        exc.code() if hasattr(exc, "code") else "UNKNOWN",
+                        exc.details() if hasattr(exc, "details") else str(exc),
+                    )
+                    raise
+                if attempt >= self.max_retries:
+                    LOGGER.error(
+                        "Streaming ASR failed permanently after %d retries. Last error: %s",
+                        self.max_retries,
+                        exc.details() if hasattr(exc, "details") else str(exc),
+                    )
+                    raise
+                delay = exponential_backoff(attempt, self.base_delay, self.max_delay)
+                LOGGER.info(
+                    "Streaming ASR connection lost (%s). Reconnecting in %.2f s "
+                    "(attempt %d/%d).",
+                    exc.code(),
+                    delay,
+                    attempt + 1,
+                    self.max_retries,
+                )
+                self._retry_count += 1
+                attempt += 1
+                time.sleep(delay)
+                # Loop continues: _buffered_request_generator will replay
+                # self._audio_buffer and then consume audio_iterator.

riva/client/auth.py

@@ -14,7 +14,7 @@ def create_channel(
     use_ssl: bool = False,
     uri: str = "localhost:50051",
     metadata: Optional[List[Tuple[str, str]]] = None,
-    options: Optional[List[Tuple[str, str]]] = [],
+    options: Optional[List[Tuple[str, Union[str, int]]]] = None,
     use_aio: Optional[bool] = False,
 ) -> grpc.Channel:
     def metadata_callback(context, callback):
@@ -61,7 +61,7 @@ def __init__(
         metadata_args: List[List[str]] = None,
         ssl_client_cert: Optional[Union[str, os.PathLike]] = None,
         ssl_client_key: Optional[Union[str, os.PathLike]] = None,
-        options: Optional[List[Tuple[str, str]]] = [],
+        options: Optional[List[Tuple[str, Union[str, int]]]] = None,
         use_aio: bool = False,
     ) -> None:
         """
@@ -82,8 +82,8 @@ def __init__(
                 Used for mutual TLS authentication. Defaults to None.
             ssl_client_key (Optional[Union[str, os.PathLike]], optional): Path to the SSL client private key file.
                 Used for mutual TLS authentication. Defaults to None.
-            options (Optional[List[Tuple[str, str]]], optional): Additional gRPC channel options.
-                Each tuple should contain (option_name, option_value). Defaults to [].
+            options (Optional[List[Tuple[str, Union[str, int]]]], optional): Additional gRPC channel options.
+                Each tuple should contain an option name and value.
             use_aio (bool, optional): Whether to use asyncio for the channel. Defaults to False.
 
         Raises:

riva/client/retry.py

@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+"""Shared retry utilities for resilient streaming ASR and TTS clients.
+
+This module provides constants and helpers for handling transient gRPC
+failures with exponential backoff. It is designed to be used by both
+:mod:`riva.client.asr` and :mod:`riva.client.tts`.
+"""
+
+import logging
+import random
+import grpc
+
+LOGGER = logging.getLogger(__name__)
+
+# gRPC status codes that are generally considered transient and safe to retry.
+RETRYABLE_GRPC_CODES = frozenset({
+    grpc.StatusCode.UNAVAILABLE,
+    grpc.StatusCode.DEADLINE_EXCEEDED,
+    grpc.StatusCode.INTERNAL,
+    grpc.StatusCode.RESOURCE_EXHAUSTED,
+    grpc.StatusCode.ABORTED,
+})
+
+# gRPC status codes that should NEVER be retried (client-side errors).
+NON_RETRYABLE_GRPC_CODES = frozenset({
+    grpc.StatusCode.INVALID_ARGUMENT,
+    grpc.StatusCode.PERMISSION_DENIED,
+    grpc.StatusCode.UNAUTHENTICATED,
+    grpc.StatusCode.NOT_FOUND,
+    grpc.StatusCode.ALREADY_EXISTS,
+    grpc.StatusCode.FAILED_PRECONDITION,
+    grpc.StatusCode.OUT_OF_RANGE,
+    grpc.StatusCode.UNIMPLEMENTED,
+})
+
+def is_retryable_grpc_error(exc: grpc.RpcError) -> bool:
+    """Return ``True`` if *exc* is a transient gRPC error that is safe to retry.
+
+    Args:
+        exc: The exception raised by a gRPC call.
+
+    Returns:
+        ``True`` if the error code is in :data:`RETRYABLE_GRPC_CODES`.
+    """
+    code = exc.code() if hasattr(exc, "code") else None
+    if code is None:
+        return False
+    return code in RETRYABLE_GRPC_CODES
+
+
+def exponential_backoff(
+    attempt: int,
+    base_delay: float = 1.0,
+    max_delay: float = 60.0,
+    jitter: bool = True,
+) -> float:
+    """Compute a sleep duration for the *attempt*-th retry.
+
+    Uses capped exponential backoff with optional full jitter to avoid
+    thundering-herd behaviour when many clients reconnect simultaneously.
+
+    Args:
+        attempt: Zero-based retry attempt number.
+        base_delay: Initial delay in seconds.
+        max_delay: Upper bound for the delay in seconds.
+        jitter: If ``True``, multiply the delay by a random factor in ``[0, 1)``.
+
+    Returns:
+        The number of seconds to sleep before the next attempt.
+    """
+    delay = min(base_delay * (2 ** attempt), max_delay)
+    if jitter:
+        delay = delay * random.random()
+    return delay