From 1fef4919bb77a4df57f86513b341879eab7bf655 Mon Sep 17 00:00:00 2001
From: Rick Chao <rchao@google.com>
Date: Tue, 18 Jun 2019 19:37:32 -0700
Subject: [PATCH] Roll back create_cluster_spec change in estimator until
 nightly picks up the updated tf.

PiperOrigin-RevId: 253917889
---
 tensorflow/python/distribute/multi_worker_test_base.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/distribute/multi_worker_test_base.py b/tensorflow/python/distribute/multi_worker_test_base.py
index ba396074366..f33f5f6e921 100644
--- a/tensorflow/python/distribute/multi_worker_test_base.py
+++ b/tensorflow/python/distribute/multi_worker_test_base.py
@@ -183,11 +183,11 @@ def create_in_process_cluster(num_workers,
       protocol='grpc')
 
 
-def create_cluster_spec(test_obj,
-                        has_chief=False,
+def create_cluster_spec(has_chief=False,
                         num_workers=1,
                         num_ps=0,
-                        has_eval=False):
+                        has_eval=False,
+                        test_obj=None):
   """Create a cluster spec with tasks with unused local ports."""
   if _portpicker_import_error:
     raise _portpicker_import_error  # pylint: disable=raising-bad-type
@@ -207,6 +207,8 @@ def create_cluster_spec(test_obj,
     if has_eval:
       cluster_spec['evaluator'] = ['localhost:%s' % pick_unused_port()]
   except portpicker.NoFreePortFoundError:
+    if test_obj is None:
+      raise
     test_obj.skipTest('Flakes in portpicker library do not represent '
                       'TensorFlow errors.')
   return cluster_spec