Function async_call_with_retries

tensorboard/util/grpc_util.py:112–184

Initiate an asynchronous call to a gRPC stub, with retry logic. This is similar to the `async_call` API, except that the call is handled asynchronously, and the completion may be handled by another thread. The caller must provide a `done_callback` argument which will handle the resu

(api_method, request, clock=None)

def async_call_with_retries(api_method, request, clock=None):
    """Initiate an asynchronous call to a gRPC stub, with retry logic.

    This is similar to the `async_call` API, except that the call is handled
    asynchronously, and the completion may be handled by another thread. The
    caller must provide a `done_callback` argument which will handle the
    result or exception rising from the gRPC completion.

    Retries are handled with jittered exponential backoff to spread out failures
    due to request spikes.

    This only supports unary-unary RPCs: i.e., no streaming on either end.

    Args:
      api_method: Callable for the API method to invoke.
      request: Request protocol buffer to pass to the API method.
      clock: an interface object supporting `time()` and `sleep()` methods
        like the standard `time` module; if not passed, uses the normal module.

    Returns:
      An `AsyncCallFuture` which will encapsulate the `grpc.Future`
      corresponding to the gRPC call which either completes successfully or
      represents the final try.
    """
    if clock is None:
        clock = time
    logger.debug("Async RPC call %s with request: %r", api_method, request)

    completion_event = threading.Event()
    async_future = AsyncCallFuture(completion_event)

    def async_call(handler):
        """Invokes the gRPC future and orchestrates it via the AsyncCallFuture."""
        future = api_method.future(
            request,
            timeout=_GRPC_DEFAULT_TIMEOUT_SECS,
            metadata=version_metadata(),
        )
        # Ensure we set the active future before invoking the done callback, to
        # avoid the case where the done callback completes immediately and
        # triggers completion event while async_future still holds the old
        # future.
        async_future._set_active_future(future)
        future.add_done_callback(handler)

    # retry_handler is the continuation of the `async_call`.  It should:
    #   * If the grpc call succeeds: trigger the `completion_event`.
    #   * If there are no more retries: trigger the `completion_event`.
    #   * Otherwise, invoke a new async_call with the same
    #     retry_handler.
    def retry_handler(future, num_attempts):
        e = future.exception()
        if e is None:
            completion_event.set()
            return
        else:
            logger.info("RPC call %s got error %s", api_method, e)
            # If unable to retry, proceed to completion.