From bf019cc5a7194639bf32957e32969c0fabf08408 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= <andras@schmelczer.dev>
Date: Wed, 6 Apr 2022 19:01:34 +0200
Subject: [PATCH] Add fastapi deployment
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: András Schmelczer <andras@schmelczer.dev>
---
 .vscode/settings.json                         |  3 ++
 example/main.py                               |  4 +-
 good_ai/src/good_ai/good_ai/__init__.py       |  4 +-
 good_ai/src/good_ai/good_ai/config.py         | 14 ++++++
 .../src/good_ai/good_ai/deploy/__init__.py    |  1 +
 .../good_ai/good_ai/deploy/process_batch.py   |  3 +-
 good_ai/src/good_ai/good_ai/deploy/serve.py   | 46 +++++++++++++++++++
 .../src/good_ai/good_ai/models/use_model.py   |  3 +-
 .../src/good_ai/good_ai/set_default_config.py |  2 +-
 .../src/good_ai/good_ai/tracing/__init__.py   |  2 -
 .../good_ai/tracing/tracing_context.py        |  4 +-
 good_ai/src/good_ai/good_ai/views/__init__.py |  3 ++
 .../good_ai/views/health_check_response.py    |  5 ++
 .../good_ai/{tracing => views}/model.py       |  0
 .../good_ai/{tracing => views}/trace.py       |  0
 good_ai/src/good_ai/open_s3/large_file.py     |  2 -
 16 files changed, 83 insertions(+), 13 deletions(-)
 create mode 100644 good_ai/src/good_ai/good_ai/config.py
 create mode 100644 good_ai/src/good_ai/good_ai/views/__init__.py
 create mode 100644 good_ai/src/good_ai/good_ai/views/health_check_response.py
 rename good_ai/src/good_ai/good_ai/{tracing => views}/model.py (100%)
 rename good_ai/src/good_ai/good_ai/{tracing => views}/trace.py (100%)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index a4135c2..03a850a 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,11 +1,14 @@
 {
   "cSpell.words": [
+    "boto",
     "botocore",
     "pydantic",
     "pyplot",
     "sklearn",
+    "starlette",
     "Tfidf",
     "tinydb",
+    "uvicorn",
     "Vectorizer",
     "xmargin"
   ],
diff --git a/example/main.py b/example/main.py
index 4621f2e..2d1def9 100644
--- a/example/main.py
+++ b/example/main.py
@@ -4,9 +4,11 @@ from random import shuffle
 from devtools import debug
 from predict_domain import predict_domain
 
-from good_ai import process_batch
+from good_ai import process_batch, serve
 
 if __name__ == "__main__":
+    serve(predict_domain)
+    
     with open(".cache/ss-data-0/s2-corpus-1583.json") as f:
         raw = json.load(f)
 
diff --git a/good_ai/src/good_ai/good_ai/__init__.py b/good_ai/src/good_ai/good_ai/__init__.py
index 2bc2318..57df559 100644
--- a/good_ai/src/good_ai/good_ai/__init__.py
+++ b/good_ai/src/good_ai/good_ai/__init__.py
@@ -1,3 +1,3 @@
-from .deploy import process_batch
+from .deploy import process_batch, serve
 from .models import save_model, use_model
-from .set_default_config import set_default_config, set_default_config_if_uninitialized
+from .set_default_config import set_default_config
diff --git a/good_ai/src/good_ai/good_ai/config.py b/good_ai/src/good_ai/good_ai/config.py
new file mode 100644
index 0000000..9e501dc
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/config.py
@@ -0,0 +1,14 @@
+import logging
+import os
+
+logger = logging.getLogger("good_ai")
+
+PRODUCTION_KEY = "production"
+
+environment = os.environ.get("ENVIRONMENT", PRODUCTION_KEY).lower()
+is_production = environment == PRODUCTION_KEY
+
+if is_production:
+    logger.info("Running in production mode ✅")
+else:
+    logger.warn("Running in development mode ‼️")
diff --git a/good_ai/src/good_ai/good_ai/deploy/__init__.py b/good_ai/src/good_ai/good_ai/deploy/__init__.py
index 5013cad..d13ff87 100644
--- a/good_ai/src/good_ai/good_ai/deploy/__init__.py
+++ b/good_ai/src/good_ai/good_ai/deploy/__init__.py
@@ -1 +1,2 @@
 from .process_batch import process_batch
+from .serve import serve
diff --git a/good_ai/src/good_ai/good_ai/deploy/process_batch.py b/good_ai/src/good_ai/good_ai/deploy/process_batch.py
index 2e76485..10256da 100644
--- a/good_ai/src/good_ai/good_ai/deploy/process_batch.py
+++ b/good_ai/src/good_ai/good_ai/deploy/process_batch.py
@@ -3,7 +3,8 @@ from typing import Any, Callable, Iterable, Optional, Sequence
 from good_ai.utilities.parallel_map import parallel_map
 
 from ..set_default_config import set_default_config_if_uninitialized
-from ..tracing import Trace, TracingContext
+from ..tracing import TracingContext
+from ..views import Trace
 
 
 def process_batch(
diff --git a/good_ai/src/good_ai/good_ai/deploy/serve.py b/good_ai/src/good_ai/good_ai/deploy/serve.py
index 8b13789..c1a7ebf 100644
--- a/good_ai/src/good_ai/good_ai/deploy/serve.py
+++ b/good_ai/src/good_ai/good_ai/deploy/serve.py
@@ -1 +1,47 @@
+from typing import Any, Callable
 
+import uvicorn
+from fastapi import FastAPI, status
+from fastapi.openapi.docs import get_swagger_ui_html
+from fastapi.responses import RedirectResponse
+from starlette.responses import HTMLResponse
+from typing_extensions import Never
+
+from ..set_default_config import set_default_config_if_uninitialized
+from ..tracing import TracingContext
+from ..views import HealthCheckResponse, Trace
+
+
+def serve(
+    function: Callable[..., Any],
+) -> Never:
+    set_default_config_if_uninitialized()
+
+    app = FastAPI(
+        title=function.__name__.capitalize().replace("_", " "),
+        description=f"REST API wrapper for interacting with the {function.__name__} function.",
+        docs_url=None,
+        redoc_url=None,
+    )
+
+    @app.get("/", include_in_schema=False)
+    def redirect_to_docs() -> RedirectResponse:
+        return RedirectResponse("/docs")
+
+    @app.get("/docs", include_in_schema=False)
+    def custom_swagger_ui_html() -> HTMLResponse:
+        return get_swagger_ui_html(openapi_url="openapi.json", title=app.title)
+
+    @app.get("/health", status_code=status.HTTP_200_OK)
+    def check_health() -> HealthCheckResponse:
+        return HealthCheckResponse(is_healthy=True)
+
+    @app.post("/score", status_code=status.HTTP_200_OK, response_model=Trace)
+    def process(input: Any) -> Trace:
+        with TracingContext() as t:
+            t.log_input(input)
+            result = function(input)
+            output = t.log_output(result)
+        return output
+
+    uvicorn.run(app, host="0.0.0.0", port=5050)
diff --git a/good_ai/src/good_ai/good_ai/models/use_model.py b/good_ai/src/good_ai/good_ai/models/use_model.py
index 680e443..189ba65 100644
--- a/good_ai/src/good_ai/good_ai/models/use_model.py
+++ b/good_ai/src/good_ai/good_ai/models/use_model.py
@@ -1,7 +1,8 @@
 from functools import wraps
 from typing import Any, Callable, Literal, Union
 
-from ..tracing import Model, TracingContext
+from ..tracing import TracingContext
+from ..views import Model
 from .load_model import load_model
 
 
diff --git a/good_ai/src/good_ai/good_ai/set_default_config.py b/good_ai/src/good_ai/good_ai/set_default_config.py
index e7cd4f0..e99ad38 100644
--- a/good_ai/src/good_ai/good_ai/set_default_config.py
+++ b/good_ai/src/good_ai/good_ai/set_default_config.py
@@ -33,7 +33,7 @@ def set_default_config(
 
     _initialized = True
 
-    logger.info(f"Defaults: configured ✅")
+    logger.info("Defaults: configured ✅")
 
 
 def _initialize_large_file(s3_config: Path) -> None:
diff --git a/good_ai/src/good_ai/good_ai/tracing/__init__.py b/good_ai/src/good_ai/good_ai/tracing/__init__.py
index 510390e..12773c1 100644
--- a/good_ai/src/good_ai/good_ai/tracing/__init__.py
+++ b/good_ai/src/good_ai/good_ai/tracing/__init__.py
@@ -1,4 +1,2 @@
-from .model import Model
 from .persistence import MongoDbDriver, PersistenceDriver, TinyDbDriver
-from .trace import Trace
 from .tracing_context import TracingContext
diff --git a/good_ai/src/good_ai/good_ai/tracing/tracing_context.py b/good_ai/src/good_ai/good_ai/tracing/tracing_context.py
index bbfda4d..f710ee3 100644
--- a/good_ai/src/good_ai/good_ai/tracing/tracing_context.py
+++ b/good_ai/src/good_ai/good_ai/tracing/tracing_context.py
@@ -5,13 +5,11 @@ from datetime import datetime
 from types import TracebackType
 from typing import Any, DefaultDict, List, Optional, Type
 
+from ..views import Model, Trace
 from .persistence import PersistenceDriver
 
 logger = logging.getLogger("good_ai")
 
-from .model import Model
-from .trace import Trace
-
 
 class TracingContext:
     persistence_driver: PersistenceDriver
diff --git a/good_ai/src/good_ai/good_ai/views/__init__.py b/good_ai/src/good_ai/good_ai/views/__init__.py
new file mode 100644
index 0000000..a367601
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/views/__init__.py
@@ -0,0 +1,3 @@
+from .health_check_response import HealthCheckResponse
+from .model import Model
+from .trace import Trace
diff --git a/good_ai/src/good_ai/good_ai/views/health_check_response.py b/good_ai/src/good_ai/good_ai/views/health_check_response.py
new file mode 100644
index 0000000..5ffa26c
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/views/health_check_response.py
@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class HealthCheckResponse(BaseModel):
+    is_healthy: bool
diff --git a/good_ai/src/good_ai/good_ai/tracing/model.py b/good_ai/src/good_ai/good_ai/views/model.py
similarity index 100%
rename from good_ai/src/good_ai/good_ai/tracing/model.py
rename to good_ai/src/good_ai/good_ai/views/model.py
diff --git a/good_ai/src/good_ai/good_ai/tracing/trace.py b/good_ai/src/good_ai/good_ai/views/trace.py
similarity index 100%
rename from good_ai/src/good_ai/good_ai/tracing/trace.py
rename to good_ai/src/good_ai/good_ai/views/trace.py
diff --git a/good_ai/src/good_ai/open_s3/large_file.py b/good_ai/src/good_ai/open_s3/large_file.py
index 59c07d9..a742487 100644
--- a/good_ai/src/good_ai/open_s3/large_file.py
+++ b/good_ai/src/good_ai/open_s3/large_file.py
@@ -96,7 +96,6 @@ class LargeFile:
         credentials.default_section
         cls.configure_credentials(**credentials[credentials.default_section])
 
-    
     @classmethod
     def configure_credentials(
         cls,
@@ -114,7 +113,6 @@ class LargeFile:
         cls.bucket_name = large_files_bucket_name
         cls.endpoint_url = aws_endpoint_url
 
-
     def __enter__(self) -> IO:
         self._file: IO[Any] = (
             tempfile.NamedTemporaryFile(