feat: add automatic recovery for streaming ASR/TTS

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, split_recovery_configuration
+
+
+LOGGER = logging.getLogger(__name__)
 
 
 def get_wav_file_parameters(input_file: Union[str, os.PathLike]) -> Dict[str, Union[int, float]]:
@@ -432,16 +439,67 @@ def streaming_response_generator(
                     config = RecognitionConfig(enable_automatic_punctuation=True)
                     streaming_config = StreamingRecognitionConfig(config, interim_results=True)
 
+                Set ``client_auto_recover:true`` in ``config.custom_configuration`` to enable
+                retry. ``client_max_retries`` and ``client_lookback_seconds`` control the
+                retry count and bounded PCM lookback buffer. These client-only keys are not
+                sent to Riva.
+
         Yields:
             :obj:`riva.client.proto.riva_asr_pb2.StreamingRecognizeResponse`: responses for audio chunks in
             :param:`audio_chunks`. You may find description of response fields in declaration of
             ``StreamingRecognizeResponse``
             message `here
             <https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/protos.html#riva-proto-riva-asr-proto>`_.
         """
-        generator = streaming_request_generator(audio_chunks, streaming_config)
-        for response in self.stub.StreamingRecognize(generator, metadata=self.auth.get_auth_metadata()):
-            yield response
+        server_configuration, auto_recover, max_retries, lookback_seconds = split_recovery_configuration(
+            streaming_config.config.custom_configuration
+        )
+        server_streaming_config = rasr.StreamingRecognitionConfig()
+        server_streaming_config.CopyFrom(streaming_config)
+        server_streaming_config.config.custom_configuration.clear()
+        server_streaming_config.config.custom_configuration.update(server_configuration)
+
+        if not auto_recover:
+            generator = streaming_request_generator(audio_chunks, server_streaming_config)
+            yield from self.stub.StreamingRecognize(generator, metadata=self.auth.get_auth_metadata())
+            return
+
+        sample_rate_hz = server_streaming_config.config.sample_rate_hertz or 16000
+        channel_count = server_streaming_config.config.audio_channel_count or 1
+        max_buffered_bytes = int(lookback_seconds * sample_rate_hz * channel_count * 2)
+        buffered_audio: Deque[bytes] = deque()
+        buffered_bytes = 0
+        audio_iterator = iter(audio_chunks)
+        attempt = 0
+
+        while True:
+            def request_generator() -> Generator[rasr.StreamingRecognizeRequest, None, None]:
+                nonlocal buffered_bytes
+                yield rasr.StreamingRecognizeRequest(streaming_config=server_streaming_config)
+                for chunk in buffered_audio:
+                    yield rasr.StreamingRecognizeRequest(audio_content=chunk)
+                for chunk in audio_iterator:
+                    buffered_audio.append(chunk)
+                    buffered_bytes += len(chunk)
+                    while buffered_audio and buffered_bytes > max_buffered_bytes:
+                        buffered_bytes -= len(buffered_audio.popleft())
+                    yield rasr.StreamingRecognizeRequest(audio_content=chunk)
+
+            try:
+                yield from self.stub.StreamingRecognize(
+                    request_generator(), metadata=self.auth.get_auth_metadata()
+                )
+                return
+            except grpc.RpcError as exc:
+                if not is_retryable_grpc_error(exc) or attempt >= max_retries:
+                    raise
+                delay = exponential_backoff(attempt)
+                LOGGER.info(
+                    "Streaming ASR connection lost (%s); retrying in %.2f seconds (attempt %d/%d).",
+                    exc.code(), delay, attempt + 1, max_retries,
+                )
+                attempt += 1
+                time.sleep(delay)
 
     def offline_recognize(
         self, audio_bytes: bytes, config: rasr.RecognitionConfig, future: bool = False

riva/client/retry.py

@@ -0,0 +1,106 @@
+# 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 random
+from typing import Dict, Mapping, Tuple
+
+import grpc
+
+CLIENT_AUTO_RECOVER = "client_auto_recover"
+CLIENT_MAX_RETRIES = "client_max_retries"
+CLIENT_LOOKBACK_SECONDS = "client_lookback_seconds"
+CLIENT_RECOVERY_KEYS = frozenset({
+    CLIENT_AUTO_RECOVER,
+    CLIENT_MAX_RETRIES,
+    CLIENT_LOOKBACK_SECONDS,
+})
+
+# 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
+
+
+def split_recovery_configuration(
+    custom_configuration: Mapping[str, str],
+) -> Tuple[Dict[str, str], bool, int, float]:
+    """Separate client-only streaming recovery options from server options.
+
+    The reserved ``client_*`` keys control retry behaviour in the Python
+    client and are intentionally not forwarded to Riva. All other key/value
+    pairs are returned unchanged for the server.
+    """
+    server_configuration = {
+        key: value for key, value in custom_configuration.items() if key not in CLIENT_RECOVERY_KEYS
+    }
+    enabled = str(custom_configuration.get(CLIENT_AUTO_RECOVER, "false")).lower() == "true"
+    max_retries = int(custom_configuration.get(CLIENT_MAX_RETRIES, "3"))
+    lookback_seconds = float(custom_configuration.get(CLIENT_LOOKBACK_SECONDS, "2.0"))
+    if max_retries < 0:
+        raise ValueError(f"{CLIENT_MAX_RETRIES} must be non-negative")
+    if lookback_seconds <= 0:
+        raise ValueError(f"{CLIENT_LOOKBACK_SECONDS} must be greater than zero")
+    return server_configuration, enabled, max_retries, lookback_seconds

riva/client/tts.py

@@ -1,16 +1,23 @@
 # SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: MIT
 
-from typing import Dict, Generator, Optional, Union, Iterable
+import logging
+import time
+from typing import Dict, Generator, Iterable, List, Optional, Union
 
+import grpc
 from grpc._channel import _MultiThreadedRendezvous
 
 import riva.client.proto.riva_tts_pb2 as rtts
 import riva.client.proto.riva_tts_pb2_grpc as rtts_srv
 from riva.client import Auth
 from riva.client.proto.riva_audio_pb2 import AudioEncoding
+from riva.client.retry import exponential_backoff, is_retryable_grpc_error, split_recovery_configuration
 import wave
 
+
+LOGGER = logging.getLogger(__name__)
+
 def parse_custom_configuration(custom_configuration: str) -> Dict[str, str]:
     """Parse a comma-separated ``key:value`` string into a dictionary.
 
@@ -124,8 +131,11 @@ def synthesize(
             if zero_shot_transcript is not None:
                 req.zero_shot_data.transcript = zero_shot_transcript
 
-        if custom_configuration:
-            for key, value in custom_configuration.items():
+        server_configuration, _, _, _ = split_recovery_configuration(
+            custom_configuration or {}
+        )
+        if server_configuration:
+            for key, value in server_configuration.items():
                 req.custom_configuration[key] = str(value)
 
         add_custom_dictionary_to_config(req, custom_dictionary)
@@ -167,7 +177,10 @@ def synthesize_online(
             zero_shot_quality: (:obj:`int`): Required quality of output audio, ranges between 1-40.
             custom_dictionary (:obj:`dict`, `optional`): Dictionary with key-value pair containing grapheme and corresponding phoneme
             custom_configuration (:obj:`Dict[str, str]`, `optional`): Free-form key/value parameters forwarded
-                to the synthesizer (e.g. ``{"exaggeration_factor": "1.5"}``). Model-specific.
+                to the synthesizer (e.g. ``{"exaggeration_factor": "1.5"}``). Set
+                ``client_auto_recover`` to ``"true"`` to enable client-side retry; use
+                ``client_max_retries`` to set the retry count. These reserved client keys are
+                not forwarded to the synthesizer.
             enable_word_time_offsets (:obj:`bool`, `optional`): If :obj:`True`, request per-word
                 start/end timestamps, returned in ``response.meta.words`` (supported by models that produce
                 word alignment, e.g. Magpie TTS).
@@ -196,8 +209,11 @@ def synthesize_online(
             req.zero_shot_data.encoding = audio_prompt_encoding
             req.zero_shot_data.quality = zero_shot_quality
 
-        if custom_configuration:
-            for key, value in custom_configuration.items():
+        server_configuration, auto_recover, max_retries, _ = split_recovery_configuration(
+            custom_configuration or {}
+        )
+        if server_configuration:
+            for key, value in server_configuration.items():
                 req.custom_configuration[key] = str(value)
 
         add_custom_dictionary_to_config(req, custom_dictionary)
@@ -217,4 +233,30 @@ def request_generator(text):
             else:
                 raise ValueError(f"Invalid text type: {type(text)}")
 
-        return self.stub.SynthesizeOnline(request_generator(text), metadata=self.auth.get_auth_metadata())
+        if not auto_recover:
+            return self.stub.SynthesizeOnline(request_generator(text), metadata=self.auth.get_auth_metadata())
+
+        def recovery_generator() -> Generator[rtts.SynthesizeSpeechResponse, None, None]:
+            segments = [text] if isinstance(text, str) else text
+            for segment in segments:
+                attempt = 0
+                while True:
+                    try:
+                        responses = self.stub.SynthesizeOnline(
+                            request_generator(segment), metadata=self.auth.get_auth_metadata()
+                        )
+                        # Do not expose a partial segment: retrying it must not duplicate audio.
+                        yield from list(responses)
+                        break
+                    except grpc.RpcError as exc:
+                        if not is_retryable_grpc_error(exc) or attempt >= max_retries:
+                            raise
+                        delay = exponential_backoff(attempt)
+                        LOGGER.info(
+                            "Streaming TTS connection lost (%s); retrying in %.2f seconds (attempt %d/%d).",
+                            exc.code(), delay, attempt + 1, max_retries,
+                        )
+                        attempt += 1
+                        time.sleep(delay)
+
+        return recovery_generator()

tests/unit/test_retry.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+import grpc
+import pytest
+
+from riva.client.retry import (
+    CLIENT_AUTO_RECOVER,
+    CLIENT_LOOKBACK_SECONDS,
+    CLIENT_MAX_RETRIES,
+    RETRYABLE_GRPC_CODES,
+    exponential_backoff,
+    is_retryable_grpc_error,
+    split_recovery_configuration,
+)
+
+
+class FakeRpcError(grpc.RpcError):
+    def __init__(self, code):
+        self._code = code
+
+    def code(self):
+        return self._code
+
+    def details(self):
+        return "fake error"
+
+
+class TestIsRetryableGrpcError:
+    def test_retryable_codes(self):
+        for code in RETRYABLE_GRPC_CODES:
+            exc = FakeRpcError(code)
+            assert is_retryable_grpc_error(exc) is True
+
+    def test_non_retryable_code(self):
+        exc = FakeRpcError(grpc.StatusCode.INVALID_ARGUMENT)
+        assert is_retryable_grpc_error(exc) is False
+
+    def test_no_code_method(self):
+        exc = Exception("plain exception")
+        assert is_retryable_grpc_error(exc) is False
+
+
+class TestExponentialBackoff:
+    def test_no_jitter_growth(self):
+        assert exponential_backoff(0, base_delay=1.0, jitter=False) == 1.0
+        assert exponential_backoff(1, base_delay=1.0, jitter=False) == 2.0
+        assert exponential_backoff(2, base_delay=1.0, jitter=False) == 4.0
+
+    def test_max_delay_cap(self):
+        assert exponential_backoff(10, base_delay=1.0, max_delay=8.0, jitter=False) == 8.0
+
+    def test_jitter_reduces_delay(self):
+        for _ in range(20):
+            d = exponential_backoff(2, base_delay=1.0, jitter=True)
+            assert 0.0 <= d < 4.0
+
+
+class TestRecoveryConfiguration:
+    def test_client_options_are_not_forwarded_to_riva(self):
+        server_configuration, enabled, max_retries, lookback_seconds = split_recovery_configuration({
+            "exaggeration_factor": "1.5",
+            CLIENT_AUTO_RECOVER: "true",
+            CLIENT_MAX_RETRIES: "4",
+            CLIENT_LOOKBACK_SECONDS: "3.5",
+        })
+
+        assert server_configuration == {"exaggeration_factor": "1.5"}
+        assert enabled is True
+        assert max_retries == 4
+        assert lookback_seconds == 3.5