changes to support in_process

Author: Bryannah Hernandez

src/sagemaker/serve/builder/model_builder.py

@@ -36,6 +36,7 @@
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.mode.sagemaker_endpoint_mode import SageMakerEndpointMode
 from sagemaker.serve.mode.local_container_mode import LocalContainerMode
+from sagemaker.serve.mode.in_process_mode import InProcessMode
 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
@@ -410,7 +411,7 @@ def _prepare_for_mode(
             )
             self.env_vars.update(env_vars_sagemaker)
             return self.s3_upload_path, env_vars_sagemaker
-        if self.mode == Mode.LOCAL_CONTAINER:
+        elif self.mode == Mode.LOCAL_CONTAINER:
             # init the LocalContainerMode object
             self.modes[str(Mode.LOCAL_CONTAINER)] = LocalContainerMode(
                 inference_spec=self.inference_spec,
@@ -422,9 +423,21 @@ def _prepare_for_mode(
             )
             self.modes[str(Mode.LOCAL_CONTAINER)].prepare()
             return None
+        elif self.mode == Mode.IN_PROCESS:
+            # init the InProcessMode object
+            self.modes[str(Mode.IN_PROCESS)] = InProcessMode(
+                inference_spec=self.inference_spec,
+                schema_builder=self.schema_builder,
+                session=self.sagemaker_session,
+                model_path=self.model_path,
+                env_vars=self.env_vars,
+                model_server=self.model_server,
+            )
+            self.modes[str(Mode.IN_PROCESS)].prepare()
+            return None
 
         raise ValueError(
-            "Please specify mode in: %s, %s" % (Mode.LOCAL_CONTAINER, Mode.SAGEMAKER_ENDPOINT)
+            "Please specify mode in: %s, %s" % (Mode.LOCAL_CONTAINER, Mode.SAGEMAKER_ENDPOINT, Mode.IN_PROCESS)
         )
 
     def _get_client_translators(self):
@@ -796,9 +809,10 @@ def _initialize_for_mlflow(self, artifact_path: str) -> None:
         self.dependencies.update({"requirements": mlflow_model_dependency_path})
 
     # Model Builder is a class to build the model for deployment.
-    # It supports two modes of deployment
+    # It supports two* modes of deployment
     # 1/ SageMaker Endpoint
     # 2/ Local launch with container
+    # 3/ In process mode with Transformers server in beta release
     def build(  # pylint: disable=R0911
         self,
         mode: Type[Mode] = None,
@@ -897,7 +911,10 @@ def build(  # pylint: disable=R0911
     def _build_validations(self):
         """Validations needed for model server overrides, or auto-detection or fallback"""
         if self.mode == Mode.IN_PROCESS and self.model_server is not ModelServer.MMS:
-            raise ValueError("IN_PROCESS mode is not supported yet!")
+            raise ValueError(
+                "IN_PROCESS mode is not supported yet for model server. It is "
+                "supported for MMS/Transformers server in beta release"
+            )
 
         if self.inference_spec and self.model:
             raise ValueError("Can only set one of the following: model, inference_spec.")

src/sagemaker/serve/builder/transformers_builder.py

@@ -37,7 +37,10 @@
 )
 from sagemaker.serve.detector.pickler import save_pkl
 from sagemaker.serve.utils.optimize_utils import _is_optimized
-from sagemaker.serve.utils.predictors import TransformersLocalModePredictor
+from sagemaker.serve.utils.predictors import (
+    TransformersLocalModePredictor,
+    TransformersInProcessModePredictor,
+)
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
@@ -232,7 +235,7 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
         if self.mode == Mode.IN_PROCESS:
             timeout = kwargs.get("model_data_download_timeout")
 
-            predictor = TransformersLocalModePredictor(
+            predictor = TransformersInProcessModePredictor(
                 self.modes[str(Mode.IN_PROCESS)], serializer, deserializer
             )
 
@@ -242,7 +245,6 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
                 None,
                 predictor,
                 self.pysdk_model.env,
-                jumpstart=False,
             )
             return predictor
 

src/sagemaker/serve/mode/in_process_mode.py

@@ -1,9 +1,8 @@
-"""Module that defines the LocalContainerMode class"""
+"""Module that defines the InProcessMode class"""
 
 from __future__ import absolute_import
 from pathlib import Path
 import logging
-from datetime import datetime, timedelta
 from typing import Dict, Type
 import base64
 import time
@@ -13,7 +12,6 @@
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
-from sagemaker.serve.utils.logging_agent import pull_logs
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.utils.exceptions import LocalDeepPingException
 from sagemaker.serve.model_server.multi_model_server.server import InProcessMultiModelServer
@@ -32,7 +30,7 @@
 class InProcessMode(
     InProcessMultiModelServer,
 ):
-    """A class that holds methods to deploy model to a container in local environment"""
+    """A class that holds methods to deploy model to a container in process environment"""
 
     def __init__(
         self,
@@ -83,9 +81,9 @@ def create_server(
     ):
         """Placeholder docstring"""
 
-        self._pull_image(image=image)
+        # self._pull_image(image=image)
 
-        self.destroy_server()
+        # self.destroy_server()
 
         logger.info("Waiting for model server %s to start up...", self.model_server)
 
@@ -97,28 +95,10 @@ def create_server(
                 secret_key=secret_key,
                 env_vars=env_vars if env_vars else self.env_vars,
             )
+            logger.info("Starting PING")
             self._ping_container = self._multi_model_server_deep_ping
 
-        # allow some time for container to be ready
-        time.sleep(10)
-
-        log_generator = self.container.logs(follow=True, stream=True)
-        time_limit = datetime.now() + timedelta(seconds=container_timeout_seconds)
-        healthy = False
         while True:
-            now = datetime.now()
-            final_pull = now > time_limit
-            pull_logs(
-                (x.decode("UTF-8").rstrip() for x in log_generator),
-                log_generator.close,
-                datetime.now() + timedelta(seconds=_PING_HEALTH_CHECK_INTERVAL_SEC),
-                now > time_limit,
-            )
-
-            if final_pull:
-                break
-
-            # allow some time for container to be ready
             time.sleep(10)
 
             healthy, response = self._ping_container(predictor)
@@ -128,37 +108,4 @@ def create_server(
 
         if not healthy:
             raise LocalDeepPingException(_PING_HEALTH_CHECK_FAIL_MSG)
-
-    def destroy_server(self):
-        """Placeholder docstring"""
-        if self.container:
-            try:
-                logger.debug("Stopping currently running container...")
-                self.container.kill()
-            except docker.errors.APIError as exc:
-                if exc.response.status_code < 400 or exc.response.status_code > 499:
-                    raise Exception("Error encountered when cleaning up local container") from exc
-            self.container = None
-
-    def _pull_image(self, image: str):
-        """Placeholder docstring"""
-        try:
-            encoded_token = (
-                self.ecr.get_authorization_token()
-                .get("authorizationData")[0]
-                .get("authorizationToken")
-            )
-            decoded_token = base64.b64decode(encoded_token).decode("utf-8")
-            username, password = decoded_token.split(":")
-            ecr_uri = image.split("/")[0]
-            login_command = ["docker", "login", "-u", username, "-p", password, ecr_uri]
-            subprocess.run(login_command, check=True, capture_output=True)
-        except subprocess.CalledProcessError as e:
-            logger.warning("Unable to login to ecr: %s", e)
-
-        self.client = docker.from_env()
-        try:
-            logger.info("Pulling image %s from repository...", image)
-            self.client.images.pull(image)
-        except docker.errors.NotFound as e:
-            raise ValueError("Could not find remote image to pull") from e
+        

src/sagemaker/serve/model_server/multi_model_server/server.py

@@ -68,6 +68,23 @@ def _invoke_multi_model_server_serving(self, request: object, content_type: str,
             raise Exception("Unable to send request to the local container server") from e
 
         return (True, response)
+    
+    def _multi_model_server_deep_ping(self, predictor: PredictorBase):
+        """Placeholder docstring"""
+        response = None
+        logger.debug("AM I HERE? PING PING")
+
+        # 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 LocalMultiModelServer:
     """Local Multi Model server instance"""

src/sagemaker/serve/utils/exceptions.py

@@ -1,4 +1,4 @@
-"""Placeholder Docstring"""
+"""Exceptions used across different model builder invocations"""
 
 from __future__ import absolute_import
 
@@ -22,6 +22,16 @@ class LocalDeepPingException(ModelBuilderException):
 
     def __init__(self, message):
         super().__init__(message=message)
+        
+
+class InProcessDeepPingException(ModelBuilderException):
+    """Raise when in process model serving does not pass the deep ping check"""
+
+    fmt = "Error Message: {message}"
+    model_builder_error_code = 1
+
+    def __init__(self, message):
+        super().__init__(message=message)
 
 
 class LocalModelOutOfMemoryException(ModelBuilderException):

src/sagemaker/serve/utils/predictors.py

@@ -6,6 +6,7 @@
 
 from sagemaker import Session
 from sagemaker.serve.mode.local_container_mode import LocalContainerMode
+from sagemaker.serve.mode.in_process_mode import InProcessMode
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
 from sagemaker.serializers import IdentitySerializer, JSONSerializer
 from sagemaker.deserializers import BytesDeserializer, JSONDeserializer
@@ -209,6 +210,49 @@ def delete_predictor(self):
         self._mode_obj.destroy_server()
 
 
+class TransformersInProcessModePredictor(PredictorBase):
+    """Lightweight Transformers predictor for local deployment"""
+
+    def __init__(
+        self,
+        mode_obj: Type[InProcessMode],
+        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_multi_model_server_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 TeiLocalModePredictor(PredictorBase):
     """Lightweight Tei predictor for local deployment in IN_PROCESS and LOCAL_CONTAINER modes"""