Merge branch 'master' into feat/jumpstartmodel-attach

Author: evakravi

CHANGELOG.md

@@ -1,5 +1,35 @@
 # Changelog
 
+## v2.221.0 (2024-05-20)
+
+### Features
+
+ * onboard tei image config to pysdk
+
+### Bug Fixes and Other Changes
+
+ * JS Model with non-TGI/non-DJL deployment failure
+ * cover tei with image_uris.retrieve API
+ * Add more debuging
+ * model builder limited container support for endpoint mode.
+ * Image URI should take precedence for HF models
+
+## v2.220.0 (2024-05-15)
+
+### Features
+
+ * AutoGluon 1.1.0 image_uris update
+ * add new images for HF TGI release
+ * Add telemetry support for mlflow models
+
+### Bug Fixes and Other Changes
+
+ * add debug logs to workflow container dist creation
+ * model builder race condition on sagemaker session
+ * Add tensorflow_serving support for mlflow models and enable lineage tracking for mlflow models
+ * update image_uri_configs  05-09-2024 07:17:41 PST
+ * skip flakey tests pending investigation
+
 ## v2.219.0 (2024-05-08)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.219.1.dev0
+2.221.1.dev0

src/sagemaker/huggingface/llm_utils.py

@@ -65,6 +65,13 @@ def get_huggingface_llm_image_uri(
             image_scope="inference",
             inference_tool="neuronx",
         )
+    if backend == "huggingface-tei":
+        return image_uris.retrieve(
+            "huggingface-tei",
+            region=region,
+            version=version,
+            image_scope="inference",
+        )
     if backend == "lmi":
         version = version or "0.24.0"
         return image_uris.retrieve(framework="djl-deepspeed", region=region, version=version)

src/sagemaker/image_uri_config/huggingface-tei.json

@@ -0,0 +1,59 @@
+{
+    "inference": {
+        "processors": [
+            "gpu"
+        ],
+        "version_aliases": {
+            "1.2": "1.2.3"
+        },
+        "versions": {
+            "1.2.3": {
+                "py_versions": [
+                    "py310"
+                ],
+                "registries": {
+                    "af-south-1": "510948584623",
+                    "ap-east-1": "651117190479",
+                    "ap-northeast-1": "354813040037",
+                    "ap-northeast-2": "366743142698",
+                    "ap-northeast-3": "867004704886",
+                    "ap-south-1": "720646828776",
+                    "ap-south-2": "628508329040",
+                    "ap-southeast-1": "121021644041",
+                    "ap-southeast-2": "783357654285",
+                    "ap-southeast-3": "951798379941",
+                    "ap-southeast-4": "106583098589",
+                    "ca-central-1": "341280168497",
+                    "ca-west-1": "190319476487",
+                    "cn-north-1": "450853457545",
+                    "cn-northwest-1": "451049120500",
+                    "eu-central-1": "492215442770",
+                    "eu-central-2": "680994064768",
+                    "eu-north-1": "662702820516",
+                    "eu-south-1": "978288397137",
+                    "eu-south-2": "104374241257",
+                    "eu-west-1": "141502667606",
+                    "eu-west-2": "764974769150",
+                    "eu-west-3": "659782779980",
+                    "il-central-1": "898809789911",
+                    "me-central-1": "272398656194",
+                    "me-south-1": "801668240914",
+                    "sa-east-1": "737474898029",
+                    "us-east-1": "683313688378",
+                    "us-east-2": "257758044811",
+                    "us-gov-east-1": "237065988967",
+                    "us-gov-west-1": "414596584902",
+                    "us-iso-east-1": "833128469047",
+                    "us-isob-east-1": "281123927165",
+                    "us-west-1": "746614075791",
+                    "us-west-2": "246618743249"
+                },
+                "tag_prefix": "2.0.1-tei1.2.3",
+                "repository": "tei",
+                "container_version": {
+                    "gpu": "cu122-ubuntu22.04"
+                }
+            }
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -37,6 +37,7 @@
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
 HUGGING_FACE_LLM_FRAMEWORK = "huggingface-llm"
+HUGGING_FACE_TEI_FRAMEWORK = "huggingface-tei"
 HUGGING_FACE_LLM_NEURONX_FRAMEWORK = "huggingface-llm-neuronx"
 XGBOOST_FRAMEWORK = "xgboost"
 SKLEARN_FRAMEWORK = "sklearn"
@@ -477,6 +478,7 @@ def _validate_version_and_set_if_needed(version, config, framework):
     if version is None and framework in [
         DATA_WRANGLER_FRAMEWORK,
         HUGGING_FACE_LLM_FRAMEWORK,
+        HUGGING_FACE_TEI_FRAMEWORK,
         HUGGING_FACE_LLM_NEURONX_FRAMEWORK,
         STABILITYAI_FRAMEWORK,
     ]:

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -23,6 +23,7 @@
 from sagemaker import model_uris
 from sagemaker.serve.model_server.djl_serving.prepare import prepare_djl_js_resources
 from sagemaker.serve.model_server.djl_serving.utils import _get_admissible_tensor_parallel_degrees
+from sagemaker.serve.model_server.multi_model_server.prepare import prepare_mms_js_resources
 from sagemaker.serve.model_server.tgi.prepare import prepare_tgi_js_resources, _create_dir_structure
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.exceptions import (
@@ -35,6 +36,7 @@
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
     TgiLocalModePredictor,
+    TransformersLocalModePredictor,
 )
 from sagemaker.serve.utils.local_hardware import (
     _get_nb_instance,
@@ -90,6 +92,7 @@ def __init__(self):
         self.existing_properties = None
         self.prepared_for_tgi = None
         self.prepared_for_djl = None
+        self.prepared_for_mms = None
         self.schema_builder = None
         self.nb_instance_type = None
         self.ram_usage_model_load = None
@@ -137,7 +140,11 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
 
             if overwrite_mode == Mode.SAGEMAKER_ENDPOINT:
                 self.mode = self.pysdk_model.mode = Mode.SAGEMAKER_ENDPOINT
-                if not hasattr(self, "prepared_for_djl") or not hasattr(self, "prepared_for_tgi"):
+                if (
+                    not hasattr(self, "prepared_for_djl")
+                    or not hasattr(self, "prepared_for_tgi")
+                    or not hasattr(self, "prepared_for_mms")
+                ):
                     self.pysdk_model.model_data, env = self._prepare_for_mode()
             elif overwrite_mode == Mode.LOCAL_CONTAINER:
                 self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
@@ -160,6 +167,13 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                         dependencies=self.dependencies,
                         model_data=self.pysdk_model.model_data,
                     )
+                elif not hasattr(self, "prepared_for_mms"):
+                    self.js_model_config, self.prepared_for_mms = prepare_mms_js_resources(
+                        model_path=self.model_path,
+                        js_id=self.model,
+                        dependencies=self.dependencies,
+                        model_data=self.pysdk_model.model_data,
+                    )
 
                 self._prepare_for_mode()
                 env = {}
@@ -179,6 +193,10 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                 predictor = TgiLocalModePredictor(
                     self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
                 )
+            elif self.model_server == ModelServer.MMS:
+                predictor = TransformersLocalModePredictor(
+                    self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
+                )
 
             ram_usage_before = _get_ram_usage_mb()
             self.modes[str(Mode.LOCAL_CONTAINER)].create_server(
@@ -254,6 +272,24 @@ def _build_for_tgi_jumpstart(self):
 
         self.pysdk_model.env.update(env)
 
+    def _build_for_mms_jumpstart(self):
+        """Placeholder docstring"""
+
+        env = {}
+        if self.mode == Mode.LOCAL_CONTAINER:
+            if not hasattr(self, "prepared_for_mms"):
+                self.js_model_config, self.prepared_for_mms = prepare_mms_js_resources(
+                    model_path=self.model_path,
+                    js_id=self.model,
+                    dependencies=self.dependencies,
+                    model_data=self.pysdk_model.model_data,
+                )
+            self._prepare_for_mode()
+        elif self.mode == Mode.SAGEMAKER_ENDPOINT and hasattr(self, "prepared_for_mms"):
+            self.pysdk_model.model_data, env = self._prepare_for_mode()
+
+        self.pysdk_model.env.update(env)
+
     def _tune_for_js(self, sharded_supported: bool, max_tuning_duration: int = 1800):
         """Tune for Jumpstart Models in Local Mode.
 
@@ -264,7 +300,7 @@ def _tune_for_js(self, sharded_supported: bool, max_tuning_duration: int = 1800)
         returns:
             Tuned Model.
         """
-        if self.mode != Mode.LOCAL_CONTAINER:
+        if self.mode == Mode.SAGEMAKER_ENDPOINT:
             logger.warning(
                 "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
             )
@@ -438,7 +474,6 @@ def _build_for_jumpstart(self):
         self.jumpstart = True
 
         pysdk_model = self._create_pre_trained_js_model()
-
         image_uri = pysdk_model.image_uri
 
         logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
@@ -451,7 +486,6 @@ def _build_for_jumpstart(self):
         if "djl-inference" in image_uri:
             logger.info("Building for DJL JumpStart Model ID...")
             self.model_server = ModelServer.DJL_SERVING
-
             self.pysdk_model = pysdk_model
             self.image_uri = self.pysdk_model.image_uri
 
@@ -461,16 +495,23 @@ def _build_for_jumpstart(self):
         elif "tgi-inference" in image_uri:
             logger.info("Building for TGI JumpStart Model ID...")
             self.model_server = ModelServer.TGI
-
             self.pysdk_model = pysdk_model
             self.image_uri = self.pysdk_model.image_uri
 
             self._build_for_tgi_jumpstart()
 
             self.pysdk_model.tune = self.tune_for_tgi_jumpstart
-        else:
+        elif "huggingface-pytorch-inference:" in image_uri:
+            logger.info("Building for MMS JumpStart Model ID...")
+            self.model_server = ModelServer.MMS
+            self.pysdk_model = pysdk_model
+            self.image_uri = self.pysdk_model.image_uri
+
+            self._build_for_mms_jumpstart()
+        elif self.mode != Mode.SAGEMAKER_ENDPOINT:
             raise ValueError(
-                "JumpStart Model ID was not packaged with djl-inference or tgi-inference container."
+                "JumpStart Model ID was not packaged "
+                "with djl-inference, tgi-inference, or mms-inference container."
             )
 
         return self.pysdk_model

src/sagemaker/serve/builder/model_builder.py

@@ -36,6 +36,7 @@
 from sagemaker.serve.detector.pickler import save_pkl, save_xgboost
 from sagemaker.serve.builder.serve_settings import _ServeSettings
 from sagemaker.serve.builder.djl_builder import DJL
+from sagemaker.serve.builder.tei_builder import TEI
 from sagemaker.serve.builder.tgi_builder import TGI
 from sagemaker.serve.builder.jumpstart_builder import JumpStart
 from sagemaker.serve.builder.transformers_builder import Transformers
@@ -95,9 +96,9 @@
 }
 
 
-# pylint: disable=attribute-defined-outside-init, disable=E1101, disable=R0901
+# pylint: disable=attribute-defined-outside-init, disable=E1101, disable=R0901, disable=R1705
 @dataclass
-class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing):
+class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing, TEI):
     """Class that builds a deployable model.
 
     Args:
@@ -753,7 +754,7 @@ def build(  # pylint: disable=R0911
                 model_task = self.model_metadata.get("HF_TASK")
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()
-            if self._is_djl():  # pylint: disable=R1705
+            if self._is_djl():
                 return self._build_for_djl()
             else:
                 hf_model_md = get_huggingface_model_metadata(
@@ -764,8 +765,10 @@ def build(  # pylint: disable=R0911
                     model_task = hf_model_md.get("pipeline_tag")
                 if self.schema_builder is None and model_task is not None:
                     self._hf_schema_builder_init(model_task)
-                if model_task == "text-generation":  # pylint: disable=R1705
+                if model_task == "text-generation":
                     return self._build_for_tgi()
+                if model_task == "sentence-similarity":
+                    return self._build_for_tei()
                 elif self._can_fit_on_single_gpu():
                     return self._build_for_transformers()
                 elif (