fix: tensorflow-2.0 library code changes (#247)

Author: chuyang-deng

setup.py

@@ -57,7 +57,7 @@ def read_version():
                       'pandas', 'Pillow', 'h5py'],
     extras_require={
         'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist', 'mock',
-                 'sagemaker==1.19.1', 'tensorflow<2.0', 'docker-compose', 'botocore>=1.12.140'],
+                 'sagemaker==1.19.1', 'tensorflow', 'docker-compose', 'botocore>=1.12.140'],
         'benchmark': ['click']
     },
 )

src/sagemaker_tensorflow_container/training.py

@@ -96,9 +96,9 @@ def _run_ps(env, cluster):
     # 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})
+    no_gpu_config = tf.compat.v1.ConfigProto(device_count={'GPU': 0})
 
-    server = tf.train.Server(
+    server = tf.distribute.Server(
         cluster_spec, job_name='ps', task_index=task_index, config=no_gpu_config
     )
 

test/unit/test_training.py

@@ -102,7 +102,7 @@ def test_train_horovod(run_module, single_machine_training_env):
 @pytest.mark.skipif(sys.version_info.major != 3,
                     reason="Skip this for python 2 because of dict key order mismatch")
 @patch('tensorflow.train.ClusterSpec')
-@patch('tensorflow.train.Server')
+@patch('tensorflow.distribute.Server')
 @patch('sagemaker_containers.beta.framework.entry_point.run')
 @patch('multiprocessing.Process', lambda target: target())
 @patch('time.sleep', MagicMock())
@@ -114,7 +114,7 @@ def test_train_distributed_master(run, tf_server, cluster_spec, distributed_trai
                                      'ps': ['host1:2223', 'host2:2223']})
 
     tf_server.assert_called_with(
-        cluster_spec(), job_name='ps', task_index=0, config=tf.ConfigProto(device_count={'GPU': 0})
+        cluster_spec(), job_name='ps', task_index=0, config=tf.compat.v1.ConfigProto(device_count={'GPU': 0})
     )
     tf_server().join.assert_called_with()
 
@@ -132,7 +132,7 @@ def test_train_distributed_master(run, tf_server, cluster_spec, distributed_trai
 @pytest.mark.skipif(sys.version_info.major != 3,
                     reason="Skip this for python 2 because of dict key order mismatch")
 @patch('tensorflow.train.ClusterSpec')
-@patch('tensorflow.train.Server')
+@patch('tensorflow.distribute.Server')
 @patch('sagemaker_containers.beta.framework.entry_point.run')
 @patch('multiprocessing.Process', lambda target: target())
 @patch('time.sleep', MagicMock())
@@ -146,7 +146,7 @@ def test_train_distributed_worker(run, tf_server, cluster_spec, distributed_trai
                                      'ps': ['host1:2223', 'host2:2223']})
 
     tf_server.assert_called_with(
-        cluster_spec(), job_name='ps', task_index=1, config=tf.ConfigProto(device_count={'GPU': 0})
+        cluster_spec(), job_name='ps', task_index=1, config=tf.compat.v1.ConfigProto(device_count={'GPU': 0})
     )
     tf_server().join.assert_called_with()