feature: add EIA support for TFS 1.15.0 and 2.0.0 (#1436)

Co-authored-by: Deng <[email protected]>

Author: chuyang-deng

README.rst

@@ -215,7 +215,7 @@ By using TensorFlow SageMaker Estimators, you can train and host TensorFlow mode
 
 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.0``, ``1.15.0``, ``1.15.2``, ``2.0.0``, ``2.0.1``, ``2.1.0``.
 
-Supported versions of TensorFlow for Elastic Inference: ``1.11.0``, ``1.12.0``, ``1.13.1``, ``1.14.0``.
+Supported versions of TensorFlow for Elastic Inference: ``1.11.0``, ``1.12.0``, ``1.13.1``, ``1.14.0``, ``1.15.0``, ``2.0.0``.
 
 We recommend that you use the latest supported version, because that's where we focus most of our development efforts.
 

src/sagemaker/tensorflow/README.rst

@@ -20,7 +20,7 @@ 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.14``.
+Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``, ``1.14``, ``1.15``, ``2.0``.
 
 Supported versions of TensorFlow for Inferentia: ``1.15.0``.
 

src/sagemaker/tensorflow/serving.py

@@ -130,7 +130,7 @@ class Model(sagemaker.model.FrameworkModel):
         logging.ERROR: "error",
         logging.CRITICAL: "crit",
     }
-    LATEST_EIA_VERSION = [1, 14]
+    LATEST_EIA_VERSION = [2, 0]
 
     def __init__(
         self,

tests/conftest.py

@@ -65,7 +65,7 @@ def pytest_addoption(parser):
     )
     parser.addoption("--sklearn-full-version", action="store", default=SKLEARN_VERSION)
     parser.addoption("--tf-full-version", action="store")
-    parser.addoption("--ei-tf-full-version", action="store", default=TensorFlow.LATEST_VERSION)
+    parser.addoption("--ei-tf-full-version", action="store")
     parser.addoption("--xgboost-full-version", action="store", default=SKLEARN_VERSION)
 
 
@@ -268,9 +268,13 @@ def tf_full_version(request):
         return tf_version
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="module", params=["1.15.0", "2.0.0"])
 def ei_tf_full_version(request):
-    return request.config.getoption("--ei-tf-full-version")
+    tf_ei_version = request.config.getoption("--ei-tf-full-version")
+    if tf_ei_version is None:
+        return request.param
+    else:
+        tf_ei_version
 
 
 @pytest.fixture(scope="session")

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, cpu_instance_type):
+def tfs_predictor_with_accelerator(sagemaker_session, ei_tf_full_version, cpu_instance_type):
     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, cpu_insta
         model = Model(
             model_data=model_data,
             role="SageMakerRole",
-            framework_version="1.14",
+            framework_version=ei_tf_full_version,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(

tests/unit/test_fw_utils.py

@@ -302,13 +302,13 @@ def test_tf_eia_images():
         "us-west-2",
         "tensorflow-serving",
         "ml.m4.xlarge",
-        "1.14.0",
+        "2.0.0",
         "py3",
         accelerator_type="ml.eia1.medium",
     )
     assert (
         image_uri
-        == "{}.dkr.ecr.us-west-2.amazonaws.com/tensorflow-inference-eia:1.14.0-cpu".format(
+        == "{}.dkr.ecr.us-west-2.amazonaws.com/tensorflow-inference-eia:2.0.0-cpu".format(
             fw_utils.ASIMOV_PROD_ACCOUNT
         )
     )

tests/unit/test_tfs.py

@@ -115,7 +115,7 @@ def test_tfs_model_image_accelerator_not_supported(sagemaker_session):
     model = Model(
         "s3://some/data.tar.gz",
         role=ROLE,
-        framework_version="1.15",
+        framework_version="2.1",
         sagemaker_session=sagemaker_session,
     )
 
@@ -130,7 +130,7 @@ def test_tfs_model_image_accelerator_not_supported(sagemaker_session):
             initial_instance_count=1,
         )
 
-    assert str(e.value) == "The TensorFlow version 1.15 doesn't support EIA."
+    assert str(e.value) == "The TensorFlow version 2.1 doesn't support EIA."
 
 
 def test_tfs_model_with_log_level(sagemaker_session, tf_version):