這兩天琢磨了下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()
現在整個流程應該是比較清晰了。
