feat: InferenceSpec support for MMS and testing

Author: Bryannah Hernandez

src/sagemaker/serve/builder/model_builder.py

@@ -881,8 +881,8 @@ def _build_for_model_server(self):  # pylint: disable=R0911, R1710
         if self.model_metadata:
             mlflow_path = self.model_metadata.get(MLFLOW_MODEL_PATH)
 
-        if not self.model and not mlflow_path:
-            raise ValueError("Missing required parameter `model` or 'ml_flow' path")
+        if not self.model and not mlflow_path and not self.inference_spec:
+            raise ValueError("Missing required parameter `model` or 'ml_flow' path or inf_spec")
 
         if self.model_server == ModelServer.TORCHSERVE:
             return self._build_for_torchserve()

src/sagemaker/serve/builder/transformers_builder.py

@@ -72,6 +72,7 @@ def __init__(self):
         self.pytorch_version = None
         self.instance_type = None
         self.schema_builder = None
+        self.inference_spec = None
 
     @abstractmethod
     def _prepare_for_mode(self):
@@ -109,7 +110,7 @@ def _get_hf_metadata_create_model(self) -> Type[Model]:
         """
 
         hf_model_md = get_huggingface_model_metadata(
-            self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
+            self.env_vars.get("HF_MODEL_ID"), self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
         )
         hf_config = image_uris.config_for_framework("huggingface").get("inference")
         config = hf_config["versions"]
@@ -246,18 +247,22 @@ def _build_transformers_env(self):
 
         _create_dir_structure(self.model_path)
         if not hasattr(self, "pysdk_model"):
-            self.env_vars.update({"HF_MODEL_ID": self.model})
+
+            if self.inference_spec is not None:
+                self.env_vars.update({"HF_MODEL_ID": self.inference_spec.get_model()})
+            else:
+                self.env_vars.update({"HF_MODEL_ID": self.model})
 
             logger.info(self.env_vars)
 
             # TODO: Move to a helper function
             if hasattr(self.env_vars, "HF_API_TOKEN"):
                 self.hf_model_config = _get_model_config_properties_from_hf(
-                    self.model, self.env_vars.get("HF_API_TOKEN")
+                    self.env_vars.get("HF_MODEL_ID"), self.env_vars.get("HF_API_TOKEN")
                 )
             else:
                 self.hf_model_config = _get_model_config_properties_from_hf(
-                    self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
+                    self.env_vars.get("HF_MODEL_ID"), self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
                 )
 
         self.pysdk_model = self._create_transformers_model()

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

@@ -1,24 +1,17 @@
 """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
 
 
@@ -32,20 +25,12 @@ def model_fn(model_dir):
 
     serve_path = Path(__file__).parent.joinpath("serve.pkl")
     with open(str(serve_path), mode="rb") as file:
-        global inference_spec, native_model, schema_builder
+        global inference_spec, 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:
+            
+    if inference_spec:
         return partial(inference_spec.invoke, model=inference_spec.load(model_dir))
 
 

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

@@ -43,7 +43,6 @@ def _start_serving(
                     "mode": "rw",
                 },
             },
-
             environment={
                 "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
                 "SAGEMAKER_PROGRAM": "inference.py",
@@ -88,12 +87,10 @@ 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)
@@ -132,7 +129,6 @@ def _upload_server_artifacts(
             "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(),
         }