feat: Enable galactus integ tests

jiapinw · jiapinw · commit 566f17bb6c35 · 2024-01-16T12:09:00.000-08:00
diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
@@ -39,3 +39,6 @@ tritonclient[http]<2.37.0
 onnx==1.14.1
 # tf2onnx==1.15.1
 nbformat>=5.9,<6
+torch==2.0.1
+torchvision==0.15.2
+torchaudio==2.0.2
diff --git a/tests/integ/sagemaker/serve/constants.py b/tests/integ/sagemaker/serve/constants.py
@@ -12,24 +12,24 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-# import os
+import os
 import platform
 
-# from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR
 
-# SERVE_IN_PROCESS_TIMEOUT = 5
-# SERVE_MODEL_PACKAGE_TIMEOUT = 10
-# SERVE_LOCAL_CONTAINER_TIMEOUT = 10
+SERVE_IN_PROCESS_TIMEOUT = 5
+SERVE_MODEL_PACKAGE_TIMEOUT = 10
+SERVE_LOCAL_CONTAINER_TIMEOUT = 10
 SERVE_SAGEMAKER_ENDPOINT_TIMEOUT = 15
-# SERVE_SAVE_TIMEOUT = 2
+SERVE_SAVE_TIMEOUT = 2
 
-# NOT_RUNNING_ON_PY38 = platform.python_version_tuple()[1] != "8"
+NOT_RUNNING_ON_PY38 = platform.python_version_tuple()[1] != "8"
 NOT_RUNNING_ON_PY310 = platform.python_version_tuple()[1] != "10"
-# NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE = os.getenv("TEST_OWNER") != "INF_EXP_DEV"
+NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE = os.getenv("TEST_OWNER") != "INF_EXP_DEV"
 
-# XGB_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "xgboost")
-# PYTORCH_SQUEEZENET_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "pytorch")
-# TF_EFFICIENT_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "tensorflow")
-# HF_DIR = os.path.join(DATA_DIR, "serve_resources", "hf")
+XGB_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "xgboost")
+PYTORCH_SQUEEZENET_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "pytorch")
+TF_EFFICIENT_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "tensorflow")
+HF_DIR = os.path.join(DATA_DIR, "serve_resources", "hf")
 
-# BYOC_IMAGE_URI_TEMPLATE = "661407751302.dkr.ecr.{}.amazonaws.com/byoc-integ-test-images:{}"
+BYOC_IMAGE_URI_TEMPLATE = "661407751302.dkr.ecr.{}.amazonaws.com/byoc-integ-test-images:{}"
diff --git a/tests/integ/sagemaker/serve/test_serve_pt_happy.py b/tests/integ/sagemaker/serve/test_serve_pt_happy.py
@@ -10,112 +10,141 @@
 # 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 pytest
-# import torch
-# from PIL import Image
-# import os
-
-# from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
-# from sagemaker.serve.builder.schema_builder import SchemaBuilder
-# from sagemaker.serve.spec.inference_spec import InferenceSpec
-# from torchvision.transforms import transforms
-# from torchvision.models.squeezenet import squeezenet1_1
-
-# from tests.integ.sagemaker.serve.constants import (
-#     PYTORCH_SQUEEZENET_RESOURCE_DIR,
-#     SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
-#     NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE,
-#     NOT_RUNNING_ON_PY310,
-# )
-# from tests.integ.timeout import timeout
-# from tests.integ.utils import cleanup_model_resources
-# import logging
-
-# logger = logging.getLogger(__name__)
-
-# ROLE_NAME = "Admin"
-
-# GH_USER_NAME = os.getenv("GH_USER_NAME")
-# GH_ACCESS_TOKEN = os.getenv("GH_ACCESS_TOKEN")
-
-
-# @pytest.fixture
-# def pt_dependencies():
-#     return {
-#         "auto": True,
-#         "custom": [
-#             "boto3==1.26.*",
-#             "botocore==1.29.*",
-#             "s3transfer==0.6.*",
-#             (
-#                 f"git+https://{GH_USER_NAME}:{GH_ACCESS_TOKEN}@github.com"
-#                 "/aws/sagemaker-python-sdk-staging.git@inference-experience-dev"
-#             ),
-#         ],
-#     }
-
-
-# @pytest.fixture
-# def test_image():
-#     return Image.open(str(os.path.join(PYTORCH_SQUEEZENET_RESOURCE_DIR, "zidane.jpeg")))
-
-
-# @pytest.fixture
-# def squeezenet_inference_spec():
-#     class MySqueezeNetModel(InferenceSpec):
-#         def __init__(self) -> None:
-#             super().__init__()
-#             self.transform = transforms.Compose(
-#                 [
-#                     transforms.Resize(256),
-#                     transforms.CenterCrop(224),
-#                     transforms.ToTensor(),
-#                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-#                 ]
-#             )
-
-#         def invoke(self, input_object: object, model: object):
-#             # transform
-#             image_tensor = self.transform(input_object)
-#             input_batch = image_tensor.unsqueeze(0)
-#             # invoke
-#             with torch.no_grad():
-#                 output = model(input_batch)
-#             return output
-
-#         def load(self, model_dir: str):
-#             model = squeezenet1_1()
-#             model.load_state_dict(torch.load(model_dir + "/model.pth"))
-#             model.eval()
-#             return model
-
-#     return MySqueezeNetModel()
-
-
-# @pytest.fixture
-# def squeezenet_schema():
-#     input_image = Image.open(os.path.join(PYTORCH_SQUEEZENET_RESOURCE_DIR, "zidane.jpeg"))
-#     output_tensor = torch.rand(3, 4)
-#     return SchemaBuilder(sample_input=input_image, sample_output=output_tensor)
-
-
-# @pytest.fixture
-# def model_builder_inference_spec_schema_builder(
-#     squeezenet_inference_spec, squeezenet_schema, pt_dependencies
-# ):
-#     return ModelBuilder(
-#         model_path=PYTORCH_SQUEEZENET_RESOURCE_DIR,
-#         inference_spec=squeezenet_inference_spec,
-#         schema_builder=squeezenet_schema,
-#         dependencies=pt_dependencies,
-#     )
-
-
-# @pytest.fixture
-# def model_builder(request):
-#     return request.getfixturevalue(request.param)
+from __future__ import absolute_import
+
+import pytest
+import torch
+from PIL import Image
+import os
+import io
+import numpy as np
+
+from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
+from sagemaker.serve.builder.schema_builder import SchemaBuilder, CustomPayloadTranslator
+from sagemaker.serve.spec.inference_spec import InferenceSpec
+from torchvision.transforms import transforms
+from torchvision.models.squeezenet import squeezenet1_1
+
+from tests.integ.sagemaker.serve.constants import (
+    PYTORCH_SQUEEZENET_RESOURCE_DIR,
+    SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
+    NOT_RUNNING_ON_PY310,
+)
+from tests.integ.timeout import timeout
+from tests.integ.utils import cleanup_model_resources
+import logging
+
+logger = logging.getLogger(__name__)
+
+ROLE_NAME = "SageMakerRole"
+
+
+@pytest.fixture
+def test_image():
+    return Image.open(str(os.path.join(PYTORCH_SQUEEZENET_RESOURCE_DIR, "zidane.jpeg")))
+
+
+@pytest.fixture
+def squeezenet_inference_spec():
+    class MySqueezeNetModel(InferenceSpec):
+        def __init__(self) -> None:
+            super().__init__()
+
+        def invoke(self, input_object: object, model: object):
+            with torch.no_grad():
+                output = model(input_object)
+            return output
+
+        def load(self, model_dir: str):
+            model = squeezenet1_1()
+            model.load_state_dict(torch.load(model_dir + "/model.pth"))
+            model.eval()
+            return model
+
+    return MySqueezeNetModel()
+
+
+@pytest.fixture
+def custom_request_translator():
+    # request translator
+    class MyRequestTranslator(CustomPayloadTranslator):
+        def __init__(self):
+            super().__init__()
+            # Define image transformation
+            self.transform = transforms.Compose(
+                [
+                    transforms.Resize(256),
+                    transforms.CenterCrop(224),
+                    transforms.ToTensor(),
+                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ]
+            )
+
+        # This function converts the payload to bytes - happens on client side
+        def serialize_payload_to_bytes(self, payload: object) -> bytes:
+            # converts an image to bytes
+            image_tensor = self.transform(payload)
+            input_batch = image_tensor.unsqueeze(0)
+            input_ndarray = input_batch.numpy()
+            return self._convert_numpy_to_bytes(input_ndarray)
+
+        # This function converts the bytes to payload - happens on server side
+        def deserialize_payload_from_stream(self, stream) -> torch.Tensor:
+            # convert payload back to torch.Tensor
+            np_array = np.load(io.BytesIO(stream.read()))
+            return torch.from_numpy(np_array)
+
+        def _convert_numpy_to_bytes(self, np_array: np.ndarray) -> bytes:
+            buffer = io.BytesIO()
+            np.save(buffer, np_array)
+            return buffer.getvalue()
+
+    return MyRequestTranslator()
+
+
+@pytest.fixture
+def custom_response_translator():
+    # response translator
+    class MyResponseTranslator(CustomPayloadTranslator):
+        # This function converts the payload to bytes - happens on server side
+        def serialize_payload_to_bytes(self, payload: torch.Tensor) -> bytes:
+            return self._convert_numpy_to_bytes(payload.numpy())
+
+        # This function converts the bytes to payload - happens on client side
+        def deserialize_payload_from_stream(self, stream) -> object:
+            return torch.from_numpy(np.load(io.BytesIO(stream.read())))
+
+        def _convert_numpy_to_bytes(self, np_array: np.ndarray) -> bytes:
+            buffer = io.BytesIO()
+            np.save(buffer, np_array)
+            return buffer.getvalue()
+
+    return MyResponseTranslator()
+
+
+@pytest.fixture
+def squeezenet_schema(custom_request_translator, custom_response_translator):
+    input_image = Image.open(os.path.join(PYTORCH_SQUEEZENET_RESOURCE_DIR, "zidane.jpeg"))
+    output_tensor = torch.rand(3, 4)
+    return SchemaBuilder(
+        sample_input=input_image,
+        sample_output=output_tensor,
+        input_translator=custom_request_translator,
+        output_translator=custom_response_translator,
+    )
+
+@pytest.fixture
+def model_builder_inference_spec_schema_builder(squeezenet_inference_spec, squeezenet_schema):
+    return ModelBuilder(
+        model_path=PYTORCH_SQUEEZENET_RESOURCE_DIR,
+        inference_spec=squeezenet_inference_spec,
+        schema_builder=squeezenet_schema,
+    )
+
+
+@pytest.fixture
+def model_builder(request):
+    return request.getfixturevalue(request.param)
 
 
 # @pytest.mark.skipif(
@@ -149,45 +178,45 @@
 #                 ), f"{caught_ex} was thrown when running pytorch squeezenet local container test"
 
 
-# @pytest.mark.skipif(
-#     NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE or NOT_RUNNING_ON_PY310,
-#     reason="The goal of these test are to test the serving components of our feature",
-# )
-# @pytest.mark.parametrize(
-#     "model_builder", ["model_builder_inference_spec_schema_builder"], indirect=True
-# )
-# def test_happy_pytorch_sagemaker_endpoint(
-#     sagemaker_session, model_builder, cpu_instance_type, test_image
-# ):
-#     logger.info("Running in SAGEMAKER_ENDPOINT mode...")
-#     caught_ex = None
-
-#     iam_client = sagemaker_session.boto_session.client("iam")
-#     role_arn = iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"]
-
-#     model = model_builder.build(
-#         mode=Mode.SAGEMAKER_ENDPOINT, role_arn=role_arn, sagemaker_session=sagemaker_session
-#     )
-
-#     with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
-#         try:
-#             logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
-#             predictor = model.deploy(instance_type=cpu_instance_type, initial_instance_count=1)
-#             logger.info("Endpoint successfully deployed.")
-#             predictor.predict(test_image)
-#         except Exception as e:
-#             caught_ex = e
-#         finally:
-#             cleanup_model_resources(
-#                 sagemaker_session=model_builder.sagemaker_session,
-#                 model_name=model.name,
-#                 endpoint_name=model.endpoint_name,
-#             )
-#             if caught_ex:
-#                 logger.exception(caught_ex)
-#                 assert (
-#                     False
-#                 ), f"{caught_ex} was thrown when running pytorch squeezenet sagemaker endpoint test"
+@pytest.mark.skipif(
+    NOT_RUNNING_ON_PY310,  # or NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE,
+    reason="The goal of these test are to test the serving components of our feature",
+)
+@pytest.mark.parametrize(
+    "model_builder", ["model_builder_inference_spec_schema_builder"], indirect=True
+)
+def test_happy_pytorch_sagemaker_endpoint(
+    sagemaker_session, model_builder, cpu_instance_type, test_image
+):
+    logger.info("Running in SAGEMAKER_ENDPOINT mode...")
+    caught_ex = None
+
+    iam_client = sagemaker_session.boto_session.client("iam")
+    role_arn = iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"]
+
+    model = model_builder.build(
+        mode=Mode.SAGEMAKER_ENDPOINT, role_arn=role_arn, sagemaker_session=sagemaker_session
+    )
+
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        try:
+            logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
+            predictor = model.deploy(instance_type=cpu_instance_type, initial_instance_count=1)
+            logger.info("Endpoint successfully deployed.")
+            predictor.predict(test_image)
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=model_builder.sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                logger.exception(caught_ex)
+                assert (
+                    False
+                ), f"{caught_ex} was thrown when running pytorch squeezenet sagemaker endpoint test"
 
 
 # @pytest.mark.skipif(
