feature: support for TensorFlow 1.14

Author: mvsusp

README.rst

@@ -189,7 +189,7 @@ TensorFlow SageMaker Estimators
 
 By using TensorFlow SageMaker Estimators, you can train and host TensorFlow models on Amazon SageMaker.
 
-Supported versions of TensorFlow: ``1.4.1``, ``1.5.0``, ``1.6.0``, ``1.7.0``, ``1.8.0``, ``1.9.0``, ``1.10.0``, ``1.11.0``, ``1.12.0``, ``1.13.1``.
+Supported versions of TensorFlow: ``1.4.1``, ``1.5.0``, ``1.6.0``, ``1.7.0``, ``1.8.0``, ``1.9.0``, ``1.10.0``, ``1.11.0``, ``1.12.0``, ``1.13.1``, ``1.14``.
 
 Supported versions of TensorFlow for Elastic Inference: ``1.11.0``, ``1.12.0``, ``1.13.0``
 

doc/using_tf.rst

@@ -8,6 +8,9 @@ models on SageMaker Hosting.
 
 For general information about using the SageMaker Python SDK, see :ref:`overview:Using the SageMaker Python SDK`.
 
+.. warning::
+    The TensorFlow estimator is available only for Python 3, starting by the TensorFlow version 1.14.
+
 .. warning::
     We have added a new format of your TensorFlow training script with TensorFlow version 1.11.
     This new way gives the user script more flexibility.

src/sagemaker/tensorflow/estimator.py

@@ -195,17 +195,18 @@ class TensorFlow(Framework):
 
     __framework_name__ = "tensorflow"
 
-    LATEST_VERSION = "1.13"
+    LATEST_VERSION = "1.14"
     """The latest version of TensorFlow included in the SageMaker pre-built Docker images."""
 
     _LOWEST_SCRIPT_MODE_ONLY_VERSION = [1, 13]
+    _LOWEST_PYTHON_2_ONLY_VERSION = [1, 14]
 
     def __init__(
         self,
         training_steps=None,
         evaluation_steps=None,
         checkpoint_path=None,
-        py_version="py2",
+        py_version=None,
         framework_version=None,
         model_dir=None,
         requirements_file="",
@@ -279,6 +280,9 @@ def __init__(
             logger.warning(fw.empty_framework_version_warning(TF_VERSION, self.LATEST_VERSION))
         self.framework_version = framework_version or TF_VERSION
 
+        if not py_version:
+            py_version = "py3" if self._only_python_3_supported() else "py2"
+
         super(TensorFlow, self).__init__(image_name=image_name, **kwargs)
         self.checkpoint_path = checkpoint_path
 
@@ -337,6 +341,13 @@ def _validate_args(
                     )
                 )
 
+        if py_version == "py2" and self._only_python_3_supported():
+            msg = (
+                "Python 2 containers are only available until TensorFlow version 1.13.1. "
+                "Please use a Python 3 container."
+            )
+            raise AttributeError(msg)
+
         if (not self._script_mode_enabled()) and self._only_script_mode_supported():
             logger.warning(
                 "Legacy mode is deprecated in versions 1.13 and higher. Using script mode instead."
@@ -349,6 +360,12 @@ def _only_script_mode_supported(self):
             int(s) for s in self.framework_version.split(".")
         ] >= self._LOWEST_SCRIPT_MODE_ONLY_VERSION
 
+    def _only_python_3_supported(self):
+        """Placeholder docstring"""
+        return [
+            int(s) for s in self.framework_version.split(".")
+        ] >= self._LOWEST_PYTHON_2_ONLY_VERSION
+
     def _validate_requirements_file(self, requirements_file):
         """Placeholder docstring"""
         if not requirements_file:

tests/unit/test_fw_utils.py

@@ -137,6 +137,13 @@ def test_create_image_uri_gov_cloud():
 
 
 def test_create_image_uri_merged():
+    image_uri = fw_utils.create_image_uri(
+        "us-west-2", "tensorflow-scriptmode", "ml.p3.2xlarge", "1.14", "py3"
+    )
+    assert (
+        image_uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:1.14-gpu-py3"
+    )
+
     image_uri = fw_utils.create_image_uri(
         "us-west-2", "tensorflow-scriptmode", "ml.p3.2xlarge", "1.13.1", "py3"
     )

tests/unit/test_tf_estimator.py

@@ -25,7 +25,6 @@
 from sagemaker.tensorflow import defaults, TensorFlow, TensorFlowModel, TensorFlowPredictor
 import sagemaker.tensorflow.estimator as tfe
 
-
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
 SCRIPT_FILE = "dummy_script.py"
 SCRIPT_PATH = os.path.join(DATA_DIR, SCRIPT_FILE)
@@ -956,6 +955,30 @@ def test_script_mode_deprecated_args(sagemaker_session):
     ) in str(e.value)
 
 
+def test_py2_version_deprecated(sagemaker_session):
+    with pytest.raises(AttributeError) as e:
+        _build_tf(sagemaker_session=sagemaker_session, framework_version="1.14", py_version="py2")
+
+    msg = "Python 2 containers are only available until TensorFlow version 1.13.1. Please use a Python 3 container."
+    assert msg in str(e.value)
+
+
+def test_py3_is_default_version_after_tf1_14(sagemaker_session):
+    estimator = _build_tf(sagemaker_session=sagemaker_session, framework_version="1.14")
+
+    assert estimator.py_version == "py3"
+
+
+def test_py3_is_default_version_before_tf1_14(sagemaker_session):
+    estimator = _build_tf(sagemaker_session=sagemaker_session, framework_version="1.13")
+
+    assert estimator.py_version == "py2"
+
+    estimator = _build_tf(sagemaker_session=sagemaker_session, framework_version="1.10")
+
+    assert estimator.py_version == "py2"
+
+
 def test_legacy_mode_deprecated(sagemaker_session):
     tf = _build_tf(
         sagemaker_session=sagemaker_session,