message: Gracefully handle workers without heartbeat support enabled.
diff: @@ -28,6 +28,7 @@ from tensorflow.core.protobuf import config_pb2
  from tensorflow.core.util import event_pb2
  from tensorflow.python.client import session as session_lib
  from tensorflow.python.framework import dtypes
  +from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.platform import tf_logging as logging
  @@ -78,6 +79,15 @@ class WorkerHeartbeatManager(object):
  return WorkerHeartbeatManager(session, devices, heartbeat_ops,
  request_placeholder)
  + def heartbeat_supported(self):
  + """Returns True if heartbeat operations are supported on all workers."""
  + try:
  + # Send ping to verify worker has heartbeat support.
  + self.ping()
  + return True
  + except errors.InvalidArgumentError as _:
  + return False
  +
  def configure(self, message):
  """Configure heartbeat manager for all devices.
  @@ -106,7 +116,7 @@ class WorkerHeartbeatManager(object):
  event_pb2.WorkerHeartbeatResponse.FromString(res_pb)
  for res_pb in results
  ]
  - logging.info('Results: %s', parsed_results)
  + logging.debug('Ping results: %s', parsed_results)
  return parsed_results
  def lame_workers(self):
  @@ -189,7 +199,9 @@ class WatchdogManager(threading.Thread):
  self._running = False
  self._graph = ops.Graph()
  self._session = session_lib.Session(
  - target=session.sess_str, graph=self._graph)
  + target=session.sess_str,
  + graph=self._graph,
  + )
  with self._graph.as_default():
  if devices is None:
  @@ -249,6 +261,7 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
  self._graph = ops.Graph()
  self._workers = None
  self._session = None
  + self._heartbeat_supported = False
  def after_create_session(self, training_session, coord): # pylint: disable=unused-argument
  # N.B. We have to pull the global step here to avoid it being unavailable
  @@ -264,10 +277,16 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
  target=training_session.sess_str, graph=self._graph)
  self._workers = WorkerHeartbeatManager.from_devices(
  self._session, all_worker_devices(self._session))
  -
  + self._heartbeat_supported = self._workers.heartbeat_supported()
  + if self._heartbeat_supported:
  self._workers.configure(
  event_pb2.WorkerHeartbeatRequest(
  shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR))
  + else:
  + logging.warn(
  + 'Worker heartbeats not supported by all workers. No failure '
  + 'handling will be enabled.'
  + )
  def saver(self):
  if self._saver:
  @@ -286,6 +305,9 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
  def after_run(self, run_context, run_values):
  del run_values
  + if not self._heartbeat_supported:
  + return
  +
  lame_workers = self._workers.lame_workers()
  if lame_workers:
  logging.info('ShutdownHook: lame workers found: %s', lame_workers)
  