minor changes

Author: Bryannah Hernandez

src/sagemaker/serve/app.py

@@ -1,75 +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("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"
-    )
+    logger.error("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",
-    )
-
-    @app.get("/")
-    def read_root():
+    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"""
-        return {"Hello": "World"}
+        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("/generate")
-    async def generate_text(prompt: Request):
-        """Placeholder docstring"""
-        str_prompt = await prompt.json()
+        self.server = uvicorn.Server(config)
+        self.host = config.host
+        self.port = config.port
 
-        generated_text = generator(
-            str_prompt, max_length=30, num_return_sequences=5, truncation=True
-        )
-        return generated_text[0]["generated_text"]
+    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()
 
-    @app.post("/post")
-    def post(payload: dict):
+    def stop_server(self):
+        """Destroys the uvicorn server."""
+        # TODO: Implement me.
+
+    def _start_run_async_in_thread(self):
         """Placeholder docstring"""
-        return payload
+        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):
+        """Placeholder docstring"""
+        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,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,13 @@
 
 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
@@ -26,78 +27,57 @@
 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
 
-        time.sleep(10)
+        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)
 
-    def _invoke_multi_model_server_serving(self, request: object, content_type: str, accept: str):
+        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"""
-        time.sleep(2)
-        background_tasks = set()
-        task = asyncio.create_task(self.generate_connect())
-        background_tasks.add(task)
-        task.add_done_callback(background_tasks.discard)
-        return task.result()
+        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"""
-        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
         try:
             response = predictor.predict(self.schema_builder.sample_input)
-            return True, response
+            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 False, 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)
-
-        data = await reader.read()
-        logger.info("Response from server")
-        logger.info(data)
-        writer.close()
-        await writer.wait_closed()
-        return data
-
-    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",
-        )
 
-        data = await reader.read()
-        logger.info("Response from server")
-        logger.info(data)
-        writer.close()
-        await writer.wait_closed()
-        return data
+        return healthy, response
 
 
 class LocalMultiModelServer:

tests/unit/sagemaker/serve/mode/test_in_process_mode.py

@@ -18,7 +18,7 @@
 from sagemaker.serve.mode.in_process_mode import InProcessMode
 from sagemaker.serve import SchemaBuilder
 from sagemaker.serve.utils.types import ModelServer
-from sagemaker.serve.utils.exceptions import LocalDeepPingException
+from sagemaker.serve.utils.exceptions import InProcessDeepPingException
 
 
 mock_prompt = "Hello, I'm a language model,"
@@ -163,4 +163,4 @@ def test_create_server_ex(
         in_process_mode._multi_model_server_deep_ping = mock_multi_model_server_deep_ping
         in_process_mode._start_serving = mock_start_serving
 
-        self.assertRaises(LocalDeepPingException, in_process_mode.create_server, mock_predictor)
+        self.assertRaises(InProcessDeepPingException, in_process_mode.create_server, mock_predictor)