Allow running PS nodes on the spark driver

[winston-zillow]

The current behavior is to start the tensorflow PS nodes on an spark executor. This would mean wasting GPUs available on that node. These changes allow one to start the PS nodes in the driver while the workers are started on the spark executors

[anfeng]

@winston-zillow please address the conflict

[leewyang]

@winston-zillow Finally had time to take a more detailed look at this in my environments (Spark Standalone, Hadoop/YARN).

In my setup, I saw the following:

1. It looks like the Spark job isn't stopping cleanly (I'm assuming due to the PS thread).
2. I had to explicitly set --cluster_size equal to --num-executors plus one for the PS node.
3. I had to run a TensorFlow/CPU build on the driver, while running a TensorFlow/GPU build on the executors. Not sure if this is well supported by TF, but it worked.

Have you seen similar issues in your env? If not, can you describe your setup?

[winston-zillow]

@leewyang my spark job was able to completed and that TFonSpark joint successfully. I was on hadoop/YARN in EMR env and use the python codes to start the spark jobs. I haven't tried the spark submit. is that what you use? will try again to see if I have any issue.

[leewyang]

Yes, we use a dedicated Hadoop/YARN cluster with spark-submit.

[winston-zillow]

@leewyang I fixed the problem of the driver node not terminating. Also this seems to work only in TENSORFLOW mode, so I put in a check.

I ran it successfully in Yarn/EMR environment

export SPARK_HOME=/usr/lib/spark
export HADOOP_HDFS_HOME=/usr/lib/hadoop-hdfs
export PYTHONFAULTHANDLER=true

# note: --cluster-size is 1 + num. executor
${SPARK_HOME}/bin/spark-submit \
--master yarn \
--deploy-mode client \
--num-executors 2 \
--executor-memory 8G \
--executor-cores 1 \
--py-files tensorflowonspark.zip,TensorFlowOnSpark/examples/mnist/tf/mnist_dist.py \
--conf spark.dynamicAllocation.enabled=false \
--conf spark.yarn.maxAppAttempts=1 \
--conf spark.yarn.executor.memoryOverhead=16G \
--conf spark.executorEnv.LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/lib/jvm/java/jre/lib/amd64/server:/usr/lib/jvm/java/jre/lib/amd64:/tmp/libhdfs \
--conf spark.executorEnv.HADOOP_HDFS_HOME=$HADOOP_HDFS_HOME \
TensorFlowOnSpark/examples/mnist/tf/mnist_spark.py \
--steps 10 \
--images /user/hadoop/mnist_data/train/images \
--labels /user/hadoop/mnist_data/train/labels \
--format csv \
--mode train \
--model mnist_model \
--driver_ps_nodes True \
--cluster_size 3

Console log at driver:

2018-01-19 00:13:05,432 INFO (MainThread-26109) Shutting down cluster
2018-01-19 00:13:05,633 INFO (Thread-3-26109) Got msg: None
2018-01-19 00:13:05,633 INFO (Thread-3-26109) Terminating PS
2018-01-19T00:13:10.837377 ===== Stop
18/01/19 00:13:10 INFO SparkContext: Invoking stop() from shutdown hook
18/01/19 00:13:10 INFO SparkUI: Stopped Spark web UI at http://172.30.0.203:4040
18/01/19 00:13:11 INFO YarnClientSchedulerBackend: Interrupting monitor thread
18/01/19 00:13:11 INFO YarnClientSchedulerBackend: Shutting down all executors
18/01/19 00:13:11 INFO YarnSchedulerBackend$YarnDriverEndpoint: Asking each executor to shut down
18/01/19 00:13:11 INFO SchedulerExtensionServices: Stopping SchedulerExtensionServices
(serviceOption=None,
 services=List(),
 started=false)
18/01/19 00:13:11 INFO YarnClientSchedulerBackend: Stopped
18/01/19 00:13:11 INFO MapOutputTrackerMasterEndpoint: MapOutputTrackerMasterEndpoint stopped!
18/01/19 00:13:11 INFO MemoryStore: MemoryStore cleared
18/01/19 00:13:11 INFO BlockManager: BlockManager stopped
18/01/19 00:13:11 INFO BlockManagerMaster: BlockManagerMaster stopped
18/01/19 00:13:11 INFO OutputCommitCoordinator$OutputCommitCoordinatorEndpoint: OutputCommitCoordinator stopped!
18/01/19 00:13:11 INFO SparkContext: Successfully stopped SparkContext
18/01/19 00:13:11 INFO ShutdownHookManager: Shutdown hook called
18/01/19 00:13:11 INFO ShutdownHookManager: Deleting directory /mnt/tmp/spark-0f01312e-b879-4714-870a-6a9a06279ecd/pyspark-e61b56a4-9c2f-4080-9cd0-7bbf8fdf3a93
18/01/19 00:13:11 INFO ShutdownHookManager: Deleting directory /mnt/tmp/spark-0f01312e-b879-4714-870a-6a9a06279ecd
-bash-4.2$ 

[leewyang]

This looks good. I was able to run successfully in my environment, however I have one comment for something that tripped me up during testing.

[leewyang]

So previously, we had: num_workers + num_ps = cluster_size, where cluster_size == num_executors.
With the --driver_ps_nodes option, this is now a bit different, since num_workers == num_executors.

Can you add something like you will need to set cluster_size = num_executors + num_ps?

[winston-zillow]

OK, will do

[leewyang]

Looks good. Thank you for your contribution.