Force parameter server to run on CPU (#143)

Author: icywang86rui

src/sagemaker_tensorflow_container/training.py

@@ -93,8 +93,14 @@ def _run_ps(env, cluster):
 
     cluster_spec = tf.train.ClusterSpec(cluster)
     task_index = env.hosts.index(env.current_host)
-
-    server = tf.train.Server(cluster_spec, job_name='ps', task_index=task_index)
+    # Force parameter server to run on cpu. Running multiple TensorFlow processes on the same
+    # GPU is not safe:
+    # https://stackoverflow.com/questions/46145100/is-it-unsafe-to-run-multiple-tensorflow-processes-on-the-same-gpu
+    no_gpu_config = tf.ConfigProto(device_count={'GPU': 0})
+
+    server = tf.train.Server(
+        cluster_spec, job_name='ps', task_index=task_index, config=no_gpu_config
+    )
 
     threading.Thread(target=lambda: server.join()).start()
 

test/unit/test_training.py

@@ -17,6 +17,7 @@
 
 from mock import MagicMock, patch
 import pytest
+import tensorflow as tf
 
 from sagemaker_tensorflow_container import training
 
@@ -98,7 +99,9 @@ def test_train_distributed_master(run, tf_server, cluster_spec, distributed_trai
                                      'master': ['host1:2222'],
                                      'ps': ['host1:2223', 'host2:2223']})
 
-    tf_server.assert_called_with(cluster_spec(), job_name='ps', task_index=0)
+    tf_server.assert_called_with(
+        cluster_spec(), job_name='ps', task_index=0, config=tf.ConfigProto(device_count={'GPU': 0})
+    )
     tf_server().join.assert_called_with()
 
     tf_config = '{"cluster": {' \
@@ -128,7 +131,9 @@ def test_train_distributed_worker(run, tf_server, cluster_spec, distributed_trai
                                      'master': ['host1:2222'],
                                      'ps': ['host1:2223', 'host2:2223']})
 
-    tf_server.assert_called_with(cluster_spec(), job_name='ps', task_index=1)
+    tf_server.assert_called_with(
+        cluster_spec(), job_name='ps', task_index=1, config=tf.ConfigProto(device_count={'GPU': 0})
+    )
     tf_server().join.assert_called_with()
 
     tf_config = '{"cluster": {' \