feature: Add support for TF-2.0.0.

README.rst

@@ -193,7 +193,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``, ``1.14``.
+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.``, ``1.15.0``, ``2.0.0``.
 
 Supported versions of TensorFlow for Elastic Inference: ``1.11.0``, ``1.12.0``, ``1.13.1``, ``1.14``.
 

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.0``, ``1.12.0``.
+Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``, ``1.14``.
 
 
 *****************************

src/sagemaker/tensorflow/README.rst

@@ -20,52 +20,52 @@ Documentation of the previous Legacy Mode versions: `1.4.1 <https://github.com/a
 | You can find the Legacy Mode README `here <https://github.com/aws/sagemaker-python-sdk/tree/v1.12.0/src/sagemaker/tensorflow#tensorflow-sagemaker-estimators-and-models>`_. |
 +-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 
-Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``, ``1.15``.
+Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``, ``1.14``.
 
 For information about using TensorFlow with the SageMaker Python SDK, see https://sagemaker.readthedocs.io/en/stable/using_tf.html.
 
 SageMaker TensorFlow Docker containers
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-The containers include the following Python packages:
+The latest containers include the following Python packages:
 
-+--------------------------------+---------------+
-| Dependencies                   | Script Mode   |
-+--------------------------------+---------------+
-| awscli                         | 1.16.130      |
-+--------------------------------+---------------+
-| boto3                          | Latest        |
-+--------------------------------+---------------+
-| botocore                       | Latest        |
-+--------------------------------+---------------+
-| h5py                           | 2.9.0         |
-+--------------------------------+---------------+
-| keras                          | 2.2.4         |
-+--------------------------------+---------------+
-| mpi4py                         | 3.0.1         |
-+--------------------------------+---------------+
-| numpy                          | 1.16.2        |
-+--------------------------------+---------------+
-| pandas                         | 0.24.2        |
-+--------------------------------+---------------+
-| pip                            | 19.1.1        |
-+--------------------------------+---------------+
-| Pillow                         | 5.4.1         |
-+--------------------------------+---------------+
-| Python                         | 2.7 or 3.6    |
-+--------------------------------+---------------+
-| requests                       | 2.21.0        |
-+--------------------------------+---------------+
-| sagemaker-containers           | >=2.4.6       |
-+--------------------------------+---------------+
-| sagemaker-tensorflow-container | Latest        |
-+--------------------------------+---------------+
-| scipy                          | 1.2.1         |
-+--------------------------------+---------------+
-| sklean                         | 0.20.3        |
-+--------------------------------+---------------+
-| tensorflow                     | 1.15.0        |
-+--------------------------------+---------------+
++--------------------------------+---------------+---------------+
+| Dependencies                   | TF 1.15       | TF 2.0        |
++--------------------------------+---------------+---------------+
+| awscli                         | 1.16.296      | 1.16.303      |
++--------------------------------+---------------+---------------+
+| boto3                          | Latest        | Latest        |
++--------------------------------+---------------+---------------+
+| botocore                       | Latest        | Latest        |
++--------------------------------+---------------+---------------+
+| h5py                           | 2.9.0         | 2.10.0        |
++--------------------------------+---------------+---------------+
+| horovod                        | 0.18.2        | 0.18.2        |
++--------------------------------+---------------+---------------+
+| keras                          | 2.3.1         | 2.3.1         |
++--------------------------------+---------------+---------------+
+| mpi4py                         | 3.0.2         | 3.0.3         |
++--------------------------------+---------------+---------------+
+| numpy                          | 1.17.4        | 1.17.4        |
++--------------------------------+---------------+---------------+
+| pandas                         | 0.24.2        | 0.25.3        |
++--------------------------------+---------------+---------------+
+| pip                            | 19.3.1        | 19.3.1        |
++--------------------------------+---------------+---------------+
+| Pillow                         | 6.1.0         | 6.2.1         |
++--------------------------------+---------------+---------------+
+| Python                         | 2.7 or 3.6    | 2.7 or 3.6    |
++--------------------------------+---------------+---------------+
+| requests                       | 2.22.0        | 2.22.0        |
++--------------------------------+---------------+---------------+
+| sagemaker-containers           | >=2.6.1       | >=2.6.2       |
++--------------------------------+---------------+---------------+
+| sagemaker-tensorflow-container | 1.15.0.1.0.0  | 2.0.0.1.0.0   |
++--------------------------------+---------------+---------------+
+| scipy                          | 1.2.2         | 1.2.2         |
++--------------------------------+---------------+---------------+
+| tensorflow                     | 1.15.0        | 2.0.0         |
++--------------------------------+---------------+---------------+
 
 Script Mode TensorFlow Docker images support both Python 2.7 and Python 3.6. The Docker images extend Ubuntu 16.04.
 

src/sagemaker/tensorflow/estimator.py

@@ -197,13 +197,15 @@ class TensorFlow(Framework):
 
     __framework_name__ = "tensorflow"
 
-    LATEST_VERSION = "1.15"
+    LATEST_VERSION = "2.0.0"
     """The latest version of TensorFlow included in the SageMaker pre-built Docker images."""
 
+    _LATEST_1X_VERSION = "1.15.0"
+
     _LOWEST_SCRIPT_MODE_ONLY_VERSION = [1, 13]
-    # 1.15.0 still supports py2
+    # 2.0.0 still supports py2
     # we will need to update this version number if future versions still support py2
-    _HIGHEST_PYTHON_2_VERSION = [1, 15, 0]
+    _HIGHEST_PYTHON_2_VERSION = [2, 0, 0]
 
     def __init__(
         self,

tests/conftest.py

@@ -56,7 +56,7 @@ def pytest_addoption(parser):
         "--rl-ray-full-version", action="store", default=RLEstimator.RAY_LATEST_VERSION
     )
     parser.addoption("--sklearn-full-version", action="store", default=SKLEARN_VERSION)
-    parser.addoption("--tf-full-version", action="store", default=TensorFlow.LATEST_VERSION)
+    parser.addoption("--tf-full-version", action="store")
     parser.addoption("--ei-tf-full-version", action="store", default=TensorFlow.LATEST_VERSION)
     parser.addoption("--xgboost-full-version", action="store", default=SKLEARN_VERSION)
 
@@ -246,9 +246,13 @@ def sklearn_full_version(request):
     return request.config.getoption("--sklearn-full-version")
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="module", params=[TensorFlow._LATEST_1X_VERSION, TensorFlow.LATEST_VERSION])
 def tf_full_version(request):
-    return request.config.getoption("--tf-full-version")
+    tf_version = request.config.getoption("--tf-full-version")
+    if tf_version is None:
+        return request.param
+    else:
+        return tf_version
 
 
 @pytest.fixture(scope="module")

tests/integ/test_tf_script_mode.py

@@ -39,7 +39,7 @@
 TAGS = [{"Key": "some-key", "Value": "some-value"}]
 
 
-def test_mnist_with_checkpoint_config(sagemaker_session, instance_type):
+def test_mnist_with_checkpoint_config(sagemaker_session, instance_type, tf_full_version):
     checkpoint_s3_uri = "s3://{}/checkpoints/tf-{}".format(
         sagemaker_session.default_bucket(), sagemaker_timestamp()
     )
@@ -51,7 +51,7 @@ def test_mnist_with_checkpoint_config(sagemaker_session, instance_type):
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,
         script_mode=True,
-        framework_version=TensorFlow.LATEST_VERSION,
+        framework_version=tf_full_version,
         py_version=tests.integ.PYTHON_VERSION,
         metric_definitions=[{"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}],
         checkpoint_s3_uri=checkpoint_s3_uri,
@@ -82,7 +82,7 @@ def test_mnist_with_checkpoint_config(sagemaker_session, instance_type):
     assert actual_training_checkpoint_config == expected_training_checkpoint_config
 
 
-def test_server_side_encryption(sagemaker_session):
+def test_server_side_encryption(sagemaker_session, tf_full_version):
     boto_session = sagemaker_session.boto_session
     with kms_utils.bucket_with_encryption(boto_session, ROLE) as (bucket_with_kms, kms_key):
         output_path = os.path.join(
@@ -97,7 +97,7 @@ def test_server_side_encryption(sagemaker_session):
             train_instance_type="ml.c5.xlarge",
             sagemaker_session=sagemaker_session,
             script_mode=True,
-            framework_version=TensorFlow.LATEST_VERSION,
+            framework_version=tf_full_version,
             py_version=tests.integ.PYTHON_VERSION,
             code_location=output_path,
             output_path=output_path,
@@ -125,7 +125,7 @@ def test_server_side_encryption(sagemaker_session):
 
 
 @pytest.mark.canary_quick
-def test_mnist_distributed(sagemaker_session, instance_type):
+def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version):
     estimator = TensorFlow(
         entry_point=SCRIPT,
         role=ROLE,
@@ -134,7 +134,7 @@ def test_mnist_distributed(sagemaker_session, instance_type):
         sagemaker_session=sagemaker_session,
         py_version=tests.integ.PYTHON_VERSION,
         script_mode=True,
-        framework_version=TensorFlow.LATEST_VERSION,
+        framework_version=tf_full_version,
         distributions=PARAMETER_SERVER_DISTRIBUTION,
     )
     inputs = estimator.sagemaker_session.upload_data(
@@ -159,6 +159,7 @@ def test_mnist_async(sagemaker_session, cpu_instance_type):
         py_version=tests.integ.PYTHON_VERSION,
         sagemaker_session=sagemaker_session,
         script_mode=True,
+        # testing py-sdk functionality, no need to run against all TF versions
         framework_version=TensorFlow.LATEST_VERSION,
         tags=TAGS,
     )
@@ -191,7 +192,7 @@ def test_mnist_async(sagemaker_session, cpu_instance_type):
         _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name)
 
 
-def test_deploy_with_input_handlers(sagemaker_session, instance_type):
+def test_deploy_with_input_handlers(sagemaker_session, instance_type, tf_full_version):
     estimator = TensorFlow(
         entry_point="training.py",
         source_dir=TFS_RESOURCE_PATH,
@@ -201,7 +202,7 @@ def test_deploy_with_input_handlers(sagemaker_session, instance_type):
         py_version=tests.integ.PYTHON_VERSION,
         sagemaker_session=sagemaker_session,
         script_mode=True,
-        framework_version=TensorFlow.LATEST_VERSION,
+        framework_version=tf_full_version,
         tags=TAGS,
     )
 

tests/integ/test_tuner.py

@@ -590,7 +590,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):
+def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type, tf_full_version):
     resource_path = os.path.join(DATA_DIR, "tensorflow_mnist")
     script_path = os.path.join(resource_path, "mnist.py")
 
@@ -602,7 +602,7 @@ def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type):
         script_mode=True,
         sagemaker_session=sagemaker_session,
         py_version=PYTHON_VERSION,
-        framework_version=TensorFlow.LATEST_VERSION,
+        framework_version=tf_full_version,
     )
 
     hyperparameter_ranges = {"epochs": IntegerParameter(1, 2)}

tests/unit/test_tf_estimator.py

@@ -982,7 +982,7 @@ def test_script_mode_deprecated_args(sagemaker_session):
 
 def test_py2_version_deprecated(sagemaker_session):
     with pytest.raises(AttributeError) as e:
-        _build_tf(sagemaker_session=sagemaker_session, framework_version="1.15.1", py_version="py2")
+        _build_tf(sagemaker_session=sagemaker_session, framework_version="2.0.1", py_version="py2")
 
     msg = "Python 2 containers are only available until January 1st, 2020. Please use a Python 3 container."
     assert msg in str(e.value)
@@ -993,6 +993,10 @@ def test_py2_version_is_not_deprecated(sagemaker_session):
         sagemaker_session=sagemaker_session, framework_version="1.15.0", py_version="py2"
     )
     assert estimator.py_version == "py2"
+    estimator = _build_tf(
+        sagemaker_session=sagemaker_session, framework_version="2.0.0", py_version="py2"
+    )
+    assert estimator.py_version == "py2"
 
 
 def test_py3_is_default_version_before_tf1_14(sagemaker_session):