前言

這兩天琢磨了下spark-deep-learning和spark-sklearn兩個(gè)項(xiàng)目，但是感覺都不盡人如意。在training時(shí)，都需要把數(shù)據(jù)broadcast到各個(gè)節(jié)點(diǎn)進(jìn)行并行訓(xùn)練，基本就失去實(shí)用價(jià)值了（tranning數(shù)據(jù)都會(huì)大于單節(jié)點(diǎn)內(nèi)存的好么），而且spark-deep-learning目前還沒有實(shí)現(xiàn)和tf cluster的結(jié)合。所以這個(gè)時(shí)候轉(zhuǎn)向了開源已久的yahoo的TensorFlowOnSpark項(xiàng)目。簡單了過了下他的源碼，大致理清楚了原理，這里算是記錄下來，也希望能幫到讀者。

TensorFlowOnSpark 代碼運(yùn)行剖析

從項(xiàng)目中打開examples/mnist/spark/mnist_spark/mnist_dist.py，

第一步通過pyspark創(chuàng)建SparkContext,這個(gè)過程其實(shí)就啟動(dòng)了Spark cluster，至于如何通過python啟動(dòng)spark 并且進(jìn)行相互通訊，具體可以參考我這篇文章：PySpark如何設(shè)置worker的python命令。

sc = SparkContext(conf=SparkConf().setAppName("mnist_spark"))
executors = sc._conf.get("spark.executor.instances")

第二步是接受一些命令行參數(shù)，這個(gè)我就不貼了。

第三步是使用標(biāo)準(zhǔn)的pyspark API 從HDFS獲取圖片數(shù)據(jù),構(gòu)成一個(gè)dataframe/rdd：

dataRDD = images.map(lambda x: toNumpy(str(x[0])))

接著就是開始進(jìn)入正題，啟動(dòng)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函數(shù)則包含了你的tensorflow業(yè)務(wù)代碼，在這個(gè)示例里就是minist的模型代碼，模型代碼具體細(xì)節(jié)代碼我們會(huì)晚點(diǎn)說。我們先看看TFCluster.run方法：

    cluster_template = {}
    cluster_template['ps'] = range(num_ps)
    cluster_template['worker'] = range(num_ps, num_executors)

上面是確定parameter server和worker的數(shù)目，這兩個(gè)概念是和tf相關(guān)的。
接著會(huì)啟動(dòng)一個(gè)Server：

    server = reservation.Server(num_executors)
    server_addr = server.start()

在driver端啟動(dòng)一個(gè)Server,主要是為了監(jiān)聽待會(huì)spark executor端啟動(dòng)的tf worker，進(jìn)行協(xié)調(diào)。

 # 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
    }

上面的代碼獲取完整的啟動(dòng)tf cluster所需要的信息。建議大家可以去google下如何手動(dòng)配置tf cluster,然后就能更深入理解TensorFlowOnSpark是如何預(yù)先收集好哪些參數(shù)。

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")

上面的第一段代碼其實(shí)是為了確保啟動(dòng)cluster_size個(gè)task,每個(gè)task對應(yīng)一個(gè)partition，每個(gè)partition其實(shí)只有一個(gè)元素，就是worker的編號(hào)。通過對partition進(jìn)行foreatch來啟動(dòng)對應(yīng)的tf worker(包含ps)。倒數(shù)第二行代碼我們又看到了，前面的那個(gè)server了，它會(huì)阻塞代碼往下執(zhí)行，直到所有tf worker都啟動(dòng)為止。
到這里我們也可以看到，一個(gè)spark executor可能會(huì)啟動(dòng)多個(gè)tf worker。

現(xiàn)在我們進(jìn)入 TFSparkNode.run看看，這里面包含了具體如何啟動(dòng)tf worker的邏輯，記得這些代碼已經(jīng)在executor執(zhí)行了。

def run(fn, tf_args, cluster_meta, tensorboard, queues, background):
    """
    Wraps the TensorFlow main function in a Spark mapPartitions-compatible function.
    """
    def _mapfn(iter):

首先定義了一個(gè)函數(shù)_mapfn，他的參數(shù)是一個(gè)iter,這個(gè)iter 沒啥用，就是前面的worker編號(hào)，只有一個(gè)元素。該函數(shù)里主要作用其實(shí)就是啟動(dòng)tf worker(PS)的，并且運(yùn)行用戶的代碼的：

client = reservation.Client(cluster_meta['server_addr'])
        cluster_info = client.get_reservations()

啟動(dòng)的過程中會(huì)啟動(dòng)一個(gè)client,連接我們前面說的Server,報(bào)告自己成功啟動(dòng)了。

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))

這里會(huì)判斷是ps還是worker。如果是后臺(tái)運(yùn)行，則通過multiprocessing.Process直接運(yùn)行我們前年提到的mnist_dist.map_fun方法，而mnist_dist.map_fun其實(shí)包含了tf session的邏輯代碼。當(dāng)然這個(gè)時(shí)候模型雖然啟動(dòng)了，但是因?yàn)樵讷@取數(shù)據(jù)時(shí)使用了queue.get(block=True) 時(shí)，這個(gè)時(shí)候還沒有數(shù)據(jù)進(jìn)來，所以會(huì)被阻塞住。值得注意的是，這里的代碼會(huì)發(fā)送給spark起的python worker里執(zhí)行。

在獲得cluster對象后，我們就可以調(diào)用train方法做真實(shí)的訓(xùn)練了，本質(zhì)上就是開始喂數(shù)據(jù)：

if args.mode == "train":
  cluster.train(dataRDD, args.epochs)

進(jìn)入 cluster.train看下，會(huì)進(jìn)入如下代碼：

unionRDD.foreachPartition(TFSparkNode.train(self.cluster_info, self.cluster_meta, qname))

這里會(huì)把數(shù)據(jù)按partition的方式喂給每個(gè)TF worker(通過調(diào)用train方法):

def _train(iter):
     queue = mgr.get_queue(qname)
      ....
      for item in iter:
                count += 1
                queue.put(item, block=True)
      ....
       queue.join()

這里會(huì)拿到tf的queue,然后通過iter(也就是實(shí)際的spark rdd包含的訓(xùn)練數(shù)據(jù))往里面放，如果放滿了就會(huì)阻塞。

直至，大致流程就完成了?，F(xiàn)在我們回過頭來看我們的業(yè)務(wù)代碼mnist_dist.map_fun,該方法其實(shí)是在每個(gè)tf worker上執(zhí)行的：

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則停止在這，否則執(zhí)行構(gòu)建模型的工作。在with tf.device.. 里面就是開始定義模型什么的了，標(biāo)準(zhǔn)的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)

當(dāng)然，在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提供了一個(gè)便捷的獲取批量數(shù)據(jù)的方式，讓你不用操心queue的事情。
在訓(xùn)練達(dá)到必要的數(shù)目后，你可以停止訓(xùn)練：

 if sv.should_stop() or step >= args.steps:
        tf_feed.terminate()

現(xiàn)在整個(gè)流程應(yīng)該是比較清晰了。