Fix tensorboard thread after fitting (#229)

* Fix tensorboard unit tests

* Join tensorboard thread after fitting

Author: yangaws

CHANGELOG.rst

@@ -6,7 +6,7 @@ CHANGELOG
 ========
 * feature: Allow Local Serving of Models in S3
 * enhancement: Allow option for ``HyperparameterTuner`` to not include estimator metadata in job
-
+* bug-fix: Estimators: Join tensorboard thread after fitting
 
 1.4.2
 =====

src/sagemaker/tensorflow/estimator.py

@@ -19,6 +19,7 @@
 import subprocess
 import tempfile
 import threading
+import time
 
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import create_image_uri, framework_name_from_image, framework_version_from_tag
@@ -234,7 +235,10 @@ def fit_super():
                 tensorboard.start()
                 fit_super()
             finally:
+                # sleep 20 secs for tensorboard start up if fit() quits instantly
+                time.sleep(20)
                 tensorboard.event.set()
+                tensorboard.join()
         else:
             fit_super()
 

tests/unit/test_tf_estimator.py

@@ -286,17 +286,20 @@ def test_run_tensorboard_locally_without_awscli_binary(time, strftime, popen, ca
                                'following command: \n pip install awscli'
 
 
+@patch('sagemaker.tensorflow.estimator.Tensorboard._sync_directories')
 @patch('tempfile.mkdtemp', return_value='/my/temp/folder')
+@patch('shutil.rmtree')
 @patch('os.access', return_value=True)
 @patch('subprocess.call')
 @patch('subprocess.Popen')
 @patch('time.strftime', return_value=TIMESTAMP)
 @patch('time.time', return_value=TIME)
-@pytest.mark.skip(reason="this test fails sometimes and it needs further investigation")
-def test_run_tensorboard_locally(time, strftime, popen, call, access, sagemaker_session):
+def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session):
     tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                     train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)
 
+    popen().poll.return_value = None
+
     tf.fit(inputs='s3://mybucket/train', run_tensorboard_locally=True)
 
     popen.assert_called_with(['tensorboard', '--logdir', '/my/temp/folder', '--host', 'localhost', '--port', '6006'],
@@ -305,19 +308,21 @@ def test_run_tensorboard_locally(time, strftime, popen, call, access, sagemaker_
                              )
 
 
+@patch('sagemaker.tensorflow.estimator.Tensorboard._sync_directories')
 @patch('tempfile.mkdtemp', return_value='/my/temp/folder')
+@patch('shutil.rmtree')
 @patch('socket.socket')
 @patch('os.access', return_value=True)
 @patch('subprocess.call')
 @patch('subprocess.Popen')
 @patch('time.strftime', return_value=TIMESTAMP)
 @patch('time.time', return_value=TIME)
-@pytest.mark.skip(reason="this test fails sometimes and it needs further investigation")
-def test_run_tensorboard_locally_port_in_use(time, strftime, popen, call, access, socket, sagemaker_session):
+def test_run_tensorboard_locally_port_in_use(time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync,
+                                             sagemaker_session):
     tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                     train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)
 
-    popen().poll.side_effect = [True, False]
+    popen().poll.side_effect = [-1, None]
 
     tf.fit(inputs='s3://mybucket/train', run_tensorboard_locally=True)