feature:support custom workflow deployment in ModelBuilder using SMD image. (#1661)

* feature:support custom workflow deployment in ModelBuilder using SMD inference image.

* Rename test case and pass session.

* Address PR comments.

* Tweak resource cleanup logic in integ test.

* Fixing CodeBuild integ test failures.

* Renamed integ test.

* Remove unused integ test, restore once GA.

---------

Co-authored-by: Joseph Zhang <[email protected]>

Author: 2 people

src/sagemaker/image_uri_config/sagemaker-distribution.json

@@ -0,0 +1,37 @@
+{
+    "processors": ["cpu", "gpu"],
+    "scope": ["inference"],
+    "version_aliases": {
+            "3.0": "3.0.0"
+        },
+    "versions": {
+        "3.0.0": {
+            "registries": {
+                "us-east-1": "885854791233",
+                "us-east-2": "137914896644",
+                "us-west-1": "053634841547",
+                "us-west-2": "542918446943",
+                "af-south-1": "238384257742",
+                "ap-east-1": "523751269255",
+                "ap-south-1": "245090515133",
+                "ap-northeast-2": "064688005998",
+                "ap-southeast-1": "022667117163",
+                "ap-southeast-2": "648430277019",
+                "ap-northeast-1": "010972774902",
+                "ca-central-1": "481561238223",
+                "eu-central-1": "545423591354",
+                "eu-west-1": "819792524951",
+                "eu-west-2": "021081402939",
+                "eu-west-3": "856416204555",
+                "eu-north-1": "175620155138",
+                "eu-south-1": "810671768855",
+                "sa-east-1": "567556641782",
+                "ap-northeast-3": "564864627153",
+                "ap-southeast-3": "370607712162",
+                "me-south-1": "523774347010",
+                "me-central-1": "358593528301"
+            },
+            "repository": "sagemaker-distribution-prod"
+        }
+    }
+}

src/sagemaker/serve/builder/model_builder.py

@@ -16,17 +16,21 @@
 from sagemaker.serve.model_server.djl_serving.server import SageMakerDjlServing
 from sagemaker.serve.model_server.tgi.server import SageMakerTgiServing
 from sagemaker.serve.model_server.multi_model_server.server import SageMakerMultiModelServer
+from sagemaker.serve.model_server.smd.server import SageMakerSmdServer
+
 
 logger = logging.getLogger(__name__)
 
 
+# pylint: disable=R0901
 class SageMakerEndpointMode(
     SageMakerTorchServe,
     SageMakerTritonServer,
     SageMakerDjlServing,
     SageMakerTgiServing,
     SageMakerMultiModelServer,
     SageMakerTensorflowServing,
+    SageMakerSmdServer,
 ):
     """Holds the required method to deploy a model to a SageMaker Endpoint"""
 
@@ -144,6 +148,16 @@ def prepare(
                 should_upload_artifacts=should_upload_artifacts,
             )
 
+        if self.model_server == ModelServer.SMD:
+            upload_artifacts = self._upload_smd_artifacts(
+                model_path=model_path,
+                sagemaker_session=sagemaker_session,
+                secret_key=secret_key,
+                s3_model_data_url=s3_model_data_url,
+                image=image,
+                should_upload_artifacts=True,
+            )
+
         if upload_artifacts or isinstance(self.model_server, ModelServer):
             return upload_artifacts
 

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -0,0 +1,72 @@
+# Copyright 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.
+"""This module is for SageMaker inference.py."""
+
+from __future__ import absolute_import
+import asyncio
+import os
+import platform
+import cloudpickle
+import logging
+from pathlib import Path
+from sagemaker.serve.validations.check_integrity import perform_integrity_check
+
+logger = LOGGER = logging.getLogger("sagemaker")
+
+
+def initialize_custom_orchestrator():
+    """Initializes the custom orchestrator."""
+    code_dir = os.getenv("SAGEMAKER_INFERENCE_CODE_DIRECTORY", None)
+    serve_path = Path(code_dir).joinpath("serve.pkl")
+    with open(str(serve_path), mode="rb") as pkl_file:
+        return cloudpickle.load(pkl_file)
+
+
+def _run_preflight_diagnostics():
+    _py_vs_parity_check()
+    _pickle_file_integrity_check()
+
+
+def _py_vs_parity_check():
+    container_py_vs = platform.python_version()
+    local_py_vs = os.getenv("LOCAL_PYTHON")
+
+    if not local_py_vs or container_py_vs.split(".")[1] != local_py_vs.split(".")[1]:
+        logger.warning(
+            f"The local python version {local_py_vs} differs from the python version "
+            f"{container_py_vs} on the container. Please align the two to avoid unexpected behavior"
+        )
+
+
+def _pickle_file_integrity_check():
+    with open("/opt/ml/model/code/serve.pkl", "rb") as f:
+        buffer = f.read()
+
+    metadata_path = Path("/opt/ml/model/code/metadata.json")
+    perform_integrity_check(buffer=buffer, metadata_path=metadata_path)
+
+
+_run_preflight_diagnostics()
+custom_orchestrator, _ = initialize_custom_orchestrator()
+
+
+async def handler(request):
+    """Custom service entry point function.
+
+    :param request: raw input from request
+    :return: outputs to be send back to client
+    """
+    if asyncio.iscoroutinefunction(custom_orchestrator.handle):
+        return await custom_orchestrator.handle(request)
+    else:
+        return custom_orchestrator.handle(request)

src/sagemaker/serve/model_server/smd/custom_execution_inference.py

@@ -0,0 +1,74 @@
+"""Summary of MyModule.
+
+Extended discussion of my module.
+"""
+
+from __future__ import absolute_import
+import os
+from pathlib import Path
+import shutil
+from typing import List
+
+from sagemaker.serve.spec.inference_spec import InferenceSpec
+from sagemaker.serve.detector.dependency_manager import capture_dependencies
+from sagemaker.serve.validations.check_integrity import (
+    generate_secret_key,
+    compute_hash,
+)
+from sagemaker.remote_function.core.serialization import _MetaData
+from sagemaker.serve.spec.inference_base import CustomOrchestrator, AsyncCustomOrchestrator
+
+
+def prepare_for_smd(
+    model_path: str,
+    shared_libs: List[str],
+    dependencies: dict,
+    inference_spec: InferenceSpec = None,
+) -> str:
+    """Prepares artifacts for SageMaker model deployment.
+
+    Args:to
+        model_path (str) : Argument
+        shared_libs (List[]) : Argument
+        dependencies (dict) : Argument
+        inference_spec (InferenceSpec, optional) : Argument
+            (default is None)
+
+    Returns:
+        ( str ) :
+
+    """
+    model_path = Path(model_path)
+    if not model_path.exists():
+        model_path.mkdir()
+    elif not model_path.is_dir():
+        raise Exception("model_dir is not a valid directory")
+
+    if inference_spec and isinstance(inference_spec, InferenceSpec):
+        inference_spec.prepare(str(model_path))
+
+    code_dir = model_path.joinpath("code")
+    code_dir.mkdir(exist_ok=True)
+
+    if inference_spec and isinstance(inference_spec, (CustomOrchestrator, AsyncCustomOrchestrator)):
+        shutil.copy2(Path(__file__).parent.joinpath("custom_execution_inference.py"), code_dir)
+        os.rename(
+            str(code_dir.joinpath("custom_execution_inference.py")),
+            str(code_dir.joinpath("inference.py")),
+        )
+
+    shared_libs_dir = model_path.joinpath("shared_libs")
+    shared_libs_dir.mkdir(exist_ok=True)
+    for shared_lib in shared_libs:
+        shutil.copy2(Path(shared_lib), shared_libs_dir)
+
+    capture_dependencies(dependencies=dependencies, work_dir=code_dir)
+
+    secret_key = generate_secret_key()
+    with open(str(code_dir.joinpath("serve.pkl")), "rb") as f:
+        buffer = f.read()
+    hash_value = compute_hash(buffer=buffer, secret_key=secret_key)
+    with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
+        metadata.write(_MetaData(hash_value).to_json())
+
+    return secret_key

src/sagemaker/serve/model_server/smd/prepare.py

@@ -0,0 +1,59 @@
+"""Module for SMD Server"""
+
+from __future__ import absolute_import
+
+import logging
+import platform
+from sagemaker.serve.utils.optimize_utils import _is_s3_uri
+from sagemaker.session import Session
+from sagemaker.s3_utils import determine_bucket_and_prefix, parse_s3_url
+from sagemaker import fw_utils
+from sagemaker.serve.utils.uploader import upload
+
+logger = logging.getLogger(__name__)
+
+
+class SageMakerSmdServer:
+    """Placeholder docstring"""
+
+    def _upload_smd_artifacts(
+        self,
+        model_path: str,
+        sagemaker_session: Session,
+        secret_key: str,
+        s3_model_data_url: str = None,
+        image: str = None,
+        should_upload_artifacts: bool = False,
+    ):
+        """Tar the model artifact and upload to S3 bucket, then prepare for the environment variables"""
+        s3_upload_path = None
+        if _is_s3_uri(model_path):
+            s3_upload_path = model_path
+        elif should_upload_artifacts:
+            if s3_model_data_url:
+                bucket, key_prefix = parse_s3_url(url=s3_model_data_url)
+            else:
+                bucket, key_prefix = None, None
+
+            code_key_prefix = fw_utils.model_code_key_prefix(key_prefix, None, image)
+
+            bucket, code_key_prefix = determine_bucket_and_prefix(
+                bucket=bucket, key_prefix=code_key_prefix, sagemaker_session=sagemaker_session
+            )
+
+            logger.debug(
+                "Uploading the model resources to bucket=%s, key_prefix=%s.",
+                bucket,
+                code_key_prefix,
+            )
+            s3_upload_path = upload(sagemaker_session, model_path, bucket, code_key_prefix)
+            logger.debug("Model resources uploaded to: %s", s3_upload_path)
+
+        env_vars = {
+            "SAGEMAKER_INFERENCE_CODE_DIRECTORY": "/opt/ml/model/code",
+            "SAGEMAKER_INFERENCE_CODE": "inference.handler",
+            "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
+            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+            "LOCAL_PYTHON": platform.python_version(),
+        }
+        return s3_upload_path, env_vars

src/sagemaker/serve/model_server/smd/server.py

@@ -0,0 +1,40 @@
+# Copyright 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.
+"""Holds templated classes to enable users to provide custom inference scripting capabilities"""
+from __future__ import absolute_import
+from abc import ABC, abstractmethod
+
+
+class CustomOrchestrator(ABC):
+    """Templated class to standardize sync entrypoint-based inference scripts"""
+
+    @property
+    def client(self):
+        """Boto3 SageMaker runtime client to use with custom orchestrator"""
+        from boto3 import Session
+
+        return Session().client("sagemaker-runtime")
+
+    @abstractmethod
+    def handle(self, data, context=None):
+        """Abstract class for defining an entrypoint for the model server"""
+        return NotImplemented
+
+
+class AsyncCustomOrchestrator(ABC):
+    """Templated class to standardize async entrypoint-based inference scripts"""
+
+    @abstractmethod
+    async def handle(self, data, context=None):
+        """Abstract class for defining an aynchronous entrypoint for the model server"""
+        return NotImplemented

src/sagemaker/serve/spec/inference_base.py

@@ -64,6 +64,7 @@
     str(ModelServer.TRITON): 5,
     str(ModelServer.TGI): 6,
     str(ModelServer.TEI): 7,
+    str(ModelServer.SMD): 8,
 }
 
 MLFLOW_MODEL_PATH_CODE = {

src/sagemaker/serve/utils/telemetry_logger.py

@@ -19,6 +19,7 @@ def __str__(self):
     TRITON = 5
     TGI = 6
     TEI = 7
+    SMD = 8
 
 
 class HardwareType(Enum):

src/sagemaker/serve/utils/types.py

@@ -25,6 +25,7 @@
 
 PYTHON_VERSION_IS_NOT_38 = platform.python_version_tuple()[1] != "8"
 PYTHON_VERSION_IS_NOT_310 = platform.python_version_tuple()[1] != "10"
+PYTHON_VERSION_IS_NOT_312 = platform.python_version_tuple()[1] != "12"
 
 XGB_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "xgboost")
 PYTORCH_SQUEEZENET_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "pytorch")