class_docstr: Preemption and error handler for synchronous training.
  
  Note: This API only supports use with
  `tf.distribute.MultiWorkerMirroredStrategy` and `tf.distribute.TPUStrategy`.
  
  A `PreemptionCheckpointHandler` coordinates all workers to save a checkpoint
  upon receiving a preemption signal. It also helps disseminate application
  error messages accurately among the cluster. When a
  `PreemptionCheckpointHandler` object is created, it restores values from
  the latest checkpoint file if any exists.
  
  Right after the initialization, the object starts to watch out for termination
  signal for any member in the cluster. If receiving a signal, the next time the
  worker executes `PreemptionCheckpointHandler.run`, the
  `PreemptionCheckpointHandler` will align all workers to save a checkpoint.
  Then, if an `exit_fn` is configured via
  `tf.distribute.experimental.TerminationConfig`, it will be invoked. Otherwise,
  the process will simply exit and later the platform should restart it.
  
  Note: We advise users of `tf.distribute.MultiWorkerMirroredStrategy` who
  choose to configure their
  own `exit_fn` in `tf.distribute.experimental.TerminationConfig` to include a
  `sys.exit(CODE_OR_MESSAGE)` in the `exit_fn` so that after the restart, all
  workers can initialize communication services correctly. For users of
  `tf.distribute.TPUStrategy`, if they do not wish to do a cluster restart but
  would like an in-process restart (i.e., keep the coordinator alive and re-do
  the steps to connect to cluster, initialize TPU system, and make the
  `TPUStrategy` object), they could configure the `exit_fn` to a no-op.
  
  For users of `tf.distribute.MultiWorkerMirroredStrategy`, the core API is
  `PreemptionCheckpointHandler.run`:
  
  ```python
  strategy = tf.distribute.MultiWorkerMirroredStrategy()
  
  trained_epoch = tf.Variable(initial_value=tf.constant(0, dtype=tf.dtypes.int64), name='epoch')
  step_in_epoch = tf.Variable(initial_value=tf.constant(0, dtype=tf.dtypes.int64), name='step_in_epoch')
  
  with strategy.scope():
    dataset, model, optimizer = ...
  
    checkpoint = tf.train.Checkpoint(optimizer=optimizer,
                                     model=model,
                                     trained_epoch=trained_epoch,
                                     step_in_epoch=step_in_epoch)
  
    preemption_checkpoint_handler = tf.distribute.experimental.PreemptionCheckpointHandler(cluster_resolver, checkpoint, checkpoint_dir)
  
  while trained_epoch.numpy() < NUM_EPOCH:
  
    while step_in_epoch.numpy() < STEPS_PER_EPOCH:
  
      # distributed_train_function contains a call to strategy.run.
      loss += preemption_checkpoint_handler.run(distributed_train_function, args=(next(iterator),))
      # For users of MultiWorkerMirroredStrategy, usually
      # STEPS_PER_TRAIN_FUNCTION = 1.
      step_in_epoch.assign_add(STEPS_PER_TRAIN_FUNCTION)
      ...
  
    epoch.assign_add(1)
    step_in_epoch.assign(0)
  ```
  
  For users of `tf.distribute.TPUStrategy`, the core APIs are
  `PreemptionCheckpointHandler.run` and
  `PreemptionCheckpointHandler.watch_preemption_scope`:
  
  ```python
  
  strategy = tf.distribute.TPUStrategy(tpu_cluster_resolver)
  
  # Rest of TPU init omitted, see documentation for TPUSTrategy.
  
  with preemption_checkpoint_handler.watch_preemption_scope():
    while trained_epoch.numpy() < NUM_EPOCH:
  
      while step_in_epoch.numpy() < STEPS_PER_EPOCH:
  
        # distributed_train_function contains a call to strategy.run.
        loss += preemption_checkpoint_handler.run(distributed_train_function, args=(next(iterator),))
  
        # For users of TPUStrategy, usually STEPS_PER_TRAIN_FUNCTION >> 1 since
        # clustering multiple steps within a tf.function amortizes the overhead
        # of launching a multi-device function on TPU Pod.
        step_in_epoch.assign_add(STEPS_PER_TRAIN_FUNCTION)
        ...
  
      epoch.assign_add(1)
      step_in_epoch.assign(0)
  ```
  
  Not all interruptions come with advance notice so that the
  `PreemptionCheckpointHandler` can handle them, e.g., those caused by hardware
  failure. For a user who saves checkpoints for these cases themselves outside
  the `PreemptionCheckpointHandler`, if they are using a
  `tf.train.CheckpointManager`, pass it as the
  `checkpoint_or_checkpoint_manager` argument to the
  `PreemptionCheckpointHandler`. If they do not have a
  `tf.train.CheckpointManager` but are directly working with
  `tf.train.Checkpoint`, we advise saving the checkpoints in the directory
  that's passed as the `checkpoint_dir` argument. In this way, at the program
  beginning, `PreemptionCheckpointHandler` can restore the latest checkpoint
  from the directory, no matter it's saved by the user themselves or saved by
  the `PreemptionCheckpointHandler` before preemption happens.
  
  **A note on the platform:**
  
  `PreemptionCheckpointHandler` can only handle the kind of termination with
  advance notice. For now, the API recognizes the termination signal for CPU,
  GPU, and TPU on Google Borg and CPU and GPU on the Google Cloud Platform. In
  these cases, `PreemptionCheckpointHandler` will automatically adopt the
  correct preemption/maintenance notification detection mechanism. Users of
  other platforms can configure a detection monitoring behavior through the
  `tf.distribute.experimental.TerminationConfig`. Customization for the exit
  behavior and grace period length could also be done here.
class_docstr_tok: ["Preemption","and","error","handler","for","synchronous","training",".","Note",":","This","API","only","supports","use","with","`","tf.distribute.MultiWorkerMirroredStrategy","`","and","`","tf.distribute.TPUStrategy","`",".","A","`","PreemptionCheckpointHandler","`","coordinates","all","workers","to","save","a","checkpoint","upon","receiving","a","preemption","signal",".","It","also","helps","disseminate","application","error","messages","accurately","among","the","cluster",".","When","a","`","PreemptionCheckpointHandler","`","object","is","created",",","it","restores","values","from","the","latest","checkpoint","file","if","any","exists",".","Right","after","the","initialization",",","the","object","starts","to","watch","out","for","termination","signal","for","any","member","in","the","cluster",".","If","receiving","a","signal",",","the","next","time","the","worker","executes","`","PreemptionCheckpointHandler.run","`",",","the","`","PreemptionCheckpointHandler","`","will","align","all","workers","to","save","a","checkpoint",".","Then",",","if","an","`","exit_fn","`","is","configured","via","`","tf.distribute.experimental.TerminationConfig","`",",","it","will","be","invoked",".","Otherwise",",","the","process","will","simply","exit","and","later","the","platform","should","restart","it",".","Note",":","We","advise","users","of","`","tf.distribute.MultiWorkerMirroredStrategy","`","who","choose","to","configure","their","own","`","exit_fn","`","in","`","tf.distribute.experimental.TerminationConfig","`","to","include","a","`","sys.exit","(","CODE_OR_MESSAGE",")","`","in","the","`","exit_fn","`","so","that","after","the","restart",",","all","workers","can","initialize","communication","services","correctly",".","For","users","of","`","tf.distribute.TPUStrategy","`",",","if","they","do","not","wish","to","do","a","cluster","restart","but","would","like","an","in-process","restart","(","i.e.",",","keep","the","coordinator","alive","and","re-do","the","steps","to","connect","to","cluster",",","initialize","TPU","system",",","and","make","the","`","TPUStrategy","`","object",")",",","they","could","configure","the","`","exit_fn","`","to","a","no-op",".","For","users","of","`","tf.distribute.MultiWorkerMirroredStrategy","`",",","the","core","API","is","`","PreemptionCheckpointHandler.run","`",":","``","`","python","strategy","=","tf.distribute.MultiWorkerMirroredStrategy","(",")","trained_epoch","=","tf.Variable","(","initial_value=tf.constant","(","0",",","dtype=tf.dtypes.int64",")",",","name='epoch","'",")","step_in_epoch","=","tf.Variable","(","initial_value=tf.constant","(","0",",","dtype=tf.dtypes.int64",")",",","name='step_in_epoch","'",")","with","strategy.scope","(",")",":","dataset",",","model",",","optimizer","=","...","checkpoint","=","tf.train.Checkpoint","(","optimizer=optimizer",",","model=model",",","trained_epoch=trained_epoch",",","step_in_epoch=step_in_epoch",")","preemption_checkpoint_handler","=","tf.distribute.experimental.PreemptionCheckpointHandler","(","cluster_resolver",",","checkpoint",",","checkpoint_dir",")","while","trained_epoch.numpy","(",")","<","NUM_EPOCH",":","while","step_in_epoch.numpy","(",")","<","STEPS_PER_EPOCH",":","#","distributed_train_function","contains","a","call","to","strategy.run",".","loss","+=","preemption_checkpoint_handler.run","(","distributed_train_function",",","args=","(","next","(","iterator",")",",",")",")","#","For","users","of","MultiWorkerMirroredStrategy",",","usually","#","STEPS_PER_TRAIN_FUNCTION","=","1.","step_in_epoch.assign_add","(","STEPS_PER_TRAIN_FUNCTION",")","...","epoch.assign_add","(","1",")","step_in_epoch.assign","(","0",")","``","`","For","users","of","`","tf.distribute.TPUStrategy","`",",","the","core","APIs","are","`","PreemptionCheckpointHandler.run","`","and","`","PreemptionCheckpointHandler.watch_preemption_scope","`",":","``","`","python","strategy","=","tf.distribute.TPUStrategy","(","tpu_cluster_resolver",")","#","Rest","of","TPU","init","omitted",",","see","documentation","for","TPUSTrategy",".","with","preemption_checkpoint_handler.watch_preemption_scope","(",")",":","while","trained_epoch.numpy","(",")","<","NUM_EPOCH",":","while","step_in_epoch.numpy","(",")","<","STEPS_PER_EPOCH",":","#","distributed_train_function","contains","a","call","to","strategy.run",".","loss","+=","preemption_checkpoint_handler.run","(","distributed_train_function",",","args=","(","next","(","iterator",")",",",")",")","#","For","users","of","TPUStrategy",",","usually","STEPS_PER_TRAIN_FUNCTION",">",">","1","since","#","clustering","multiple","steps","within","a","tf.function","amortizes","the","overhead","#","of","launching","a","multi-device","function","on","TPU","Pod",".","step_in_epoch.assign_add","(","STEPS_PER_TRAIN_FUNCTION",")","...","epoch.assign_add","(","1",")","step_in_epoch.assign","(","0",")","``","`","Not","all","interruptions","come","with","advance","notice","so","that","the","`","PreemptionCheckpointHandler","`","can","handle","them",",","e.g.",",","those","caused","by","hardware","failure",".","For","a","user","who","saves","checkpoints","for","these","cases","themselves","outside","the","`","PreemptionCheckpointHandler","`",",","if","they","are","using","a","`","tf.train.CheckpointManager","`",",","pass","it","as","the","`","checkpoint_or_checkpoint_manager","`","argument","to","the","`","PreemptionCheckpointHandler","`",".","If","they","do","not","have","a","`","tf.train.CheckpointManager","`","but","are","directly","working","with","`","tf.train.Checkpoint","`",",","we","advise","saving","the","checkpoints","in","the","directory","that","'s","passed","as","the","`","checkpoint_dir","`","argument",".","In","this","way",",","at","the","program","beginning",",","`","PreemptionCheckpointHandler","`","can","restore","the","latest","checkpoint","from","the","directory",",","no","matter","it","'s","saved","by","the","user","themselves","or","saved","by","the","`","PreemptionCheckpointHandler","`","before","preemption","happens",".","*","*","A","note","on","the","platform",":","*","*","`","PreemptionCheckpointHandler","`","can","only","handle","the","kind","of","termination","with","advance","notice",".","For","now",",","the","API","recognizes","the","termination","signal","for","CPU",",","GPU",",","and","TPU","on","Google","Borg","and","CPU","and","GPU","on","the","Google","Cloud","Platform",".","In","these","cases",",","`","PreemptionCheckpointHandler","`","will","automatically","adopt","the","correct","preemption\/maintenance","notification","detection","mechanism",".","Users","of","other","platforms","can","configure","a","detection","monitoring","behavior","through","the","`","tf.distribute.experimental.TerminationConfig","`",".","Customization","for","the","exit","behavior","and","grace","period","length","could","also","be","done","here","."]
docstr: Makes configurations and start watchers for MS, MWMS, or OneDevice.
docstr_tok: ["Makes","configurations","and","start","watchers","for","MS",",","MWMS",",","or","OneDevice","."]
code: def _initialize_for_mirrored_and_multi_worker_mirrored(self):
      """Makes configurations and start watchers for MS, MWMS, or OneDevice."""
      if (
          not self._cluster_resolver
          or not self._cluster_resolver.cluster_spec().jobs
      ):
          # For MirroredStrategy, OneDeviceStrategy, and local-mode
          # MultiWorkerMirroredStrategy, an empty cluster spec is passed, and
          # coordination service is not enabled nor is it needed (since
          # it's used for cross-worker communication). Thus we will directly name
          # the worker id and is_chief properties and also skip the
          # uploading/reading from coordination service logic.
          self._local_mode = True
          self._id_in_cluster = "single_worker"
          self._is_chief = True
      else:
          self._local_mode = False
          self._id_in_cluster = str(
              multi_worker_util.id_in_cluster(
                  self._cluster_resolver.cluster_spec(),
                  self._cluster_resolver.task_type,
                  self._cluster_resolver.task_id,
              )
          )
          self._is_chief = multi_worker_util.is_chief(
              cluster_spec=self._cluster_resolver.cluster_spec(),
              task_type=self._cluster_resolver.task_type,
              task_id=self._cluster_resolver.task_id,
          )
      # The number of calls to `PreemptionCheckpointHandler.run` when the latest
      # checkpoint was saved.
      self._checkpointed_runs = variables.Variable(
          initial_value=constant_op.constant(0, dtype=dtypes.int64),
          trainable=False,
          name=_ITERATION_VARIABLE,
      )
  
      self._maybe_create_checkpoint_manager()
  
      if not hasattr(
          self._write_checkpoint_manager._checkpoint,  # pylint: disable=protected-access
          _ITERATION_VARIABLE,
      ):
          setattr(
              self._write_checkpoint_manager._checkpoint,
              _ITERATION_VARIABLE,  # pylint: disable=protected-access
              self._checkpointed_runs,
          )
  
      if not hasattr(
          self._read_checkpoint_manager._checkpoint,  # pylint: disable=protected-access
          _ITERATION_VARIABLE,
      ):
          setattr(
              self._read_checkpoint_manager._checkpoint,
              _ITERATION_VARIABLE,  # pylint: disable=protected-access
              self._checkpointed_runs,
          )
  
      self._read_checkpoint_manager.restore_or_initialize()
  
      # grace period countdown. Set to True for all workers once they finish
      # timing saving a checkpoint. Once entering this phase, new
      # preemption/maintenance notice will not be handled, since the whole cluster
      # goes down as the worker who first initiates the grace period goes down.
      self._final_checkpoint_countdown = False
  
      self._estimated_run_time = 0
  
      # An internal step counter that's restored to checkpointed_iterations when
      # training is restored. It increments by one every time
      # `PreemptionCheckpointHandler.run` is called. Note that in this case, the
      # user must pass a single-step training function to
      # `PreemptionCheckpointHandler.run` instead of a multiple-step one.
      self._run_counter = self._checkpointed_runs.numpy()
  
      # The worker itself has received preeption signal.
      self._received_own_sigterm = threading.Event()
  
      # Some member (could be oneself) has received preemption signal, and the
      # step number to save a checkpoint has been aligned.
      self._received_checkpoint_step = threading.Event()
  
      distribute_lib.distribution_strategy_input_api_counter.get_cell(
          self._platform_device.name, "PreemptionCheckpointHandler"
      ).increase_by(1)
  
      if not self._local_mode:
          # When training is interrupted, we explicitly call the cleanup methods for
          # the thread watching for local worker's termination signal and the thread
          # watching for clusterwise information before we save a checkpoint and
          # exit. In the final chapter of the training where no interruption is
          # encountered, we rely on __del__ to clean up. However, there is no
          # guarantee when or whether __del__ is executed, thus we make the threads
          # daemon to avoid it preventing program from exit.
          self._cluster_wise_termination_watcher_thread = (
              threading.Thread(
                  target=self._watch_step_to_save_key,
                  name="PeerTerminationWatcher-%s"
                  % self._id_in_cluster,
                  daemon=True,
              )
          )
          logging.info("Start watcher for peer's signal.")
          self._cluster_wise_termination_watcher_thread.start()
  
      else:
          self._cluster_wise_termination_watcher_thread = None
  
      self._poll_termination_signal_thread = None
  
      if self._termination_watcher_fn:
          self._start_polling_for_termination_signal()
      else:
          self._start_watching_for_signal()
  
imports: import os
  import signal
  import sys
  import threading
  import time
  from tensorflow.core.distributed_runtime.preemption import gen_check_preemption_op
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_context
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import multi_worker_util
  from tensorflow.python.distribute.failure_handling import failure_handling_util
  from tensorflow.python.eager import context
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.lib.io import file_io
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import gfile
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.util import tf_contextlib
  from tensorflow.python.util.deprecation import deprecated
  from tensorflow.python.util.tf_export import tf_export
  from tensorflow.tools.docs import doc_controls
label_desc: Use image node_id 3 for calling the PreemptionCheckpointHandler obj's underlying member method code with example usage: obj._initialize_for_mirrored_and_multi_worker_mirrored() without return types