feat: FastAPI integration for In_Process Mode (2/2)

requirements/extras/huggingface_requirements.txt

@@ -1,2 +1,5 @@
 accelerate>=0.24.1,<=0.27.0
 sagemaker_schema_inference_artifacts>=0.0.5
+uvicorn>=0.30.1
+fastapi>=0.111.0
+nest-asyncio

requirements/extras/test_requirements.txt

@@ -40,3 +40,6 @@ schema==0.7.5
 tensorflow>=2.1,<=2.16
 mlflow>=2.12.2,<2.13
 huggingface_hub>=0.23.4
+uvicorn>=0.30.1
+fastapi>=0.111.0
+nest-asyncio

src/sagemaker/serve/app.py

@@ -0,0 +1,79 @@
+"""FastAPI requests"""
+
+from __future__ import absolute_import
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+try:
+    from fastapi import FastAPI, Request
+
+    app = FastAPI(
+        title="Transformers In Process Server",
+        version="1.0",
+        description="A simple server",
+    )
+
+    @app.get("/")
+    def read_root():
+        """Placeholder docstring"""
+        return {"Hello": "World"}
+
+    @app.get("/generate")
+    async def generate_text(prompt: Request):
+        """Placeholder docstring"""
+        logger.info("Generating Text....")
+
+        str_prompt = await prompt.json()
+
+        logger.info(str_prompt)
+
+        generated_text = generator(
+            str_prompt, max_length=30, num_return_sequences=5, truncation=True
+        )
+        return generated_text[0]["generated_text"]
+
+except ImportError:
+    logger.error("To enable in_process mode for Transformers install fastapi from HuggingFace hub")
+
+
+try:
+    from transformers import pipeline
+
+    generator = pipeline("text-generation", model="gpt2")
+
+except ImportError:
+    logger.error(
+        "To enable in_process mode for Transformers install transformers from HuggingFace hub"
+    )
+
+try:
+    import uvicorn
+
+except ImportError:
+    logger.error("To enable in_process mode for Transformers install uvicorn from HuggingFace hub")
+
+
+@app.post("/post")
+def post(payload: dict):
+    """Placeholder docstring"""
+    return payload
+
+
+async def main():
+    """Running server locally with uvicorn"""
+    logger.info("Running")
+    config = uvicorn.Config(
+        "sagemaker.app:app",
+        host="127.0.0.1",
+        port=9007,
+        log_level="info",
+        loop="asyncio",
+        reload=True,
+        workers=3,
+        use_colors=True,
+    )
+    server = uvicorn.Server(config)
+    logger.info("Waiting for a connection...")
+    await server.serve()

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,22 @@ 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, %s"
+            % (Mode.LOCAL_CONTAINER, Mode.SAGEMAKER_ENDPOINT, Mode.IN_PROCESS)
         )
 
     def _get_client_translators(self):
@@ -606,6 +620,9 @@ def _overwrite_mode_in_deploy(self, overwrite_mode: str):
         elif overwrite_mode == Mode.LOCAL_CONTAINER:
             self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
             self._prepare_for_mode()
+        elif overwrite_mode == Mode.IN_PROCESS:
+            self.mode = self.pysdk_model.mode = Mode.IN_PROCESS
+            self._prepare_for_mode()
         else:
             raise ValueError("Mode %s is not supported!" % overwrite_mode)
 
@@ -795,9 +812,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 three 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,
@@ -895,8 +913,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:
-            raise ValueError("IN_PROCESS mode is not supported yet!")
+        if self.mode == Mode.IN_PROCESS and self.model_server is not ModelServer.MMS:
+            raise ValueError(
+                "IN_PROCESS mode is only 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

@@ -36,7 +36,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
@@ -47,6 +50,7 @@
 
 logger = logging.getLogger(__name__)
 DEFAULT_TIMEOUT = 1800
+LOCAL_MODES = [Mode.LOCAL_CONTAINER, Mode.IN_PROCESS]
 
 
 """Retrieves images for different libraries - Pytorch, TensorFlow from HuggingFace hub
@@ -228,6 +232,18 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
             )
             return predictor
 
+        if self.mode == Mode.IN_PROCESS:
+            timeout = kwargs.get("model_data_download_timeout")
+
+            predictor = TransformersInProcessModePredictor(
+                self.modes[str(Mode.IN_PROCESS)], serializer, deserializer
+            )
+
+            self.modes[str(Mode.IN_PROCESS)].create_server(
+                predictor,
+            )
+            return predictor
+
         self._set_instance(kwargs)
 
         if "mode" in kwargs:
@@ -293,7 +309,7 @@ def _build_transformers_env(self):
 
         self.pysdk_model = self._create_transformers_model()
 
-        if self.mode == Mode.LOCAL_CONTAINER:
+        if self.mode in LOCAL_MODES:
             self._prepare_for_mode()
 
         return self.pysdk_model

src/sagemaker/serve/mode/in_process_mode.py

@@ -0,0 +1,90 @@
+"""Module that defines the InProcessMode class"""
+
+from __future__ import absolute_import
+from pathlib import Path
+import logging
+from typing import Dict, Type
+import time
+from datetime import datetime, timedelta
+
+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.types import ModelServer
+from sagemaker.serve.utils.exceptions import LocalDeepPingException
+from sagemaker.serve.model_server.multi_model_server.server import InProcessMultiModelServer
+from sagemaker.session import Session
+
+logger = logging.getLogger(__name__)
+
+_PING_HEALTH_CHECK_FAIL_MSG = (
+    "Ping health check did not pass. "
+    + "Please increase container_timeout_seconds or review your inference code."
+)
+
+
+class InProcessMode(
+    InProcessMultiModelServer,
+):
+    """A class that holds methods to deploy model to a container in process environment"""
+
+    def __init__(
+        self,
+        model_server: ModelServer,
+        inference_spec: Type[InferenceSpec],
+        schema_builder: Type[SchemaBuilder],
+        session: Session,
+        model_path: str = None,
+        env_vars: Dict = None,
+    ):
+        # pylint: disable=bad-super-call
+        super().__init__()
+
+        self.inference_spec = inference_spec
+        self.model_path = model_path
+        self.env_vars = env_vars
+        self.session = session
+        self.schema_builder = schema_builder
+        self.model_server = model_server
+        self._ping_container = None
+
+    def load(self, model_path: str = None):
+        """Loads model path, checks that path exists"""
+        path = Path(model_path if model_path else self.model_path)
+        if not path.exists():
+            raise ValueError("model_path does not exist")
+        if not path.is_dir():
+            raise ValueError("model_path is not a valid directory")
+
+        return self.inference_spec.load(str(path))
+
+    def prepare(self):
+        """Prepares the server"""
+
+    def create_server(
+        self,
+        predictor: PredictorBase,
+    ):
+        """Creating the server and checking ping health."""
+        logger.info("Waiting for model server %s to start up...", self.model_server)
+
+        if self.model_server == ModelServer.MMS:
+            self._start_serving()
+            self._ping_container = self._multi_model_server_deep_ping
+
+        time_limit = datetime.now() + timedelta(seconds=5)
+        while self._ping_container is not None:
+            final_pull = datetime.now() > time_limit
+
+            if final_pull:
+                break
+
+            time.sleep(10)
+
+            healthy, response = self._ping_container(predictor)
+            if healthy:
+                logger.debug("Ping health check has passed. Returned %s", str(response))
+                break
+
+        if not healthy:
+            raise LocalDeepPingException(_PING_HEALTH_CHECK_FAIL_MSG)

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

@@ -2,9 +2,11 @@
 
 from __future__ import absolute_import
 
+import asyncio
 import requests
 import logging
 import platform
+import time
 from pathlib import Path
 from sagemaker import Session, fw_utils
 from sagemaker.serve.utils.exceptions import LocalModelInvocationException
@@ -13,13 +15,79 @@
 from sagemaker.s3 import S3Uploader
 from sagemaker.local.utils import get_docker_host
 from sagemaker.serve.utils.optimize_utils import _is_s3_uri
+from sagemaker.serve.app import main
 
 MODE_DIR_BINDING = "/opt/ml/model/"
 _DEFAULT_ENV_VARS = {}
 
 logger = logging.getLogger(__name__)
 
 
+class InProcessMultiModelServer:
+    """In Process Mode Multi Model server instance"""
+
+    def _start_serving(self):
+        """Initializes the start of the server"""
+        background_tasks = set()
+        task = asyncio.create_task(main())
+        background_tasks.add(task)
+        task.add_done_callback(background_tasks.discard)
+
+        time.sleep(10)
+
+    def _invoke_multi_model_server_serving(self, request: object, content_type: str, accept: str):
+        """Placeholder docstring"""
+        background_tasks = set()
+        task = asyncio.create_task(self.generate_connect())
+        background_tasks.add(task)
+        task.add_done_callback(background_tasks.discard)
+
+    def _multi_model_server_deep_ping(self, predictor: PredictorBase):
+        """Sends a deep ping to ensure prediction"""
+        background_tasks = set()
+        task = asyncio.create_task(self.tcp_connect())
+        background_tasks.add(task)
+        task.add_done_callback(background_tasks.discard)
+        response = None
+        return True, response
+
+    async def generate_connect(self):
+        """Writes the lines in bytes for server"""
+        reader, writer = await asyncio.open_connection("127.0.0.1", 9007)
+        a = (
+            b"GET /generate HTTP/1.1\r\nHost: 127.0.0.1:9007\r\nUser-Agent: "
+            b"python-requests/2.31.0\r\nAccept-Encoding: gzip, deflate, br\r\nAccept: */*\r\nConnection: ",
+            "keep-alive\r\nContent-Length: 33\r\nContent-Type: application/json\r\n\r\n",
+        )
+        b = b'"\\"Hello, I\'m a language model\\""'
+        list = [a, b]
+        writer.writelines(list)
+        logger.debug(writer.get_extra_info("peername"))
+        logger.debug(writer.transport)
+
+        data = await reader.read()
+        logger.info("Response from server")
+        logger.info(data)
+        writer.close()
+        await writer.wait_closed()
+
+    async def tcp_connect(self):
+        """Writes the lines in bytes for server"""
+        reader, writer = await asyncio.open_connection("127.0.0.1", 9007)
+        writer.write(
+            b"GET / HTTP/1.1\r\nHost: 127.0.0.1:9007\r\nUser-Agent: python-requests/2.32.3\r\nAccept-Encoding: gzip, ",
+            "deflate, br\r\nAccept: */*\r\nConnection: keep-alive\r\n\r\n",
+        )
+        logger.debug(writer.get_extra_info("peername"))
+        logger.debug(writer.transport)
+
+        data = await reader.read()
+        logger.info("Response from server")
+        logger.info(data)
+        writer.close()
+        await writer.wait_closed()
+
+
 class LocalMultiModelServer:
     """Local Multi Model server instance"""