diff options
| author |  Frank Chen <frankchn@google.com> | 2017-10-03 21:35:54 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-10-03 21:40:26 -0700 |
| commit | bfaaefa9ecbbbc797f5af60f3d87f6a3c3ac7a09 (patch) | |
| tree | 956e75283a697a915d7ee3273297200412a3ca6a /tensorflow/contrib/cluster_resolver | |
| parent | c31c118a350f4b7010de41fc60a640f2f68e110e (diff) | |
Update APIs for TPU Cluster Resolver to remove the custom API definition and instead use a standard definition file stored in GCS.
PiperOrigin-RevId: 170960877
Diffstat (limited to 'tensorflow/contrib/cluster_resolver')
| -rw-r--r-- | tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py | 15 |
1 files changed, 7 insertions, 8 deletions
diff --git a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py index ceb583abe0..d76ddf8c65 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py @@ -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 |