message: An internal change.
diff: @@ -27,7 +27,6 @@ import enum
  import functools
  import os
  import re
  -import sys
  import threading
  import time
  import weakref
  @@ -1291,20 +1290,6 @@ class ClusterCoordinator(object):
  return nest.map_structure(_maybe_fetch, val)
  -# pylint: disable=missing-function-docstring
  -@contextlib.contextmanager
  -def handle_parameter_server_failure():
  - try:
  - yield
  - except errors.UnavailableError as e: # pylint: disable=broad-except
  - restart_exit_code = os.environ.get("TF_CLIENT_NON_FATAL_RESTART_EXIT_CODE",
  - None)
  - if restart_exit_code is not None:
  - sys.exit(int(restart_exit_code))
  - else:
  - raise
  -
  -
  class _PerWorkerDistributedDataset(object):
  """Represents worker-distributed datasets created from dataset function."""
  