def response_thread_fn():
"""Consumes response iter and exposes the value on corresponding Event."""
try:
logging.debug('Response thread: blocking for next response')
for response in response_iter:
if response.WhichOneof('response') is None:
# TODO(b/175927125): We currently pass an empty response in some
# error cases and pass a GRPC error back via the ServicerContext in
# some others. Unify this error-passing.
raise execution_context.RetryableError(
'Unknown error on the service side.')
logging.debug(
'Response thread: processing response of type %s, seq_no %s',
response.WhichOneof('response'),
response.sequence_number)
# Get the corresponding response Event
response_event = self._response_event_dict[
response.sequence_number]
# Attach the response as an attribute on the Event
response_event.response = response
response_event.set()
# Set the event indicating the stream has been closed
self._stream_closed_event.set()
except Exception as error: # pylint: disable=broad-except
logging.exception('Error calling remote executor: %s', error)
if _is_retryable_grpc_error(error):
logging.exception('gRPC error is retryable')
error = execution_context.RetryableError(error)
with self._response_event_lock:
self._stream_error = error
for _, response_event in self._response_event_dict.items():
if not response_event.isSet():
response_event.response = error
response_event.set()
self._stream_closed_event.set()
def test_is_retryable_error(self):
retryable_error = execution_context.RetryableError()
self.assertTrue(execution_context._is_retryable_error(retryable_error))
self.assertFalse(execution_context._is_retryable_error(TypeError()))
self.assertFalse(execution_context._is_retryable_error(1))
self.assertFalse(execution_context._is_retryable_error('a'))
self.assertFalse(execution_context._is_retryable_error(None))
def _configure_remote_workers(num_clients, remote_executors):
""""Configures `num_clients` across `remote_executors`."""
loop, must_close_loop = _get_event_loop()
available_executors = [ex for ex in remote_executors if ex.is_ready]
logging.info('%s TFF workers available out of a total of %s.',
len(available_executors), len(remote_executors))
if not available_executors:
raise execution_context.RetryableError(
'No workers are ready; try again to reconnect.')
try:
remaining_clients = num_clients
live_workers = []
for ex_idx, ex in enumerate(available_executors):
remaining_executors = len(available_executors) - ex_idx
num_clients_to_host = remaining_clients // remaining_executors
remaining_clients -= num_clients_to_host
if num_clients_to_host > 0:
_configure_remote_executor(
ex, {placements.CLIENTS: num_clients_to_host}, loop)
live_workers.append(ex)
finally:
if must_close_loop:
loop.stop()
loop.close()
return [
_wrap_executor_in_threading_stack(e, can_resolve_references=False)
for e in live_workers
]
def response_thread_fn():
"""Consumes response iter and exposes the value on corresponding Event."""
try:
logging.debug('Response thread: blocking for next response')
for response in response_iter:
logging.debug(
'Response thread: processing response of type %s, seq_no %s',
response.WhichOneof('response'),
response.sequence_number)
# Get the corresponding response Event
response_event = self._response_event_dict[
response.sequence_number]
# Attach the response as an attribute on the Event
response_event.response = response
response_event.set()
# Set the event indicating the stream has been closed
self._stream_closed_event.set()
except grpc.RpcError as error:
logging.exception('Error calling remote executor: %s', error)
if _is_retryable_grpc_error(error): # pytype: disable=attribute-error
logging.info('gRPC error is retryable')
error = execution_context.RetryableError(error)
# Set all response events to errors. This is heavy-handed and
# potentially can be scaled back.
for _, response_event in self._response_event_dict.items():
if not response_event.isSet():
response_event.response = error
response_event.set()
self._stream_closed_event.set()
def response_thread_fn():
"""Consumes response iter and exposes the value on corresponding Event."""
try:
logging.debug('Response thread: blocking for next response')
for response in response_iter:
logging.debug(
'Response thread: processing response of type %s, seq_no %s',
response.WhichOneof('response'),
response.sequence_number)
# Get the corresponding response Event
response_event = self._response_event_dict[
response.sequence_number]
# Attach the response as an attribute on the Event
response_event.response = response
response_event.set()
# Set the event indicating the stream has been closed
self._stream_closed_event.set()
except Exception as error: # pylint: disable=broad-except
logging.exception('Error calling remote executor: %s', error)
if _is_retryable_grpc_error(error):
logging.exception('gRPC error is retryable')
error = execution_context.RetryableError(error)
with self._response_event_lock:
self._stream_error = error
for _, response_event in self._response_event_dict.items():
if not response_event.isSet():
response_event.response = error
response_event.set()
self._stream_closed_event.set()
def wrapped_fn():
try:
return func()
except grpc.RpcError as e:
if _is_retryable_grpc_error(e):
logging.info('Received retryable gRPC error: %s', e)
raise execution_context.RetryableError(e)
else:
raise e
def _request(rpc_func, request):
with tracing.wrap_rpc_in_trace_context():
try:
return rpc_func(request)
except grpc.RpcError as e:
if _is_retryable_grpc_error(e):
logging.info('Received retryable gRPC error: %s', e)
raise execution_context.RetryableError(e)
else:
raise
