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 for using the handler with TPUStrategy.
docstr_tok: ["Makes","configurations","for","using","the","handler","with","TPUStrategy","."]
code: def _initialize_for_tpu_strategy(self):
      """Makes configurations for using the handler with TPUStrategy."""
      self._is_chief = True
      self._poll_termination_signal_thread = None
      self._cluster_wise_termination_watcher_thread = None
      self._maybe_create_checkpoint_manager()
      self._read_checkpoint_manager.restore_or_initialize()
      self._run_counter = 0
  
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 2 for calling the PreemptionCheckpointHandler obj's underlying member method code with example usage: obj._initialize_for_tpu_strategy() without return types