Add model saving warning at end of training (#171)

* Add model saving warning at end of training

* Add warning when no model artifact is found
* Add warning if model is not saved in the SavedModel bundle format

* Combine logging messages

* Enforce psutil version

* Remove pinned version of sagemaker-containers and install this package last

Author: icywang86rui

docker/1.12.0/Dockerfile.cpu

@@ -76,14 +76,13 @@ COPY $framework_support_installable .
 
 RUN pip install --no-cache-dir -U \
         keras==2.2.4 \
-        sagemaker-containers==2.4.2 \
-        $framework_support_installable \
         "sagemaker-tensorflow>=1.12,<1.13" && \
     # Let's install TensorFlow separately in the end to avoid
     # the library version to be overwritten
     pip install --force-reinstall --no-cache-dir -U \
         tensorflow-1.12.0-py2.py3-none-any.whl \
         horovod && \
+    pip install --no-cache-dir -U $framework_support_installable && \
     rm -f tensorflow-1.12.0-py2.py3-none-any.whl && \
     rm -f $framework_support_installable && \
     pip uninstall -y --no-cache-dir \

setup.py

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

src/sagemaker_tensorflow_container/training.py

@@ -28,6 +28,7 @@
 logger = logging.getLogger(__name__)
 
 SAGEMAKER_PARAMETER_SERVER_ENABLED = 'sagemaker_parameter_server_enabled'
+MODEL_DIR = '/opt/ml/model'
 
 
 def _is_host_master(hosts, current_host):
@@ -159,6 +160,33 @@ def train(env):
                                   runner=runner_type)
 
 
+def _log_model_missing_warning(model_dir):
+    pb_file_exists = False
+    file_exists = False
+    for dirpath, dirnames, filenames in os.walk(model_dir):
+        if filenames:
+            file_exists = True
+        for f in filenames:
+            if 'saved_model.pb' in f or 'saved_model.pbtxt' in f:
+                pb_file_exists = True
+                path, direct_parent_dir = os.path.split(dirpath)
+                if not str.isdigit(direct_parent_dir):
+                    logger.warn('Your model will NOT be servable with SageMaker TensorFlow Serving containers.'
+                                'The SavedModel bundle is under directory \"{}\", not a numeric name.'
+                                .format(direct_parent_dir))
+
+    if not file_exists:
+        logger.warn('No model artifact is saved under path {}.'
+                    ' Your training job will not save any model files to S3.\n'
+                    'For details of how to construct your training script see:\n'
+                    'https://github.com/aws/sagemaker-python-sdk/tree/master/src/sagemaker/tensorflow#adapting-your-local-tensorflow-script' # noqa
+                    .format(model_dir))
+    elif not pb_file_exists:
+        logger.warn('Your model will NOT be servable with SageMaker TensorFlow Serving container.'
+                    'The model artifact was not saved in the TensorFlow SavedModel directory structure:\n'
+                    'https://www.tensorflow.org/guide/saved_model#structure_of_a_savedmodel_directory')
+
+
 def main():
     """Training entry point
     """
@@ -167,3 +195,4 @@ def main():
     s3_utils.configure(env.hyperparameters.get('model_dir'), os.environ.get('SAGEMAKER_REGION'))
     logger.setLevel(env.log_level)
     train(env)
+    _log_model_missing_warning(MODEL_DIR)

test/resources/test_dir_correct_model/12345/saved_model.pb

@@ -41,6 +41,7 @@
 PS_TASK_2 = {'index': 1, 'type': 'ps'}
 MODEL_DIR = 's3://bucket/prefix'
 REGION = 'us-west-2'
+RESOURCE_PATH = os.path.join(os.path.dirname(__file__), '..', 'resources')
 
 
 @pytest.fixture
@@ -200,18 +201,54 @@ def test_build_tf_config_error():
     assert 'Cannot have a ps task if there are no parameter servers in the cluster' in str(error)
 
 
+@patch('sagemaker_tensorflow_container.training.logger')
+def test_log_model_missing_warning_no_model(logger):
+    path = os.path.join(RESOURCE_PATH, 'test_dir_empty')
+    if not os.path.exists(path):
+        os.mkdir(path)
+    training._log_model_missing_warning(path)
+    logger.warn.assert_called_with('No model artifact is saved under path {}.'
+                                   ' Your training job will not save any model files to S3.\n'
+                                   'For details of how to construct your training script see:\n'
+                                   'https://github.com/aws/sagemaker-python-sdk/tree/master/src/sagemaker/tensorflow#adapting-your-local-tensorflow-script' # noqa
+                                   .format(path))
+
+
+@patch('sagemaker_tensorflow_container.training.logger')
+def test_log_model_missing_warning_wrong_format(logger):
+    training._log_model_missing_warning(os.path.join(RESOURCE_PATH, 'test_dir_wrong_model'))
+    logger.warn.assert_called_with('Your model will NOT be servable with SageMaker TensorFlow Serving container.'
+                                   'The model artifact was not saved in the TensorFlow '
+                                   'SavedModel directory structure:\n'
+                                   'https://www.tensorflow.org/guide/saved_model#structure_of_a_savedmodel_directory')
+
+
+@patch('sagemaker_tensorflow_container.training.logger')
+def test_log_model_missing_warning_wrong_parent_dir(logger):
+    training._log_model_missing_warning(os.path.join(RESOURCE_PATH, 'test_dir_wrong_parent_dir'))
+    logger.warn.assert_called_with('Your model will NOT be servable with SageMaker TensorFlow Serving containers.'
+                                   'The SavedModel bundle is under directory \"{}\", not a numeric name.'
+                                   .format('not-digit'))
+
+
+@patch('sagemaker_tensorflow_container.training.logger')
+def test_log_model_missing_warning_correct(logger):
+    training._log_model_missing_warning(os.path.join(RESOURCE_PATH, 'test_dir_correct_model'))
+    logger.warn.assert_not_called()
+
+
+@patch('sagemaker_tensorflow_container.training.logger')
 @patch('sagemaker_tensorflow_container.training.train')
 @patch('logging.Logger.setLevel')
 @patch('sagemaker_containers.beta.framework.training_env')
 @patch('sagemaker_containers.beta.framework.env.read_hyperparameters', return_value={})
 @patch('sagemaker_tensorflow_container.s3_utils.configure')
 def test_main(configure_s3_env, read_hyperparameters, training_env,
-              set_level, train, single_machine_training_env):
+              set_level, train, logger, single_machine_training_env):
     training_env.return_value = single_machine_training_env
     os.environ['SAGEMAKER_REGION'] = REGION
     training.main()
     read_hyperparameters.assert_called_once_with()
     training_env.assert_called_once_with(hyperparameters={})
-    set_level.assert_called_once_with(LOG_LEVEL)
     train.assert_called_once_with(single_machine_training_env)
     configure_s3_env.assert_called_once()

test/resources/test_dir_wrong_model/fake_model.h5

@@ -12,7 +12,7 @@ python =
   3.6: py36, flake8
 
 [flake8]
-max-line-length = 100
+max-line-length = 120
 exclude =
     build/
     .git