# 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.
# ==============================================================================
"""Test that we can connect to a real Cloud TPU."""

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

from absl import flags
from absl.testing import absltest

from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
from tensorflow.python.eager import context
from tensorflow.python.eager import remote
from tensorflow.python.tpu import tpu_strategy_util

FLAGS = flags.FLAGS
flags.DEFINE_string('tpu', '', 'Name of TPU to connect to.')
flags.DEFINE_string('project', None, 'Name of GCP project with TPU.')
flags.DEFINE_string('zone', None, 'Name of GCP zone with TPU.')

EXPECTED_DEVICES_PRE_CONNECT = [
    '/job:localhost/replica:0/task:0/device:CPU:0',
    '/job:localhost/replica:0/task:0/device:XLA_CPU:0'
]
EXPECTED_DEVICES_AFTER_CONNECT = [
    '/job:localhost/replica:0/task:0/device:CPU:0',
    '/job:localhost/replica:0/task:0/device:XLA_CPU:0',
    '/job:worker/replica:0/task:0/device:CPU:0',
    '/job:worker/replica:0/task:0/device:XLA_CPU:0',
    '/job:worker/replica:0/task:0/device:TPU_SYSTEM:0',
    '/job:worker/replica:0/task:0/device:TPU:0',
    '/job:worker/replica:0/task:0/device:TPU:1',
    '/job:worker/replica:0/task:0/device:TPU:2',
    '/job:worker/replica:0/task:0/device:TPU:3',
    '/job:worker/replica:0/task:0/device:TPU:4',
    '/job:worker/replica:0/task:0/device:TPU:5',
    '/job:worker/replica:0/task:0/device:TPU:6',
    '/job:worker/replica:0/task:0/device:TPU:7',
]


class RemoteCloudTPUTest(absltest.TestCase):
  """Test that we can connect to a real Cloud TPU."""

  def test_connect(self):
    self.assertCountEqual(
        EXPECTED_DEVICES_PRE_CONNECT,
        context.list_devices())

    resolver = tpu_cluster_resolver.TPUClusterResolver(
        tpu=FLAGS.tpu, zone=FLAGS.zone, project=FLAGS.project
    )
    remote.connect_to_cluster(resolver)

    self.assertCountEqual(
        EXPECTED_DEVICES_AFTER_CONNECT,
        context.list_devices())

    tpu_strategy_util.initialize_tpu_system(resolver)

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

[((30, 0, 30, 60), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(30, 20, 30, 25): '"""tpu"""', (30, 27, 30, 29): '""""""', (30, 31, 30, 59): '"""Name of TPU to connect to."""'}, {}), "('tpu', '', 'Name of TPU to connect to.')", False, 'from absl import flags\n'), ((31, 0, 31, 69), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(31, 20, 31, 29): '"""project"""', (31, 31, 31, 35): 'None', (31, 37, 31, 68): '"""Name of GCP project with TPU."""'}, {}), "('project', None, 'Name of GCP project with TPU.')", False, 'from absl import flags\n'), ((32, 0, 32, 63), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(32, 20, 32, 26): '"""zone"""', (32, 28, 32, 32): 'None', (32, 34, 32, 62): '"""Name of GCP zone with TPU."""'}, {}), "('zone', None, 'Name of GCP zone with TPU.')", False, 'from absl import flags\n'), ((75, 2, 75, 17), 'absl.testing.absltest.main', 'absltest.main', ({}, {}), '()', False, 'from absl.testing import absltest\n'), ((63, 15, 65, 5), 'tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver.TPUClusterResolver', 'tpu_cluster_resolver.TPUClusterResolver', (), '', False, 'from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver\n'), ((66, 4, 66, 39), 'tensorflow.python.eager.remote.connect_to_cluster', 'remote.connect_to_cluster', ({(66, 30, 66, 38): 'resolver'}, {}), '(resolver)', False, 'from tensorflow.python.eager import remote\n'), ((72, 4, 72, 53), 'tensorflow.python.tpu.tpu_strategy_util.initialize_tpu_system', 'tpu_strategy_util.initialize_tpu_system', ({(72, 44, 72, 52): 'resolver'}, {}), '(resolver)', False, 'from tensorflow.python.tpu import tpu_strategy_util\n'), ((61, 8, 61, 30), 'tensorflow.python.eager.context.list_devices', 'context.list_devices', ({}, {}), '()', False, 'from tensorflow.python.eager import context\n'), ((70, 8, 70, 30), 'tensorflow.python.eager.context.list_devices', 'context.list_devices', ({}, {}), '()', False, 'from tensorflow.python.eager import context\n')]