# Copyright 2017 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.
# ==============================================================================
"""Implementation of Cluster Resolvers for Cloud TPUs."""

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

import collections
import os
import re

from six.moves import urllib
from six.moves.urllib.error import URLError
from six.moves.urllib.request import Request
from six.moves.urllib.request import urlopen

from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import get_accelerator_devices
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import server_lib
from tensorflow.python.util import compat
from tensorflow.python.util.tf_export import tf_export

_GOOGLE_API_CLIENT_INSTALLED = True
try:
  from googleapiclient import discovery  # pylint: disable=g-import-not-at-top
  from oauth2client.client import GoogleCredentials  # pylint: disable=g-import-not-at-top
except ImportError:
  _GOOGLE_API_CLIENT_INSTALLED = False

_GKE_ENV_VARIABLE = 'KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'
_ENDPOINTS_SEPARATOR = ','
_DEFAULT_ENV_VARIABLE = 'TPU_NAME'
_DISCOVERY_SERVICE_URL_ENV_VARIABLE = 'TPU_API_DISCOVERY_URL'

_TPU_DEVICE_REGEX = re.compile(
    r'.*task:(?P<host_id>\d+)/.*device:TPU:(?P<core_id>\d+)$')
_TPU_CONN_RETRIES = 120

_GCE_METADATA_ENDPOINT = 'http://metadata.google.internal'

DeviceDetails = collections.namedtuple(
    'DeviceDetails', ['device_map', 'total_cores'])


def is_running_in_gce():
  """Checks for GCE presence by attempting to query the metadata service."""
  try:
    req = Request(
        '%s/computeMetadata/v1' % _GCE_METADATA_ENDPOINT,
        headers={'Metadata-Flavor': 'Google'})
    resp = urllib.request.urlopen(req, timeout=1)
    info = resp.info()
    if 'Metadata-Flavor' in info and info['Metadata-Flavor'] == 'Google':
      return True
  except URLError:
    pass
  return False


@tf_export('distribute.cluster_resolver.TPUClusterResolver')
class TPUClusterResolver(ClusterResolver):
  """Cluster Resolver for Google Cloud TPUs.

  This is an implementation of cluster resolvers for the Google Cloud TPU
  service. As Cloud TPUs are in alpha, you will need to specify a API definition
  file for this to consume, in addition to a list of Cloud TPUs in your Google
  Cloud Platform project.

  TPUClusterResolver supports the following distinct environments:
  Google Compute Engine
  Google Kubernetes Engine
  Google internal
  """

  def _tpu_service(self):
    """Creates a new Cloud TPU API object.

    This works around an issue where the underlying HTTP connection sometimes
    times out when the script has been running for too long. Other methods in
    this object call this method to get a new API object whenever they need
    to communicate with the Cloud API.

    Returns:
      A Google Cloud TPU API object.
    """
    if self._service:
      return self._service

    credentials = self._credentials
    if credentials is None or credentials == 'default':
      credentials = GoogleCredentials.get_application_default()

    if self._discovery_url:
      return discovery.build(
          'tpu',
          'v1',
          credentials=credentials,
          discoveryServiceUrl=self._discovery_url,
          cache_discovery=False)
    else:
      return discovery.build(
          'tpu', 'v1', credentials=credentials, cache_discovery=False)

  def _request_compute_metadata(self, path):
    req = Request('%s/computeMetadata/v1/%s' % (_GCE_METADATA_ENDPOINT, path),
                  headers={'Metadata-Flavor': 'Google'})
    resp = urlopen(req)
    return compat.as_bytes(resp.read())

  def _is_google_environment(self):
    return (
        self._tpu == compat.as_bytes('') or
        self._tpu == compat.as_bytes('local') or
        self._tpu.startswith(compat.as_bytes('localhost:')) or
        self._tpu.startswith(compat.as_bytes('/bns')) or
        self._tpu.startswith(compat.as_bytes('uptc://')))

  def _should_resolve(self):
    if isinstance(self._should_resolve_override, bool):
      return self._should_resolve_override
    else:
      return not (self._tpu.startswith(compat.as_bytes('grpc://')) or
                  self._is_google_environment())

  @staticmethod
  def _get_device_dict_and_cores(devices):
    """Returns a dict of hosts to cores and total cores given devices names.

    Returns a namedtuple with two attributes:
      device_map: A map of host_ids to a list of core_ids.
      total_cores: The total number of cores within the TPU system.

    Args:
      devices: A list of devices returned by session.list_devices()
    """
    device_map = collections.defaultdict(list)
    num_cores = 0
    for device in devices:
      match = _TPU_DEVICE_REGEX.match(device.name)
      if match:
        host_id = match.group('host_id')
        core_id = match.group('core_id')
        device_map[host_id].append(core_id)
        num_cores += 1
    return DeviceDetails(device_map, num_cores)

  @staticmethod
  def _verify_and_return_same_core_count(device_dict):
    """Verifies that every device in device_dict has the same # of cores."""
    num_cores_per_host_set = (
        {len(core_ids) for core_ids in device_dict.values()})
    if len(num_cores_per_host_set) != 1:
      raise RuntimeError('TPU cores on each device is not the same. This '
                         'should never happen. Devices: {}'.format(device_dict))
    return num_cores_per_host_set.pop()

  @staticmethod
  def _in_gke():
    """When running in GKE, the environment variable will be set."""
    return _GKE_ENV_VARIABLE in os.environ

  @staticmethod
  def _gke_endpoints():
    return os.environ[_GKE_ENV_VARIABLE]

  @staticmethod
  def _env_var_fallback():
    if _DEFAULT_ENV_VARIABLE in os.environ:
      return os.environ[_DEFAULT_ENV_VARIABLE]
    return None

  @staticmethod
  def _environment_discovery_url():
    return os.environ.get(_DISCOVERY_SERVICE_URL_ENV_VARIABLE)

  def __init__(self,
               tpu=None,
               zone=None,
               project=None,
               job_name='worker',
               coordinator_name=None,
               coordinator_address=None,
               credentials='default',
               service=None,
               discovery_url=None):
    """Creates a new TPUClusterResolver object.

    The ClusterResolver will then use the parameters to query the Cloud TPU APIs
    for the IP addresses and ports of each Cloud TPU listed.

    Args:
      tpu: A string corresponding to the TPU to use. If the string is an empty
        string, the string 'local', or a string that begins with 'grpc://' or
          '/bns', then it is assumed to not correspond with a Cloud TPU and will
          instead be passed as the session master and no ClusterSpec propagation
          will be done. In the future, this may also support a list of strings
          when multiple Cloud TPUs are used.
      zone: Zone where the TPUs are located. If omitted or empty, we will assume
        that the zone of the TPU is the same as the zone of the GCE VM, which we
        will try to discover from the GCE metadata service.
      project: Name of the GCP project containing Cloud TPUs. If omitted or
        empty, we will try to discover the project name of the GCE VM from the
        GCE metadata service.
      job_name: Name of the TensorFlow job the TPUs belong to.
      coordinator_name: The name to use for the coordinator. Set to None if the
        coordinator should not be included in the computed ClusterSpec.
      coordinator_address: The address of the coordinator (typically an ip:port
        pair). If set to None, a TF server will be started. If coordinator_name
        is None, a TF server will not be started even if coordinator_address is
        None.
      credentials: GCE Credentials. If None, then we use default credentials
        from the oauth2client
      service: The GCE API object returned by the googleapiclient.discovery
        function. If you specify a custom service object, then the credentials
        parameter will be ignored.
      discovery_url: A URL template that points to the location of the discovery
        service. It should have two parameters {api} and {apiVersion} that when
        filled in produce an absolute URL to the discovery document for that
        service. The environment variable 'TPU_API_DISCOVERY_URL' will override
        this.

    Raises:
      ImportError: If the googleapiclient is not installed.
      ValueError: If no TPUs are specified.
      RuntimeError: If an empty TPU name is specified and this is running in a
        Google Cloud environment.
    """
    if isinstance(tpu, list):
      if not tpu:
        raise ValueError('At least one TPU must be specified.')
      if len(tpu) != 1:
        raise NotImplementedError(
            'Using multiple TPUs in a single session is not yet implemented')
      tpu = tpu[0]

    in_gke = self._in_gke()
    # When using GKE with Cloud TPUs, the env variable will be set.
    if tpu is None:
      if in_gke:
        tpu = self._gke_endpoints()
      else:
        tpu = self._env_var_fallback()

    if tpu is None:
      raise ValueError('Please provide a TPU Name to connect to.')

    self._tpu = compat.as_bytes(tpu)  # self._tpu is always bytes

    # If we are running in Cloud and don't specify a TPU name
    if is_running_in_gce() and not self._tpu:
      raise RuntimeError('You need to specify a TPU Name if you are running in '
                         'the Google Cloud environment.')

    # By default the task_type is 'worker` and the task_id is 0 (which is the
    # first worker in the task).
    self.task_type = job_name
    self.task_id = 0

    if self._is_google_environment():
      self._environment = 'google'
      self.rpc_layer = None

      # TODO(rsopher): remove this logic when possible
      if self._tpu and self._tpu.startswith(compat.as_bytes('/bns')):
        bns_and_port = self._tpu.rsplit(compat.as_bytes(':'), 1)
        if len(bns_and_port) == 2:
          try:
            int(bns_and_port[1])
          except ValueError:
            # Leave named ports.
            pass
          else:
            # Strip numerical ports.
            self._tpu = bns_and_port[0]
    else:
      self._environment = ''
      self.rpc_layer = 'grpc'

    # Setting this overrides the return value of self._should_resolve()
    self._should_resolve_override = None

    # We strip out the protocol if it is included, and override the
    # shouldResolve function to never resolve. We are adding the protocol back
    # in later in self.master().
    if self.rpc_layer is not None and tpu.startswith(self.rpc_layer + '://'):
      tpu = tpu[len(self.rpc_layer + '://'):]
      self._tpu = compat.as_bytes(tpu)  # self._tpu is always bytes
      self._should_resolve_override = False

    # Whether we should actually attempt to contact Cloud APIs
    should_resolve = self._should_resolve()

    # We error out if we are in a non-Cloud environment which cannot talk to the
    # Cloud APIs using the standard class and a special object is not passed in.
    self._service = service
    if (self._service is None and should_resolve and
        not _GOOGLE_API_CLIENT_INSTALLED):
      raise ImportError('googleapiclient and oauth2client must be installed '
                        'before using the TPU cluster resolver. Execute: '
                        '`pip install --upgrade google-api-python-client` '
                        'and `pip install --upgrade oauth2client` to '
                        'install with pip.')

    # We save user-passed credentials, unless the user didn't pass in anything.
    self._credentials = credentials
    if (credentials == 'default' and should_resolve and
        _GOOGLE_API_CLIENT_INSTALLED):
      self._credentials = None

    # Automatically detect project and zone if unspecified.
    if not project and should_resolve:
      project = compat.as_str(
          self._request_compute_metadata('project/project-id'))
    if not zone and should_resolve:
      zone_path = compat.as_str(self._request_compute_metadata('instance/zone'))
      zone = zone_path.split('/')[-1]
    self._project = project
    self._zone = zone

    self._discovery_url = self._environment_discovery_url() or discovery_url

    self._coordinator_name = coordinator_name
    if (coordinator_name and not coordinator_address and
        (should_resolve or in_gke)):
      self._start_local_server()
    else:
      self._coordinator_address = coordinator_address

  def master(self, task_type=None, task_id=None, rpc_layer=None):
    """Get the Master string to be used for the session.

    In the normal case, this returns the grpc path (grpc://1.2.3.4:8470) of
    first instance in the ClusterSpec returned by the cluster_spec function.

    If a non-TPU name is used when constructing a TPUClusterResolver, that will
    be returned instead (e.g. If the tpus argument's value when constructing
    this TPUClusterResolver was 'grpc://10.240.1.2:8470',
    'grpc://10.240.1.2:8470' will be returned).

    Args:
      task_type: (Optional, string) The type of the TensorFlow task of the
        master.
      task_id: (Optional, integer) The index of the TensorFlow task of the
        master.
      rpc_layer: (Optional, string) The RPC protocol TensorFlow should use to
        communicate with TPUs.

    Returns:
      string, the connection string to use when creating a session.

    Raises:
      ValueError: If none of the TPUs specified exists.
    """
    if self._should_resolve():
      # We are going to communicate with the Cloud TPU APIs to get a Cluster.
      cluster_spec = self.cluster_spec()
      if task_type is not None and task_id is not None:
        # task_type and task_id is from the function parameter
        master = cluster_spec.task_address(task_type, task_id)
      elif self.task_type is not None and self.task_id is not None:
        # task_type and task_id is from the object
        master = cluster_spec.task_address(self.task_type, self.task_id)
      else:
        # by default we take the first item in the cluster with the right name
        job_tasks = cluster_spec.job_tasks(self.task_type)
        if not job_tasks:
          raise ValueError('No TPUs with the specified names exist.')
        master = job_tasks[0]
    else:
      if isinstance(self._tpu, (bytes, bytearray)):
        master = compat.as_text(self._tpu).split(_ENDPOINTS_SEPARATOR)[0]
      else:
        master = self._tpu.split(_ENDPOINTS_SEPARATOR)[0]
    return format_master_url(master, rpc_layer or self.rpc_layer)

  def get_master(self):
    return self.master()

  def get_job_name(self):
    if ops.executing_eagerly_outside_functions() or self._should_resolve(
    ) or is_running_in_gce():
      return self.task_type

  def cluster_spec(self):
    """Returns a ClusterSpec object based on the latest TPU information.

    We retrieve the information from the GCE APIs every time this method is
    called.

    Returns:
      A ClusterSpec containing host information returned from Cloud TPUs,
      or None.

    Raises:
      RuntimeError: If the provided TPU is not healthy.
    """
    ############################################################################
    # There are 5 potential cases this code must handle:
    #  1. [Normal case.] We should resolve the TPU name to a set of tasks, and
    #      a. Create a ClusterSpec that includes the coordinator job
    #      b. Create a ClusterSpec without the coordinator job.
    #  2. [GKE / No API Access.] We should not resolve the TPU name to a set of
    #     tasks and
    #      a. Create a ClusterSpec with the coordinator
    #      b. Create a ClusterSpec without the coordinator
    #  3. [Other (legacy non-gRPC).] We should return None.
    ############################################################################

    if self._should_resolve():
      # Case 1.
      response = self._fetch_cloud_tpu_metadata()  # pylint: disable=protected-access

      if 'state' in response and response['state'] != 'READY':
        raise RuntimeError('TPU "%s" is not yet ready; state: "%s"' %
                           (compat.as_text(self._tpu), response['state']))

      if 'networkEndpoints' in response:
        worker_list = [
            '%s:%s' % (endpoint['ipAddress'], endpoint['port'])
            for endpoint in response['networkEndpoints']
        ]
      else:
        # Fall back to the deprecated response format
        instance_url = '%s:%s' % (response['ipAddress'], response['port'])
        worker_list = [instance_url]

      cluster_spec = {self.task_type: worker_list}
    else:
      is_eager = ops.executing_eagerly_outside_functions()
      if self.rpc_layer is None and not is_eager:
        # Case 3.
        return None
      # Case 2.
      tpus = []
      for tpu in compat.as_text(self._tpu).split(_ENDPOINTS_SEPARATOR):
        # We are working around the fact that GKE environment variable that is
        # supplied to us has the protocol string embedded in it, but we want
        # to strip it out for the ClusterSpec.
        if (self.rpc_layer is not None and
            tpu.startswith(self.rpc_layer + '://')):
          tpus.append(tpu[len(self.rpc_layer + '://'):])
        else:
          tpus.append(tpu)
      cluster_spec = {self.task_type: tpus}

    if self._coordinator_address:
      # {1, 2}.a
      cluster_spec[self._coordinator_name] = [self._coordinator_address]

    return server_lib.ClusterSpec(cluster_spec)

  def _fetch_cloud_tpu_metadata(self):
    """Returns the TPU metadata object from the TPU Get API call."""
    try:
      full_name = 'projects/%s/locations/%s/nodes/%s' % (
          self._project, self._zone, compat.as_text(self._tpu))
      service = self._tpu_service()
      request = service.projects().locations().nodes().get(name=full_name)
      return request.execute()
    except Exception as e:
      raise ValueError("Could not lookup TPU metadata from name '%s'. Please "
                       "doublecheck the tpu argument in the TPUClusterResolver "
                       "constructor. Exception: %s" % (self._tpu, e))

  def num_accelerators(self,
                       task_type=None,
                       task_id=None,
                       config_proto=None):
    """Returns the number of TPU cores per worker.

    Connects to the master and list all the devices present in the master,
    and counts them up. Also verifies that the device counts per host in the
    cluster is the same before returning the number of TPU cores per host.

    Args:
      task_type: Unused.
      task_id: Unused.
      config_proto: Used to create a connection to a TPU master in order to
        retrieve the system metadata.

    Raises:
      RuntimeError: If we cannot talk to a TPU worker after retrying or if the
        number of TPU devices per host is different.
    """
    retry_count = 1
    # TODO(b/120564445): Replace with standard library for retries.
    while True:
      try:
        device_details = TPUClusterResolver._get_device_dict_and_cores(
            get_accelerator_devices(self.master(), config_proto=config_proto))
        break
      except errors.DeadlineExceededError:
        error_message = ('Failed to connect to master. The TPU might not be '
                         'ready (e.g. still scheduling) or the master '
                         'address is incorrect: got (%s)' % self.master())
        if retry_count <= _TPU_CONN_RETRIES:
          logging.warning(error_message)
          logging.warning('Retrying (%d/%d)...', retry_count, _TPU_CONN_RETRIES)
          retry_count += 1
        else:
          raise RuntimeError(error_message)

    if device_details.total_cores:
      return {'TPU': TPUClusterResolver._verify_and_return_same_core_count(
          device_details.device_map)}
    return {'TPU': 0}

  @property
  def environment(self):
    """Returns the current environment which TensorFlow is running in."""
    return self._environment

  def _start_local_server(self):
    address = compat.as_text(
        self._request_compute_metadata('instance/network-interfaces/0/ip'))
    self._server = server_lib.Server(
        {
            'local': ['0.0.0.0:0']
        }, protocol='grpc', config=None, start=True)
    # self._server.target is of the form: grpc://ipaddress:port
    target = compat.as_bytes(self._server.target)
    splits = target.split(compat.as_bytes(':'))
    assert len(splits) == 3, self._server.target
    assert splits[0] == compat.as_bytes('grpc'), self._server.target
    self._coordinator_port = compat.as_text(splits[2])
    self._coordinator_address = '%s:%s' % (
        address, compat.as_text(self._coordinator_port))

  def __deepcopy__(self, memo):
    # TODO(b/73668574): Remove this once RunConfig avoids performing deepcopy.
    return self
