diff --git a/riva/client/asr.py b/riva/client/asr.py index be9a362..82eb43c 100644 --- a/riva/client/asr.py +++ b/riva/client/asr.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: MIT import io +import logging import os import sys import time @@ -11,8 +12,10 @@ 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 @@ -20,6 +23,10 @@ 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,6 +439,11 @@ 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 @@ -439,9 +451,55 @@ def streaming_response_generator( message `here `_. """ - 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 diff --git a/riva/client/retry.py b/riva/client/retry.py new file mode 100644 index 0000000..15f1f4a --- /dev/null +++ b/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 diff --git a/riva/client/tts.py b/riva/client/tts.py index c057956..f43c85f 100644 --- a/riva/client/tts.py +++ b/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()) \ No newline at end of file + 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() diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py new file mode 100644 index 0000000..9a1a4c0 --- /dev/null +++ b/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