feature: support for TensorFlow 1.14 (#967)

Author: mvsusp

README.rst

@@ -189,9 +189,9 @@ 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``
+Supported versions of TensorFlow for Elastic Inference: ``1.11.0``, ``1.12.0``, ``1.13.1``
 
 We recommend that you use the latest supported version, because that's where we focus most of our development efforts.
 

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:

src/sagemaker/tensorflow/serving.py

@@ -131,6 +131,7 @@ class Model(sagemaker.model.FrameworkModel):
         logging.ERROR: "error",
         logging.CRITICAL: "crit",
     }
+    LATEST_EIA_VERSION = [1, 13]
 
     def __init__(
         self,
@@ -176,6 +177,37 @@ def __init__(
         self._framework_version = framework_version
         self._container_log_level = container_log_level
 
+    def deploy(
+        self,
+        initial_instance_count,
+        instance_type,
+        accelerator_type=None,
+        endpoint_name=None,
+        update_endpoint=False,
+        tags=None,
+        kms_key=None,
+        wait=True,
+    ):
+
+        if accelerator_type and not self._eia_supported():
+            msg = "The TensorFlow version %s doesn't support EIA." % self._framework_version
+
+            raise AttributeError(msg)
+        return super(Model, self).deploy(
+            initial_instance_count,
+            instance_type,
+            accelerator_type,
+            endpoint_name,
+            update_endpoint,
+            tags,
+            kms_key,
+            wait,
+        )
+
+    def _eia_supported(self):
+        """Return true if TF version is EIA enabled"""
+        return [int(s) for s in self._framework_version.split(".")][:2] <= self.LATEST_EIA_VERSION
+
     def prepare_container_def(self, instance_type, accelerator_type=None):
         """
         Args:

tests/integ/test_tfs.py

@@ -111,7 +111,7 @@ def tfs_predictor_with_model_and_entry_point_and_dependencies(
 
 
 @pytest.fixture(scope="module")
-def tfs_predictor_with_accelerator(sagemaker_session, tf_full_version):
+def tfs_predictor_with_accelerator(sagemaker_session):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
         path=os.path.join(tests.integ.DATA_DIR, "tensorflow-serving-test-model.tar.gz"),
@@ -121,7 +121,7 @@ def tfs_predictor_with_accelerator(sagemaker_session, tf_full_version):
         model = Model(
             model_data=model_data,
             role="SageMakerRole",
-            framework_version=tf_full_version,
+            framework_version="1.13",
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(

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,

tests/unit/test_tfs.py

@@ -96,6 +96,41 @@ def test_tfs_model_image_accelerator(sagemaker_session, tf_version):
     assert isinstance(predictor, Predictor)
 
 
+def test_tfs_model_image_accelerator_not_supported(sagemaker_session):
+    model = Model(
+        "s3://some/data.tar.gz",
+        role=ROLE,
+        framework_version="1.13.1",
+        sagemaker_session=sagemaker_session,
+    )
+
+    # assert error is not raised
+
+    model.deploy(
+        instance_type="ml.c4.xlarge", initial_instance_count=1, accelerator_type="ml.eia1.medium"
+    )
+
+    model = Model(
+        "s3://some/data.tar.gz",
+        role=ROLE,
+        framework_version="1.14",
+        sagemaker_session=sagemaker_session,
+    )
+
+    # assert error is not raised
+
+    model.deploy(instance_type="ml.c4.xlarge", initial_instance_count=1)
+
+    with pytest.raises(AttributeError) as e:
+        model.deploy(
+            instance_type="ml.c4.xlarge",
+            accelerator_type="ml.eia1.medium",
+            initial_instance_count=1,
+        )
+
+    assert str(e.value) == "The TensorFlow version 1.14 doesn't support EIA."
+
+
 def test_tfs_model_with_log_level(sagemaker_session, tf_version):
     model = Model(
         "s3://some/data.tar.gz",

tox.ini

@@ -33,7 +33,7 @@ ignore =
     FI15,
     FI16,
     FI17,
-    FI18,
+    FI18,  # __future__ import "annotations" missing -> check only Python 3.7 compatible
     FI50,
     FI51,
     FI52,