Address pr comments

Author: icywang86rui

src/sagemaker/tensorflow/estimator.py

@@ -185,6 +185,8 @@ def __init__(self, training_steps=None, evaluation_steps=None, checkpoint_path=N
             py_version (str): Python version you want to use for executing your model training code (default: 'py2').
             framework_version (str): TensorFlow version you want to use for executing your model training code.
                 List of supported versions https://github.com/aws/sagemaker-python-sdk#tensorflow-sagemaker-estimators
+            model_dir (str): S3 location where the checkpoint data and models can be exported to during training
+                (default: None). If not specified a default S3 URI will be generated.
             requirements_file (str): Path to a ``requirements.txt`` file (default: ''). The path should be within and
                 relative to ``source_dir``. Details on the format can be found in the
                 `Pip User Guide <https://pip.pypa.io/en/stable/reference/pip_install/#requirements-file-format>`_.
@@ -194,6 +196,10 @@ def __init__(self, training_steps=None, evaluation_steps=None, checkpoint_path=N
                     Examples:
                         123.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:1.0
                         custom-image:latest.
+            script_mode (bool): If set to True will the estimator will use the Script Mode containers (default: False).
+                This will be ignored if py_version is set to 'py3'.
+            distribution (dict): A dictionary with information on how to run distributed training
+                (default: None).
             **kwargs: Additional kwargs passed to the Framework constructor.
         """
         if framework_version is None:
@@ -207,7 +213,7 @@ def __init__(self, training_steps=None, evaluation_steps=None, checkpoint_path=N
         self.evaluation_steps = evaluation_steps
         self.model_dir = model_dir
         self.script_mode = script_mode
-        self.distributions = distributions
+        self.distributions = distributions or {}
 
         self._validate_args(py_version=py_version, script_mode=script_mode, framework_version=framework_version,
                             training_steps=training_steps, evaluation_steps=evaluation_steps,
@@ -283,12 +289,11 @@ def fit_super():
         if run_tensorboard_locally and wait is False:
             raise ValueError("Tensorboard is not supported with async fit")
 
-        if run_tensorboard_locally:
-
-            if self.script_mode_enabled():
+        if self._script_mode_enabled():
+            if run_tensorboard_locally:
                 LOGGER.warning(_SCRIPT_MODE_TENSORBOARD_WARNING.format(self.model_dir))
-                return
-
+            fit_super()
+        elif run_tensorboard_locally:
             tensorboard = Tensorboard(self)
             tensorboard.validate_requirements()
 
@@ -371,12 +376,9 @@ def create_model(self, model_server_workers=None, role=None,
         """
 
         role = role or self.role
-        if endpoint_type == 'tensorflow-serving':
+        if endpoint_type == 'tensorflow-serving' or self._script_mode_enabled():
             return self._create_tfs_model(role=role, vpc_config_override=vpc_config_override)
 
-        if self.script_mode_enabled():
-            raise ValueError(_SCRIPT_MODE_SERVING_ERROR_MSG)
-
         return self._create_default_model(model_server_workers=model_server_workers, role=role,
                                           vpc_config_override=vpc_config_override)
 
@@ -408,17 +410,14 @@ def hyperparameters(self):
         """Return hyperparameters used by your custom TensorFlow code during model training."""
         hyperparameters = super(TensorFlow, self).hyperparameters()
 
-        if not self.checkpoint_path:
-            self.checkpoint_path = self._default_s3_path('checkpoints')
+        self.checkpoint_path = self.checkpoint_path or self._default_s3_path('checkpoints')
 
-        if self.script_mode_enabled():
-            if not self.model_dir:
-                self.model_dir = self._default_s3_path('model')
+        if self._script_mode_enabled():
+            self.model_dir = self.model_dir or self._default_s3_path('model')
             additional_hyperparameters = {'model_dir': self.model_dir}
-            if self.distributions:
-                if 'parameter_server' in self.distributions:
-                    enabled = self.distributions['parameter_server'].get('enabled', False)
-                    additional_hyperparameters[self.LAUNCH_PS_ENV_NAME] = enabled
+            if 'parameter_server' in self.distributions:
+                enabled = self.distributions['parameter_server'].get('enabled', False)
+                additional_hyperparameters[self.LAUNCH_PS_ENV_NAME] = enabled
         else:
             additional_hyperparameters = {'checkpoint_path': self.checkpoint_path,
                                           'training_steps': self.training_steps,
@@ -435,15 +434,15 @@ def _default_s3_path(self, directory):
         else:
             return os.path.join(self.output_path, self._current_job_name, directory)
 
-    def script_mode_enabled(self):
+    def _script_mode_enabled(self):
         return self.py_version == 'py3' or self.script_mode
 
     def train_image(self):
         if self.image_name:
             return self.image_name
 
-        if self.script_mode_enabled():
+        if self._script_mode_enabled():
             return fw.create_image_uri(self.sagemaker_session.boto_region_name, _SCRIPT_MODE,
                                        self.train_instance_type, self.framework_version, self.py_version)
-        else:
-            return super(TensorFlow, self).train_image()
+
+        return super(TensorFlow, self).train_image()

tests/data/tensorflow_mnist/mnist.py

@@ -10,19 +10,16 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-from __future__ import absolute_import
+from __future__ import absolute_import, division, print_function
 
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-import tensorflow as tf
-import os
-import json
 import argparse
-from tensorflow.python.platform import tf_logging
+import json
 import logging as _logging
+import numpy as np
+import os
 import sys as _sys
+import tensorflow as tf
+from tensorflow.python.platform import tf_logging
 
 tf.logging.set_verbosity(tf.logging.DEBUG)
 _handler = _logging.StreamHandler(_sys.stdout)
@@ -137,11 +134,11 @@ def _parse_args():
     # hyperparameters sent by the client are passed as command-line arguments to the script.
     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('--train', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
-    parser.add_argument('--hosts', type=list, default=json.loads(os.environ['SM_HOSTS']))
-    parser.add_argument('--current-host', type=str, default=os.environ['SM_CURRENT_HOST'])
+    parser.add_argument('--output-data-dir', type=str, default=os.environ.get('SM_OUTPUT_DATA_DIR'))
+    parser.add_argument('--model_dir', type=str)
+    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAINING'))
+    parser.add_argument('--hosts', type=list, default=json.loads(os.environ.get('SM_HOSTS')))
+    parser.add_argument('--current-host', type=str, default=os.environ.get('SM_CURRENT_HOST'))
 
     return parser.parse_known_args()
 

tests/integ/test_tf_script_mode.py

@@ -56,7 +56,8 @@ def test_mnist_distributed(sagemaker_session, instance_type):
                            train_instance_count=2,
                            train_instance_type=instance_type,
                            sagemaker_session=sagemaker_session,
-                           py_version='py3',
+                           py_version=integ.PYTHON_VERSION,
+                           script_mode=True,
                            framework_version='1.11',
                            distributions=DISTRIBUTION_ENABLED,
                            base_job_name='test-tf-sm-mnist')

tests/unit/test_tf_estimator.py

@@ -690,22 +690,19 @@ def test_script_mode_deprecated_args(sagemaker_session):
 
 def test_script_mode_enabled(sagemaker_session):
     tf = _build_tf(sagemaker_session=sagemaker_session, py_version='py3')
-    assert tf.script_mode_enabled() is True
+    assert tf._script_mode_enabled() is True
 
     tf = _build_tf(sagemaker_session=sagemaker_session, script_mode=True)
-    assert tf.script_mode_enabled() is True
+    assert tf._script_mode_enabled() is True
 
     tf = _build_tf(sagemaker_session=sagemaker_session)
-    assert tf.script_mode_enabled() is False
+    assert tf._script_mode_enabled() is False
 
 
 @patch('sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model')
 def test_script_mode_create_model(create_tfs_model, sagemaker_session):
     tf = _build_tf(sagemaker_session=sagemaker_session, py_version='py3')
-    with pytest.raises(ValueError) as e:
-        tf.create_model()
-    assert tfe._SCRIPT_MODE_SERVING_ERROR_MSG in str(e)
-    tf.create_model(endpoint_type='tensorflow-serving')
+    tf.create_model()
     create_tfs_model.assert_called_once()