Feat: Add TEI support for ModelBuilder (#4694)

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Add TEI Serving

* Notebook testing

* Notebook testing

* Notebook testing

* Refactoring

* Refactoring

* UT

* UT

* Refactoring

* Test coverage

* Refactoring

* Refactoring

---------

Co-authored-by: Jonathan Makunga <[email protected]>

Author: makungaj1

src/sagemaker/serve/builder/model_builder.py

@@ -169,7 +169,7 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
             in order for model builder to build the artifacts correctly (according
             to the model server). Possible values for this argument are
             ``TORCHSERVE``, ``MMS``, ``TENSORFLOW_SERVING``, ``DJL_SERVING``,
-            ``TRITON``, and``TGI``.
+            ``TRITON``,``TGI``, and ``TEI``.
         model_metadata (Optional[Dict[str, Any]): Dictionary used to override model metadata.
             Currently, ``HF_TASK`` is overridable for HuggingFace model. HF_TASK should be set for
             new models without task metadata in the Hub, adding unsupported task types will throw

src/sagemaker/serve/builder/tei_builder.py

@@ -25,7 +25,7 @@
     _get_nb_instance,
 )
 from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
-from sagemaker.serve.utils.predictors import TgiLocalModePredictor
+from sagemaker.serve.utils.predictors import TeiLocalModePredictor
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
@@ -74,16 +74,16 @@ def _prepare_for_mode(self):
     def _get_client_translators(self):
         """Placeholder docstring"""
 
-    def _set_to_tgi(self):
+    def _set_to_tei(self):
         """Placeholder docstring"""
-        if self.model_server != ModelServer.TGI:
+        if self.model_server != ModelServer.TEI:
             messaging = (
                 "HuggingFace Model ID support on model server: "
                 f"{self.model_server} is not currently supported. "
-                f"Defaulting to {ModelServer.TGI}"
+                f"Defaulting to {ModelServer.TEI}"
             )
             logger.warning(messaging)
-            self.model_server = ModelServer.TGI
+            self.model_server = ModelServer.TEI
 
     def _create_tei_model(self, **kwargs) -> Type[Model]:
         """Placeholder docstring"""
@@ -142,7 +142,7 @@ def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
         if self.mode == Mode.LOCAL_CONTAINER:
             timeout = kwargs.get("model_data_download_timeout")
 
-            predictor = TgiLocalModePredictor(
+            predictor = TeiLocalModePredictor(
                 self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
             )
 
@@ -180,7 +180,9 @@ def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
 
-        if not self.nb_instance_type and "instance_type" not in kwargs:
+        if self.nb_instance_type and "instance_type" not in kwargs:
+            kwargs.update({"instance_type": self.nb_instance_type})
+        elif not self.nb_instance_type and "instance_type" not in kwargs:
             raise ValueError(
                 "Instance type must be provided when deploying " "to SageMaker Endpoint mode."
             )
@@ -216,7 +218,7 @@ def _build_for_tei(self):
         """Placeholder docstring"""
         self.secret_key = None
 
-        self._set_to_tgi()
+        self._set_to_tei()
 
         self.pysdk_model = self._build_for_hf_tei()
         return self.pysdk_model

src/sagemaker/serve/mode/local_container_mode.py

@@ -21,6 +21,7 @@
 from sagemaker.serve.model_server.djl_serving.server import LocalDJLServing
 from sagemaker.serve.model_server.triton.server import LocalTritonServer
 from sagemaker.serve.model_server.tgi.server import LocalTgiServing
+from sagemaker.serve.model_server.tei.server import LocalTeiServing
 from sagemaker.serve.model_server.multi_model_server.server import LocalMultiModelServer
 from sagemaker.session import Session
 
@@ -69,6 +70,7 @@ def __init__(
         self.container = None
         self.secret_key = None
         self._ping_container = None
+        self._invoke_serving = None
 
     def load(self, model_path: str = None):
         """Placeholder docstring"""
@@ -156,6 +158,19 @@ def create_server(
                 env_vars=env_vars if env_vars else self.env_vars,
             )
             self._ping_container = self._tensorflow_serving_deep_ping
+        elif self.model_server == ModelServer.TEI:
+            tei_serving = LocalTeiServing()
+            tei_serving._start_tei_serving(
+                client=self.client,
+                image=image,
+                model_path=model_path if model_path else self.model_path,
+                secret_key=secret_key,
+                env_vars=env_vars if env_vars else self.env_vars,
+            )
+            tei_serving.schema_builder = self.schema_builder
+            self.container = tei_serving.container
+            self._ping_container = tei_serving._tei_deep_ping
+            self._invoke_serving = tei_serving._invoke_tei_serving
 
         # allow some time for container to be ready
         time.sleep(10)

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -6,6 +6,7 @@
 import logging
 from typing import Type
 
+from sagemaker.serve.model_server.tei.server import SageMakerTeiServing
 from sagemaker.serve.model_server.tensorflow_serving.server import SageMakerTensorflowServing
 from sagemaker.session import Session
 from sagemaker.serve.utils.types import ModelServer
@@ -37,6 +38,8 @@ def __init__(self, inference_spec: Type[InferenceSpec], model_server: ModelServe
         self.inference_spec = inference_spec
         self.model_server = model_server
 
+        self._tei_serving = SageMakerTeiServing()
+
     def load(self, model_path: str):
         """Placeholder docstring"""
         path = Path(model_path)
@@ -66,8 +69,9 @@ def prepare(
                 + "session to be created or supply `sagemaker_session` into @serve.invoke."
             ) from e
 
+        upload_artifacts = None
         if self.model_server == ModelServer.TORCHSERVE:
-            return self._upload_torchserve_artifacts(
+            upload_artifacts = self._upload_torchserve_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -76,7 +80,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.TRITON:
-            return self._upload_triton_artifacts(
+            upload_artifacts = self._upload_triton_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
@@ -85,15 +89,15 @@ def prepare(
             )
 
         if self.model_server == ModelServer.DJL_SERVING:
-            return self._upload_djl_artifacts(
+            upload_artifacts = self._upload_djl_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
             )
 
         if self.model_server == ModelServer.TGI:
-            return self._upload_tgi_artifacts(
+            upload_artifacts = self._upload_tgi_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
@@ -102,20 +106,31 @@ def prepare(
             )
 
         if self.model_server == ModelServer.MMS:
-            return self._upload_server_artifacts(
+            upload_artifacts = self._upload_server_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
             )
 
         if self.model_server == ModelServer.TENSORFLOW_SERVING:
-            return self._upload_tensorflow_serving_artifacts(
+            upload_artifacts = self._upload_tensorflow_serving_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 secret_key=secret_key,
                 s3_model_data_url=s3_model_data_url,
                 image=image,
             )
 
+        if self.model_server == ModelServer.TEI:
+            upload_artifacts = self._tei_serving._upload_tei_artifacts(
+                model_path=model_path,
+                sagemaker_session=sagemaker_session,
+                s3_model_data_url=s3_model_data_url,
+                image=image,
+            )
+
+        if upload_artifacts:
+            return upload_artifacts
+
         raise ValueError("%s model server is not supported" % self.model_server)

src/sagemaker/serve/model_server/tei/__init__.py

@@ -0,0 +1,160 @@
+"""Module for Local TEI Serving"""
+
+from __future__ import absolute_import
+
+import requests
+import logging
+from pathlib import Path
+from docker.types import DeviceRequest
+from sagemaker import Session, fw_utils
+from sagemaker.serve.utils.exceptions import LocalModelInvocationException
+from sagemaker.base_predictor import PredictorBase
+from sagemaker.s3_utils import determine_bucket_and_prefix, parse_s3_url, s3_path_join
+from sagemaker.s3 import S3Uploader
+from sagemaker.local.utils import get_docker_host
+
+
+MODE_DIR_BINDING = "/opt/ml/model/"
+_SHM_SIZE = "2G"
+_DEFAULT_ENV_VARS = {
+    "TRANSFORMERS_CACHE": "/opt/ml/model/",
+    "HUGGINGFACE_HUB_CACHE": "/opt/ml/model/",
+}
+
+logger = logging.getLogger(__name__)
+
+
+class LocalTeiServing:
+    """LocalTeiServing class"""
+
+    def _start_tei_serving(
+        self, client: object, image: str, model_path: str, secret_key: str, env_vars: dict
+    ):
+        """Starts a local tei serving container.
+
+        Args:
+            client: Docker client
+            image: Image to use
+            model_path: Path to the model
+            secret_key: Secret key to use for authentication
+            env_vars: Environment variables to set
+        """
+        if env_vars and secret_key:
+            env_vars["SAGEMAKER_SERVE_SECRET_KEY"] = secret_key
+
+        self.container = client.containers.run(
+            image,
+            shm_size=_SHM_SIZE,
+            device_requests=[DeviceRequest(count=-1, capabilities=[["gpu"]])],
+            network_mode="host",
+            detach=True,
+            auto_remove=True,
+            volumes={
+                Path(model_path).joinpath("code"): {
+                    "bind": MODE_DIR_BINDING,
+                    "mode": "rw",
+                },
+            },
+            environment=_update_env_vars(env_vars),
+        )
+
+    def _invoke_tei_serving(self, request: object, content_type: str, accept: str):
+        """Invokes a local tei serving container.
+
+        Args:
+            request: Request to send
+            content_type: Content type to use
+            accept: Accept to use
+        """
+        try:
+            response = requests.post(
+                f"http://{get_docker_host()}:8080/invocations",
+                data=request,
+                headers={"Content-Type": content_type, "Accept": accept},
+                timeout=600,
+            )
+            response.raise_for_status()
+            return response.content
+        except Exception as e:
+            raise Exception("Unable to send request to the local container server") from e
+
+    def _tei_deep_ping(self, predictor: PredictorBase):
+        """Checks if the local tei serving container is up and running.
+
+        If the container is not up and running, it will raise an exception.
+        """
+        response = None
+        try:
+            response = predictor.predict(self.schema_builder.sample_input)
+            return (True, response)
+            # pylint: disable=broad-except
+        except Exception as e:
+            if "422 Client Error: Unprocessable Entity for url" in str(e):
+                raise LocalModelInvocationException(str(e))
+            return (False, response)
+
+        return (True, response)
+
+
+class SageMakerTeiServing:
+    """SageMakerTeiServing class"""
+
+    def _upload_tei_artifacts(
+        self,
+        model_path: str,
+        sagemaker_session: Session,
+        s3_model_data_url: str = None,
+        image: str = None,
+        env_vars: dict = None,
+    ):
+        """Uploads the model artifacts to S3.
+
+        Args:
+            model_path: Path to the model
+            sagemaker_session: SageMaker session
+            s3_model_data_url: S3 model data URL
+            image: Image to use
+            env_vars: Environment variables to set
+        """
+        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
+        )
+
+        code_dir = Path(model_path).joinpath("code")
+
+        s3_location = s3_path_join("s3://", bucket, code_key_prefix, "code")
+
+        logger.debug("Uploading TEI Model Resources uncompressed to: %s", s3_location)
+
+        model_data_url = S3Uploader.upload(
+            str(code_dir),
+            s3_location,
+            None,
+            sagemaker_session,
+        )
+
+        model_data = {
+            "S3DataSource": {
+                "CompressionType": "None",
+                "S3DataType": "S3Prefix",
+                "S3Uri": model_data_url + "/",
+            }
+        }
+
+        return (model_data, _update_env_vars(env_vars))
+
+
+def _update_env_vars(env_vars: dict) -> dict:
+    """Placeholder docstring"""
+    updated_env_vars = {}
+    updated_env_vars.update(_DEFAULT_ENV_VARS)
+    if env_vars:
+        updated_env_vars.update(env_vars)
+    return updated_env_vars