experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Job name should be picked based on the cluster_spec

Browse Source 
PiperOrigin-RevId: 211833041
This commit is contained in:
Sourabh Bajaj 2018-09-06 10:50:35 -07:00 committed by TensorFlower Gardener
parent b9310932ce
commit 9638524520
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -229,6 +229,10 @@ class TPUClusterResolver(ClusterResolver):
def get_master(self): def get_master(self):
return self.master() return self.master()
def get_job_name(self):
if self._shouldResolve():
return self._job_name
def cluster_spec(self): def cluster_spec(self):
"""Returns a ClusterSpec object based on the latest TPU information. """Returns a ClusterSpec object based on the latest TPU information.

3
tensorflow/contrib/distribute/python/tpu_strategy.py
View File

@ -310,7 +310,8 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy):
def get_host_cpu_device(self, host_id): def get_host_cpu_device(self, host_id):
if self._tpu_cluster_resolver.get_master() in ('', 'local'): if self._tpu_cluster_resolver.get_master() in ('', 'local'):
return '/replica:0/task:0/device:CPU:0' return '/replica:0/task:0/device:CPU:0'
return '/job:tpu_worker/task:%d/device:CPU:0' % (host_id,) job_name = self._tpu_cluster_resolver.get_job_name() or 'tpu_worker'
return '/job:%s/task:%d/device:CPU:0' % (job_name, host_id)
def configure(self, def configure(self,
session_config=None, session_config=None,