message: Add retry in MWMS check health
  in case check health fails due to transient issues
diff: @@ -193,6 +193,8 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
  # TODO(b/151232436): now the inital barrier may hang in a rare case, so we
  # need a finite timeout.
  _check_health_initial_timeout = 1200
  + # Times to retry before considering the peer is down.
  + _check_health_retry_limit = 3
  def __init__(self,
  container_strategy,
  @@ -639,23 +641,39 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
  while True:
  if self._check_health_thread_should_stop.is_set():
  return
  - try:
  for job in self._cluster_spec.jobs:
  for task_id in range(self._cluster_spec.num_tasks(job)):
  - context.context().check_collective_ops_peer_health(
  - "/job:{}/replica:0/task:{}".format(job, task_id))
  - except (errors.UnavailableError, errors.FailedPreconditionError) as e:
  - # TODO(b/151232436): Always raise UnavailableError when a peer fails.
  - # Now there could be many kinds of errors:
  + peer = "/job:{}/replica:0/task:{}".format(job, task_id)
  + attempts = 0
  + while True:
  + attempts += 1
  + try:
  + context.context().check_collective_ops_peer_health(peer)
  + # If check_collective_ops_peer_health doesn't raise an Exception,
  + # the peer is healthy.
  + break
  + except (errors.UnavailableError,
  + errors.FailedPreconditionError) as e:
  + # TODO(b/151232436): Always raise UnavailableError when a peer
  + # fails. Now there could be many kinds of errors:
  # - Unavailable: when the peer is not reachable, e.g. it's down.
  # - FailedPrecondition: when the peer has restarted.
  - logging.error("Cluster check alive failed, aborting collectives")
  + if attempts < self._check_health_retry_limit:
  + logging.warning("%s seems down, retrying %d/%d", peer, attempts,
  + self._check_health_retry_limit)
  + continue
  + logging.error(
  + "Cluster check alive failed, %s is down, "
  + "aborting collectives: %s", peer, e)
  context.context().abort_collective_ops(
  - errors.UNAVAILABLE, "cluster check alive failed: %s" % e)
  + errors.UNAVAILABLE,
  + "cluster check alive failed, {} is down".format(peer))
  + return
  except Exception as e: # pylint: disable=broad-except
  - logging.exception("Unexpected exception in check alive.")
  + logging.error("Unexpected exception in check alive: %s", e)
  context.context().abort_collective_ops(
  - errors.INTERNAL, "unexecpted exception in check alive: %s" % e)
  + errors.INTERNAL,
  + "unexecpted exception in check alive: %s" % e)
  return
  time.sleep(self._check_health_interval)
  