# 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 that the system configuration methods work properly."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import parameterized

from tensorflow.core.protobuf import cluster_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.core.protobuf import tensorflow_server_pb2
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
from tensorflow.python.util import compat


def reset_eager(fn):
  def wrapper(*args, **kwargs):
    try:
      return fn(*args, **kwargs)
    finally:
      # Reset the context.
      context._context = None
      ops.enable_eager_execution_internal()
      assert context._context is not None

  return wrapper


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

  @test_util.run_gpu_only
  @reset_eager
  def testDevicePolicy(self):
    self.assertEqual(context.DEVICE_PLACEMENT_SILENT,
                     context.context().device_policy)

    # If no op has been executed we should be able to set the device policy as
    # well as any init-time configs.
    config.set_intra_op_parallelism_threads(1)
    config.set_device_policy('silent')
    config.set_intra_op_parallelism_threads(2)

    context.ensure_initialized()

    def copy_tensor(dtype=dtypes.int32):
      cpu_tensor = constant_op.constant(1, dtype=dtype)
      gpu_tensor = cpu_tensor.gpu()
      self.assertAllEqual(cpu_tensor + gpu_tensor, 2.0)

    config.set_device_policy('silent')
    self.assertEqual(config.get_device_policy(), 'silent')
    self.assertEqual(context.DEVICE_PLACEMENT_SILENT,
                     context.context().device_policy)
    copy_tensor()

    config.set_device_policy('silent_for_int32')
    self.assertEqual(config.get_device_policy(), 'silent_for_int32')
    self.assertEqual(context.DEVICE_PLACEMENT_SILENT_FOR_INT32,
                     context.context().device_policy)
    with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                 'Tensors on conflicting devices'):
      copy_tensor(dtypes.float32)
    copy_tensor()

    config.set_device_policy('warn')
    self.assertEqual(config.get_device_policy(), 'warn')
    self.assertEqual(context.DEVICE_PLACEMENT_WARN,
                     context.context().device_policy)
    copy_tensor()

    config.set_device_policy('explicit')
    self.assertEqual(config.get_device_policy(), 'explicit')
    self.assertEqual(context.DEVICE_PLACEMENT_EXPLICIT,
                     context.context().device_policy)
    with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                 'Tensors on conflicting devices'):
      copy_tensor()

    config.set_device_policy(None)
    self.assertEqual(config.get_device_policy(), 'silent')

  @reset_eager
  def testExecutionMode(self):
    self.assertTrue(config.get_synchronous_execution())
    self.assertEqual(context.SYNC, context.context().execution_mode)

    # If no op has been executed we should be able to set the execution mode as
    # well as any init-time configs.
    config.set_intra_op_parallelism_threads(1)
    config.set_synchronous_execution(False)
    config.set_intra_op_parallelism_threads(2)

    config.set_synchronous_execution(True)
    self.assertTrue(config.get_synchronous_execution())
    self.assertEqual(context.SYNC, context.context().execution_mode)
    config.set_synchronous_execution(False)
    self.assertFalse(config.get_synchronous_execution())
    self.assertEqual(context.ASYNC, context.context().execution_mode)

  @reset_eager
  def testIntraOpParallelismThreads(self):
    config.set_intra_op_parallelism_threads(10)
    self.assertEqual(
        config.get_intra_op_parallelism_threads(),
        context.context().intra_op_parallelism_threads)

    context.ensure_initialized()

    with self.assertRaises(RuntimeError):
      config.set_intra_op_parallelism_threads(1)

    config.set_intra_op_parallelism_threads(10)

  @reset_eager
  def testInterOpParallelismThreads(self):
    config.set_inter_op_parallelism_threads(10)
    self.assertEqual(
        config.get_inter_op_parallelism_threads(),
        context.context().inter_op_parallelism_threads)

    context.ensure_initialized()

    with self.assertRaises(RuntimeError):
      config.set_inter_op_parallelism_threads(1)

    config.set_inter_op_parallelism_threads(10)

  @test_util.run_gpu_only
  @reset_eager
  def testSoftPlacement(self):
    if context.executing_eagerly():
      self.assertTrue(config.get_soft_device_placement())
    else:
      self.assertFalse(config.get_soft_device_placement())

    @def_function.function
    def mod():
      with ops.device('/device:GPU:0'):
        a = constant_op.constant(1.0)
        b = constant_op.constant(1.0)
        return math_ops.mod(a, b)

    config.set_soft_device_placement(True)
    self.assertEqual(config.get_soft_device_placement(), True)
    self.assertEqual(
        config.get_soft_device_placement(),
        context.context().soft_device_placement)

    # Since soft placement is enabled, the mod operation should work with CPU
    mod()

    config.set_soft_device_placement(False)
    self.assertEqual(config.get_soft_device_placement(), False)
    self.assertEqual(
        config.get_soft_device_placement(),
        context.context().soft_device_placement)

    # Since soft placement is disabled, the mod operation should fail on GPU
    with self.assertRaises(errors.InvalidArgumentError):
      mod()

  @reset_eager
  def testLogDevicePlacement(self):
    self.assertFalse(context.get_log_device_placement())

    context.set_log_device_placement(True)
    self.assertEqual(context.get_log_device_placement(), True)
    self.assertEqual(
        context.get_log_device_placement(),
        context.context().log_device_placement)

    context.set_log_device_placement(False)
    self.assertEqual(context.get_log_device_placement(), False)
    self.assertEqual(
        context.get_log_device_placement(),
        context.context().log_device_placement)

    context.ensure_initialized()

    with self.assertRaises(RuntimeError):
      context.set_log_device_placement(True)

    # If the setting the device placement is a no-op, do not throw a runtime
    # exception.
    context.set_log_device_placement(False)

  @test_util.run_gpu_only
  @reset_eager
  def testJit(self):
    self.assertEqual(config.get_optimizer_jit(), False)

    # the following function should cause Op fusion to occur. However, there is
    # unfortunately no straightforward way to ensure this. We will just have to
    # settle for creating a test that can trigger JIT.
    @def_function.function
    def fun(a, b):
      c = a * b
      d = c + a
      return d

    a = constant_op.constant([2., 2.])
    b = constant_op.constant([2., 2.])

    self.evaluate(fun(a, b))

    config.set_optimizer_jit(True)
    self.assertEqual(config.get_optimizer_jit(), True)
    self.assertEqual(config.get_optimizer_jit(),
                     context.context().optimizer_jit)

    self.evaluate(fun(a, b))

    config.set_optimizer_jit(False)
    self.assertEqual(config.get_optimizer_jit(), False)
    self.assertEqual(config.get_optimizer_jit(),
                     context.context().optimizer_jit)

    self.evaluate(fun(a, b))

  @parameterized.named_parameters(
      ('LayoutOptimizer', 'layout_optimizer'),
      ('ConstantFolding', 'constant_folding'),
      ('ShapeOptimization', 'shape_optimization'),
      ('Remapping', 'remapping'),
      ('ArithmeticOptimization', 'arithmetic_optimization'),
      ('DependencyOptimization', 'dependency_optimization'),
      ('LoopOptimization', 'loop_optimization'),
      ('FunctionOptimization', 'function_optimization'),
      ('DebugStripper', 'debug_stripper'),
      ('ScopedAllocatorOptimization', 'scoped_allocator_optimization'),
      ('ImplementationSelector', 'implementation_selector'),
      ('AutoMixedPrecision', 'auto_mixed_precision'))
  @reset_eager
  def testOptimizerToggleOption(self, field):
    # TODO(b/128531235): Improve testing of option
    options = config.get_optimizer_experimental_options()
    self.assertIsNone(options.get(field))

    config.set_optimizer_experimental_options({field: True})
    options[field] = True
    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
    self.assertDictEqual(
        context.context().get_optimizer_experimental_options(), options)

    config.set_optimizer_experimental_options({field: False})
    options[field] = False
    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
    self.assertDictEqual(
        context.context().get_optimizer_experimental_options(), options)

  @parameterized.named_parameters(
      ('DisableModelPruning', 'disable_model_pruning'),
      ('DisableMetaOptimizer', 'disable_meta_optimizer'))
  @reset_eager
  def testOptimizerBoolOption(self, field):
    # TODO(b/128531235): Improve testing of option
    options = config.get_optimizer_experimental_options()
    self.assertFalse(options.get(field))

    config.set_optimizer_experimental_options({field: True})
    options[field] = True
    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
    self.assertDictEqual(
        context.context().get_optimizer_experimental_options(), options)

    config.set_optimizer_experimental_options({field: False})
    options[field] = False
    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
    self.assertDictEqual(
        context.context().get_optimizer_experimental_options(), options)

  @test_util.run_gpu_only
  @reset_eager
  def testOptimizerToggleOptionPinToHost(self):
    options = config.get_optimizer_experimental_options()
    self.assertIsNone(options.get('pin_to_host_optimization'))

    @def_function.function
    def fun():
      op = test_ops.device_placement_op()
      return op

    # Force optimizer to run for all graphs
    config.set_optimizer_experimental_options({'min_graph_nodes': -1})
    options['min_graph_nodes'] = -1

    # Since pin to host is disabled, the operation should go on GPU
    gpu = self.evaluate(fun())
    self.assertIn(compat.as_bytes('GPU'), gpu)

    config.set_optimizer_experimental_options(
        {'pin_to_host_optimization': True})
    options['pin_to_host_optimization'] = True
    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
    self.assertDictEqual(
        context.context().get_optimizer_experimental_options(), options)

    # Since pin to host is enabled, the operation should go on CPU
    cpu = self.evaluate(fun())
    self.assertIn(compat.as_bytes('CPU'), cpu)

    config.set_optimizer_experimental_options(
        {'pin_to_host_optimization': False})
    options['pin_to_host_optimization'] = False
    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
    self.assertDictEqual(
        context.context().get_optimizer_experimental_options(), options)

    # Since pin to host is disabled again, the operation should go on GPU
    gpu2 = self.evaluate(fun())
    self.assertIn(compat.as_bytes('GPU'), gpu2)


class DeviceTest(test.TestCase):

  @reset_eager
  def testPhysicalDevices(self):
    cpus = config.list_physical_devices('CPU')
    self.assertGreater(len(cpus), 0)
    if test_util.is_gpu_available():
      gpus = config.list_physical_devices('GPU')
      self.assertGreater(len(gpus), 0)

  @reset_eager
  def testCpuMultiple(self):
    cpus = config.list_physical_devices('CPU')
    self.assertEqual(len(cpus), 1)

    config.set_virtual_device_configuration(cpus[0], [
        context.VirtualDeviceConfiguration(),
        context.VirtualDeviceConfiguration()
    ])

    context.ensure_initialized()

    vcpus = config.list_logical_devices('CPU')
    self.assertEqual(len(vcpus), 2)

    with ops.device('/device:CPU:0'):
      a = constant_op.constant(1.0)
      self.evaluate(a)
    with ops.device('/device:CPU:1'):
      b = constant_op.constant(1.0)
      self.evaluate(b)
    with self.assertRaisesRegexp(RuntimeError, 'unknown device'):
      with ops.device('/device:CPU:2'):
        c = constant_op.constant(1.0)
        self.evaluate(c)

    # Ensure we can place ops on each of the device names
    for vcpu in vcpus:
      with ops.device(vcpu.name):
        d = constant_op.constant(1.0)
        self.evaluate(d)

    # Modifying the CPU configuration is not supported
    with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
      config.set_virtual_device_configuration(cpus[0], [
          context.VirtualDeviceConfiguration(),
          context.VirtualDeviceConfiguration(),
          context.VirtualDeviceConfiguration()
      ])

    # Setting the same CPU configuration is fine
    config.set_virtual_device_configuration(cpus[0], [
        context.VirtualDeviceConfiguration(),
        context.VirtualDeviceConfiguration()
    ])

  @test_util.run_gpu_only
  @reset_eager
  def testGpuNone(self):
    gpus = config.list_physical_devices('GPU')
    self.assertGreater(len(gpus), 0)

    cpus = config.list_physical_devices('CPU')
    self.assertEqual(len(cpus), 1)

    self.assertEqual(len(config.get_visible_devices('CPU')), 1)
    self.assertGreater(len(config.get_visible_devices('GPU')), 0)
    config.set_visible_devices(cpus[0])
    self.assertEqual(len(config.get_visible_devices('CPU')), 1)
    self.assertEqual(len(config.get_visible_devices('GPU')), 0)

    with self.assertRaisesRegexp(RuntimeError, 'unknown device'):
      with ops.device('/device:GPU:0'):
        a = constant_op.constant(1.0)
        self.evaluate(a)

    # Modifying the visible devices is not supported
    with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
      config.set_visible_devices(gpus)

    # Setting the same visible devices is fine
    config.set_visible_devices(cpus[0])

  @reset_eager
  def testGpuMultiple(self):
    gpus = config.list_physical_devices('GPU')
    if len(gpus) < 2:
      self.skipTest('Need at least 2 GPUs')

    context.ensure_initialized()

    for i in range(0, len(gpus)):
      with ops.device('/device:GPU:' + str(i)):
        a = constant_op.constant(1.0)
        self.evaluate(a)

    with self.assertRaisesRegexp(RuntimeError, 'unknown device'):
      with ops.device('/device:GPU:' + str(len(gpus))):
        a = constant_op.constant(1.0)
        self.evaluate(a)

  @test_util.run_gpu_only
  @reset_eager
  def testVirtualGpu(self):
    gpus = config.list_physical_devices('GPU')
    self.assertNotEqual(len(gpus), 0)

    self.assertIsNone(config.get_virtual_device_configuration(gpus[-1]))
    config.set_virtual_device_configuration(gpus[-1], [
        context.VirtualDeviceConfiguration(memory_limit=10),
        context.VirtualDeviceConfiguration(memory_limit=10)
    ])
    self.assertEqual(len(config.get_virtual_device_configuration(gpus[-1])), 2)

    logical_gpus = config.list_logical_devices('GPU')
    self.assertTrue(len(logical_gpus), len(gpus) + 1)
    for i in range(0, len(logical_gpus)):
      with ops.device('/device:GPU:' + str(i)):
        a = constant_op.constant(1.0)
        self.evaluate(a)

    with self.assertRaisesRegexp(RuntimeError, 'unknown device'):
      with ops.device('/device:GPU:' + str(len(logical_gpus))):
        a = constant_op.constant(1.0)
        self.evaluate(a)

    # Modifying the GPU configuration is not supported
    with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
      config.set_virtual_device_configuration(gpus[-1], [
          context.VirtualDeviceConfiguration(memory_limit=20),
          context.VirtualDeviceConfiguration(memory_limit=20)
      ])

    with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
      config.set_virtual_device_configuration(gpus[-1], [
          context.VirtualDeviceConfiguration(memory_limit=10),
          context.VirtualDeviceConfiguration(memory_limit=10),
          context.VirtualDeviceConfiguration(memory_limit=10)
      ])

    # Setting the same GPU configuration is fine
    config.set_virtual_device_configuration(gpus[-1], [
        context.VirtualDeviceConfiguration(memory_limit=10),
        context.VirtualDeviceConfiguration(memory_limit=10)
    ])

  @test_util.run_gpu_only
  @reset_eager
  def testGpuGrowth(self):
    gpus = config.list_physical_devices('GPU')
    self.assertNotEqual(len(gpus), 0)

    self.assertIsNone(config.get_memory_growth(gpus[-1]))
    for gpu in gpus:
      config.set_memory_growth(gpu, True)

    c = context.context().config
    self.assertTrue(c.gpu_options.allow_growth)

    logical_gpus = config.list_logical_devices('GPU')
    self.assertTrue(len(logical_gpus), len(gpus))

    # Modifying the GPU configuration is not supported
    with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
      for gpu in gpus:
        config.set_memory_growth(gpu, False)

    # Setting the same GPU configuration is fine
    for gpu in gpus:
      config.set_memory_growth(gpu, True)

  @test_util.run_gpu_only
  @reset_eager
  def testGpuInvalidConfig(self):
    gpus = config.list_physical_devices('GPU')
    self.assertNotEqual(len(gpus), 0)

    if len(gpus) > 1:
      # Assert if other GPUs were not configured
      config.set_memory_growth(gpus[0], True)
      with self.assertRaisesRegexp(ValueError, 'cannot differ'):
        c = context.context().config

      # If we limit visibility to GPU 0, growth is fine
      config.set_visible_devices(gpus[0], 'GPU')
      c = context.context().config
      self.assertTrue(c.gpu_options.allow_growth)

      # Default setting for second GPU is False and works if we set visibility
      config.set_visible_devices(gpus[1], 'GPU')
      c = context.context().config
      self.assertFalse(c.gpu_options.allow_growth)

      # Growth now fails because all the GPUs are visible and not the same
      config.set_visible_devices(gpus, 'GPU')
      with self.assertRaisesRegexp(ValueError, 'cannot differ'):
        c = context.context().config

    for gpu in gpus:
      config.set_memory_growth(gpu, True)

    c = context.context().config
    self.assertTrue(c.gpu_options.allow_growth)

    with self.assertRaisesRegexp(ValueError, 'memory limit'):
      config.set_virtual_device_configuration(gpus[-1], [
          context.VirtualDeviceConfiguration(),
          context.VirtualDeviceConfiguration()
      ])

    self.assertIsNone(config.get_virtual_device_configuration(gpus[-1]))
    config.set_virtual_device_configuration(gpus[-1], [
        context.VirtualDeviceConfiguration(memory_limit=10),
        context.VirtualDeviceConfiguration(memory_limit=10)
    ])

    c = context.context().config
    self.assertFalse(c.gpu_options.allow_growth)

    with self.assertRaisesRegexp(ValueError, 'virtual devices'):
      config.set_memory_growth(gpus[-1], False)

  @test_util.run_gpu_only
  @reset_eager
  def testRemote(self):
    gpus = config.list_logical_devices('GPU')
    self.assertNotEqual(len(gpus), 0)

    context.ensure_initialized()

    gpus = config.list_logical_devices('GPU')
    self.assertNotEqual(len(gpus), 0)
    for gpu in gpus:
      self.assertIsNotNone(gpu.name)

    context.ensure_initialized()

    job_name = 'test'
    cluster_def = cluster_pb2.ClusterDef()
    job_def = cluster_def.job.add()
    job_def.name = job_name
    job_def.tasks[0] = 'localhost:0'

    server_def = tensorflow_server_pb2.ServerDef(
        cluster=cluster_def, job_name=job_name, task_index=0, protocol='grpc')

    context.set_server_def(server_def)

    gpus = config.list_logical_devices('GPU')
    for gpu in gpus:
      self.assertIsNotNone(gpu.name)

  @reset_eager
  def testV1CompatibilityDummyInivisibleDeviceList(self):
    gpus = config.list_physical_devices('GPU')
    if gpus:
      self.skipTest('Test requires no GPUs')

    # Ensure GPU options left untouched on CPU only environments
    context.context()._physical_devices = None
    context.context()._config = config_pb2.ConfigProto(
        gpu_options=config_pb2.GPUOptions(visible_device_list='0'))
    new_config = context.context().config
    self.assertEqual(new_config.gpu_options.visible_device_list, '0')

  @test_util.run_gpu_only
  @reset_eager
  def testV1Compatibility(self):
    # Ensure we set 1 CPU by default
    context.context()._config = config_pb2.ConfigProto()
    new_config = context.context().config
    self.assertEqual(new_config.device_count['CPU'], 1)
    context.context()._physical_devices = None

    # Ensure CPU is split
    context.context()._config = config_pb2.ConfigProto(device_count={'CPU': 2})
    new_config = context.context().config
    self.assertEqual(new_config.device_count['CPU'], 2)
    context.context()._physical_devices = None

    # Handle empty visible device list
    context.context()._config = config_pb2.ConfigProto(
        gpu_options=config_pb2.GPUOptions(visible_device_list=''))
    gpus = config.list_physical_devices('GPU')
    gpu_count = len(gpus)
    new_config = context.context().config
    self.assertEqual(new_config.gpu_options.visible_device_list,
                     ','.join(str(i) for i in range(len(gpus))))
    context.context()._physical_devices = None

    # Handle invalid visible device list
    context.context()._config = config_pb2.ConfigProto(
        gpu_options=config_pb2.GPUOptions(visible_device_list=str(gpu_count)))
    with self.assertRaisesRegexp(ValueError, 'Invalid visible device index'):
      gpus = config.list_physical_devices('GPU')
      new_config = context.context().config
    context.context()._physical_devices = None

    # Handle single visible device list
    context.context()._config = config_pb2.ConfigProto(
        gpu_options=config_pb2.GPUOptions(visible_device_list=str(gpu_count-1)))
    gpus = config.list_physical_devices('GPU')
    new_config = context.context().config
    self.assertEqual(new_config.gpu_options.visible_device_list,
                     str(gpu_count-1))
    context.context()._physical_devices = None

  def testConfigureCollectiveOps(self):
    context.context().configure_collective_ops(
        collective_leader='/job:worker/replica:0/task:0',
        scoped_allocator_enabled_ops=('CollectiveReduce',),
        use_nccl_communication=False,
        device_filters=['/job:worker/task:1'])
    new_config = context.context().config

    # 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)


if __name__ == '__main__':
  ops.enable_eager_execution()
  test.main()
