Disable GPU for parameter process (#121)

Author: icywang86rui

setup.py

@@ -49,7 +49,7 @@ def read(fname):
         'Programming Language :: Python :: 3.6',
     ],
 
-    install_requires=['sagemaker-containers>==2.3.1', 'numpy', 'scipy', 'sklearn',
+    install_requires=['sagemaker-containers>=2.3.4', 'numpy', 'scipy', 'sklearn',
                       'pandas', 'Pillow', 'h5py'],
     extras_require={
         'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist', 'mock',

src/sagemaker_tensorflow_container/training.py

@@ -99,6 +99,9 @@ def _env_vars_with_tf_config(env, ps_task):
 
 def _run_ps(env):
     env_vars = _env_vars_with_tf_config(env, ps_task=True)
+    # Parameter server processes should always run on CPU. Sets CUDA_VISIBLE_DEVICES to '-1' forces
+    # TensorFlow to use CPU.
+    env_vars['CUDA_VISIBLE_DEVICES'] = json.dumps(-1)
     framework.entry_point.run(env.module_dir, env.user_entry_point,
                               env.to_cmd_args(), env_vars, wait=False)