parameters: (self,
            distributed_train_function,
            *args,
            **kwargs)
return_statement: return result
docstring: Runs a training function with error and preemption handling.
  
      This function handles the preemption signal from any peer in the cluster by
      saving the training progress and exiting gracefully. (Specifically, when
      running on Borg, it exits with a special code so that the cluster
      automatically restarts the training after the down worker is back.) It will
      also propagate any program error encountered during execution of
      `distributed_train_function` to all workers so that they can raise the same
      error.
  
      The `distributed_train_function` argument should be a distributed train
      function (i.e., containing a call to `tf.distribute.Strategy.run`). For
      `tf.distribute.MultiWorkerMirroredStrategy` users, we recommend passing in a
      single-step `distributed_train_function` to
      `CoordinatedCheckpointManager.run` so that the checkpoint can be saved in
      time in case a preemption signal or maintenance notice is sent.
  
      Besides the preemption and error handling part,
      `CoordinatedCheckpointManager.run(distributed_train_function, *args,
      **kwargs)` has the same effect and output as
      `distributed_train_function(*args, **kwargs)`. `distributed_train_function`
      can return either some or no result. The following is a shortened example:
  
      ```python
  
      @tf.function
      def distributed_train_step(iterator):
        # A distributed single-step training function.
  
        def step_fn(inputs):
          # A per-replica single-step training function.
          x, y = inputs
          ...
          return loss
  
        per_replica_losses = strategy.run(step_fn, args=(next(iterator),))
        return strategy.reduce(
            tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
  
      for epoch in range(coordinated_checkpoint_manager.total_runs //
                         STEPS_PER_EPOCH, EPOCHS_TO_RUN):
        iterator = iter(multi_worker_dataset)
        total_loss = 0.0
        num_batches = 0
  
        for step in range(coordinated_checkpoint_manager.total_runs %
                          STEPS_PER_EPOCH, STEPS_PER_EPOCH):
          total_loss += coordinated_checkpoint_manager.run(distributed_train_step)
          num_batches += 1
  
        train_loss = total_loss / num_batches
        print('Epoch: %d, train_loss: %f.' %(epoch.numpy(), train_loss))
  
        train_accuracy.reset_states()
      ```
  
      Args:
        distributed_train_function: A (single-step) distributed training function.
        *args: args for `distributed_train_function`.
        **kwargs: kwargs for `distributed_train_function`.
  
      Raises:
        Program error encountered by any member in the cluster encounters one
        while executing the `distributed_train_function`, or any error from the
        program error propagation process.
  
      Returns:
        Result of running the `distributed_train_function`.
docstring_summary: Runs a training function with error and preemption handling.
docstring_tokens: Runs a training function with error and preemption handling .
function: def run(self,
            distributed_train_function,
            *args,
            **kwargs):
      """Runs a training function with error and preemption handling.
  
      This function handles the preemption signal from any peer in the cluster by
      saving the training progress and exiting gracefully. (Specifically, when
      running on Borg, it exits with a special code so that the cluster
      automatically restarts the training after the down worker is back.) It will
      also propagate any program error encountered during execution of
      `distributed_train_function` to all workers so that they can raise the same
      error.
  
      The `distributed_train_function` argument should be a distributed train
      function (i.e., containing a call to `tf.distribute.Strategy.run`). For
      `tf.distribute.MultiWorkerMirroredStrategy` users, we recommend passing in a
      single-step `distributed_train_function` to
      `CoordinatedCheckpointManager.run` so that the checkpoint can be saved in
      time in case a preemption signal or maintenance notice is sent.
  
      Besides the preemption and error handling part,
      `CoordinatedCheckpointManager.run(distributed_train_function, *args,
      **kwargs)` has the same effect and output as
      `distributed_train_function(*args, **kwargs)`. `distributed_train_function`
      can return either some or no result. The following is a shortened example:
  
      ```python
  
      @tf.function
      def distributed_train_step(iterator):
        # A distributed single-step training function.
  
        def step_fn(inputs):
          # A per-replica single-step training function.
          x, y = inputs
          ...
          return loss
  
        per_replica_losses = strategy.run(step_fn, args=(next(iterator),))
        return strategy.reduce(
            tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
  
      for epoch in range(coordinated_checkpoint_manager.total_runs //
                         STEPS_PER_EPOCH, EPOCHS_TO_RUN):
        iterator = iter(multi_worker_dataset)
        total_loss = 0.0
        num_batches = 0
  
        for step in range(coordinated_checkpoint_manager.total_runs %
                          STEPS_PER_EPOCH, STEPS_PER_EPOCH):
          total_loss += coordinated_checkpoint_manager.run(distributed_train_step)
          num_batches += 1
  
        train_loss = total_loss / num_batches
        print('Epoch: %d, train_loss: %f.' %(epoch.numpy(), train_loss))
  
        train_accuracy.reset_states()
      ```
  
      Args:
        distributed_train_function: A (single-step) distributed training function.
        *args: args for `distributed_train_function`.
        **kwargs: kwargs for `distributed_train_function`.
  
      Raises:
        Program error encountered by any member in the cluster encounters one
        while executing the `distributed_train_function`, or any error from the
        program error propagation process.
  
      Returns:
        Result of running the `distributed_train_function`.
      """
      # TODO(wxinyi): after we support use with TPUStrategy, we should expand the
      # API doc to state that `distributed_train_function` does not need to be a
      # single-step training function, since a multi-step host-training loop is
      # the dominant use case for TPU user. Besides, passing in a multi-step
      # `distributed_train_function` will require the user to track their own
      # training steps.
      try:
        self._checkpoint_if_preempted()
        result = distributed_train_function(*args, **kwargs)
        self._run_counter += 1
  
      except errors.OpError as e:
        logging.info('Propagating error to cluster: %r: %s', e, e)
        try:
          context.context().report_error_to_cluster(e.error_code, e.message)
        except Exception as ex:  # pylint: disable=broad-except
          logging.info('Ignoring error during error propagation: %r:%s', ex, ex)
        raise
  
      return result
function_tokens: def run ( self , distributed_train_function , * args , * * kwargs ) : # TODO(wxinyi): after we support use with TPUStrategy, we should expand the # API doc to state that `distributed_train_function` does not need to be a # single-step training function, since a multi-step host-training loop is # the dominant use case for TPU user. Besides, passing in a multi-step # `distributed_train_function` will require the user to track their own # training steps. try : self . _checkpoint_if_preempted ( ) result = distributed_train_function ( * args , * * kwargs ) self . _run_counter += 1 except errors . OpError as e : logging . info ( 'Propagating error to cluster: %r: %s' , e , e ) try : context . context ( ) . report_error_to_cluster ( e . error_code , e . message ) except Exception as ex : # pylint: disable=broad-except logging . info ( 'Ignoring error during error propagation: %r:%s' , ex , ex ) raise return result