Notebook testing

Author: Jonathan Makunga

src/sagemaker/serve/mode/local_container_mode.py

@@ -226,7 +226,6 @@ def _pull_image(self, image: str):
             logger.warning("Unable to login to ecr: %s", e)
 
         self.client = docker.from_env()
-        print(self.client.images.list())
         try:
             logger.info("Pulling image %s from repository...", image)
             self.client.images.pull(image)

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
@@ -26,6 +27,7 @@ class SageMakerEndpointMode(
     SageMakerTgiServing,
     SageMakerMultiModelServer,
     SageMakerTensorflowServing,
+    SageMakerTeiServing,
 ):
     """Holds the required method to deploy a model to a SageMaker Endpoint"""
 
@@ -118,4 +120,12 @@ def prepare(
                 image=image,
             )
 
+        if self.model_server == ModelServer.TEI:
+            return self._upload_tei_artifacts(
+                model_path=model_path,
+                sagemaker_session=sagemaker_session,
+                s3_model_data_url=s3_model_data_url,
+                image=image,
+            )
+
         raise ValueError("%s model server is not supported" % self.model_server)