Merge pull request #1 from makungaj1/patch-1

Update in_process_mode.py

Author: makungaj1

src/sagemaker/serve/app.py

@@ -1,75 +1,101 @@
 """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("To enable in_process mode for Transformers install uvicorn from HuggingFace hub")
+    logger.error("Unable to import uvicorn, check if uvicorn is installed.")
 
 
 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"
+        "Unable to import transformers, check if transformers is installed."
     )
 
 
 try:
-    from fastapi import FastAPI, Request
-
-    app = FastAPI(
-        title="Transformers In Process Server",
-        version="1.0",
-        description="A simple server",
-    )
+    from fastapi import FastAPI, Request, APIRouter
+except ImportError:
+    logger.error("Unable to import fastapi, check if fastapi is installed.")
+
+
+class InProcessServer:
+
+    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")
+
+        @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):
+        _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,
+        )
 
-    @app.get("/")
-    def read_root():
-        """Placeholder docstring"""
-        return {"Hello": "World"}
+        self.server = uvicorn.Server(config)
+        self.host = config.host
+        self.port = config.port
 
-    @app.get("/generate")
-    async def generate_text(prompt: Request):
-        """Placeholder docstring"""
-        str_prompt = await prompt.json()
+    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()
 
-        generated_text = generator(
-            str_prompt, max_length=30, num_return_sequences=5, truncation=True
-        )
-        return generated_text[0]["generated_text"]
+    def stop_server(self):
+        """Destroys the uvicorn server."""
+        # TODO: Implement me.
 
-    @app.post("/post")
-    def post(payload: dict):
-        """Placeholder docstring"""
-        return payload
+    def _start_run_async_in_thread(self):
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(self._serve())
 
-except ImportError:
-    logger.error("To enable in_process mode for Transformers install fastapi from HuggingFace hub")
-
-
-async def main():
-    """Running server locally with uvicorn"""
-    config = uvicorn.Config(
-        "sagemaker.serve.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()
+    async def _serve(self):
+        await self.server.serve()

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,22 +70,30 @@ 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()
 
-        time_limit = datetime.now() + timedelta(seconds=5)
-        while self._ping_container is not None:
-            final_pull = datetime.now() > time_limit
+        # allow some time for server to be ready.
+        time.sleep(1)
 
+        count = 1
+        time_limit = datetime.now() + timedelta(seconds=20)
+        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)
+            count += 1
             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):
+        self._stop_serving()

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

@@ -2,14 +2,15 @@
 
 from __future__ import absolute_import
 
-import asyncio
+import json
+
 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
+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,71 +26,52 @@
 class InProcessMultiModelServer:
     """In Process Mode Multi Model server instance"""
 
-    def __init__(self):
-        from sagemaker.serve.app import main
-
-        self._main = main
-
     def _start_serving(self):
         """Initializes the start of the server"""
-        background_tasks = set()
-        task = asyncio.create_task(self._main())
-        background_tasks.add(task)
-        task.add_done_callback(background_tasks.discard)
+        from sagemaker.serve.app import InProcessServer
+        if hasattr(self, "inference_spec"):
+            model_id = self.inference_spec.get_model()
+        else:
+            model_id = None
+        self.server = InProcessServer(model_id=model_id)
 
-        time.sleep(10)
+        self.server.start_server()
 
-    def _invoke_multi_model_server_serving(self, request: object, content_type: str, accept: str):
+    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"""
-        background_tasks = set()
-        task = asyncio.create_task(self.generate_connect())
-        background_tasks.add(task)
-        task.add_done_callback(background_tasks.discard)
+        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 not "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"""
-        background_tasks = set()
-        task = asyncio.create_task(self.tcp_connect())
-        background_tasks.add(task)
-        task.add_done_callback(background_tasks.discard)
+        healthy = False
         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()
+        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: