OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-07 00:20:20 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Update APIs for TPU Cluster Resolver to remove the custom API definition and instead use a standard definition file stored in GCS.

Browse Source 
PiperOrigin-RevId: 170960877
This commit is contained in:
Frank Chen 2017-10-03 21:35:54 -07:00 committed by TensorFlower Gardener
parent c31c118a35
commit bfaaefa9ec
1 changed files with 7 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py
View File

@ -39,7 +39,6 @@ class TPUClusterResolver(ClusterResolver):
""" """
def __init__(self, def __init__(self,
api_definition,
project, project,
zone, zone,
tpu_names, tpu_names,
@ -52,8 +51,6 @@ class TPUClusterResolver(ClusterResolver):
for the IP addresses and ports of each Cloud TPU listed. for the IP addresses and ports of each Cloud TPU listed.
Args: Args:
api_definition: (Alpha only) A copy of the JSON API definitions for
Cloud TPUs. This will be removed once Cloud TPU enters beta.
project: Name of the GCP project containing Cloud TPUs project: Name of the GCP project containing Cloud TPUs
zone: Zone where the TPUs are located zone: Zone where the TPUs are located
tpu_names: A list of names of the target Cloud TPUs. tpu_names: A list of names of the target Cloud TPUs.
@ -83,11 +80,13 @@ class TPUClusterResolver(ClusterResolver):
raise ImportError('googleapiclient must be installed before using the ' raise ImportError('googleapiclient must be installed before using the '
'TPU cluster resolver') 'TPU cluster resolver')
# TODO(frankchn): Remove once Cloud TPU API Definitions are public and # TODO(b/67375680): Remove custom URL once TPU APIs are finalized
# replace with discovery.build('tpu', 'v1') self._service = discovery.build(
self._service = discovery.build_from_document( 'tpu',
api_definition, 'v1',
credentials=self._credentials) credentials=self._credentials,
discoveryServiceUrl='https://storage.googleapis.com'
'/tpu-api-definition/v1alpha1.json')
else: else:
self._service = service self._service = service