# Lint as: python3
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Module for `CoordinatedCheckpointManager`.

This is currently under development and the API is subject to change.
"""
import os
import signal
import sys
import threading
import time

from tensorflow.python.distribute import multi_worker_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.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import checkpoint_management

_PREEMPTION_KEY = 'TERMINATED_WORKER'
_RUN_COUNT_KEY = 'RUN_TO_CHECKPOINT'
_ACKNOWLEDGE_KEY = 'RECEIVED_SIGNAL'
_ITERATION_VARIABLE = 'checkpointed_runs'
# TODO(wxinyi): Move the hardcoded 42 to an internal module.
_RESTARTABLE_EXIT_CODE = 42


def _mwms_write_checkpoint_dir(checkpoint_dir, task_type, task_id,
                               cluster_spec):
  """Returns checkpoint_dir for chief and a temp dir for any other worker."""
  dirpath = os.path.dirname(checkpoint_dir)
  base = os.path.basename(checkpoint_dir)
  if not multi_worker_util.is_chief(
      cluster_spec=cluster_spec, task_type=task_type, task_id=task_id):
    base_dirpath = 'workertemp_' + str(task_id)
    dirpath = os.path.join(dirpath, base_dirpath)
    gfile.MakeDirs(dirpath)
  return os.path.join(dirpath, base)


class CoordinatedCheckpointManager(object):
  """Preemption and error handler for synchronous training.

  The API helps coordinate all workers to save a checkpoint upon receiving a
  preemption signal and helps propagate accurate error messages during training.
  When the program recovers from preemption, the checkpoint that is passed to
  initialize a `CoordinatedCheckpointManager` object will be loaded
  automatically.

  Right after the initialization, a thread starts to watch out for a preemption
  signal for any member in the cluster, but the signal will only be handled
  (which includes aligning the step to save a checkpoint, saving a checkpoint,
  and exiting with a platform recognized restart code) after entering a
  `CoordinatedCheckpointManager.run` call.

  Example usage:
  ```python
  strategy = tf.distribute.MultiWorkerMirroredStrategy()

  with strategy.scope():
    dataset, model, optimizer = ...

    fh_checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)

    coordinated_checkpoint_manager = tf.distribute.CoordinatedCheckpointManager(
        cluster_resolver, fh_checkpoint, checkpoint_directory)


    # `coordinated_checkpoint_manager.total_runs` will be restored to its
    # checkpointed value when training is restored after interruption.
    for epoch in range(coordinated_checkpoint_manager.total_runs //
                       STEPS_PER_EPOCH, num_epochs):
      for step in range(coordinated_checkpoint_manager.total_runs %
                        STEPS_PER_EPOCH, num_steps):
        # distributed_train_step is a function wrapped by strategy.run
        loss += coordinated_checkpoint_manager.run(distributed_train_step,
                                                   args=(next(dataset),))
  ```

  `CoordinatedCheckpointManager` will create a CheckpointManager to manage the
  checkpoint and only one CheckpointManager should be active in a particular
  directory at a time. Thus, if the user would like to save a checkpoint for
  purpose other than fault tolerance, e.g., for evaluation, they should save it
  in a directory different from the one passed to a
  `CoordinatedCheckpointManager`.

  This API targets multi-client distributed training, and right now only
  `tf.distribute.MultiWorkerMirroredStrategy` is supported.
  """

  def __init__(self, cluster_resolver, checkpoint, checkpoint_dir):
    """Creates the failure handler.

    Args:
      cluster_resolver: a `tf.distribute.cluster_resolver.ClusterResolver`. You
        may also get it through the `cluster_resolver` attribute of the
        strategy in use.
      checkpoint: a `tf.train.Checkpoint` that will be saved upon preemption and
        loaded upon restart by the `CoordinatedCheckpointManager` API
        automatically.
      checkpoint_dir: a directory for the `CoordinatedCheckpointManager` to play
        with checkpoints. `CoordinatedCheckpointManager` will create a
        `tf.train.CheckpointManager` to manage the passed-in `checkpoint`. Since
        only one `tf.train.CheckpointManager` should be active in a particular