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

[((1145, 1205), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""tpu"""', '""""""', '"""Name of TPU to connect to."""'], {}), "('tpu', '', 'Name of TPU to connect to.')\n", (1164, 1205), False, 'from absl import flags\n'), ((1206, 1275), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""project"""', 'None', '"""Name of GCP project with TPU."""'], {}), "('project', None, 'Name of GCP project with TPU.')\n", (1225, 1275), False, 'from absl import flags\n'), ((1276, 1339), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""zone"""', 'None', '"""Name of GCP zone with TPU."""'], {}), "('zone', None, 'Name of GCP zone with TPU.')\n", (1295, 1339), False, 'from absl import flags\n'), ((2753, 2768), 'absl.testing.absltest.main', 'absltest.main', ([], {}), '()\n', (2766, 2768), False, 'from absl.testing import absltest\n'), ((2419, 2517), 'tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver.TPUClusterResolver', 'tpu_cluster_resolver.TPUClusterResolver', ([], {'tpu': 'FLAGS.tpu', 'zone': 'FLAGS.zone', 'project': 'FLAGS.project'}), '(tpu=FLAGS.tpu, zone=FLAGS.zone,\n    project=FLAGS.project)\n', (2458, 2517), False, 'from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver\n'), ((2532, 2567), 'tensorflow.python.eager.remote.connect_to_cluster', 'remote.connect_to_cluster', (['resolver'], {}), '(resolver)\n', (2557, 2567), False, 'from tensorflow.python.eager import remote\n'), ((2673, 2722), 'tensorflow.python.tpu.tpu_strategy_util.initialize_tpu_system', 'tpu_strategy_util.initialize_tpu_system', (['resolver'], {}), '(resolver)\n', (2712, 2722), False, 'from tensorflow.python.tpu import tpu_strategy_util\n'), ((2379, 2401), 'tensorflow.python.eager.context.list_devices', 'context.list_devices', ([], {}), '()\n', (2399, 2401), False, 'from tensorflow.python.eager import context\n'), ((2644, 2666), 'tensorflow.python.eager.context.list_devices', 'context.list_devices', ([], {}), '()\n', (2664, 2666), False, 'from tensorflow.python.eager import context\n')]