chore: import upstream snapshot with attribution

tensorflow/python/tpu/client/BUILD

@@ -0,0 +1,44 @@
+# Cloud TPU Client.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "client",
+    srcs = [
+        "client.py",
+        "version.py",
+    ],
+    strict_deps = True,
+    deps = ["@absl_py//absl/flags"],
+)
+
+py_library(
+    name = "client_lib",
+    srcs = [
+        "__init__.py",
+    ],
+    strict_deps = True,
+    deps = [":client"],
+)
+
+tf_py_strict_test(
+    name = "client_py_test",
+    size = "small",
+    srcs = ["client_test.py"],
+    grpc_enabled = True,
+    main = "client_test.py",
+    deps = [
+        ":client",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/flags",
+    ],
+)

tensorflow/python/tpu/client/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+# =============================================================================
+"""Cloud TPU Client."""
+
+from .version import __version__
+
+from tensorflow.python.tpu.client.client import Client

tensorflow/python/tpu/client/client.py

@@ -0,0 +1,438 @@
+# 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.
+# ==============================================================================
+"""Cloud TPU Client."""
+
+from concurrent import futures
+import datetime
+import json
+import logging
+import os
+import time
+import urllib
+
+from absl import flags
+
+_GOOGLE_API_CLIENT_INSTALLED = True
+try:
+  from googleapiclient import discovery  # pylint: disable=g-import-not-at-top
+  from oauth2client import client  # pylint: disable=g-import-not-at-top
+except ImportError:
+  _GOOGLE_API_CLIENT_INSTALLED = False
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_bool('runtime_oom_exit', True,
+                  'Exit the script when the TPU runtime is OOM.')
+flags.DEFINE_bool('hbm_oom_exit', True,
+                  'Exit the script when the TPU HBM is OOM.')
+
+_GKE_ENV_VARIABLE = 'KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'
+_DEFAULT_TPUCONFIG_VARIABLE = 'TPU_CONFIG'
+_ENDPOINTS_SEPARATOR = ','
+_DEFAULT_ENV_VARIABLE = 'TPU_NAME'
+_DISCOVERY_SERVICE_URL_ENV_VARIABLE = 'TPU_API_DISCOVERY_URL'
+_GCE_METADATA_URL_ENV_VARIABLE = 'GCE_METADATA_IP'
+_GCE_METADATA_ENDPOINT_ENV_VARIABLE = 'GCE_METADATA_HOST'
+_DEFAULT_ENDPOINT_PORT = '8470'
+_OOM_EVENT_COOL_TIME_SEC = 90
+_VERSION_SWITCHER_ENDPOINT = 'http://{}:8475/requestversion'
+
+
+def _utcnow():
+  """A wrapper function around datetime.datetime.utcnow.
+
+  This function is created for unit testing purpose. It's not easy to do
+  StubOutWithMock with datetime.datetime package.
+
+  Returns:
+    datetime.datetime
+  """
+  return datetime.datetime.utcnow()
+
+
+def _environment_discovery_url():
+  return os.environ.get(_DISCOVERY_SERVICE_URL_ENV_VARIABLE)
+
+
+def _gce_metadata_endpoint():
+  endpoint = os.environ.get(_GCE_METADATA_ENDPOINT_ENV_VARIABLE)
+  if not endpoint:
+    endpoint = os.environ.get(
+        _GCE_METADATA_URL_ENV_VARIABLE, 'metadata.google.internal'
+    )
+  return 'http://' + endpoint
+
+
+def _request_compute_metadata(path):
+  req = urllib.request.Request(
+      '%s/computeMetadata/v1/%s' % (_gce_metadata_endpoint(), path),
+      headers={'Metadata-Flavor': 'Google'})
+  resp = urllib.request.urlopen(req)
+  return _as_text(resp.read())
+
+
+def _environment_var_to_network_endpoints(endpoints):
+  """Yields a dict with ip address and port."""
+  for endpoint in endpoints.split(','):
+    grpc_prefix = 'grpc://'
+    if endpoint.startswith(grpc_prefix):
+      endpoint = endpoint.split(grpc_prefix)[1]
+    parts = endpoint.split(':')
+    ip_address = parts[0]
+    port = _DEFAULT_ENDPOINT_PORT
+    if len(parts) > 1:
+      port = parts[1]
+    yield {
+        'ipAddress': ip_address,
+        'port': port
+    }
+
+
+def _get_tpu_node_config():
+  tpu_config_env = os.environ.get(_DEFAULT_TPUCONFIG_VARIABLE)
+  if tpu_config_env:
+    return json.loads(tpu_config_env)
+  return None
+
+
+def _get_tpu_name(tpu):
+  if tpu:
+    return tpu
+
+  for e in [_GKE_ENV_VARIABLE, _DEFAULT_ENV_VARIABLE]:
+    if e in os.environ:
+      return os.environ[e]
+  return None
+
+
+def _as_text(s):
+  if isinstance(s, bytes):
+    return s.decode('utf-8')
+  return s
+
+
+class Client:
+  """Client for working with the Cloud TPU API.
+
+  This client is intended to be used for resolving tpu name to ip addresses.
+
+  It's recommended to use this library as a contextlib to utilize all
+  functionality.
+  """
+
+  def __init__(self,
+               tpu=None,
+               zone=None,
+               project=None,
+               credentials='default',
+               service=None,
+               discovery_url=None):
+    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]
+
+    tpu = _get_tpu_name(tpu)
+
+    if tpu is None:
+      tpu_node_config = _get_tpu_node_config()
+      if tpu_node_config:
+        tpu = tpu_node_config.get('tpu_node_name')
+        project = project or tpu_node_config.get('project')
+        zone = zone or tpu_node_config.get('zone')
+      else:
+        raise ValueError('Please provide a TPU Name to connect to.')
+
+    self._tpu = _as_text(tpu)
+
+    self._use_api = not self._tpu.startswith('grpc://')
+    self._service = service
+
+    self._credentials = None
+    self._project = None
+    self._zone = None
+    self._discovery_url = None
+    if self._use_api:
+      if credentials != 'default':
+        self._credentials = credentials
+      # Automatically detect project and zone if unspecified.
+      if project:
+        self._project = _as_text(project)
+      else:
+        self._project = _request_compute_metadata('project/project-id')
+      if zone:
+        self._zone = _as_text(zone)
+      else:
+        zone_path = _request_compute_metadata('instance/zone')
+        self._zone = zone_path.split('/')[-1]
+      self._discovery_url = _environment_discovery_url() or discovery_url
+
+  def _symptom_msg(self, msg):
+    """Return the structured Symptom message."""
+    return 'Symptom: ' + msg
+
+  def _oom_event(self, symptoms):
+    """Check if a runtime OOM event is reported."""
+    if not symptoms:
+      return False
+    for symptom in reversed(symptoms):
+      if symptom['symptomType'] != 'OUT_OF_MEMORY':
+        continue
+      oom_datetime_str = symptom['createTime'].split('.')[0]
+      oom_datetime = datetime.datetime.strptime(oom_datetime_str,
+                                                '%Y-%m-%dT%H:%M:%S')
+      time_diff = _utcnow() - oom_datetime
+      if time_diff < datetime.timedelta(seconds=_OOM_EVENT_COOL_TIME_SEC):
+        logging.warning(
+            self._symptom_msg(
+                'a recent runtime OOM has occurred ~{} seconds ago. The model '
+                'script will terminate automatically. To prevent future OOM '
+                'events, please consider reducing the model size. To disable this '
+                'behavior, set flag --runtime_oom_exit=false when starting the '
+                'script.'.format(time_diff.seconds)))
+        return True
+    return False
+
+  def _hbm_oom_event(self, symptoms):
+    """Check if a HBM OOM event is reported."""
+    if not symptoms:
+      return False
+    for symptom in reversed(symptoms):
+      if symptom['symptomType'] != 'HBM_OUT_OF_MEMORY':
+        continue
+      oom_datetime_str = symptom['createTime'].split('.')[0]
+      oom_datetime = datetime.datetime.strptime(oom_datetime_str,
+                                                '%Y-%m-%dT%H:%M:%S')
+      time_diff = _utcnow() - oom_datetime
+      if time_diff < datetime.timedelta(seconds=_OOM_EVENT_COOL_TIME_SEC):
+        logging.warning(
+            self._symptom_msg(
+                'a recent HBM OOM has occurred ~{} seconds ago. The model '
+                'script will terminate automatically. To prevent future HBM OOM '
+                'events, please consider reducing the model size. To disable this '
+                'behavior, set flag --hbm_oom_exit=false when starting the '
+                'script.'.format(time_diff.seconds)))
+        return True
+    return False
+
+  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.
+
+    Raises:
+      RuntimeError: If the dependent Python packages are missing.
+
+    Returns:
+      A Google Cloud TPU API object.
+    """
+    if self._service:
+      return self._service
+
+    if not _GOOGLE_API_CLIENT_INSTALLED:
+      raise RuntimeError('Missing runtime dependency on the Google API client. '
+                         'Run `pip install cloud-tpu-client` to fix.')
+
+    credentials = self._credentials
+    if credentials is None or credentials == 'default':
+      credentials = client.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 _full_name(self):
+    """Returns the full Cloud name for this TPU."""
+    return 'projects/%s/locations/%s/nodes/%s' % (
+        self._project, self._zone, self._tpu)
+
+  def _fetch_cloud_tpu_metadata(self):
+    """Returns the TPU metadata object from the TPU Get API call."""
+    service = self._tpu_service()
+    try:
+      r = service.projects().locations().nodes().get(name=self._full_name())
+      return r.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 _get_tpu_property(self, key):
+    if self._use_api:
+      metadata = self._fetch_cloud_tpu_metadata()
+      return metadata.get(key)
+
+    return None
+
+  def __enter__(self):
+    self._open = True
+
+  def __exit__(self, type, value, traceback):  # pylint: disable=redefined-builtin
+    del type, value, traceback
+
+  def recoverable(self):
+    """Returns true if the TPU is in a state where training should eventually resume.
+
+    If false the TPU is in a unrecoverable state and should be recreated.
+    """
+    state = self.state()
+    symptoms = self.symptoms()
+    if state and state in ['TERMINATED', 'PREEMPTED']:
+      return False
+    elif FLAGS.runtime_oom_exit and self._oom_event(symptoms):
+      return False
+    elif FLAGS.hbm_oom_exit and self._hbm_oom_event(symptoms):
+      return False
+    return True
+
+  def symptoms(self):
+    """Return Cloud TPU Symptoms of the TPU."""
+    return self._get_tpu_property('symptoms')
+
+  def state(self):
+    """Return state of the TPU."""
+    return self._get_tpu_property('state')
+
+  def health(self):
+    """Return health of the TPU."""
+    return self._get_tpu_property('health')
+
+  def runtime_version(self):
+    """Return runtime version of the TPU."""
+
+    if not self._use_api:
+      # Fallback on getting version directly from TPU.
+      url = _VERSION_SWITCHER_ENDPOINT.format(
+          self.network_endpoints()[0]['ipAddress'])
+      try:
+        req = urllib.request.Request(url)
+        resp = urllib.request.urlopen(req)
+        version_details = json.loads(resp.read())
+        return version_details.get('currentVersion')
+      except urllib.error.HTTPError as e:
+        status_code = e.code
+        if status_code == 404:
+          return None
+        else:
+          raise e
+    return self._get_tpu_property('tensorflowVersion')
+
+  def accelerator_type(self):
+    """Return accelerator type of the TPU."""
+    return self._get_tpu_property('acceleratorType')
+
+  def api_available(self):
+    """Return if the Cloud TPU API is available, if not certain features will not work."""
+    return self._use_api
+
+  def name(self):
+    """Return the name of the tpu, or the ip address if name is not provided."""
+    return self._tpu
+
+  def get_local_ip(self):
+    """Return the local ip address of the Google Cloud VM the workload is running on."""
+    return _request_compute_metadata('instance/network-interfaces/0/ip')
+
+  def network_endpoints(self):
+    """Return a list of tpu endpoints."""
+    if not self._use_api:
+      return list(_environment_var_to_network_endpoints(self._tpu))
+    response = self._fetch_cloud_tpu_metadata()
+
+    if response.get('state') != 'READY':
+      raise RuntimeError('TPU "%s" is not yet ready; state: "%s"' %
+                         (self._tpu, response.get('state')))
+    if 'networkEndpoints' in response:
+      return response['networkEndpoints']
+    else:
+      return [{'ipAddress': response['ipAddress'], 'port': response['port']}]
+
+  def wait_for_healthy(self, timeout_s=1200, interval=30):
+    """Wait for TPU to become healthy or raise error if timeout reached.
+
+    Args:
+      timeout_s (int): The timeout in seconds for waiting TPU to become healthy.
+      interval (int): The interval in seconds to poll the TPU for health.
+
+    Raises:
+      RuntimeError: If the TPU doesn't become healthy by the timeout.
+    """
+    timeout = time.time() + timeout_s
+    while self.health() != 'HEALTHY':
+      logging.warning(
+          ('Waiting for TPU "%s" with state "%s" '
+           'and health "%s" to become healthy'),
+          self.name(), self.state(), self.health())
+      if time.time() + interval > timeout:
+        raise RuntimeError(
+            'Timed out waiting for TPU "%s" to become healthy' % self.name())
+      time.sleep(interval)
+
+    logging.warning('TPU "%s" is healthy.', self.name())
+
+  def configure_tpu_version(self, version, restart_type='always'):
+    """Configure TPU software version.
+
+    Args:
+      version (string): Version of software to configure the TPU with.
+      restart_type (string): Restart behaviour when switching versions,
+        defaults to always restart. Options are 'always', 'ifNeeded'.
+
+    """
+
+    def configure_worker(worker):
+      """Configure individual TPU worker.
+
+      Args:
+        worker: A dict with the field ipAddress where the configure request will
+          be sent.
+      """
+      ip_address = worker['ipAddress']
+      url = (_VERSION_SWITCHER_ENDPOINT + '/{}?restartType={}').format(
+          ip_address, version, restart_type)
+      req = urllib.request.Request(url, data=b'')
+      try:
+        urllib.request.urlopen(req)
+      except urllib.error.HTTPError as e:
+        status_code = e.code
+        if status_code == 404:
+          raise Exception(
+              'Tensorflow version {} is not available on Cloud TPU, '
+              'try a previous nightly version or refer to '
+              'https://cloud.google.com/tpu/docs/release-notes for '
+              'the latest official version.'.format(version))
+        else:
+          raise Exception('Failed to configure worker {}'.format(ip_address))
+
+    workers = self.network_endpoints()
+
+    with futures.ThreadPoolExecutor(max_workers=len(workers)) as executor:
+      results = executor.map(configure_worker, workers)
+      for result in results:
+        if result:
+          result.result()

tensorflow/python/tpu/client/client_test.py

@@ -0,0 +1,863 @@
+# 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 for cloud tpu client."""
+
+import datetime
+import json
+import os
+import time
+import urllib
+
+from absl import flags
+
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.client import client
+
+FLAGS = flags.FLAGS
+
+mock = test.mock
+
+_UTCNOW_STR = '2000-01-01T00:30:00'
+
+
+def mock_utcnow():
+  return datetime.datetime.strptime(_UTCNOW_STR, '%Y-%m-%dT%H:%M:%S')
+
+
+def mock_request_compute_metadata(path):
+  if path == 'project/project-id':
+    return 'test-project'
+  elif path == 'instance/zone':
+    return 'projects/test-project/locations/us-central1-c'
+  elif path == 'instance/network-interfaces/0/ip':
+    return '10.128.1.2'
+  return ''
+
+
+class MockRequestClass:
+
+  def __init__(self, name, tpu_map):
+    self._name = name
+    self._tpu_map = tpu_map
+
+  def execute(self):
+    if self._name in self._tpu_map:
+      tpu_dict = self._tpu_map[self._name].copy()
+      if isinstance(tpu_dict.get('health'), list):
+        # Do extraction of health list to a single health string based on time.
+        time_now = time.time()
+        health_now = tpu_dict.get('health')[time_now]
+        tpu_dict['health'] = health_now
+      return tpu_dict
+    else:
+      raise KeyError('Resource %s was not found' % self._name)
+
+
+class MockNodeClass:
+
+  def __init__(self, tpu_map):
+    self._tpu_map = tpu_map
+
+  def get(self, name):
+    return MockRequestClass(name, self._tpu_map)
+
+
+class CloudTpuClientTest(test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if 'TPU_API_DISCOVERY_URL' in os.environ:
+      del os.environ['TPU_API_DISCOVERY_URL']
+    if 'TPU_NAME' in os.environ:
+      del os.environ['TPU_NAME']
+    self._time_now = 0
+    self.addCleanup(mock.patch.stopall)
+
+  def _mock_time(self, *args, **kwargs):
+    return self._time_now
+
+  def _mock_sleep(self, secs):
+    self._time_now += secs
+
+  def mock_service_client(self, tpu_map=None):
+    if tpu_map is None:
+      tpu_map = {}
+
+    mock_locations = mock.MagicMock()
+    mock_locations.nodes.return_value = MockNodeClass(tpu_map)
+
+    mock_project = mock.MagicMock()
+    mock_project.locations.return_value = mock_locations
+
+    mock_client = mock.MagicMock()
+    mock_client.projects.return_value = mock_project
+    return mock_client
+
+  def testEnvironmentDiscoveryUrl(self):
+    os.environ['TPU_API_DISCOVERY_URL'] = 'https://{api}.internal/{apiVersion}'
+    self.assertEqual('https://{api}.internal/{apiVersion}',
+                     (client._environment_discovery_url()))
+
+  def testEnvironmentGCEDefault(self):
+    self.assertEqual(
+        'http://metadata.google.internal', client._gce_metadata_endpoint()
+    )
+
+  @mock.patch.dict(os.environ, {'GCE_METADATA_IP': '1.2.3.4'})
+  def testEnvironmentGCEIPOverride(self):
+    self.assertEqual('http://1.2.3.4', client._gce_metadata_endpoint())
+
+  @mock.patch.dict(os.environ, {'GCE_METADATA_HOST': 'foo.bar'})
+  def testEnvironmentGCEHostOverride(self):
+    self.assertEqual('http://foo.bar', client._gce_metadata_endpoint())
+
+  def testEnvironmentVarToNetworkEndpointsSingleIp(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '1234'}],
+        list(client._environment_var_to_network_endpoints(
+            '1.2.3.4:1234')))
+
+  def testEnvironmentVarToNetworkEndpointsSingleGrpcAddress(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'}],
+        list(
+            client._environment_var_to_network_endpoints(
+                'grpc://1.2.3.4:2000')))
+
+  def testEnvironmentVarToNetworkEndpointsMultipleIps(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'},
+         {'ipAddress': '5.6.7.8', 'port': '1234'}],
+        list(
+            client._environment_var_to_network_endpoints(
+                '1.2.3.4:2000,5.6.7.8:1234')))
+
+  def testEnvironmentVarToNetworkEndpointsMultipleGrpcAddresses(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'},
+         {'ipAddress': '5.6.7.8', 'port': '1234'}],
+        list(client._environment_var_to_network_endpoints(
+            'grpc://1.2.3.4:2000,grpc://5.6.7.8:1234')))
+
+  def testEnvironmentVarToNetworkEndpointsMissingPortAndMixed(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'},
+         {'ipAddress': '5.6.7.8', 'port': '8470'}],
+        list(client._environment_var_to_network_endpoints(
+            '1.2.3.4:2000,grpc://5.6.7.8')))
+
+  def testInitializeNoArguments(self):
+    with self.assertRaisesRegex(
+        ValueError, 'Please provide a TPU Name to connect to.'):
+      client.Client()
+
+  def testInitializeMultiElementTpuArray(self):
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        'Using multiple TPUs in a single session is not yet implemented'):
+      client.Client(tpu=['multiple', 'elements'])
+
+  def assertClientContains(self, c):
+    self.assertEqual('tpu_name', c._tpu)
+    self.assertEqual(True, c._use_api)
+    self.assertIsNone(c._credentials)
+    self.assertEqual('test-project', c._project)
+    self.assertEqual('us-central1-c', c._zone)
+    self.assertIsNone(c._discovery_url)
+    self.assertEqual([{
+        'ipAddress': '10.1.2.3',
+        'port': '8470'
+    }], c.network_endpoints())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testNetworkEndpointsNotReadyWithApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertRaisesRegex(
+        RuntimeError, 'TPU .* is not yet ready; state: "None"',
+        c.network_endpoints)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeNoArgumentsWithEnvironmentVariable(self):
+    os.environ['TPU_NAME'] = 'tpu_name'
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY',
+        }
+    }
+    c = client.Client(
+        service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertClientContains(c)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeNoArgumentsWithTPUEnvironmentVariableTPUConfig(self):
+    os.environ['TPU_CONFIG'] = json.dumps({
+        'project': 'test-project',
+        'zone': 'us-central1-c',
+        'tpu_node_name': 'tpu_name',
+    })
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY',
+        }
+    }
+    c = client.Client(service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertClientContains(c)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeTpuName(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertClientContains(c)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeIpAddress(self):
+    c = client.Client(tpu='grpc://1.2.3.4:8470')
+    self.assertEqual('grpc://1.2.3.4:8470', c._tpu)
+    self.assertEqual(False, c._use_api)
+    self.assertIsNone(c._service)
+    self.assertIsNone(c._credentials)
+    self.assertIsNone(c._project)
+    self.assertIsNone(c._zone)
+    self.assertIsNone(c._discovery_url)
+    self.assertEqual([{
+        'ipAddress': '1.2.3.4',
+        'port': '8470'
+    }], c.network_endpoints())
+
+  def testInitializeWithoutMetadata(self):
+    c = client.Client(
+        tpu='tpu_name', project='project', zone='zone')
+    self.assertEqual('tpu_name', c._tpu)
+    self.assertEqual(True, c._use_api)
+    self.assertIsNone(c._service)
+    self.assertIsNone(c._credentials)
+    self.assertEqual('project', c._project)
+    self.assertEqual('zone', c._zone)
+    self.assertIsNone(c._discovery_url)
+
+  def testRecoverableNoApiAccess(self):
+    c = client.Client(tpu='grpc://1.2.3.4:8470')
+    self.assertEqual(True, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRecoverableNoState(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual(True, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRecoverableReady(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual(True, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRecoverablePreempted(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual(False, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableOOM(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:00.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:10.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:20.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True)
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableOOMDisabled(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    FLAGS.runtime_oom_exit = False
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+    FLAGS.runtime_oom_exit = True
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableOOMNoAPI(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='grpc://1.2.3.4:8470',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableHBMOOM(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:00.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:10.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:20.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True)
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableHBMOOMDisabled(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    FLAGS.hbm_oom_exit = False
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+    FLAGS.hbm_oom_exit = True
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableHBMOOMNoAPI(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='grpc://1.2.3.4:8470',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testHealthApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+            'health': 'HEALTHY',
+            'acceleratorType': 'v3-8',
+            'tensorflowVersion': 'nightly',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual('HEALTHY', c.health())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRuntimeVersionApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+            'health': 'HEALTHY',
+            'acceleratorType': 'v3-8',
+            'tensorflowVersion': 'nightly',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual('nightly', c.runtime_version())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testAcceleratorTypeApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+            'health': 'HEALTHY',
+            'acceleratorType': 'v3-8',
+            'tensorflowVersion': 'nightly',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual('v3-8', c.accelerator_type())
+
+  def testHandlesByteStrings(self):
+    self.assertEqual(
+        client.Client(
+            tpu='tpu_name', zone='zone', project='project')._full_name(),
+        client.Client(
+            tpu=b'tpu_name', zone=b'zone', project=b'project')._full_name(),
+    )
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testWaitForHealthy(self):
+    time_mock = mock.patch.object(time, 'time', autospec=True).start()
+    time_mock.side_effect = self._mock_time
+    sleep_mock = mock.patch.object(time, 'sleep', autospec=True).start()
+    sleep_mock.side_effect = self._mock_sleep
+
+    health_timeseries = (['UNHEALTHY_MAINTENANCE']*30 + ['TIMEOUT']*10
+                         + [None]*20 + ['HEALTHY']*30)
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': health_timeseries,
+        },
+    }
+
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+
+    # Doesn't throw RuntimeError as TPU becomes HEALTHY before timeout
+    timeout = 80
+    interval = 5
+    return_time = 60
+    c.wait_for_healthy(timeout_s=timeout, interval=interval)
+    self.assertEqual(time.time(), return_time)
+    self.assertEqual(sleep_mock.call_count, return_time/interval)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testWaitForHealthyRaisesError(self):
+    time_mock = mock.patch.object(time, 'time', autospec=True).start()
+    time_mock.side_effect = self._mock_time
+    sleep_mock = mock.patch.object(time, 'sleep', autospec=True).start()
+    sleep_mock.side_effect = self._mock_sleep
+
+    # Mock timeseries where takes longer than timeout.
+    health_timeseries = ['UNHEALTHY_MAINTENANCE']*50 + ['TIMEOUT']*50
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': health_timeseries,
+        },
+    }
+
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+
+    # Doesn't throw RuntimeError as TPU becomes HEALTHY before timeout
+    with self.assertRaisesRegex(
+        RuntimeError,
+        'Timed out waiting for TPU .* to become healthy'):
+      c.wait_for_healthy(timeout_s=80, interval=5)
+
+  def baseConfigureTpuVersion(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'state':
+                'READY',
+            'networkEndpoints': [
+                {
+                    'ipAddress': '1.2.3.4'
+                },
+                {
+                    'ipAddress': '5.6.7.8'
+                },
+            ]
+        }
+    }
+    return client.Client(
+        tpu='tpu_name',
+        project='test-project',
+        zone='us-central1-c',
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+  @mock.patch.object(urllib.request, 'urlopen')
+  def testConfigureTpuVersion(self, urlopen):
+    c = self.baseConfigureTpuVersion()
+    c.configure_tpu_version('1.15')
+    paths = [call[0][0].full_url for call in urlopen.call_args_list]
+    self.assertCountEqual([
+        'http://1.2.3.4:8475/requestversion/1.15?restartType=always',
+        'http://5.6.7.8:8475/requestversion/1.15?restartType=always'
+    ], sorted(paths))
+
+  @mock.patch.object(urllib.request, 'urlopen')
+  def testConfigureTpuVersionRestartIfneeded(self, urlopen):
+    c = self.baseConfigureTpuVersion()
+    c.configure_tpu_version('1.15', restart_type='ifNeeded')
+    paths = [call[0][0].full_url for call in urlopen.call_args_list]
+    self.assertCountEqual([
+        'http://1.2.3.4:8475/requestversion/1.15?restartType=ifNeeded',
+        'http://5.6.7.8:8475/requestversion/1.15?restartType=ifNeeded'
+    ], sorted(paths))
+
+  @mock.patch.object(urllib.request, 'urlopen')
+  def testGetTpuVersion(self, urlopen):
+    c = client.Client(
+        tpu='grpc://1.2.3.4:8470')
+    resp = mock.Mock()
+    resp.read.side_effect = ['{}', '{"currentVersion": "someVersion"}']
+    urlopen.return_value = resp
+    self.assertIsNone(c.runtime_version(), 'Missing key should be handled.')
+    self.assertEqual(
+        'someVersion', c.runtime_version(), 'Should return configured version.')
+    paths = [call[0][0].full_url for call in urlopen.call_args_list]
+    self.assertCountEqual([
+        'http://1.2.3.4:8475/requestversion',
+        'http://1.2.3.4:8475/requestversion',
+    ], sorted(paths))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/tpu/client/pip_package/BUILD

@@ -0,0 +1,20 @@
+# Description:
+#  Tools for building the Cloud TPU Client pip package.
+
+load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:private"],
+)
+
+licenses(["notice"])
+
+sh_binary(
+    name = "build_pip_package",
+    srcs = ["build_pip_package.sh"],
+    data = [
+        "setup.py",
+        "//tensorflow/python/tpu/client:client_lib",
+    ],
+)

tensorflow/python/tpu/client/pip_package/README

@@ -0,0 +1,3 @@
+Client responsible for communicating the Cloud TPU API. Released separately from tensorflow.
+
+https://pypi.org/project/cloud-tpu-client/

tensorflow/python/tpu/client/pip_package/build_pip_package.sh

@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+# 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.
+# =============================================================================
+
+
+set -e
+
+if [ "$(uname)" = "Darwin" ]; then
+  sedi="sed -i ''"
+else
+  sedi="sed -i"
+fi
+
+PACKAGE_NAME="cloud_tpu_client"
+PIP_PACKAGE="tensorflow/python/tpu/client/pip_package"
+RUNFILES="bazel-bin/tensorflow/python/tpu/client/pip_package/build_pip_package.runfiles/org_tensorflow/tensorflow/python/tpu/client"
+
+function main() {
+  if [ $# -lt 1 ] ; then
+    echo "No destination dir provided"
+    exit 1
+  fi
+
+  DEST=$1
+  TMPDIR=$(mktemp -d -t tmp.XXXXXXXXXX)
+
+  echo $(date) : "=== Using tmpdir: ${TMPDIR}"
+
+  cp ${PIP_PACKAGE}/README ${TMPDIR}
+  cp ${PIP_PACKAGE}/setup.py ${TMPDIR}
+  mkdir ${TMPDIR}/${PACKAGE_NAME}
+  cp -a ${RUNFILES}/. ${TMPDIR}/${PACKAGE_NAME}/
+
+  # Fix the import statements to reflect the copied over path.
+  find  ${TMPDIR}/${PACKAGE_NAME} -name \*.py |
+    xargs $sedi -e '
+      s/^from tensorflow.python.tpu.client/from '${PACKAGE_NAME}'/
+  '
+  echo $(ls $TMPDIR)
+
+  pushd ${TMPDIR}
+  echo $(date) : "=== Building wheel"
+  echo $(pwd)
+  python3 setup.py bdist_wheel >/dev/null
+  mkdir -p ${DEST}
+  cp dist/* ${DEST}
+  popd
+  rm -rf ${TMPDIR}
+  echo $(date) : "=== Output wheel file is in: ${DEST}"
+}
+
+main "$@"

tensorflow/python/tpu/client/pip_package/setup.py

@@ -0,0 +1,50 @@
+# 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.
+# =============================================================================
+"""Cloud TPU Client package."""
+
+from cloud_tpu_client.version import __version__
+from setuptools import find_packages
+from setuptools import setup
+
+setup(
+    name='cloud-tpu-client',
+    version=__version__.replace('-', ''),
+    description='Client for using Cloud TPUs',
+    long_description='Client for using Cloud TPUs',
+    url='https://cloud.google.com/tpu/',
+    author='Google Inc.',
+    author_email='packages@tensorflow.org',
+    packages=find_packages(),
+    classifiers=[
+        'Development Status :: 5 - Production/Stable',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.4',
+        'Programming Language :: Python :: 3.5',
+        'Programming Language :: Python :: 3.6',
+        'Topic :: Scientific/Engineering',
+        'Topic :: Scientific/Engineering :: Mathematics',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+    license='Apache 2.0',
+    keywords='tensorflow tpu',
+    install_requires=['google-api-python-client==1.8.0', 'oauth2client']
+)

tensorflow/python/tpu/client/version.py

@@ -0,0 +1,17 @@
+# 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.
+# =============================================================================
+"""Cloud TPU Client version information."""
+
+__version__ = "0.11"