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: Checkpoint if any worker has received a preemption signal.
  
  This function handles preemption signal reported by any worker in the
  cluster. The current implementation relies on the fact that all workers in a
  MultiWorkerMirroredStrategy training cluster have a step number difference
  maximum of 1.
  - If the signal comes from the worker itself (i.e., where this failure
  handler sits), the worker will notify all peers to checkpoint after they
  finish CURRENT_STEP+1 steps, where CURRENT_STEP is the step this worker has
  just finished. And the worker will wait for all peers to acknowledge that
  they have received its preemption signal and the final-step number before
  the worker proceeds on training the final step.
  - If the signal comes from another member in the cluster but NO final-step
  info is available, proceed on training, because it will be available after
  finishing the next step.
  - If the signal comes from some other member in the cluster, and final-step
  info is available, if the worker has not finished these steps yet, keep
  training; otherwise, checkpoint and exit with a cluster-recognized restart
  code.
  
  Args:
    *args: args for `tf.train.CheckpointManager.save()` to save checkpoint.
    **kwargs: kwargs for `tf.train.CheckpointManager.save()` to save.
docstr_tok: ["Checkpoint","if","any","worker","has","received","a","preemption","signal",".","This","function","handles","preemption","signal","reported","by","any","worker","in","the","cluster",".","The","current","implementation","relies","on","the","fact","that","all","workers","in","a","MultiWorkerMirroredStrategy","training","cluster","have","a","step","number","difference","maximum","of","1",".","-","If","the","signal","comes","from","the","worker","itself","(","i.e.",",","where","this","failure","handler","sits",")",",","the","worker","will","notify","all","peers","to","checkpoint","after","they","finish","CURRENT_STEP+1","steps",",","where","CURRENT_STEP","is","the","step","this","worker","has","just","finished",".","And","the","worker","will","wait","for","all","peers","to","acknowledge","that","they","have","received","its","preemption","signal","and","the","final-step","number","before","the","worker","proceeds","on","training","the","final","step",".","-","If","the","signal","comes","from","another","member","in","the","cluster","but","NO","final-step","info","is","available",",","proceed","on","training",",","because","it","will","be","available","after","finishing","the","next","step",".","-","If","the","signal","comes","from","some","other","member","in","the","cluster",",","and","final-step","info","is","available",",","if","the","worker","has","not","finished","these","steps","yet",",","keep","training",";","otherwise",",","checkpoint","and","exit","with","a","cluster-recognized","restart","code",".","Args",":","*","args",":","args","for","`","tf.train.CheckpointManager.save","(",")","`","to","save","checkpoint",".","*","*","kwargs",":","kwargs","for","`","tf.train.CheckpointManager.save","(",")","`","to","save","."]
code: def _check_preemption_and_maybe_checkpoint(self, *args, **kwargs):
      """Checkpoint if any worker has received a preemption signal.
  
      This function handles preemption signal reported by any worker in the
      cluster. The current implementation relies on the fact that all workers in a
      MultiWorkerMirroredStrategy training cluster have a step number difference
      maximum of 1.
      - If the signal comes from the worker itself (i.e., where this failure
      handler sits), the worker will notify all peers to checkpoint after they
      finish CURRENT_STEP+1 steps, where CURRENT_STEP is the step this worker has
      just finished. And the worker will wait for all peers to acknowledge that
      they have received its preemption signal and the final-step number before
      the worker proceeds on training the final step.
      - If the signal comes from another member in the cluster but NO final-step
      info is available, proceed on training, because it will be available after
      finishing the next step.
      - If the signal comes from some other member in the cluster, and final-step
      info is available, if the worker has not finished these steps yet, keep
      training; otherwise, checkpoint and exit with a cluster-recognized restart
      code.
  
      Args:
        *args: args for `tf.train.CheckpointManager.save()` to save checkpoint.
        **kwargs: kwargs for `tf.train.CheckpointManager.save()` to save.
      """
      if (
          self._platform_device
          == failure_handling_util.PlatformDevice.INTERNAL_TPU
      ):
          gen_check_preemption_op.check_preemption(
              preemption_key=PREEMPTION_KEY
          )
          return
  
      if self._final_checkpoint_countdown:
          run_count_config_key = _FINAL_RUN_COUNT_KEY
  
      else:
          run_count_config_key = _INITIAL_RUN_COUNT_KEY
  
      if self._received_checkpoint_step.is_set():
          if self._step_to_checkpoint == str(self._run_counter):
              self._save_checkpoint(*args, **kwargs)
  
              if self._time_to_exit():
                  self._stop_poll_termination_signal_thread()
                  self._stop_cluster_wise_termination_watcher_thread()
                  if (
                      self._api_made_checkpoint_manager
                      and not self._is_chief
                  ):
                      gfile.DeleteRecursively(
                          os.path.dirname(
                              self._write_checkpoint_manager.directory
                          )
                      )
                  logging.info(
                      "PreemptionCheckpointHandler: checkpoint saved. Exiting."
                  )
  
                  self._exit_fn()
  
              else:
                  logging.info(
                      "Continue training for the grace period."
                  )
                  self._final_checkpoint_countdown = True
                  self._received_checkpoint_step.clear()
  
      elif self._received_own_sigterm.is_set():
          # Only the worker who gets termination signal first among the cluster
          # will enter this branch. The following will happen in chronological
          # order:
          # 1. The worker just receives a preemption signal and enters this branch
          # for the first time. It will set a step-to-checkpoint and let the cluster
          # know.
          # 2. If there is a long grace period, it will also set
          # _final_checkpoint_countdown, so that during this grace period, it will
          # re-enter this branch to check if grace period is ending.
          # 3. If it is, set a step-to-checkpoint key again.
  
          if self._final_checkpoint_countdown:
              if self._target_time_for_termination < time.time():
                  logging.info(
                      "Grace period almost ended. Final call to save a checkpoint!"
                  )
              else:
                  return
  
          step_to_save_at = str(self._run_counter + 1)
  
          logging.info(
              "Termination caught in main thread on preempted worker"
          )
  
          if self._local_mode:
              self._step_to_checkpoint = step_to_save_at
              self._received_checkpoint_step.set()
  
          else:
              context.context().set_config_key_value(
                  run_count_config_key, step_to_save_at
              )
              logging.info(
                  "%s set to %s", run_count_config_key, step_to_save_at
              )
  
              if not self._local_mode:
                  worker_count = multi_worker_util.worker_count(
                      self._cluster_resolver.cluster_spec(),
                      self._cluster_resolver.task_type,
                  )
                  for i in range(worker_count):
                      context.context().get_config_key_value(
                          f"{_ACKNOWLEDGE_KEY}_{run_count_config_key}_{i}"
                      )
                      logging.info(
                          "Sigterm acknowledgement from replica %d received",
                          i,
                      )
  
          self._setup_countdown_if_has_grace_period_and_not_already_counting_down()
  
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 19 for calling the PreemptionCheckpointHandler obj's underlying member method code with example usage: obj._check_preemption_and_maybe_checkpoint() without return types