change: revise integration test for compilation

Author: wanyixia

tests/conftest.py

@@ -223,6 +223,31 @@ def neo_pytorch_cpu_instance_type():
     return "ml.c5.xlarge"
 
 
+@pytest.fixture(scope="module")
+def tfs_eia_latest_py_version():
+    return "py3"
+
+
+@pytest.fixture(scope="module")
+def tfs_eia_latest_version():
+    return "2.3"
+
+
+@pytest.fixture(scope="module")
+def tfs_eia_target_device():
+    return "ml_eia2"
+
+
+@pytest.fixture(scope="module")
+def tfs_eia_cpu_instance_type():
+    return "ml.c5.xlarge"
+
+
+@pytest.fixture(scope="module")
+def tfs_eia_compilation_job_name():
+    return utils.name_from_base("tfs-eia-compilation")
+
+
 @pytest.fixture(scope="module")
 def xgboost_framework_version(xgboost_version):
     if xgboost_version in ("1", "latest"):

tests/integ/test_tfs.py

@@ -138,42 +138,6 @@ def tfs_predictor_with_accelerator(
         yield predictor
 
 
-@pytest.fixture(scope="module")
-def tfs_trt_predictor_with_accelerator(
-    sagemaker_session, tensorflow_eia_latest_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"),
-        key_prefix="tensorflow-serving/compiledmodels",
-    )
-    bucket = sagemaker_session.default_bucket()
-    with tests.integ.timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
-        model = TensorFlowModel(
-            model_data=model_data,
-            role="SageMakerRole",
-            framework_version=tensorflow_eia_latest_version,
-            sagemaker_session=sagemaker_session,
-            name=endpoint_name,
-        )
-        data_shape = {"input": [1, 224, 224, 3]}
-        tfs_eia_compilation_job_name = "tfs_eia_compilation_job_name"
-        compiled_model_path = "s3://{}/{}/output".format(bucket, tfs_eia_compilation_job_name)
-        compiled_model = model.compile(
-            target_instance_family='ml_eia2',
-            input_shape=data_shape,
-            output_path=compiled_model_path,
-            role="SageMakerRole",
-            job_name=tfs_eia_compilation_job_name,
-            framework='tensorflow',
-            framework_version='2.3'
-        )
-        predictor = compiled_model.deploy(
-            1, cpu_instance_type, endpoint_name=endpoint_name, accelerator_type="ml.eia2.large"
-        )
-        yield predictor
-
-
 @pytest.mark.release
 def test_predict(tfs_predictor):
     input_data = {"instances": [1.0, 2.0, 5.0]}
@@ -196,23 +160,6 @@ def test_predict_with_accelerator(tfs_predictor_with_accelerator):
     assert expected_result == result
 
 
-@pytest.mark.skipif(
-    tests.integ.test_region() not in tests.integ.EI_SUPPORTED_REGIONS,
-    reason="EI is not supported in region {}".format(tests.integ.test_region()),
-)
-@pytest.mark.release
-def test_trt_predict_with_accelerator(tfs_predictor_with_accelerator):
-    import numpy as np
-    import matplotlib.image as mpimg
-    path = os.path.join(tests.integ.DATA_DIR, "cuteCat.jpg")
-    img = mpimg.imread(path)
-    img = np.resize(img, (299, 299, 3))
-    img = np.expand_dims(img, axis=0)
-    input_data = {"inputs": img}
-    result = tfs_trt_predictor_with_accelerator.predict(input_data)
-    print("trt predictor result is: " + result)
-
-
 @pytest.mark.local_mode
 def test_predict_with_entry_point(tfs_predictor_with_model_and_entry_point_same_tar):
     input_data = {"instances": [1.0, 2.0, 5.0]}

tests/integ/test_tfs_eia_compilation.py

@@ -0,0 +1,83 @@
+# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import tarfile
+
+import botocore.exceptions
+import os
+
+import pytest
+import sagemaker
+import sagemaker.predictor
+import sagemaker.utils
+import tests.integ
+import tests.integ.timeout
+import numpy as np
+import matplotlib.image as mpimg
+from sagemaker.deserializers import JSONDeserializer
+from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor
+from sagemaker.serializers import CSVSerializer, IdentitySerializer
+from tests.integ import (
+    DATA_DIR,
+)
+from tests.integ.timeout import timeout_and_delete_endpoint_by_name
+
+INPUT_MODEL = os.path.join(DATA_DIR, "tensorflow-serving-test-model.tar.gz")
+INFERENCE_IMAGE = os.path.join(DATA_DIR, "cuteCat.jpg")
+
+
+def test_compile_and_deploy_with_accelerator(
+    sagemaker_session,
+    tfs_eia_cpu_instance_type, 
+    tfs_eia_latest_version,
+    tfs_eia_latest_py_version,
+    tfs_eia_target_device,
+    tfs_eia_compilation_job_name
+):
+    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"),
+        key_prefix="tensorflow-serving/compiledmodels",
+    )
+    bucket = sagemaker_session.default_bucket()
+    with tests.integ.timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        model = TensorFlowModel(
+            model_data=model_data,
+            role="SageMakerRole",
+            framework_version=tfs_eia_latest_version,
+            py_version=tfs_eia_latest_py_version
+            sagemaker_session=sagemaker_session,
+            name=endpoint_name,
+        )
+        data_shape = {"input": [1, 224, 224, 3]}
+        compiled_model_path = "s3://{}/{}/output".format(bucket, tfs_eia_compilation_job_name)
+        compiled_model = model.compile(
+            target_instance_family=tfs_eia_target_device,
+            input_shape=data_shape,
+            output_path=compiled_model_path,
+            role="SageMakerRole",
+            job_name=tfs_eia_compilation_job_name,
+            framework='tensorflow',
+            framework_version=tfs_eia_latest_version
+        )
+        predictor = compiled_model.deploy(
+            1, tfs_eia_cpu_instance_type, endpoint_name=endpoint_name, accelerator_type="ml.eia2.large"
+        )
+   
+        image_path = os.path.join(tests.integ.DATA_DIR, "cuteCat.jpg")
+        img = mpimg.imread(image_path)
+        img = np.resize(img, (224, 224, 3))
+        img = np.expand_dims(img, axis=0)
+        input_data = {"inputs": img}
+        result = predictor.predict(input_data)