TensorFlowOnSpark 源码解析
前言
这两天琢磨了下spark-deep-learning和spark-sklearn两个项目，但是感觉都不尽人如意。在training时，都需要把数据broadcast到各个节点进行并行训练，基本就失去实用价值了（tranning数据都会大于单节点内存的好么），而且spark-deep-learning目前还没有实现和tf cluster的结合。所以这个时候转向了开源已久的yahoo的TensorFlowOnSpark项目。简单了过了下他的源码，大致理清楚了原理，这里算是记录下来，也希望能帮到读者。
TensorFlowOnSpark 代码运行剖析
从项目中打开examples/mnist/spark/mnist_spark/mnist_dist.py，
第一步通过pyspark创建SparkContext,这个过程其实就启动了Spark cluster，至于如何通过python启动spark 并且进行相互通讯，具体可以参考我这篇文章：PySpark如何设置worker的python命令。
sc = SparkContext(conf=SparkConf().setAppName("mnist_spark"))
executors = sc._conf.get("spark.executor.instances")第二步是接受一些命令行参数，这个我就不贴了。
第三步是使用标准的pyspark API 从HDFS获取图片数据,构成一个dataframe/rdd：
dataRDD = images.map(lambda x: toNumpy(str(x[0])))接着就是开始进入正题，启动tf cluster了：
cluster = TFCluster.run(sc, mnist_dist.map_fun, args, args.cluster_size, num_ps, args.tensorboard, TFCluster.InputMode.SPARK)TFCluster.run 里的sc 就是sparkcontext,mnist_dist.map_fun函数则包含了你的tensorflow业务代码，在这个示例里就是minist的模型代码，模型代码具体细节代码我们会晚点说。我们先看看TFCluster.run方法：
cluster_template = {}
cluster_template['ps'] = range(num_ps)
cluster_template['worker'] = range(num_ps, num_executors)上面是确定parameter server和worker的数目，这两个概念是和tf相关的。
接着会启动一个Server：
server = reservation.Server(num_executors)
server_addr = server.start()在driver端启动一个Server,主要是为了监听待会spark executor端启动的tf worker，进行协调。
# start TF nodes on all executors
logging.info("Starting TensorFlow on executors")
cluster_meta = {
'id': random.getrandbits(64),
'cluster_template': cluster_template,
'num_executors': num_executors,
'default_fs': defaultFS,
'working_dir': working_dir,
'server_addr': server_addr
}上面的代码获取完整的启动tf cluster所需要的信息。建议大家可以去google下如何手动配置tf cluster,然后就能更深入理解TensorFlowOnSpark是如何预先收集好哪些参数。
nodeRDD = sc.parallelize(range(num_executors), num_executors)
# start TF on a background thread (on Spark driver) to allow for feeding job
def _start():
nodeRDD.foreachPartition(TFSparkNode.run(map_fun,
tf_args,
cluster_meta,
tensorboard,
queues,
background=(input_mode == InputMode.SPARK)))
t = threading.Thread(target=_start)
t.start()
# wait for executors to register and start TFNodes before continuing
logging.info("Waiting for TFSparkNodes to start")
cluster_info = server.await_reservations()
logging.info("All TFSparkNodes started")上面的第一段代码其实是为了确保启动cluster_size个task,每个task对应一个partition，每个partition其实只有一个元素，就是worker的编号。通过对partition进行foreatch来启动对应的tf worker(包含ps)。倒数第二行代码我们又看到了，前面的那个server了，它会阻塞代码往下执行，直到所有tf worker都启动为止。
到这里我们也可以看到，一个spark executor可能会启动多个tf worker。
现在我们进入 TFSparkNode.run看看，这里面包含了具体如何启动tf worker的逻辑，记得这些代码已经在executor执行了。
def run(fn, tf_args, cluster_meta, tensorboard, queues, background):
"""
Wraps the TensorFlow main function in a Spark mapPartitions-compatible function.
"""
def _mapfn(iter):首先定义了一个函数_mapfn，他的参数是一个iter,这个iter 没啥用，就是前面的worker编号，只有一个元素。该函数里主要作用其实就是启动tf worker(PS)的，并且运行用户的代码的：
client = reservation.Client(cluster_meta['server_addr'])
cluster_info = client.get_reservations()启动的过程中会启动一个client,连接我们前面说的Server,报告自己成功启动了。
if job_name == 'ps' or background:
# invoke the TensorFlow main function in a background thread
logging.info("Starting TensorFlow {0}:{1} on cluster node {2} on background process".format(job_name, task_index, worker_num))
p = multiprocessing.Process(target=fn, args=(tf_args, ctx))
p.start()
# for ps nodes only, wait indefinitely in foreground thread for a "control" event (None == "stop")
if job_name == 'ps':
queue = TFSparkNode.mgr.get_queue('control')
done = False
while not done:
msg = queue.get(block=True)
logging.info("Got msg: {0}".format(msg))
if msg == None:
logging.info("Terminating PS")
TFSparkNode.mgr.set('state', 'stopped')
done = True
queue.task_done()
else:
# otherwise, just run TF function in the main executor/worker thread
logging.info("Starting TensorFlow {0}:{1} on cluster node {2} on foreground thread".format(job_name, task_index, worker_num))
fn(tf_args, ctx)
logging.info("Finished TensorFlow {0}:{1} on cluster node {2}".format(job_name, task_index, worker_num))这里会判断是ps还是worker。如果是后台运行，则通过multiprocessing.Process直接运行我们前年提到的mnist_dist.map_fun方法，而mnist_dist.map_fun其实包含了tf session的逻辑代码。当然这个时候模型虽然启动了，但是因为在获取数据时使用了queue.get(block=True) 时，这个时候还没有数据进来，所以会被阻塞住。值得注意的是，这里的代码会发送给spark起的python worker里执行。
在获得cluster对象后，我们就可以调用train方法做真实的训练了，本质上就是开始喂数据：
if args.mode == "train":
cluster.train(dataRDD, args.epochs)进入 cluster.train看下，会进入如下代码：
unionRDD.foreachPartition(TFSparkNode.train(self.cluster_info, self.cluster_meta, qname))这里会把数据按partition的方式喂给每个TF worker(通过调用train方法):
def _train(iter):
queue = mgr.get_queue(qname)
....
for item in iter:
count += 1
queue.put(item, block=True)
....
queue.join()这里会拿到tf的queue,然后通过iter(也就是实际的spark rdd包含的训练数据)往里面放，如果放满了就会阻塞。
直至，大致流程就完成了。现在我们回过头来看我们的业务代码mnist_dist.map_fun,该方法其实是在每个tf worker上执行的：
if job_name == "ps":
server.join()
elif job_name == "worker":
# Assigns ops to the local worker by default.
with tf.device(tf.train.replica_device_setter(
worker_device="/job:worker/task:%d" % task_index,
cluster=cluster)):简单的做了判定，如果是ps则停止在这，否则执行构建模型的工作。在with tf.device.. 里面就是开始定义模型什么的了，标准的tf 代码了：
# Variables of the hidden layer
hid_w = tf.Variable(tf.truncated_normal([IMAGE_PIXELS * IMAGE_PIXELS, hidden_units],
stddev=1.0 / IMAGE_PIXELS), name="hid_w")
hid_b = tf.Variable(tf.zeros([hidden_units]), name="hid_b")
tf.summary.histogram("hidden_weights", hid_w)当然，在TensorFlowOnSpark的示例代码里，使用了Supervisor：
if args.mode == "train":
sv = tf.train.Supervisor(is_chief=(task_index == 0),
logdir=logdir,
init_op=init_op,
summary_op=None,
saver=saver,
global_step=global_step,
stop_grace_secs=300,
save_model_secs=10)
with sv.managed_session(server.target) as sess:
step = 0
tf_feed = TFNode.DataFeed(ctx.mgr, args.mode == "train")
batch_xs, batch_ys = feed_dict(tf_feed.next_batch(batch_size)) TFNode.DataFeed提供了一个便捷的获取批量数据的方式，让你不用操心queue的事情。
在训练达到必要的数目后，你可以停止训练：
if sv.should_stop() or step >= args.steps:
tf_feed.terminate()现在整个流程应该是比较清晰了。