Refactoring

Author: Jonathan Makunga

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/mode/local_container_mode.py

@@ -42,7 +42,6 @@ class LocalContainerMode(
     LocalTgiServing,
     LocalMultiModelServer,
     LocalTensorflowServing,
-    LocalTeiServing,
 ):
     """A class that holds methods to deploy model to a container in local environment"""
 
@@ -72,6 +71,8 @@ def __init__(
         self.secret_key = None
         self._ping_container = None
 
+        self._tei_serving = LocalTeiServing()
+
     def load(self, model_path: str = None):
         """Placeholder docstring"""
         path = Path(model_path if model_path else self.model_path)
@@ -159,14 +160,14 @@ def create_server(
             )
             self._ping_container = self._tensorflow_serving_deep_ping
         elif self.model_server == ModelServer.TEI:
-            self._start_tei_serving(
+            self._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,
             )
-            self._ping_container = self._tei_deep_ping
+            self._ping_container = self._tei_serving._tei_deep_ping
 
         # allow some time for container to be ready
         time.sleep(10)

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -27,7 +27,6 @@ class SageMakerEndpointMode(
     SageMakerTgiServing,
     SageMakerMultiModelServer,
     SageMakerTensorflowServing,
-    SageMakerTeiServing,
 ):
     """Holds the required method to deploy a model to a SageMaker Endpoint"""
 
@@ -39,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)
@@ -122,7 +123,7 @@ def prepare(
             )
 
         if self.model_server == ModelServer.TEI:
-            upload_artifacts = self._upload_tei_artifacts(
+            upload_artifacts = self._tei_serving._upload_tei_artifacts(
                 model_path=model_path,
                 sagemaker_session=sagemaker_session,
                 s3_model_data_url=s3_model_data_url,