Add TEI Serving

Author: Jonathan Makunga

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
@@ -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
             )
 

src/sagemaker/serve/mode/local_container_mode.py

@@ -158,6 +158,15 @@ 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:
+            self._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,
+            )
+            self._ping_container = self._tei_deep_ping
 
         # allow some time for container to be ready
         time.sleep(10)

src/sagemaker/serve/utils/predictors.py

@@ -209,6 +209,49 @@ def delete_predictor(self):
         self._mode_obj.destroy_server()
 
 
+class TeiLocalModePredictor(PredictorBase):
+    """Lightweight Tei predictor for local deployment in IN_PROCESS and LOCAL_CONTAINER modes"""
+
+    def __init__(
+        self,
+        mode_obj: Type[LocalContainerMode],
+        serializer=JSONSerializer(),
+        deserializer=JSONDeserializer(),
+    ):
+        self._mode_obj = mode_obj
+        self.serializer = serializer
+        self.deserializer = deserializer
+
+    def predict(self, data):
+        """Placeholder docstring"""
+        return [
+            self.deserializer.deserialize(
+                io.BytesIO(
+                    self._mode_obj._invoke_tei_serving(
+                        self.serializer.serialize(data),
+                        self.content_type,
+                        self.deserializer.ACCEPT[0],
+                    )
+                ),
+                self.content_type,
+            )
+        ]
+
+    @property
+    def content_type(self):
+        """The MIME type of the data sent to the inference endpoint."""
+        return self.serializer.CONTENT_TYPE
+
+    @property
+    def accept(self):
+        """The content type(s) that are expected from the inference endpoint."""
+        return self.deserializer.ACCEPT
+
+    def delete_predictor(self):
+        """Shut down and remove the container that you created in LOCAL_CONTAINER mode"""
+        self._mode_obj.destroy_server()
+
+
 class TensorflowServingLocalPredictor(PredictorBase):
     """Lightweight predictor for local deployment in LOCAL_CONTAINER modes"""