Create parameter server in different thread (aws#129)

* Create parameter server in different thread
* Fixing some integ tests

Author: icywang86rui

src/sagemaker_tensorflow_container/training.py

@@ -17,14 +17,14 @@
 import multiprocessing
 import os
 import subprocess
+import threading
 import time
 
 import sagemaker_containers.beta.framework as framework
 import tensorflow as tf
 
 from sagemaker_tensorflow_container import s3_utils
 
-
 logger = logging.getLogger(__name__)
 
 SAGEMAKER_PARAMETER_SERVER_ENABLED = 'sagemaker_parameter_server_enabled'

test/integration/local/test_training.py

@@ -16,7 +16,6 @@
 import tarfile
 
 import pytest
-from sagemaker.estimator import Framework
 from sagemaker.tensorflow import TensorFlow
 
 from test.integration.utils import processor, py_version  # noqa: F401

test/resources/mnist/distributed_mnist.py

@@ -11,6 +11,7 @@
 from tensorflow.python.platform import tf_logging
 import logging as _logging
 import sys as _sys
+import json
 
 
 def cnn_model_fn(features, labels, mode):
@@ -122,9 +123,8 @@ def _parse_args():
     parser.add_argument('--epochs', type=int, default=1)
     # Data, model, and output directories
     parser.add_argument('--output-data-dir', type=str, default=os.environ['SM_OUTPUT_DATA_DIR'])
-    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
+    parser.add_argument('--model_dir', type=str, default=os.environ['SM_MODEL_DIR'])
     parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
-    parser.add_argument('--checkpoint_path', type=str, default=os.environ['SM_MODEL_DIR'])
 
     return parser.parse_known_args()
 
@@ -140,8 +140,12 @@ def _parse_args():
     eval_data, eval_labels = _load_testing_data(args.train)
 
     # Create the Estimator
+    if json.loads(os.environ['SM_TRAINING_ENV'])['additional_framework_parameters'].get('sagemaker_parameter_server_enabled'):
+        model_dir = args.model_dir
+    else:
+        model_dir = os.environ['SM_MODEL_DIR']
     mnist_classifier = tf.estimator.Estimator(
-        model_fn=cnn_model_fn, model_dir=args.checkpoint_path)
+        model_fn=cnn_model_fn, model_dir=model_dir)
 
     # Set up logging for predictions
     # Log the values in the "Softmax" tensor with label "probabilities"

test/unit/test_training.py

@@ -91,6 +91,8 @@ def test_single_machine(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('sagemaker_containers.beta.framework.entry_point.run')
 def test_train_horovod(run_module, single_machine_training_env):
     single_machine_training_env.additional_framework_parameters['sagemaker_mpi_enabled'] = True