message: Add the types of errors we observed in large-scale testing.
diff: @@ -1071,6 +1071,7 @@ class Client(object):
  yield
  except Exception as e: # pylint: disable=broad-except
  if _is_ps_failure(e):
  + logging.exception("Encountered parameter server failures!")
  raise ParameterServerFailureError(e)
  else:
  raise
  @@ -1167,7 +1168,10 @@ class _PerWorkerDistributedIterator(PerWorkerValues):
  def _is_ps_failure(error):
  """Whether the error is considered a parameter server failure."""
  - return _RPC_ERROR_FROM_PS in str(error)
  + if (_RPC_ERROR_FROM_PS in str(error) or
  + (isinstance(error, errors.InvalidArgumentError) and
  + "/job:ps" in str(error))):
  + return True
  def _is_worker_failure(error):
  @@ -1194,4 +1198,24 @@ def _is_worker_failure(error):
  # remote_handle" part.
  return True
  + # TODO(b/162541228): The following 3 types of errors are very rare and only
  + # observed in large-scale testing. The types of errors should be reduced.
  + # This error could show up when copying function inputs from remote tasks.
  + if isinstance(error, errors.InternalError):
  + if ("Failed copying input tensor" in str(error) or
  + "Unable to find a context_id" in str(error)):
  + return True
  +
  + # This could happen when the function registration fails. In the observed
  + # cases this only happens to the dataset related functions.
  + if isinstance(error, errors.NotFoundError):
  + if ("is neither a type of a primitive operation nor a name of a function "
  + "registered" in str(error)):
  + return True
  +
  + # This could happen when the iterator is no longer valid on the remote worker
  + # "Resource input tensor contains an invalid device"
  + if isinstance(error, errors.CancelledError):
  + return True
  +
  return False
  