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,
api_definition,
project,
zone,
tpu_names,
@ -52,8 +51,6 @@ class TPUClusterResolver(ClusterResolver):
for the IP addresses and ports of each Cloud TPU listed.
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
zone: Zone where the TPUs are located
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 '
'TPU cluster resolver')
# TODO(frankchn): Remove once Cloud TPU API Definitions are public and
# replace with discovery.build('tpu', 'v1')
self._service = discovery.build_from_document(
api_definition,
credentials=self._credentials)
# TODO(b/67375680): Remove custom URL once TPU APIs are finalized
self._service = discovery.build(
'tpu',
'v1',
credentials=self._credentials,
discoveryServiceUrl='https://storage.googleapis.com'
'/tpu-api-definition/v1alpha1.json')
else:
self._service = service