Fixed bug in tf.train.ClusterSpec constructor.

mrry · mrry · commit fb301a9848cb · 2016-04-15T07:45:18.000-07:00
Creating a `tf.train.ClusterSpec` from another ClusterSpec was broken, which in turn broke creating a `tf.train.Server` from a ClusterSpec. Fixes tensorflow#1961. Change: 119954117
diff --git a/tensorflow/python/training/server_lib.py b/tensorflow/python/training/server_lib.py
@@ -238,6 +238,7 @@ def __init__(self, cluster):
     elif isinstance(cluster, ClusterSpec):
       self._cluster_def = tensorflow_server_pb2.ClusterDef()
       self._cluster_def.MergeFrom(cluster.as_cluster_def())
+      self._cluster_spec = {}
       for job_def in self._cluster_def.job:
         self._cluster_spec[job_def.name] = [t for t in job_def.tasks.values()]
     else:
@@ -306,4 +307,3 @@ def _make_cluster_def(self):
           raise TypeError(
               "Task address %r must be bytes or unicode" % task_address)
         job_def.tasks[i] = task_address
-
diff --git a/tensorflow/python/training/server_lib_test.py b/tensorflow/python/training/server_lib_test.py
@@ -146,6 +146,29 @@ def testTwoJobs(self):
     cluster_spec = tf.train.ClusterSpec(cluster_def)
     self.assertProtoEquals(cluster_def, cluster_spec.as_cluster_def())
 
+  def testClusterSpec(self):
+    cluster_spec = tf.train.ClusterSpec(
+        {"ps": ["ps0:2222", "ps1:2222"],
+         "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]})
+
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
+                         tasks { key: 1 value: 'worker1:2222' }
+                         tasks { key: 2 value: 'worker2:2222' } }
+    """
+
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto, tf.train.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        tf.train.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        tf.train.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
+
 
 if __name__ == "__main__":
   tf.test.main()
