InferenceSpec support for HF

Author: Bryannah Hernandez

src/sagemaker/serve/model_server/multi_model_server/inference.py

@@ -0,0 +1,83 @@
+"""This module is for SageMaker inference.py."""
+
+from __future__ import absolute_import
+import os
+import io
+import cloudpickle
+import shutil
+import platform
+import importlib
+from pathlib import Path
+from functools import partial
+from sagemaker.serve.validations.check_integrity import perform_integrity_check
+from sagemaker.serve.spec.inference_spec import InferenceSpec
+from sagemaker.serve.detector.image_detector import _detect_framework_and_version, _get_model_base
+from sagemaker.serve.detector.pickler import load_xgboost_from_json
+import logging
+
+logger = logging.getLogger(__name__)
+
+inference_spec = None
+native_model = None
+schema_builder = None
+
+
+def model_fn(model_dir):
+    """Placeholder docstring"""
+    shared_libs_path = Path(model_dir + "/shared_libs")
+
+    if shared_libs_path.exists():
+        # before importing, place dynamic linked libraries in shared lib path
+        shutil.copytree(shared_libs_path, "/lib", dirs_exist_ok=True)
+
+    serve_path = Path(__file__).parent.joinpath("serve.pkl")
+    with open(str(serve_path), mode="rb") as file:
+        global inference_spec, native_model, schema_builder
+        obj = cloudpickle.load(file)
+        if isinstance(obj[0], InferenceSpec): 
+            inference_spec, schema_builder = obj
+        else:
+            native_model, schema_builder = obj
+    if native_model:
+        framework, _ = _detect_framework_and_version(
+            model_base=str(_get_model_base(model=native_model))
+        )
+        if framework == "pytorch":
+            native_model.eval()
+        return native_model if callable(native_model) else native_model.predict
+    elif inference_spec:
+        return partial(inference_spec.invoke, model=inference_spec.load(model_dir))
+
+
+def input_fn(input_data, content_type):
+    """Placeholder docstring"""
+    try:
+        if hasattr(schema_builder, "custom_input_translator"):
+            return schema_builder.custom_input_translator.deserialize(
+                io.BytesIO(input_data), content_type
+            )
+        else:
+            return schema_builder.input_deserializer.deserialize(
+                io.BytesIO(input_data), content_type[0]
+            )
+    except Exception as e:
+        raise Exception("Encountered error in deserialize_request.") from e
+
+
+def predict_fn(input_data, predict_callable):
+    """Placeholder docstring"""
+    return predict_callable(input_data)
+
+
+def output_fn(predictions, accept_type):
+    """Placeholder docstring"""
+    try:
+        if hasattr(schema_builder, "custom_output_translator"):
+            return schema_builder.custom_output_translator.serialize(predictions, accept_type)
+        else:
+            return schema_builder.output_serializer.serialize(predictions)
+    except Exception as e:
+        logger.error("Encountered error: %s in serialize_response." % e)
+        raise Exception("Encountered error in serialize_response.") from e
+
+

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

@@ -4,6 +4,7 @@
 
 import requests
 import logging
+import platform
 from pathlib import Path
 from sagemaker import Session, fw_utils
 from sagemaker.serve.utils.exceptions import LocalModelInvocationException
@@ -42,7 +43,14 @@ def _start_serving(
                     "mode": "rw",
                 },
             },
-            environment=_update_env_vars(env_vars),
+
+            environment={
+                "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+                "SAGEMAKER_PROGRAM": "inference.py",
+                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+                "LOCAL_PYTHON": platform.python_version(),
+                **env_vars,
+            },
         )
 
     def _invoke_multi_model_server_serving(self, request: object, content_type: str, accept: str):
@@ -80,10 +88,12 @@ class SageMakerMultiModelServer:
     def _upload_server_artifacts(
         self,
         model_path: str,
+        secret_key: str,
         sagemaker_session: Session,
         s3_model_data_url: str = None,
         image: str = None,
         env_vars: dict = None,
+        
     ):
         if s3_model_data_url:
             bucket, key_prefix = parse_s3_url(url=s3_model_data_url)
@@ -116,6 +126,16 @@ def _upload_server_artifacts(
                 "S3Uri": model_data_url + "/",
             }
         }
+
+        env_vars = {
+            "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "SAGEMAKER_REGION": sagemaker_session.boto_region_name,
+            "SAGEMAKER_CONTAINER_LOG_LEVEL": "10",
+            "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+            "LOCAL_PYTHON": platform.python_version(),
+        }
+
         return model_data, _update_env_vars(env_vars)
 
 

tests/unit/sagemaker/serve/model_server/multi_model_server/test_multi_model_server_prepare.py

@@ -12,13 +12,64 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+from pathlib import PosixPath
+import platform
 from unittest import TestCase
 from unittest.mock import Mock, patch
 
+import numpy as np
+
 from sagemaker.serve.model_server.multi_model_server.prepare import _create_dir_structure
 
+from sagemaker.serve.model_server.multi_model_server.server import (
+    LocalMultiModelServer,
+)
+
+CPU_TF_IMAGE = "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04"
+MODEL_PATH = "model_path"
+MODEL_REPO = f"{MODEL_PATH}/1"
+ENV_VAR = {"KEY": "VALUE"}
+PAYLOAD = np.random.rand(3, 4).astype(dtype=np.float32)
+DTYPE = "TYPE_FP32"
+SECRET_KEY = "secret_key"
+INFER_RESPONSE = {"outputs": [{"name": "output_name"}]}
+
 
 class MultiModelServerPrepareTests(TestCase):
+    def test_start_invoke_destroy_local_multi_model_server(self):
+        mock_container = Mock()
+        mock_docker_client = Mock()
+        mock_docker_client.containers.run.return_value = mock_container
+
+        local_multi_model_server = LocalMultiModelServer()
+        mock_schema_builder = Mock()
+        mock_schema_builder.input_serializer.serialize.return_value = PAYLOAD
+        local_multi_model_server.schema_builder = mock_schema_builder
+
+        local_multi_model_server._start_serving(
+            client=mock_docker_client,
+            model_path=MODEL_PATH,
+            secret_key=SECRET_KEY,
+            env_vars=ENV_VAR,
+            image=CPU_TF_IMAGE,
+        )
+
+        mock_docker_client.containers.run.assert_called_once_with(
+            CPU_TF_IMAGE,
+            "serve",
+            detach=True,
+            auto_remove=True,
+            network_mode="host",
+            volumes={PosixPath("model_path"): {"bind": "/opt/ml/model", "mode": "rw"}},
+            environment={
+                "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+                "SAGEMAKER_PROGRAM": "inference.py",
+                "SAGEMAKER_SERVE_SECRET_KEY": "secret_key",
+                "LOCAL_PYTHON": platform.python_version(),
+                "KEY": "VALUE",
+            },
+        )
+
     @patch("sagemaker.serve.model_server.multi_model_server.prepare._check_disk_space")
     @patch("sagemaker.serve.model_server.multi_model_server.prepare._check_docker_disk_usage")
     @patch("sagemaker.serve.model_server.multi_model_server.prepare.Path")