# Copyright 2018 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 for CollectiveAllReduceStrategy."""

import copy
import functools

from absl.testing import parameterized
import numpy as np

from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import collective_all_reduce_strategy
from tensorflow.python.distribute import collective_util
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import distribute_utils
from tensorflow.python.distribute import multi_worker_test_base
from tensorflow.python.distribute import multi_worker_util
from tensorflow.python.distribute import reduce_util
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.distribute import strategy_test_lib
from tensorflow.python.distribute import test_util
from tensorflow.python.distribute.cluster_resolver import cluster_resolver as cluster_resolver_lib
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
from tensorflow.python.distribute.v1 import input_lib as input_lib_v1
from tensorflow.python.eager import context
from tensorflow.python.framework import config as tf_config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import device as tf_device
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import gradients
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import init_ops_v2
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib


CollectiveAllReduceStrategy = (
    collective_all_reduce_strategy.CollectiveAllReduceStrategy)
CollectiveAllReduceExtended = (
    collective_all_reduce_strategy.CollectiveAllReduceExtended)
_CollectiveAllReduceStrategyExperimental = (
    collective_all_reduce_strategy._CollectiveAllReduceStrategyExperimental)


# TODO(b/231630416): Create more tests to cover the case that strategy uses
# different number of GPUs than the number of physical devices.
def create_test_objects(cluster_spec=None,
                        task_type=None,
                        task_id=None,
                        num_gpus=None,
                        num_tpus=None):
  if num_gpus is None:
    num_gpus = context.num_gpus()
  if num_tpus is None:
    num_tpus = context.context().list_physical_devices('TPU')
  if num_tpus:
    tpu_cluster_resolver.initialize_tpu_system()

  if cluster_spec and task_type and task_id is not None:
    cluster_resolver = cluster_resolver_lib.SimpleClusterResolver(
        cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec),
        task_type=task_type,
        task_id=task_id,
        num_accelerators={'GPU': num_gpus, 'TPU': num_tpus})
    target = 'grpc://' + cluster_spec[task_type][task_id]
  else:
    cluster_resolver = cluster_resolver_lib.SimpleClusterResolver(
        server_lib.ClusterSpec({}),
        num_accelerators={'GPU': num_gpus, 'TPU': num_tpus},
    )
    target = ''

  strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy(
      cluster_resolver=cluster_resolver)

  return strategy, target


class CollectiveAllReduceStrategyTestBase(
    multi_worker_test_base.MultiWorkerTestBase):

  def setUp(self):
    # We use a different key_base for each test so that collective keys won't be
    # reused.
    CollectiveAllReduceStrategy._collective_key_base += 100000
    super(CollectiveAllReduceStrategyTestBase, self).setUp()

  def _get_test_object(self,
                       task_type,
                       task_id,
                       num_gpus=0,
                       num_tpus=0,
                       use_devices_arg=False):
    strategy, target = create_test_objects(
        cluster_spec=self._cluster_spec,
        task_type=task_type,
        task_id=task_id,
        num_gpus=num_gpus,
        num_tpus=num_tpus)

    if use_devices_arg and num_gpus > 0:
      devices = ['GPU:%d' % i for i in range(num_gpus)]
      # Temporary workaround to manually set the `_extended` field before device
      # initialization is exposed as a public interface.
      strategy._extended = CollectiveAllReduceExtended(
          container_strategy=strategy,
          cluster_resolver=None,
          communication_options=collective_util.Options(),
          devices=devices)
      # Manually set the field since the workaround bypasses the base
      # contructor, resulting in the absence of this field.
      strategy._extended._retrace_functions_for_each_device = (num_gpus > 1)

    return strategy, target

  def _test_minimize_loss_graph(self, task_type, task_id, num_gpus):
    distribution, master_target = self._get_test_object(task_type, task_id,
                                                        num_gpus)
    with ops.Graph().as_default(), \
         self.cached_session(target=master_target) as sess, \
         distribution.scope():
      initializer = functools.partial(
          init_ops_v2.GlorotUniform(), (1, 1), dtype=dtypes.float32)
      kernel = variables.Variable(
          initial_value=initializer,
          name='gpu_%d/kernel' % distribution.extended._num_devices_per_worker,
          trainable=True)

      def loss_fn(x):
        y = array_ops.reshape(
            gen_math_ops.mat_mul(x, kernel), []) - constant_op.constant(1.)
        return y * y

      # TODO(yuefengz, apassos): eager.backprop.implicit_grad is not safe for
      # multiple graphs (b/111216820).
      def grad_fn(x):
        loss = loss_fn(x)
        var_list = (
            variables.trainable_variables() + ops.get_collection(
                ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
        grads = gradients.gradients(loss, var_list)
        ret = list(zip(grads, var_list))
        return ret

      def update(v, g):
        return v.assign_sub(0.05 * g, use_locking=True)

      one = constant_op.constant([[1.]])

      def step():
        """Perform one optimization step."""
        # Run forward & backward to get gradients, variables list.
        g_v = distribution.extended.call_for_each_replica(grad_fn, args=[one])
        # Update the variables using the gradients and the update() function.
        before_list = []
        after_list = []
        for g, v in g_v:
          fetched = distribution.extended.read_var(v)
          before_list.append(fetched)
          with ops.control_dependencies([fetched]):
            # TODO(yuefengz): support non-Mirrored variable as destinations.
            g = distribution.extended.reduce_to(
                reduce_util.ReduceOp.SUM, g, destinations=v)
            with ops.control_dependencies(
                distribution.extended.update(v, update, args=(g,),
                                             group=False)):
              after_list.append(distribution.extended.read_var(v))
        return before_list, after_list

      before_out, after_out = step()

      if (distribution.extended._local_device_type == 'GPU' and
          context.num_gpus() < distribution.extended._num_devices_per_worker):
        return True

      sess.run(variables.global_variables_initializer())

      for i in range(10):
        b, a = sess.run((before_out, after_out))
        if i == 0:
          before, = b
        after, = a

      error_before = abs(before - 1)
      error_after = abs(after - 1)
      # Error should go down
      self.assertLess(error_after, error_before)

  def _test_variable_initialization(self, task_type, task_id, num_gpus):
    distribution, master_target = self._get_test_object(task_type, task_id,
                                                        num_gpus)
    with ops.Graph().as_default(), \
         self.cached_session(target=master_target) as sess, \
         distribution.scope():

      def model_fn():
        x = variable_scope.get_variable(
            'x',
            shape=(2, 3),
            initializer=init_ops.random_uniform_initializer(
                1.0, 10.0, dtype=dtypes.float32))
        return array_ops.identity(x)

      x = distribution.extended.call_for_each_replica(model_fn)
      reduced_x = distribution.reduce(reduce_util.ReduceOp.MEAN, x, axis=None)
      x = distribution.experimental_local_results(x)[0]

      sess.run(variables.global_variables_initializer())

      x_value, reduced_x_value = sess.run([x, reduced_x])
      self.assertTrue(
          np.allclose(x_value, reduced_x_value, atol=1e-5),
          msg=('x_value = %r, reduced_x_value = %r' % (x_value,
                                                       reduced_x_value)))

  def _test_input_fn_iterator(self,
                              task_type,
                              task_id,
                              num_gpus,
                              input_fn,
                              expected_values,
                              test_reinitialize=True,
                              ignore_order=False,
                              use_devices_arg=False):
    distribution, master_target = self._get_test_object(
        task_type, task_id, num_gpus, use_devices_arg=use_devices_arg)
    devices = distribution.extended.worker_devices

    with ops.Graph().as_default(), \
         self.cached_session(target=master_target) as sess:
      iterator = distribution.make_input_fn_iterator(input_fn)
      sess.run(iterator.initializer)

      for expected_value in expected_values:
        next_element = iterator.get_next()
        computed_value = sess.run([distribute_utils.select_replica(
            r, next_element) for r in range(len(devices))])
        if ignore_order:
          self.assertCountEqual(list(expected_value), list(computed_value))
        else:
          self.assertEqual(list(expected_value), list(computed_value))

      with self.assertRaises(errors.OutOfRangeError):
        next_element = iterator.get_next()
        sess.run([distribute_utils.select_replica(r, next_element)
                  for r in range(len(devices))])

      # After re-initializing the iterator, should be able to iterate again.
      if test_reinitialize:
        sess.run(iterator.initializer)

        for expected_value in expected_values:
          next_element = iterator.get_next()
          computed_value = sess.run([
              distribute_utils.select_replica(r, next_element)
              for r in range(len(devices))])
          if ignore_order:
            self.assertCountEqual(list(expected_value), list(computed_value))
          else:
            self.assertEqual(list(expected_value), list(computed_value))


class DistributedCollectiveAllReduceStrategyTest(
    CollectiveAllReduceStrategyTestBase,
    strategy_test_lib.DistributionTestBase,
    parameterized.TestCase):

  @classmethod
  def setUpClass(cls):
    """Create a local cluster with 3 workers."""
    cls._cluster_spec = multi_worker_test_base.create_in_process_cluster(
        num_workers=3, num_ps=0)

  @combinations.generate(combinations.combine(mode=['graph']))
  def test_num_replicas_in_sync(self):
    distribution, _ = create_test_objects(
        cluster_spec=self._cluster_spec,
        task_type='worker',
        task_id=0,
        num_gpus=2)
    num_workers = len(self._cluster_spec.get('chief', []) +
                      self._cluster_spec.get('worker', []))
    self.assertEqual(2 * num_workers,
                     distribution.num_replicas_in_sync)

  @combinations.generate(combinations.combine(
      mode=['graph'],
      prefetch_to_device=[None, True]))
  def test_prefetch_to_device_dataset(self, prefetch_to_device):
    distribution, _ = self._get_test_object(
        task_type='worker', task_id=0, num_gpus=2)
    if prefetch_to_device is None:
      input_options = None
    else:
      input_options = distribute_lib.InputOptions(
          experimental_fetch_to_device=prefetch_to_device)
    dataset = dataset_ops.Dataset.range(100)
    dataset = dataset.batch(distribution.num_replicas_in_sync)
    dataset = distribution.experimental_distribute_dataset(
        dataset, options=input_options)
    if isinstance(dataset, input_lib_v1.DistributedDatasetV1):
      item = dataset.make_initializable_iterator().get_next()
    else:
      self.skipTest('unsupported test combination')
    device_types = {
        tf_device.DeviceSpec.from_string(tensor.device).device_type for
        tensor in item.values}
    self.assertAllEqual(list(device_types), ['GPU'])

  @combinations.generate(combinations.combine(mode=['graph']))
  def test_prefetch_to_host_dataset(self):
    distribution, _ = self._get_test_object(
        task_type='worker', task_id=0, num_gpus=2)
    input_options = distribute_lib.InputOptions(
        experimental_fetch_to_device=False)
    dataset = dataset_ops.Dataset.range(100)
    dataset = dataset.batch(distribution.num_replicas_in_sync)
    dataset = distribution.experimental_distribute_dataset(
        dataset, options=input_options)
    if isinstance(dataset, input_lib_v1.DistributedDatasetV1):
      item = dataset.make_initializable_iterator().get_next()
    else:
      self.skipTest('unsupported test combination')
    device_types = {
        tf_device.DeviceSpec.from_string(tensor.device).device_type for
        tensor in item.values}
    self.assertAllEqual(list(device_types), ['CPU'])

  @combinations.generate(
      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
  def testMinimizeLossGraph(self, required_gpus):
    self._run_between_graph_clients(self._test_minimize_loss_graph,
                                    self._cluster_spec, required_gpus)

  @combinations.generate(
      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
  def testVariableInitialization(self, required_gpus):
    self._run_between_graph_clients(
        self._test_variable_initialization,
        self._cluster_spec,
        num_gpus=required_gpus)

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=[0, 1, 2], use_dataset=[True, False]))
  def testMakeInputFnIterator(self, required_gpus, use_dataset):
    def _worker_fn(task_type, task_id, required_gpus):
      if use_dataset:
        fn = lambda: dataset_ops.Dataset.range(20)
      else:
        def fn():
          dataset = dataset_ops.Dataset.range(20)
          it = dataset_ops.make_one_shot_iterator(dataset)
          return it.get_next
      # We use CPU as the device when required_gpus = 0
      devices_per_worker = max(1, required_gpus)
      expected_values = [[i+j for j in range(devices_per_worker)]
                         for i in range(0, 20, devices_per_worker)]

      input_fn = self._input_fn_to_test_input_context(
          fn,
          expected_num_replicas_in_sync=3*devices_per_worker,
          expected_num_input_pipelines=3,
          expected_input_pipeline_id=task_id)
      self._test_input_fn_iterator(
          task_type,
          task_id,
          required_gpus,
          input_fn,
          expected_values,
          test_reinitialize=use_dataset,
          ignore_order=not use_dataset)

    self._run_between_graph_clients(_worker_fn, self._cluster_spec,
                                    required_gpus)

  @combinations.generate(combinations.combine(mode=['graph']))
  def testUpdateConfigProto(self):
    strategy, _ = self._get_test_object(
        task_type='worker', task_id=1, num_gpus=2)

    config_proto = config_pb2.ConfigProto(device_filters=['to_be_overridden'])
    rewrite_options = config_proto.graph_options.rewrite_options
    rewrite_options.scoped_allocator_opts.enable_op.append('to_be_removed')

    new_config = strategy.update_config_proto(config_proto)

    # Verify group leader
    self.assertEqual('/job:worker/replica:0/task:0',
                     new_config.experimental.collective_group_leader)

    # Verify device filters.
    self.assertEqual(['/job:worker/task:1'], new_config.device_filters)

    # Verify rewrite options.
    new_rewrite_options = new_config.graph_options.rewrite_options
    self.assertEqual(rewriter_config_pb2.RewriterConfig.ON,
                     new_rewrite_options.scoped_allocator_optimization)
    self.assertEqual(['CollectiveReduce'],
                     new_rewrite_options.scoped_allocator_opts.enable_op)


class DistributedCollectiveAllReduceStrategyTestWithChief(
    CollectiveAllReduceStrategyTestBase, parameterized.TestCase):

  @classmethod
  def setUpClass(cls):
    """Create a local cluster with 3 workers and 1 chief."""
    cls._cluster_spec = multi_worker_test_base.create_in_process_cluster(
        num_workers=3, num_ps=0, has_chief=True)

  @combinations.generate(
      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
  def testMinimizeLossGraph(self, required_gpus):
    self._run_between_graph_clients(self._test_minimize_loss_graph,
                                    self._cluster_spec, required_gpus)

  @combinations.generate(
      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
  def testVariableInitialization(self, required_gpus):
    self._run_between_graph_clients(
        self._test_variable_initialization,
        self._cluster_spec,
        num_gpus=required_gpus)


class SingleWorkerCollectiveAllReduceStrategy(
    CollectiveAllReduceStrategyTestBase, strategy_test_lib.DistributionTestBase,
    strategy_test_lib.TwoDeviceDistributionTestBase, parameterized.TestCase):

  @combinations.generate(combinations.combine(mode=['eager']))
  def testStrategyInitializationError(self):
    with self.assertRaisesRegex(
        ValueError,
        'cluster_resolver and devices cannot be set at the same time'):
      _ = collective_all_reduce_strategy.CollectiveAllReduceExtended(
          container_strategy=None,
          cluster_resolver=multi_worker_test_base.create_in_process_cluster(
              num_workers=3, num_ps=0),
          communication_options=collective_util.Options(),
          devices=['GPU:0', 'GPU:1'])

  @combinations.generate(
      combinations.combine(
          mode=['graph', 'eager'],
          required_gpus=[0, 1, 2],
          use_devices_arg=[True, False]))
  def testMinimizeLoss(self, required_gpus, use_devices_arg):
    # Collective ops doesn't support strategy with one device.
    if context.executing_eagerly():
      strategy, _ = self._get_test_object(
          None, None, required_gpus, use_devices_arg=use_devices_arg)
      self._test_minimize_loss_eager(strategy)
    else:
      self._test_minimize_loss_graph(None, None, required_gpus)

  @combinations.generate(
      combinations.combine(
          mode=['eager'], required_gpus=[1, 2], use_devices_arg=[True, False]))
  def testNumReplicasInSync(self, required_gpus, use_devices_arg):
    strategy, _ = self._get_test_object(
        None, None, required_gpus, use_devices_arg=use_devices_arg)
    self.assertEqual(required_gpus, strategy.num_replicas_in_sync)

  @combinations.generate(
      combinations.combine(
          mode=['eager'],
          required_tpus=[0, 1, 2],
          use_devices_arg=[True, False]))
  def testMinimizeLossTPU(self, required_tpus, use_devices_arg):
    strategy, _ = self._get_test_object(
        None, None, num_tpus=required_tpus, use_devices_arg=use_devices_arg)
    self._test_minimize_loss_eager(strategy)

  @combinations.generate(
      combinations.combine(
          mode=['graph', 'eager'],
          required_gpus=[0, 1, 2],
          use_devices_arg=[True, False]))
  def testCallAndMergeExceptions(self, required_gpus, use_devices_arg):
    strategy, _ = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    self._test_call_and_merge_exceptions(strategy)

  @combinations.generate(
      combinations.combine(
          mode=['graph'],
          required_gpus=2,
          use_dataset=[True, False],
          use_devices_arg=[True, False]))
  def testMakeInputFnIterator(self, required_gpus, use_dataset,
                              use_devices_arg):
    if use_dataset:
      fn = lambda: dataset_ops.Dataset.range(5 * required_gpus)
    else:
      def fn():
        dataset = dataset_ops.Dataset.range(5 * required_gpus)
        it = dataset_ops.make_one_shot_iterator(dataset)
        return it.get_next

    expected_values = [
        range(i, i + required_gpus) for i in range(0, 10, required_gpus)
    ]

    input_fn = self._input_fn_to_test_input_context(
        fn,
        expected_num_replicas_in_sync=required_gpus,
        expected_num_input_pipelines=1,
        expected_input_pipeline_id=0)
    self._test_input_fn_iterator(
        None,
        None,
        required_gpus,
        input_fn,
        expected_values,
        test_reinitialize=use_dataset,
        ignore_order=not use_dataset)

  @combinations.generate(
      combinations.combine(
          mode=['graph', 'eager'],
          required_gpus=[0, 1, 2],
          use_devices_arg=[True, False]))
  def testReduceToCpu(self, required_gpus, use_devices_arg):
    strategy, _ = self._get_test_object(
        None, None, required_gpus, use_devices_arg=use_devices_arg)
    with strategy.scope():
      result = strategy.extended.call_for_each_replica(_replica_id_f32)
      reduced = strategy.reduce(reduce_util.ReduceOp.SUM, result, axis=None)
      expected = sum(range(strategy.num_replicas_in_sync))
      self.assertEqual(expected, self.evaluate(reduced))

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
  def testAllReduceSum(self, required_gpus, use_devices_arg):
    distribution, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    with self.cached_session(target=target):
      self._test_all_reduce_sum(distribution)

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
  def testAllReduceSumGradients(self, required_gpus, use_devices_arg):
    distribution, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    with self.cached_session(target=target):
      self._test_all_reduce_sum_gradients(distribution)

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
  def testAllReduceSumGradientTape(self, required_gpus, use_devices_arg):
    distribution, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    with self.cached_session(target=target):
      self._test_all_reduce_sum_gradient_tape(distribution)

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
  def testAllReduceMean(self, required_gpus, use_devices_arg):
    distribution, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    with self.cached_session(target=target):
      self._test_all_reduce_mean(distribution)

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
  def testAllReduceMeanGradients(self, required_gpus, use_devices_arg):
    distribution, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    with self.cached_session(target=target):
      self._test_all_reduce_mean_gradients(distribution)

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
  def testAllReduceMeanGradientTape(self, required_gpus, use_devices_arg):
    distribution, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    with self.cached_session(target=target):
      self._test_all_reduce_mean_gradient_tape(distribution)

  @combinations.generate(
      combinations.combine(
          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
  def testNumpyDataset(self, required_gpus, use_devices_arg):
    strategy, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    self._test_numpy_dataset(
        strategy, session=self.cached_session(target=target))

  @combinations.generate(
      combinations.combine(
          mode=['eager'], required_gpus=2, use_devices_arg=[True, False]))
  def testReplicateDataset(self, required_gpus, use_devices_arg):
    strategy, _ = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    dataset_fn = lambda: dataset_ops.Dataset.range(10)
    expected_values = [[i, i + 1] for i in range(0, 10, 2)]
    input_fn = self._input_fn_to_test_input_context(
        dataset_fn,
        expected_num_replicas_in_sync=required_gpus,
        expected_num_input_pipelines=1,
        expected_input_pipeline_id=0)
    self._test_input_fn_iterable(strategy, input_fn, expected_values)

  @combinations.generate(
      combinations.combine(mode=['graph'], use_devices_arg=[True, False]))
  def testDeepCopy(self, use_devices_arg):
    distribution, _ = self._get_test_object(
        None, None, use_devices_arg=use_devices_arg)
    copy.deepcopy(distribution)

  @combinations.generate(
      combinations.combine(
          mode=['graph', 'eager'],
          required_gpus=[0, 1, 2],
          use_devices_arg=[True, False]))
  def testSummaryForReplicaZeroOnly(self, required_gpus, use_devices_arg):
    strategy, target = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    with self.cached_session(target=target):
      self._test_summary_for_replica_zero_only(strategy)

  @combinations.generate(
      combinations.combine(
          mode=['graph', 'eager'],
          required_gpus=[0, 1, 2],
          use_devices_arg=[True, False]))
  def testTrainableVariables(self, required_gpus, use_devices_arg):
    strategy, _ = self._get_test_object(
        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
    self._test_trainable_variable(strategy)


class LogicalDeviceTest(test.TestCase, parameterized.TestCase):

  @combinations.generate(combinations.combine(mode=['eager'], required_gpus=1))
  def testKeepLogicalDevice(self):
    gpus = tf_config.list_physical_devices('GPU')
    if len(gpus) > 1:
      self.skipTest('Skip logical device test on multi GPUs, since partial GPU '
                    'virtualization is not permitted.')
    # Cannot change logical device after the context initialization.
    context._reset_context()  # pylint: disable=protected-access
    cluster_spec = multi_worker_test_base.create_cluster_spec(
        has_chief=False, num_workers=1)
    resolver = cluster_resolver_lib.SimpleClusterResolver(
        cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec),
        task_type='worker',
        task_id=0)

    logical_gpus = len(gpus) * 2
    for i, device in enumerate(gpus):
      n = (i + 1) * logical_gpus // len(gpus) - i * logical_gpus // len(gpus)
      assert n > 0  # guaranteed if count >= len(devices)
      configs = []
      for ordinal in range(n):
        config = context.LogicalDeviceConfiguration(
            memory_limit=64,
            experimental_device_ordinal=ordinal)
        configs.append(config)

      tf_config.set_logical_device_configuration(device, configs)

    collective_all_reduce_strategy.CollectiveAllReduceStrategy(
        cluster_resolver=resolver)
    # Since we create two logical GPUs out of the last GPU, there should be one
    # more logical GPUs than physical GPUs.
    self.assertLen(tf_config.list_logical_devices('GPU'), logical_gpus)
    context._reset_context()  # pylint: disable=protected-access


@combinations.generate(
    combinations.combine(
        strategy=[
            strategy_combinations.multi_worker_mirrored_2x1_cpu,
            strategy_combinations.multi_worker_mirrored_2x1_gpu,
            strategy_combinations.multi_worker_mirrored_2x2_gpu,
            strategy_combinations.multi_worker_mirrored_2x2_gpu_no_merge_call,
        ],
        mode=['eager']))
class CollectiveAllReduceStrategyV2Test(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super().setUp()
    if context.context().list_physical_devices('TPU'):
      self.skipTest('Test not supported on TPUs')

  def test_replica_id_in_sync_group(self, strategy):

    def replica_fn():
      replica_ctx = distribute_lib.get_replica_context()
      return replica_ctx.replica_id_in_sync_group, replica_ctx._replica_id

    results = test_util.gather(strategy, strategy.run(replica_fn))
    self.assertAllEqual(list(range(strategy.extended._num_replicas_in_sync)),
                        results[0].numpy())
    self.assertAllEqual(
        list(range(len(strategy.extended.worker_devices))) *
        strategy.extended._num_workers, results[1].numpy())

  def test_deep_copy_not_allowed(self, strategy):
    # Check health is disabled in tests by default. We need to enable it for
    # this test to simulate the real world.
    strategy.extended._start_check_health_thread()
    try:
      with self.assertRaisesRegex(ValueError, 'cannot be deep copied'):
        copy.deepcopy(strategy)
      with self.assertRaisesRegex(ValueError, 'cannot be deep copied'):
        with ops.Graph().as_default():
          copy.deepcopy(strategy)
    finally:
      strategy.extended._stop_check_health_thread()


class ExperimentalCompatibilityTest(test.TestCase):

  def testIsInstance(self):
    # It's not uncommon for people to special case MultiWorkerMirroredStrategy,
    # so we need to make sure isinstance check works for combinations between
    # the experimental and non-experimental endpoints.
    strategy = CollectiveAllReduceStrategy()
    experimental_strategy = _CollectiveAllReduceStrategyExperimental()
    self.assertIsInstance(strategy, CollectiveAllReduceStrategy)
    self.assertIsInstance(strategy, _CollectiveAllReduceStrategyExperimental)
    self.assertIsInstance(experimental_strategy, CollectiveAllReduceStrategy)
    self.assertIsInstance(experimental_strategy,
                          _CollectiveAllReduceStrategyExperimental)

  def testName(self):
    # Estimator checks the __name__ to special case MultiWorkerMirroredStrategy.
    self.assertEqual(CollectiveAllReduceStrategy.__name__,
                     'CollectiveAllReduceStrategy')
    self.assertEqual(_CollectiveAllReduceStrategyExperimental.__name__,
                     'CollectiveAllReduceStrategy')


def _replica_id_f32():
  return math_ops.cast(
      distribute_lib.get_replica_context()
      .replica_id_in_sync_group, dtypes.float32)


if __name__ == '__main__':
  # TODO(b/172304955): enable logical devices.
  test_util.main(config_logical_devices=False)