feature: support TensorFlow training 2.2 (#1521)

* feature: TensorFlow 2.2 support

Author: chuyang-deng

doc/using_tf.rst

@@ -20,7 +20,7 @@ For general information about using the SageMaker Python SDK, see :ref:`overview
 
 .. contents::
 
-Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``, ``1.14``.
+Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``, ``1.14``, ``1.15``, ``2.0``.
 
 
 *****************************

src/sagemaker/tensorflow/defaults.py

@@ -18,7 +18,10 @@
 This is no longer updated so as to not break existing workflows.
 """
 
-LATEST_VERSION = "2.1.0"
+LATEST_VERSION = "2.2.0"
 """The latest version of TensorFlow included in the SageMaker pre-built Docker images."""
 
+LATEST_SERVING_VERSION = "2.1.0"
+"""The latest version of TensorFlow Serving included in the SageMaker pre-built Docker images."""
+
 LATEST_PY2_VERSION = "2.1.0"

tests/conftest.py

@@ -27,7 +27,8 @@
 from sagemaker.pytorch import PyTorch
 from sagemaker.rl import RLEstimator
 from sagemaker.sklearn.defaults import SKLEARN_VERSION
-from sagemaker.tensorflow.estimator import TensorFlow
+from sagemaker.tensorflow import TensorFlow
+from sagemaker.tensorflow.defaults import LATEST_VERSION, LATEST_SERVING_VERSION
 
 DEFAULT_REGION = "us-west-2"
 CUSTOM_BUCKET_NAME_PREFIX = "sagemaker-custom-bucket"
@@ -259,7 +260,7 @@ def sklearn_full_version(request):
     return request.config.getoption("--sklearn-full-version")
 
 
-@pytest.fixture(scope="module", params=[TensorFlow._LATEST_1X_VERSION, TensorFlow.LATEST_VERSION])
+@pytest.fixture(scope="module", params=[TensorFlow._LATEST_1X_VERSION, LATEST_VERSION])
 def tf_full_version(request):
     tf_version = request.config.getoption("--tf-full-version")
     if tf_version is None:
@@ -335,3 +336,10 @@ def pytest_generate_tests(metafunc):
 @pytest.fixture(scope="module")
 def xgboost_full_version(request):
     return request.config.getoption("--xgboost-full-version")
+
+
+@pytest.fixture(scope="module")
+def tf_serving_version(tf_full_version):
+    if tf_full_version == LATEST_VERSION:
+        return LATEST_SERVING_VERSION
+    return tf_full_version

tests/integ/test_data_capture_config.py

@@ -41,7 +41,7 @@
 
 
 def test_enabling_data_capture_on_endpoint_shows_correct_data_capture_status(
-    sagemaker_session, tf_full_version
+    sagemaker_session, tf_serving_version
 ):
     endpoint_name = unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
@@ -52,7 +52,7 @@ def test_enabling_data_capture_on_endpoint_shows_correct_data_capture_status(
         model = Model(
             model_data=model_data,
             role=ROLE,
-            framework_version=tf_full_version,
+            framework_version=tf_serving_version,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(
@@ -98,7 +98,7 @@ def test_enabling_data_capture_on_endpoint_shows_correct_data_capture_status(
 
 
 def test_disabling_data_capture_on_endpoint_shows_correct_data_capture_status(
-    sagemaker_session, tf_full_version
+    sagemaker_session, tf_serving_version
 ):
     endpoint_name = unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
@@ -109,7 +109,7 @@ def test_disabling_data_capture_on_endpoint_shows_correct_data_capture_status(
         model = Model(
             model_data=model_data,
             role=ROLE,
-            framework_version=tf_full_version,
+            framework_version=tf_serving_version,
             sagemaker_session=sagemaker_session,
         )
         destination_s3_uri = os.path.join(
@@ -184,7 +184,7 @@ def test_disabling_data_capture_on_endpoint_shows_correct_data_capture_status(
 
 
 def test_updating_data_capture_on_endpoint_shows_correct_data_capture_status(
-    sagemaker_session, tf_full_version
+    sagemaker_session, tf_serving_version
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
@@ -195,7 +195,7 @@ def test_updating_data_capture_on_endpoint_shows_correct_data_capture_status(
         model = Model(
             model_data=model_data,
             role=ROLE,
-            framework_version=tf_full_version,
+            framework_version=tf_serving_version,
             sagemaker_session=sagemaker_session,
         )
         destination_s3_uri = os.path.join(

tests/integ/test_model_monitor.py

@@ -88,7 +88,7 @@
 
 
 @pytest.fixture(scope="module")
-def predictor(sagemaker_session, tf_full_version):
+def predictor(sagemaker_session, tf_serving_version):
     endpoint_name = 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"),
@@ -100,7 +100,7 @@ def predictor(sagemaker_session, tf_full_version):
         model = Model(
             model_data=model_data,
             role=ROLE,
-            framework_version=tf_full_version,
+            framework_version=tf_serving_version,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(

tests/integ/test_tf_script_mode.py

@@ -19,6 +19,7 @@
 import pytest
 
 from sagemaker.tensorflow import TensorFlow
+from sagemaker.tensorflow.defaults import LATEST_SERVING_VERSION
 from sagemaker.utils import unique_name_from_base, sagemaker_timestamp
 
 import tests.integ
@@ -40,10 +41,8 @@
 
 
 @pytest.fixture(scope="module")
-def py_version(tf_full_version):
-    return (
-        "py37" if tf_full_version == TensorFlow._LATEST_1X_VERSION else tests.integ.PYTHON_VERSION
-    )
+def py_version(tf_full_version, tf_serving_version):
+    return "py37" if tf_full_version == tf_serving_version else tests.integ.PYTHON_VERSION
 
 
 def test_mnist_with_checkpoint_config(
@@ -61,7 +60,7 @@ def test_mnist_with_checkpoint_config(
         sagemaker_session=sagemaker_session,
         script_mode=True,
         framework_version=tf_full_version,
-        py_version=py_version,
+        py_version="py37",
         metric_definitions=[{"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}],
         checkpoint_s3_uri=checkpoint_s3_uri,
         checkpoint_local_path=checkpoint_local_path,
@@ -91,7 +90,7 @@ def test_mnist_with_checkpoint_config(
     assert actual_training_checkpoint_config == expected_training_checkpoint_config
 
 
-def test_server_side_encryption(sagemaker_session, tf_full_version, py_version):
+def test_server_side_encryption(sagemaker_session, tf_serving_version, py_version):
     with kms_utils.bucket_with_encryption(sagemaker_session, ROLE) as (bucket_with_kms, kms_key):
         output_path = os.path.join(
             bucket_with_kms, "test-server-side-encryption", time.strftime("%y%m%d-%H%M")
@@ -105,7 +104,7 @@ def test_server_side_encryption(sagemaker_session, tf_full_version, py_version):
             train_instance_type="ml.c5.xlarge",
             sagemaker_session=sagemaker_session,
             script_mode=True,
-            framework_version=tf_full_version,
+            framework_version=tf_serving_version,
             py_version=py_version,
             code_location=output_path,
             output_path=output_path,
@@ -140,7 +139,7 @@ def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version, py
         train_instance_count=2,
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,
-        py_version=py_version,
+        py_version="py37",
         script_mode=True,
         framework_version=tf_full_version,
         distributions=PARAMETER_SERVER_DISTRIBUTION,
@@ -168,7 +167,7 @@ def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, py_v
         sagemaker_session=sagemaker_session,
         script_mode=True,
         # testing py-sdk functionality, no need to run against all TF versions
-        framework_version=TensorFlow.LATEST_VERSION,
+        framework_version=LATEST_SERVING_VERSION,
         tags=TAGS,
     )
     inputs = estimator.sagemaker_session.upload_data(
@@ -200,7 +199,9 @@ def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, py_v
         _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name)
 
 
-def test_deploy_with_input_handlers(sagemaker_session, instance_type, tf_full_version, py_version):
+def test_deploy_with_input_handlers(
+    sagemaker_session, instance_type, tf_serving_version, py_version
+):
     estimator = TensorFlow(
         entry_point="training.py",
         source_dir=TFS_RESOURCE_PATH,
@@ -210,7 +211,7 @@ def test_deploy_with_input_handlers(sagemaker_session, instance_type, tf_full_ve
         py_version=py_version,
         sagemaker_session=sagemaker_session,
         script_mode=True,
-        framework_version=tf_full_version,
+        framework_version=tf_serving_version,
         tags=TAGS,
     )
 

tests/integ/test_tfs.py

@@ -27,7 +27,7 @@
 
 
 @pytest.fixture(scope="module")
-def tfs_predictor(sagemaker_session, tf_full_version):
+def tfs_predictor(sagemaker_session, tf_serving_version):
     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"),
@@ -37,7 +37,7 @@ def tfs_predictor(sagemaker_session, tf_full_version):
         model = Model(
             model_data=model_data,
             role="SageMakerRole",
-            framework_version=tf_full_version,
+            framework_version=tf_serving_version,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(1, "ml.c5.xlarge", endpoint_name=endpoint_name)
@@ -54,7 +54,7 @@ def tar_dir(directory, tmpdir):
 
 @pytest.fixture
 def tfs_predictor_with_model_and_entry_point_same_tar(
-    sagemaker_local_session, tf_full_version, tmpdir
+    sagemaker_local_session, tf_serving_version, tmpdir
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
 
@@ -65,7 +65,7 @@ def tfs_predictor_with_model_and_entry_point_same_tar(
     model = Model(
         model_data="file://" + model_tar,
         role="SageMakerRole",
-        framework_version=tf_full_version,
+        framework_version=tf_serving_version,
         sagemaker_session=sagemaker_local_session,
     )
     predictor = model.deploy(1, "local", endpoint_name=endpoint_name)
@@ -78,7 +78,7 @@ def tfs_predictor_with_model_and_entry_point_same_tar(
 
 @pytest.fixture(scope="module")
 def tfs_predictor_with_model_and_entry_point_and_dependencies(
-    sagemaker_local_session, tf_full_version
+    sagemaker_local_session, tf_serving_version
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
 
@@ -98,7 +98,7 @@ def tfs_predictor_with_model_and_entry_point_and_dependencies(
         model_data=model_data,
         role="SageMakerRole",
         dependencies=dependencies,
-        framework_version=tf_full_version,
+        framework_version=tf_serving_version,
         sagemaker_session=sagemaker_local_session,
     )
 

tests/integ/test_transformer.py

@@ -25,6 +25,7 @@
 from sagemaker.mxnet import MXNet
 from sagemaker.pytorch import PyTorchModel
 from sagemaker.tensorflow import TensorFlow
+from sagemaker.tensorflow.defaults import LATEST_SERVING_VERSION
 from sagemaker.transformer import Transformer
 from sagemaker.estimator import Estimator
 from sagemaker.utils import unique_name_from_base
@@ -351,7 +352,7 @@ def test_transform_tf_kms_network_isolation(sagemaker_session, cpu_instance_type
         role="SageMakerRole",
         train_instance_count=1,
         train_instance_type=cpu_instance_type,
-        framework_version=TensorFlow.LATEST_VERSION,
+        framework_version=LATEST_SERVING_VERSION,
         script_mode=True,
         py_version=PYTHON_VERSION,
         sagemaker_session=sagemaker_session,

tests/integ/test_tuner.py

@@ -37,6 +37,7 @@
 from sagemaker.predictor import json_deserializer
 from sagemaker.pytorch import PyTorch
 from sagemaker.tensorflow import TensorFlow
+from sagemaker.tensorflow.defaults import LATEST_VERSION
 from sagemaker.tuner import (
     IntegerParameter,
     ContinuousParameter,
@@ -51,6 +52,13 @@
 
 DATA_PATH = os.path.join(DATA_DIR, "iris", "data")
 
+PY37_SUPPORTED_FRAMEWORK_VERSION = [TensorFlow._LATEST_1X_VERSION, LATEST_VERSION]
+
+
+@pytest.fixture(scope="module")
+def py_version(tf_full_version):
+    return "py37" if tf_full_version in PY37_SUPPORTED_FRAMEWORK_VERSION else PYTHON_VERSION
+
 
 @pytest.fixture(scope="module")
 def kmeans_train_set(sagemaker_session):
@@ -590,7 +598,7 @@ def test_tuning_mxnet(sagemaker_session, mxnet_full_version, cpu_instance_type):
 
 
 @pytest.mark.canary_quick
-def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type, tf_full_version):
+def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type, tf_full_version, py_version):
     resource_path = os.path.join(DATA_DIR, "tensorflow_mnist")
     script_path = os.path.join(resource_path, "mnist.py")
 
@@ -601,7 +609,7 @@ def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type, tf_full_ver
         train_instance_type=cpu_instance_type,
         script_mode=True,
         sagemaker_session=sagemaker_session,
-        py_version=PYTHON_VERSION,
+        py_version=py_version,
         framework_version=tf_full_version,
     )