# Copyright 2021 DeepMind Technologies Limited
#
# 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 xmanager.xm.resources."""

from absl.testing import absltest
from absl.testing import parameterized

from lxm3._vendor.xmanager import xm
from lxm3._vendor.xmanager.xm import resources
from lxm3._vendor.xmanager.xm.resources import JobRequirements
from lxm3._vendor.xmanager.xm.resources import ResourceType


class ResourceDictTest(absltest.TestCase):

  def test_resource_type_by_name(self):
    self.assertEqual(ResourceType['cpu'], ResourceType.CPU)
    self.assertEqual(ResourceType['Cpu'], ResourceType.CPU)
    self.assertEqual(ResourceType['CPU'], ResourceType.CPU)
    with self.assertRaises(KeyError):
      ResourceType['UPC']  # pylint: disable=pointless-statement

  def test_resource_dict_to_string(self):
    resource_dict = resources.ResourceDict()
    resource_dict[ResourceType.V100] = 8
    resource_dict[ResourceType.CPU] = 4.2 * xm.vCPU
    resource_dict[ResourceType.MEMORY] = 16.5 * xm.GiB

    self.assertEqual(
        str(resource_dict), 'CPU: 4.2, MEMORY: 17716740096.0, V100: 8'
    )

  def test_resource_dict_from_job_requirements(self):
    requirements = JobRequirements(cpu=0.5 * xm.vCPU, memory=2 * xm.MiB, v100=8)
    resource_dict = requirements.task_requirements
    self.assertEqual(resource_dict[ResourceType.CPU], 0.5)
    self.assertEqual(resource_dict[ResourceType.MEMORY], 2097152)
    self.assertEqual(resource_dict[ResourceType.V100], 8)

  def test_job_requirements_unknown_key(self):
    with self.assertRaises(KeyError):
      JobRequirements(cpu=0.5 * xm.vCPU, upc=2)

  def test_requirements_summation(self):
    first = resources.JobRequirements(cpu=1, tpu_v2='2x2')
    second = resources.JobRequirements(ram=4 * xm.GiB, replicas=10)
    total = (
        first.replicas * first.task_requirements
        + second.replicas * second.task_requirements
    )
    self.assertEqual(total[ResourceType.CPU], 1)
    self.assertEqual(total[ResourceType.RAM], 40 * xm.GiB)
    self.assertEqual(total[ResourceType.TPU_V2], 4)


class TopologyTest(parameterized.TestCase):

  @parameterized.parameters(
      ('2', 2), ('4x4', 16), ('2x3x5', 30), ('4x4_twisted', 16)
  )
  def test_resource_type_by_name(self, topology, chip_count):
    self.assertEqual(resources.Topology(topology).chip_count, chip_count)

  def test_invalid_topology(self):
    with self.assertRaises(resources.InvalidTpuTopologyError):
      resources.Topology('euclidian')

  def test_topology_repr(self):
    self.assertEqual(repr(resources.Topology('4x4')), "xm.Topology('4x4')")

  def test_topology_eq(self):
    self.assertEqual(resources.Topology('4x4'), resources.Topology('4x4'))
    self.assertNotEqual(resources.Topology('2x2'), resources.Topology('4x4'))

    self.assertEqual(
        hash(resources.Topology('4x4')), hash(resources.Topology('4x4'))
    )


class JobRequirementsTest(parameterized.TestCase):

  def test_cpu_job(self):
    requirements = resources.JobRequirements(cpu=1.2, ram=1 * xm.GiB)
    self.assertEqual(
        requirements.task_requirements[resources.ResourceType.CPU], 1.2
    )
    self.assertEqual(
        requirements.task_requirements[resources.ResourceType.RAM], 1 * xm.GiB
    )
    self.assertIsNone(requirements.accelerator)
    self.assertIsNone(requirements.topology)
    self.assertEqual(requirements.replicas, 1)

  def test_construct_requirements(self):
    requirements = resources.JobRequirements(
        {resources.ResourceType.CPU: 4}, v100=1
    )
    task_requirements = requirements.task_requirements
    self.assertEqual(task_requirements[resources.ResourceType.CPU], 4)
    self.assertEqual(task_requirements[resources.ResourceType.V100], 1)
    self.assertEqual(requirements.replicas, 1)

  def test_resource_specified_twice(self):
    with self.assertRaises(ValueError):
      resources.JobRequirements({resources.ResourceType.CPU: 1}, cpu=2)

  def test_tpu_job(self):
    requirements = resources.JobRequirements(tpu_v3='4x4')
    self.assertEqual(requirements.accelerator, resources.ResourceType.TPU_V3)
    self.assertEqual(requirements.topology.name, '4x4')

  def test_multihost_gpu(self):
    requirements = resources.JobRequirements(v100='4x2')
    self.assertEqual(requirements.accelerator, resources.ResourceType.V100)
    self.assertEqual(requirements.topology.name, '4x2')
    self.assertEqual(requirements.replicas, 2)

  def test_multihost_gpu_with_cpu(self):
    requirements = resources.JobRequirements(v100='4x2', cpu=1)
    self.assertEqual(requirements.accelerator, resources.ResourceType.V100)
    self.assertEqual(requirements.topology.name, '4x2')
    self.assertEqual(requirements.replicas, 2)

  def test_multihost_gpu_with_replicas(self):
    with self.assertRaises(ValueError):
      resources.JobRequirements(v100='4x2', replicas=3)

  def test_location(self):
    requirements = resources.JobRequirements(location='lon_r7')
    self.assertEqual(requirements.location, 'lon_r7')

  def test_service_tier(self):
    requirements = resources.JobRequirements(
        service_tier=resources.ServiceTier.PROD
    )
    self.assertEqual(requirements.service_tier, resources.ServiceTier.PROD)

  def test_service_tier_mutable(self):
    requirements = resources.JobRequirements(
        service_tier=resources.ServiceTier.PROD
    )
    requirements.service_tier = resources.ServiceTier.BATCH
    self.assertEqual(requirements.service_tier, resources.ServiceTier.BATCH)

  def test_replicas(self):
    requirements = resources.JobRequirements(replicas=2)
    self.assertEqual(requirements.replicas, 2)

    with self.assertRaises(ValueError):
      resources.JobRequirements(replicas=2, tpu_v3='4x4')

    resources.JobRequirements(replicas=2, v100='4x2')
    resources.JobRequirements(v100='4x2')
    with self.assertRaises(ValueError):
      resources.JobRequirements(replicas=4, v100='4x2')

  def test_str(self):
    self.assertEqual(
        repr(
            resources.JobRequirements(
                cpu=1,
                location='lon_r7',
                service_tier=resources.ServiceTier.BATCH,
                replicas=2,
            )
        ),
        (
            "xm.JobRequirements(cpu=1.0, location='lon_r7',"
            ' service_tier=xm.ServiceTier.BATCH, replicas=2)'
        ),
    )

  def test_str_omits_empty_fields(self):
    self.assertEqual(
        repr(resources.JobRequirements(cpu=1)), 'xm.JobRequirements(cpu=1.0)'
    )

  def test_is_gpu_tpu_given_cpu(self):
    requirements = resources.JobRequirements(cpu=1, ram=4 * xm.GiB)
    self.assertNotIn(requirements.accelerator, xm.GpuType)
    self.assertNotIn(requirements.accelerator, xm.TpuType)

  def test_is_gpu_tpu_given_gpu(self):
    requirements = resources.JobRequirements(cpu=1, v100=4)
    self.assertIn(requirements.accelerator, xm.GpuType)
    self.assertNotIn(requirements.accelerator, xm.TpuType)

  def test_is_gpu_tpu_given_tpu(self):
    requirements = resources.JobRequirements(cpu=1, tpu_v2='2x2')
    self.assertNotIn(requirements.accelerator, xm.GpuType)
    self.assertIn(requirements.accelerator, xm.TpuType)


class EnumSubsetTest(parameterized.TestCase):

  def test_construction(self):
    self.assertEqual(resources.GpuType['V100'], resources.GpuType(17))
    with self.assertRaises(AttributeError):
      resources.GpuType['TPU_V3']  # pylint: disable=pointless-statement
    with self.assertRaises(ValueError):
      resources.GpuType(170)

  def test_equivalence(self):
    self.assertEqual(resources.ResourceType.V100, resources.GpuType.V100)

  def test_iteration(self):
    gpus = set(iter(resources.GpuType))
    self.assertIn(resources.ResourceType.V100, gpus)
    self.assertNotIn(resources.ResourceType.TPU_V3, gpus)

  def test_contains(self):
    self.assertIn(resources.ResourceType.V100, resources.GpuType)
    self.assertNotIn(resources.ResourceType.TPU_V3, resources.GpuType)


if __name__ == '__main__':
  absltest.main()
