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

* InferenceSpec support for HF

* feat: InferenceSpec support for MMS and testing

* Introduce changes for InProcess Mode

* mb_inprocess updates

* In_Process mode for TGI transformers, edits

* Remove InfSpec from branch

* changes to support in_process

* changes to get pre-checks passing

* pylint fix

* unit test, test mb

* period missing, added

* suggestions and test added

* pre-push fix

* missing an @

* fixes to test, added stubbing

* removing for fixes

* variable fixes

* init fix

* tests for in process mode

* prepush fix

* FastAPI with In_Process

* minor fixes

* putting 6 args

* server change

* changes

* format fix

* fastapi fixes

* port

* start

* changing port

* prepush

* import and unused fix

* moving files and transformers

* fix imports

* changes

* fixing modules

* placement

* dep fixes

* minor change

* fastapi predictor fix

* minor changes

* import transformers

* pylint comment

* delete local_run.sh

* format

* fix

* uvicorn fix

* fastapi

* try and except

* app

* deps comment out

* app func fix

* deps

* fix

* test fix

* non object

* comment out test for in_process

* removing unnecessary loggers

* fixing UT

* clean up loggers

* Update in_process_mode.py

* In Process Mode

* workers flags not needed when reload is set to true.

* Refactore

* wheel

* minor changes

* delete whl

* ut for app.py

* unit test fir app

* more unit tests

* py310 and higher

* delete whl file

---------

Co-authored-by: Bryannah Hernandez <[email protected]>
Co-authored-by: sage-maker <[email protected]>
Co-authored-by: Jonathan Makunga <[email protected]>
Co-authored-by: Jonathan Makunga <[email protected]>
Co-authored-by: Erick Benitez-Ramos <[email protected]>

Author: 6 people

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,100 @@
+"""FastAPI requests"""
+
+from __future__ import absolute_import
+
+import asyncio
+import logging
+import threading
+from typing import Optional
+
+
+logger = logging.getLogger(__name__)
+
+
+try:
+    import uvicorn
+except ImportError:
+    logger.error("Unable to import uvicorn, check if uvicorn is installed.")
+
+
+try:
+    from transformers import pipeline
+except ImportError:
+    logger.error("Unable to import transformers, check if transformers is installed.")
+
+
+try:
+    from fastapi import FastAPI, Request, APIRouter
+except ImportError:
+    logger.error("Unable to import fastapi, check if fastapi is installed.")
+
+
+class InProcessServer:
+    """Placeholder docstring"""
+
+    def __init__(self, model_id: Optional[str] = None, task: Optional[str] = None):
+        self._thread = None
+        self._loop = None
+        self._stop_event = asyncio.Event()
+        self._router = APIRouter()
+        self._model_id = model_id
+        self._task = task
+        self.server = None
+        self.port = None
+        self.host = None
+        # TODO: Pick up device automatically.
+        self._generator = pipeline(task, model=model_id, device="cpu")
+
+        # pylint: disable=unused-variable
+        @self._router.post("/generate")
+        async def generate_text(prompt: Request):
+            """Placeholder docstring"""
+            str_prompt = await prompt.json()
+            str_prompt = str_prompt["inputs"] if "inputs" in str_prompt else str_prompt
+
+            generated_text = self._generator(
+                str_prompt, max_length=30, num_return_sequences=1, truncation=True
+            )
+            return generated_text
+
+        self._create_server()
+
+    def _create_server(self):
+        """Placeholder docstring"""
+        app = FastAPI()
+        app.include_router(self._router)
+
+        config = uvicorn.Config(
+            app,
+            host="127.0.0.1",
+            port=9007,
+            log_level="info",
+            loop="asyncio",
+            reload=True,
+            use_colors=True,
+        )
+
+        self.server = uvicorn.Server(config)
+        self.host = config.host
+        self.port = config.port
+
+    def start_server(self):
+        """Starts the uvicorn server."""
+        if not (self._thread and self._thread.is_alive()):
+            logger.info("Waiting for a connection...")
+            self._thread = threading.Thread(target=self._start_run_async_in_thread, daemon=True)
+            self._thread.start()
+
+    def stop_server(self):
+        """Destroys the uvicorn server."""
+        # TODO: Implement me.
+
+    def _start_run_async_in_thread(self):
+        """Placeholder docstring"""
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(self._serve())
+
+    async def _serve(self):
+        """Placeholder docstring"""
+        await self.server.serve()

src/sagemaker/serve/builder/model_builder.py

@@ -812,7 +812,7 @@ 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

src/sagemaker/serve/builder/requirements_manager.py

@@ -36,7 +36,7 @@ def capture_and_install_dependencies(self, dependencies: Optional[str] = None) -
         Returns:
             file path of the existing or generated dependencies file
         """
-        _dependencies = dependencies or self._detect_conda_env_and_local_dependencies()
+        _dependencies = dependencies or self._detect_conda_env_and_local_dependencies
 
         # Dependencies specified as either req.txt or conda_env.yml
         if _dependencies.endswith(".txt"):

src/sagemaker/serve/builder/transformers_builder.py

@@ -421,6 +421,6 @@ def _create_conda_env(self):
         """Creating conda environment by running commands"""
 
         try:
-            RequirementsManager().capture_and_install_dependencies(self)
+            RequirementsManager().capture_and_install_dependencies
         except subprocess.CalledProcessError:
             print("Failed to create and activate conda environment.")

src/sagemaker/serve/mode/in_process_mode.py

@@ -1,6 +1,7 @@
 """Module that defines the InProcessMode class"""
 
 from __future__ import absolute_import
+
 from pathlib import Path
 import logging
 from typing import Dict, Type
@@ -11,7 +12,7 @@
 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.utils.exceptions import InProcessDeepPingException
 from sagemaker.serve.model_server.multi_model_server.server import InProcessMultiModelServer
 from sagemaker.session import Session
 
@@ -46,7 +47,7 @@ def __init__(
         self.session = session
         self.schema_builder = schema_builder
         self.model_server = model_server
-        self._ping_container = None
+        self._ping_local_server = None
 
     def load(self, model_path: str = None):
         """Loads model path, checks that path exists"""
@@ -69,21 +70,29 @@ def create_server(
         logger.info("Waiting for model server %s to start up...", self.model_server)
 
         if self.model_server == ModelServer.MMS:
-            self._ping_container = self._multi_model_server_deep_ping
+            self._ping_local_server = self._multi_model_server_deep_ping
+            self._start_serving()
+
+        # allow some time for server to be ready.
+        time.sleep(1)
 
         time_limit = datetime.now() + timedelta(seconds=5)
-        while self._ping_container is not None:
+        healthy = True
+        while True:
             final_pull = datetime.now() > time_limit
-
             if final_pull:
                 break
 
-            time.sleep(10)
-
-            healthy, response = self._ping_container(predictor)
+            healthy, response = self._ping_local_server(predictor)
             if healthy:
                 logger.debug("Ping health check has passed. Returned %s", str(response))
                 break
 
+        time.sleep(1)
+
         if not healthy:
-            raise LocalDeepPingException(_PING_HEALTH_CHECK_FAIL_MSG)
+            raise InProcessDeepPingException(_PING_HEALTH_CHECK_FAIL_MSG)
+
+    def destroy_server(self):
+        """Placeholder docstring"""
+        self._stop_serving()

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

@@ -2,12 +2,16 @@
 
 from __future__ import absolute_import
 
+import json
+
 import requests
 import logging
 import platform
 from pathlib import Path
+
 from sagemaker import Session, fw_utils
 from sagemaker.serve.utils.exceptions import LocalModelInvocationException
+from sagemaker.serve.utils.exceptions import InProcessDeepPingException
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.s3_utils import determine_bucket_and_prefix, parse_s3_url, s3_path_join
 from sagemaker.s3 import S3Uploader
@@ -25,16 +29,55 @@ class InProcessMultiModelServer:
 
     def _start_serving(self):
         """Initializes the start of the server"""
-        return Exception("Not implemented")
+        from sagemaker.serve.app import InProcessServer
 
-    def _invoke_multi_model_server_serving(self, request: object, content_type: str, accept: str):
-        """Invokes the MMS server by sending POST request"""
-        return Exception("Not implemented")
+        if hasattr(self, "inference_spec"):
+            model_id = self.inference_spec.get_model()
+            if not model_id:
+                raise ValueError("Model id was not provided in Inference Spec.")
+        else:
+            model_id = None
+        self.server = InProcessServer(model_id=model_id)
+
+        self.server.start_server()
+
+    def _stop_serving(self):
+        """Stops the server"""
+        self.server.stop_server()
+
+    def _invoke_multi_model_server_serving(self, request: bytes, content_type: str, accept: str):
+        """Placeholder docstring"""
+        try:
+            response = requests.post(
+                f"http://{self.server.host}:{self.server.port}/generate",
+                data=request,
+                headers={"Content-Type": content_type, "Accept": accept},
+                timeout=600,
+            )
+            response.raise_for_status()
+            if isinstance(response.content, bytes):
+                return json.loads(response.content.decode("utf-8"))
+            return response.content
+        except Exception as e:
+            if "Connection refused" in str(e):
+                raise Exception(
+                    "Unable to send request to the local server: Connection refused."
+                ) from e
+            raise Exception("Unable to send request to the local server.") from e
 
     def _multi_model_server_deep_ping(self, predictor: PredictorBase):
         """Sends a deep ping to ensure prediction"""
+        healthy = False
         response = None
-        return (True, response)
+        try:
+            response = predictor.predict(self.schema_builder.sample_input)
+            healthy = response is not None
+            # pylint: disable=broad-except
+        except Exception as e:
+            if "422 Client Error: Unprocessable Entity for url" in str(e):
+                raise InProcessDeepPingException(str(e))
+
+        return healthy, response
 
 
 class LocalMultiModelServer:

src/sagemaker/serve/utils/predictors.py

@@ -3,7 +3,7 @@
 from __future__ import absolute_import
 import io
 from typing import Type
-
+import logging
 from sagemaker import Session
 from sagemaker.serve.mode.local_container_mode import LocalContainerMode
 from sagemaker.serve.mode.in_process_mode import InProcessMode
@@ -16,6 +16,8 @@
 
 APPLICATION_X_NPY = "application/x-npy"
 
+logger = logging.getLogger(__name__)
+
 
 class TorchServeLocalPredictor(PredictorBase):
     """Lightweight predictor for local deployment in IN_PROCESS and LOCAL_CONTAINER modes"""
@@ -211,7 +213,7 @@ def delete_predictor(self):
 
 
 class TransformersInProcessModePredictor(PredictorBase):
-    """Lightweight Transformers predictor for local deployment"""
+    """Lightweight Transformers predictor for in process mode deployment"""
 
     def __init__(
         self,
@@ -225,18 +227,11 @@ def __init__(
 
     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,
-            )
-        ]
+        return self._mode_obj._invoke_multi_model_server_serving(
+            self.serializer.serialize(data),
+            self.content_type,
+            self.deserializer.ACCEPT[0],
+        )
 
     @property
     def content_type(self):

tests/unit/sagemaker/serve/builder/test_requirements_manager.py

@@ -29,7 +29,7 @@ class TestRequirementsManager(unittest.TestCase):
     @patch(
         "sagemaker.serve.builder.requirements_manager.RequirementsManager._detect_conda_env_and_local_dependencies"
     )
-    def test_capture_and_install_dependencies(
+    def test_capture_and_install_dependencies_txt(
         self,
         mock_detect_conda_env_and_local_dependencies,
         mock_install_requirements_txt,
@@ -40,8 +40,7 @@ def test_capture_and_install_dependencies(
         RequirementsManager().capture_and_install_dependencies()
         mock_install_requirements_txt.assert_called_once()
 
-        mock_detect_conda_env_and_local_dependencies.side_effect = lambda: ".yml"
-        RequirementsManager().capture_and_install_dependencies()
+        RequirementsManager().capture_and_install_dependencies("conda.yml")
         mock_update_conda_env_in_path.assert_called_once()
 
     @patch(