FastAPI with In_Process

Author: Bryannah Hernandez

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

src/sagemaker/app.py

@@ -0,0 +1,48 @@
+from __future__ import absolute_import
+
+import asyncio
+import logging
+from transformers import pipeline
+from fastapi import FastAPI
+import uvicorn
+
+logger = logging.getLogger(__name__)
+
+app = FastAPI(
+    title="Transformers In Process Server",
+    version="1.0",
+    description="A simple server",
+)
+
+
+@app.get("/")
+def read_root():
+    return {"Hello": "World"}
+
+
+@app.post("/generate")
+def generate_text(prompt: str, max_length=500, num_return_sequences=1):
+    logger.info("Generating Text....")
+
+    generated_text = generator(
+        prompt, max_length=max_length, num_return_sequences=num_return_sequences
+    )
+    return generated_text[0]["generated_text"]
+
+
+generator = pipeline("text-generation", model="gpt2")
+
+
+@app.post("/post")
+def post(prompt: str):
+    return prompt
+
+
+async def main():
+    logger.info("Running")
+    config = uvicorn.Config(
+        "sagemaker.app:app", host="0.0.0.0", port=8080, log_level="info", loop="asyncio"
+    )
+    server = uvicorn.Server(config)
+    await server.serve()
+    

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

@@ -2,6 +2,7 @@
 
 from __future__ import absolute_import
 
+import asyncio
 import requests
 import logging
 import platform
@@ -13,6 +14,8 @@
 from sagemaker.s3 import S3Uploader
 from sagemaker.local.utils import get_docker_host
 from sagemaker.serve.utils.optimize_utils import _is_s3_uri
+import time
+from sagemaker.app import main
 
 MODE_DIR_BINDING = "/opt/ml/model/"
 _DEFAULT_ENV_VARS = {}
@@ -25,11 +28,36 @@ class InProcessMultiModelServer:
 
     def _start_serving(self):
         """Initializes the start of the server"""
-        return Exception("Not implemented")
+
+        logger.info("Server started at http://0.0.0.0")
+
+        asyncio.create_task(main())
 
     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")
+
+        logger.info(request)
+        logger.info(content_type)
+        logger.info(accept)
+
+        try:
+            response = requests.post(
+                f"http://0.0.0.0:8080/generate",
+                data=request,
+                headers={"Content-Type": content_type, "Accept": accept},
+                timeout=600,
+            )
+            response.raise_for_status()
+
+            return response.content
+        except requests.exceptions.ConnectionError as e:
+            logger.debug(f"Error connecting to the server: {e}")
+        except requests.exceptions.HTTPError as e:
+            logger.debug(f"HTTP error occurred: {e}")
+        except requests.exceptions.RequestException as e:
+            logger.debug(f"An error occurred: {e}")
+        except Exception as e:
+            raise Exception("Unable to send request to the local container server") from e
 
     def _multi_model_server_deep_ping(self, predictor: PredictorBase):
         """Sends a deep ping to ensure prediction"""

src/sagemaker/serve/utils/predictors.py

@@ -3,7 +3,8 @@
 from __future__ import absolute_import
 import io
 from typing import Type
-
+import logging
+import json
 from sagemaker import Session
 from sagemaker.serve.mode.local_container_mode import LocalContainerMode
 from sagemaker.serve.mode.in_process_mode import InProcessMode
@@ -16,6 +17,7 @@
 
 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,