# Copyright 2019 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.
# ==============================================================================
"""Tests Keras multi worker fault tolerance."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import json
import os
import signal
import sys
import tempfile
import threading
from absl.testing import parameterized
from tensorflow.python.distribute import collective_all_reduce_strategy as collective_strategy
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import distribute_coordinator as dc
from tensorflow.python.distribute import mirrored_strategy
from tensorflow.python.distribute import multi_worker_test_base as test_base
from tensorflow.python.keras import backend as K
from tensorflow.python.keras import callbacks
from tensorflow.python.keras.distribute import multi_worker_testing_utils
from tensorflow.python.keras.distribute import multi_worker_training_state as training_state
from tensorflow.python.platform import test


def get_strategy_object(strategy_cls):
  if strategy_cls == mirrored_strategy.MirroredStrategy:
    return strategy_cls(mirrored_strategy.all_local_devices())
  else:
    # CollectiveAllReduceStrategy and ParameterServerStrategy.
    return strategy_cls()


class KerasMultiWorkerFaultToleranceTest(test_base.IndependentWorkerTestBase,
                                         parameterized.TestCase):

  class PreemptionAtBatchBoundarySimulatingCallback(callbacks.Callback):
    """Callback to simulate preemtion at batch boundary."""

    def on_epoch_begin(self, epoch, logs=None):
      self._current_epoch = epoch

    def on_batch_begin(self, batch, logs=None):
      if self._current_epoch == 1 and batch == 1 and not test_base.is_chief():
        # Simulate preemtion at the start of second batch of second epoch.
        raise RuntimeError('Preemption!')

    def on_batch_end(self, batch, logs=None):
      assert self._current_epoch < 1 or batch < 1

    def on_epoch_end(self, epoch, logs=None):
      assert epoch < 1

  # TODO(rchao): Add tests for checking 0th and 2nd epoch boundary.
  class PreemptionAtEpochBoundarySimulatingCallback(callbacks.Callback):
    """Callback to simulate preemtion at epoch boundary."""

    def on_epoch_begin(self, epoch, logs=None):
      if epoch == 1 and not test_base.is_chief():
        # Simulate preemtion at the start of second epoch.
        raise RuntimeError('Preemption!')

    def on_epoch_end(self, epoch, logs=None):
      assert epoch < 1

  @combinations.generate(
      combinations.combine(
          # Eager runtime unfortunately cannot be tested with multi-threading.
          # TODO(rchao): Add test to use multi-process for eager mode after
          # b/132095481 is resolved.
          mode=['graph'],
          strategy_cls=[collective_strategy.CollectiveAllReduceStrategy],
          required_gpus=[0, 1],
          file_format=['h5', 'tf'],
          preemption_callback=[
              PreemptionAtEpochBoundarySimulatingCallback,
              PreemptionAtBatchBoundarySimulatingCallback
          ],
          # FT should work regardless of `ModelCheckpoint`'s parameters.
          save_weights_only=[True, False],
          load_weights_on_restart=[True, False],
      ))
  def testFaultToleranceInSyncStrategy(self, strategy_cls, file_format,
                                       preemption_callback, save_weights_only,
                                       load_weights_on_restart):
    """Test fault-tolerance with multi-threading using sync dist-strat.

    This test simulates multi-worker training that is interrupted by a
    preemption, by having two threads, each of which represents a chief and a
    non-chief worker, where the non-chief raises an error in the middle of
    training loop. Upon excepting the error, a new thread with a new cluster
    spec is created to simulate the recovered non-chief worker. Meanwhile, the
    chief worker cannot proceed and hangs since the non-chief worker has
    crashed. To simulate a restart of the chief, a new thread has been prepared
    to run to take over chief with the help of a condition variable. It is
    expected that after the restart of both chief and non-chief workers, the
    training continues from the epoch they previously failed at. The test
    concludes by verifying the preemption-interrupted training can finish with
    the same loss and accuracy had the preemption not occurred.

    TODO(rchao): Add test to check preemption on chief (possibly using multi
    processes).

    TODO(rchao): Add test to check fault-tolerance with multiple `model.fit()`.

    Arguments:
      strategy_cls: The strategy class to use.
      file_format: `h5` or `tf`.
      preemption_callback: The callback to simulate preemption.
      save_weights_only: The argument for `model.fit()`'s `save_weights_only`.
      load_weights_on_restart: The argument for `model.fit()`'s
        `load_weights_on_restart`.
    """

    def _independent_worker_fn(*args, **kwargs):  # pylint: disable=unused-argument
      with test.mock.patch.object(dc, '_run_std_server',
                                  self._make_mock_run_std_server()):
        # `before_restart` is True for the threads that represent the original
        # chief and non-chief worker, and False for threads that represent the
        # restarted chief and non-chief workers.
        before_restart = kwargs['before_restart']

        # Model building under strategy scope. Following is the code we expect
        # the user runs on every worker.
        strategy = get_strategy_object(strategy_cls)
        batch_size = 64
        steps = 3
        train_ds, _ = multi_worker_testing_utils.mnist_synthetic_dataset(
            batch_size, steps)

        with strategy.scope():
          model = multi_worker_testing_utils.get_mnist_model((28, 28, 1))

        # Function to start a new thread. This will be called twice in the
        # following code: one represents the restart of the non-chief, and one
        # represents the restart of the chief as a result of the restart of the
        # non-chief (so the training can continue in sync).
        def start_new_thread(new_chief):
          new_thread_tf_config = json.loads(os.environ['TF_CONFIG'])

          # Update the ports in new chief and new worker threads.
          new_thread_tf_config['cluster']['worker'] = kwargs['reserved_ports']

          # Since both new chief and new worker threads are started from the
          # worker thread, we need to overwrite the tf config task index.
          new_thread_tf_config['task']['index'] = 0 if new_chief else 1
          return self._run_task_in_thread(
              task_fn=_independent_worker_fn,
              cluster_spec=None,
              task_type=None,
              task_id=None,
              tf_config=new_thread_tf_config,
              before_restart=False,
              new_chief=new_chief)

        try:

          class CkptSavedEpochAssertingCallback(callbacks.Callback):

            def __init__(self, test_obj):
              super(CkptSavedEpochAssertingCallback, self).__init__()
              self.test_obj = test_obj

            def on_epoch_begin(self, epoch, logs=None):
              # `_ckpt_saved_epoch` attribute is set at the end of every epoch.
              self.test_obj.assertEqual(
                  K.eval(self.model._ckpt_saved_epoch) ==
                  training_state.CKPT_SAVED_EPOCH_UNUSED_VALUE, epoch == 0)

          callbacks_list = [
              callbacks.ModelCheckpoint(
                  filepath=saving_filepath,
                  save_weights_only=save_weights_only,
                  load_weights_on_restart=load_weights_on_restart),
              CkptSavedEpochAssertingCallback(self)
          ]
          if before_restart:
            callbacks_list.append(preemption_callback())