From 03d770b78d4cb799ce7945adcbc8ac10fe6f4d38 Mon Sep 17 00:00:00 2001 From: Brennan Saeta Date: Thu, 10 May 2018 17:32:40 -0700 Subject: [TPU]: If the $TPU_NAME env var is set, fallback to that. PiperOrigin-RevId: 196196939 --- .../python/training/tpu_cluster_resolver.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) (limited to 'tensorflow/contrib/cluster_resolver') 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 1403483d28..8ede28602f 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py @@ -36,6 +36,7 @@ except ImportError: _GKE_ENV_VARIABLE = 'KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS' +_DEFAULT_ENV_VARIABLE = 'TPU_NAME' class TPUClusterResolver(ClusterResolver): @@ -70,6 +71,12 @@ class TPUClusterResolver(ClusterResolver): def _gkeMaster(): return os.environ[_GKE_ENV_VARIABLE].split(',')[0] + @staticmethod + def _envVarFallback(): + if _DEFAULT_ENV_VARIABLE in os.environ: + return os.environ[_DEFAULT_ENV_VARIABLE] + return None + def __init__(self, tpu=None, zone=None, @@ -123,8 +130,11 @@ class TPUClusterResolver(ClusterResolver): in_gke = self._inGke() # When using GKE with Cloud TPUs, the env variable will be set. - if tpu is None and in_gke: - tpu = self._gkeMaster() + if tpu is None: + if in_gke: + tpu = self._gkeMaster() + else: + tpu = self._envVarFallback() self._tpu = compat.as_bytes(tpu) # self._tpu is always bytes self._job_name = job_name -- cgit v1.2.3