Distributing graphs across several machines in Distributed Tensorflow - tensorflow

Question

I am currently working on a project using Distributed Tensorflow. My goal is to run several independent graphs across several different machines.
As an example, I want to do something like this (assume that the server is open on each machine)
import tensorflow as tf
a = tf.constant(3)
b = tf.constant(2)
x = tf.mul(a,b) # To be run on "grpc://www.example0.com:2222"
y = tf.mul(a,b) # To be run on "grpc://www.example1.com:2222"
z = tf.mul(a,b) # To be run on "grpc://www.example2.com:2222"
with tf.Session() as sess:
sess.run([x,y,z]) # Ops x,y,z are run on different machines in parallel
My current attempt at this is shown in the following code. However, this code runs the sessions in serial, but I want them to be executed in a parallel distributed manner.
import tensorflow as tf
a = tf.constant(3)
b = tf.constant(2)
x = tf.mul(a,b) # To be run on "grpc://www.example0.com:2222"
y = tf.mul(a,b) # To be run on "grpc://www.example1.com:2222"
z = tf.mul(a,b) # To be run on "grpc://www.example2.com:2222"
with tf.Session("grpc://www.example0.com:2222") as sess:
sess.run(x)
with tf.Session("grpc://www.example1.com:2222") as sess:
sess.run(y)
with tf.Session("grpc://www.example2.com:2222") as sess:
sess.run(z)
While reading the documentation about Distributed Tensorflow, I found that tf.device allows me to set which CPU or GPU to run Tensorflow Ops on. Is there something similar that allows me to set the session target to specify which machine will run which op? Or is there another way of distributing Tensorflow Ops?

Answer 1

I'm currently struggling with this myself. The following is mostly cribbed from the tensorflow distributed how-to guide.
You can pin ops to a job/task using tf.device:
clusterspec = \
{ "worker":
[ "www.example0.com:2222"
, "www.example1.com:2222"
, "www.example2.com:2222"
]
, "master":
[ "localhost:2222" ]
}
cluster = tf.ClusterSpec(clusterspec)
a = tf.constant(3)
b = tf.constant(2)
# pin 'x' to www.example0.com
with tf.device("/job:worker/task:0"):
x = tf.mul(a, b)
# pin 'y' to www.example1.com
with tf.device("/job:worker/task:1"):
y = tf.mul(a, b)
server = tf.train.Server(cluster, job_name="master", task_index=0)
with tf.Session(server.target) as sess:
# run the ops
print(sess.run([x, y]))
However, at least for me, this only works if all the worker processes are on the same machine as the master. Otherwise, it hangs at sess.run.
This turned out to be a problem with the use of localhost in the cluster specification. If you share the same cluster specification between servers, don't use localhost; instead, use the IP address or hostname of the computer that you think localhost refers to. In the case of the above example, suppose that you're running the master script on www.master.com. You have two options:
1. One clusterspec per server using localhost
On each server, localhost refers to the machine running the server.
# on www.example0.com
clusterspec = \
{ "worker":
[ "localhost:2222"
, "www.example1.com:2222"
, "www.example2.com:2222"
]
, "master":
[ "www.master.com:2222" ]
}
cluster = tf.ClusterSpec(clusterspec)
server = tf.train.Server(cluster, job_name="worker", task_index=0)
server.join()
# on www.example1.com
clusterspec = \
{ "worker":
[ "www.example0.com:2222"
, "localhost:2222"
, "www.example2.com:2222"
]
, "master":
[ "www.master.com:2222" ]
}
cluster = tf.ClusterSpec(clusterspec)
server = tf.train.Server(cluster, job_name="worker", task_index=1)
server.join()
# on www.example2.com
clusterspec = \
{ "worker":
[ "www.example0.com:2222"
, "www.example1.com:2222"
, "localhost:2222"
]
, "master":
[ "www.master.com:2222" ]
}
cluster = tf.ClusterSpec(clusterspec)
server = tf.train.Server(cluster, job_name="worker", task_index=2)
server.join()
# on www.master.com
clusterspec = \
{ "worker":
[ "www.example0.com:2222"
, "www.example1.com:2222"
, "www.example2.com:2222"
]
, "master":
[ "localhost:2222" ]
}
cluster = tf.ClusterSpec(clusterspec)
a = tf.constant(3)
b = tf.constant(2)
with tf.device("/job:worker/task:0"):
x = tf.mul(a, b)
with tf.device("/job:worker/task:1"):
y = tf.mul(a, b)
server = tf.train.Server(cluster, job_name="master", task_index=0)
with tf.Session(server.target) as sess:
print(sess.run([x, y]))
2. Shared clusterspec
One cluster specification, using IP addresses / domain names that can all be seen from every node.
Saved in clusterspec.json:
{ "worker":
[ "www.example0.com:2222"
, "www.example1.com:2222"
, "www.example2.com:2222"
]
, "master":
[ "www.master.com:2222" ]
}
Then on each worker:
import json
with open('clusterspec.json', 'r') as f:
clusterspec = json.load(f)
cluster = tf.ClusterSpec(clusterspec)
server = tf.train.Server(cluster, job_name="worker", task_index=<INDEX OF TASK>)
Then on the master:
import json
with open('clusterspec.json', 'r') as f:
clusterspec = json.load(f)
cluster = tf.ClusterSpec(clusterspec)
a = tf.constant(3)
b = tf.constant(2)
with tf.device("/job:worker/task:0"):
x = tf.mul(a, b)
with tf.device("/job:worker/task:1"):
y = tf.mul(a, b)
server = tf.train.Server(cluster, job_name="master", task_index=0)
with tf.Session(server.target) as sess:
print(sess.run([x, y]))