From 09137e146c16c94693b96b1f5a2cb0e229e672b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= <andras@schmelczer.dev>
Date: Sat, 9 Apr 2022 18:53:57 +0200
Subject: [PATCH] Refactor by adding context
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: András Schmelczer <andras@schmelczer.dev>
---
 example/predict_domain.py                     |  1 +
 good_ai/src/good_ai/good_ai/__init__.py       |  4 +-
 good_ai/src/good_ai/good_ai/config.py         | 13 ----
 .../src/good_ai/good_ai/context/__init__.py   |  2 +
 .../src/good_ai/good_ai/context/context.py    | 12 +++
 .../good_ai/good_ai/context/get_context.py    | 76 +++++++++++++++++++
 .../src/good_ai/good_ai/deploy/__init__.py    |  1 +
 .../good_ai/deploy/create_fastapi_app.py      | 18 ++---
 .../good_ai/good_ai/deploy/process_batch.py   |  3 -
 .../good_ai/good_ai/deploy/process_single.py  | 12 +++
 good_ai/src/good_ai/good_ai/deploy/serve.py   | 17 ++++-
 .../good_ai/metrics/create_dash_app.py        |  5 +-
 .../src/good_ai/good_ai/models/load_model.py  |  4 +-
 .../src/good_ai/good_ai/models/save_model.py  |  5 +-
 .../{tracing => }/persistence/__init__.py     |  0
 .../persistence/mongodb_driver.py             |  0
 .../persistence/persistence_driver.py         |  0
 .../persistence/tinydb_driver.py              |  0
 .../src/good_ai/good_ai/set_default_config.py | 55 --------------
 .../src/good_ai/good_ai/tracing/__init__.py   |  1 -
 .../good_ai/tracing/tracing_context.py        |  5 +-
 21 files changed, 134 insertions(+), 100 deletions(-)
 delete mode 100644 good_ai/src/good_ai/good_ai/config.py
 create mode 100644 good_ai/src/good_ai/good_ai/context/__init__.py
 create mode 100644 good_ai/src/good_ai/good_ai/context/context.py
 create mode 100644 good_ai/src/good_ai/good_ai/context/get_context.py
 create mode 100644 good_ai/src/good_ai/good_ai/deploy/process_single.py
 rename good_ai/src/good_ai/good_ai/{tracing => }/persistence/__init__.py (100%)
 rename good_ai/src/good_ai/good_ai/{tracing => }/persistence/mongodb_driver.py (100%)
 rename good_ai/src/good_ai/good_ai/{tracing => }/persistence/persistence_driver.py (100%)
 rename good_ai/src/good_ai/good_ai/{tracing => }/persistence/tinydb_driver.py (100%)
 delete mode 100644 good_ai/src/good_ai/good_ai/set_default_config.py

diff --git a/example/predict_domain.py b/example/predict_domain.py
index e2ace96..e6d6938 100644
--- a/example/predict_domain.py
+++ b/example/predict_domain.py
@@ -15,6 +15,7 @@ def predict_domain(
     text: str, model: Pipeline, cut_off_probability: float = 0.2
 ) -> List[DomainPrediction]:
     assert 0 <= cut_off_probability <= 1
+    
     cleaned = clean(text, convert_to_ascii=True)
     text = re.sub(r"[^a-zA-Z0-9]", " ", cleaned)
 
diff --git a/good_ai/src/good_ai/good_ai/__init__.py b/good_ai/src/good_ai/good_ai/__init__.py
index 57df559..bc171fa 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, serve
+from .context import set_default_config
+from .deploy import process_batch, process_single, serve
 from .models import save_model, use_model
-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
deleted file mode 100644
index 85e4be0..0000000
--- a/good_ai/src/good_ai/good_ai/config.py
+++ /dev/null
@@ -1,13 +0,0 @@
-METRICS_PATH = "/metrics"
-
-# 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/context/__init__.py b/good_ai/src/good_ai/good_ai/context/__init__.py
new file mode 100644
index 0000000..58cd719
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/context/__init__.py
@@ -0,0 +1,2 @@
+from .context import Context
+from .get_context import get_context, set_default_config
diff --git a/good_ai/src/good_ai/good_ai/context/context.py b/good_ai/src/good_ai/good_ai/context/context.py
new file mode 100644
index 0000000..2358759
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/context/context.py
@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+
+from ..persistence import PersistenceDriver
+
+
+class Context(BaseModel):
+    metrics_path: str
+    persistence: PersistenceDriver
+    is_production: bool
+
+    class Config:
+        arbitrary_types_allowed = True
diff --git a/good_ai/src/good_ai/good_ai/context/get_context.py b/good_ai/src/good_ai/good_ai/context/get_context.py
new file mode 100644
index 0000000..b44a54b
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/context/get_context.py
@@ -0,0 +1,76 @@
+import logging
+import os
+import random
+from pathlib import Path
+from typing import Optional, cast
+
+from good_ai.open_s3 import LargeFile
+
+from ..persistence import PersistenceDriver, TinyDbDriver
+from .context import Context
+
+logger = logging.getLogger("good_ai")
+
+_context: Optional[Context] = None
+PRODUCTION_KEY = "production"
+
+
+def get_context() -> Context:
+    if _context is None:
+        set_default_config()
+
+    return cast(Context, _context)
+
+
+def set_default_config(
+    log_level: int = logging.INFO,
+    s3_config: Path = Path("s3.ini"),
+    seed: int = 42,
+    persistence_driver: PersistenceDriver = TinyDbDriver(Path("tracing_database.json")),
+    is_production_mode_override: Optional[bool] = None,
+) -> None:
+    global _context
+    logging.basicConfig(level=log_level)
+
+    is_production = _is_in_production_mode(override=is_production_mode_override)
+    _initialize_large_file(s3_config)
+    _set_seed(seed)
+
+    _context = Context(
+        metrics_path="/metrics",
+        persistence=persistence_driver,
+        is_production=is_production,
+    )
+
+    logger.info("Defaults: configured ✅")
+
+
+def _is_in_production_mode(override: Optional[bool]) -> bool:
+    environment = os.environ.get("ENVIRONMENT", PRODUCTION_KEY).lower()
+    is_production = environment == PRODUCTION_KEY if override is None else override
+
+    if is_production:
+        logger.info("Running in production mode ✅")
+    else:
+        logger.warn("Running in development mode ‼️")
+
+    return is_production
+
+
+def _initialize_large_file(s3_config: Path) -> None:
+    if s3_config.exists():
+        LargeFile.configure_credentials_from_file(s3_config)
+    else:
+        logger.info(
+            f"Provided S3 config ({s3_config.resolve()}) not found, skipping LargeFile initialisation"
+        )
+
+
+def _set_seed(seed: int) -> None:
+    random.seed(seed)
+    try:
+        import numpy
+
+        numpy.random.seed(seed + 1)
+    except ImportError:
+        pass
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 d13ff87..c274dbd 100644
--- a/good_ai/src/good_ai/good_ai/deploy/__init__.py
+++ b/good_ai/src/good_ai/good_ai/deploy/__init__.py
@@ -1,2 +1,3 @@
 from .process_batch import process_batch
+from .process_single import process_single
 from .serve import serve
diff --git a/good_ai/src/good_ai/good_ai/deploy/create_fastapi_app.py b/good_ai/src/good_ai/good_ai/deploy/create_fastapi_app.py
index 0767d16..71e6cc9 100644
--- a/good_ai/src/good_ai/good_ai/deploy/create_fastapi_app.py
+++ b/good_ai/src/good_ai/good_ai/deploy/create_fastapi_app.py
@@ -1,16 +1,12 @@
 from fastapi import FastAPI, status
-from fastapi.middleware.wsgi import WSGIMiddleware
 from fastapi.openapi.docs import get_swagger_ui_html
-from fastapi.responses import RedirectResponse
 from starlette.responses import HTMLResponse
 
-from ..config import METRICS_PATH
 from ..helper import snake_case_to_text
-from ..metrics import create_dash_app
 from ..views import HealthCheckResponse
 
 
-def create_fastapi_app(function_name: str) -> FastAPI:
+def create_fastapi_app(function_name: str, disable_docs: bool) -> FastAPI:
     app = FastAPI(
         title=snake_case_to_text(function_name),
         description=f"REST API wrapper for interacting with the '{function_name}' function.",
@@ -18,15 +14,11 @@ def create_fastapi_app(function_name: str) -> FastAPI:
         redoc_url=None,
     )
 
-    @app.get("/", include_in_schema=False)
-    def redirect_to_entrypoint() -> RedirectResponse:
-        return RedirectResponse("/metrics")
+    if not disable_docs:
 
-    app.mount(METRICS_PATH, WSGIMiddleware(create_dash_app(function_name)))
-
-    @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("/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:
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 10256da..9dc0f3a 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
@@ -2,7 +2,6 @@ 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 TracingContext
 from ..views import Trace
 
@@ -12,8 +11,6 @@ def process_batch(
     batch: Iterable[Any],
     concurrency: Optional[int] = None,
 ) -> Sequence[Trace]:
-    set_default_config_if_uninitialized()
-
     def inner(input: Any) -> Trace:
         with TracingContext() as t:
             t.log_input(input)
diff --git a/good_ai/src/good_ai/good_ai/deploy/process_single.py b/good_ai/src/good_ai/good_ai/deploy/process_single.py
new file mode 100644
index 0000000..19d8346
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/deploy/process_single.py
@@ -0,0 +1,12 @@
+from typing import Any, Callable
+
+from ..tracing import TracingContext
+from ..views import Trace
+
+
+def process_single(function: Callable[..., Any], input_value: Any) -> Trace:
+    with TracingContext() as t:
+        t.log_input(input_value)
+        result = function(input_value)
+        output = t.log_output(result)
+    return output
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 ddfee06..3ffde5c 100644
--- a/good_ai/src/good_ai/good_ai/deploy/serve.py
+++ b/good_ai/src/good_ai/good_ai/deploy/serve.py
@@ -2,21 +2,32 @@ from typing import Any, Callable
 
 import uvicorn
 from fastapi import status
+from fastapi.middleware.wsgi import WSGIMiddleware
+from fastapi.responses import RedirectResponse
 from typing_extensions import Never
 
 from good_ai.good_ai.deploy.create_fastapi_app import create_fastapi_app
 
-from ..set_default_config import set_default_config_if_uninitialized
+from ..context import get_context
+from ..metrics import create_dash_app
 from ..tracing import TracingContext
 from ..views import Trace
 
 
 def serve(
     function: Callable[..., Any],
+    disable_docs: bool = False,
+    disable_metrics: bool = False,
 ) -> Never:
-    set_default_config_if_uninitialized()
+    app = create_fastapi_app(function.__name__, disable_docs=disable_docs)
 
-    app = create_fastapi_app(function.__name__)
+    if not disable_metrics:
+        dash_app = create_dash_app(function.__name__)
+        app.mount(get_context().metrics_path, WSGIMiddleware(dash_app))
+
+        @app.get("/", include_in_schema=False)
+        def redirect_to_entrypoint() -> RedirectResponse:
+            return RedirectResponse("/metrics")
 
     @app.post("/score", status_code=status.HTTP_200_OK, response_model=Trace)
     def process(input: Any) -> Trace:
diff --git a/good_ai/src/good_ai/good_ai/metrics/create_dash_app.py b/good_ai/src/good_ai/good_ai/metrics/create_dash_app.py
index a261888..9c116a3 100644
--- a/good_ai/src/good_ai/good_ai/metrics/create_dash_app.py
+++ b/good_ai/src/good_ai/good_ai/metrics/create_dash_app.py
@@ -3,12 +3,13 @@ import plotly.express as px
 from dash import Dash, dcc, html
 from flask import Flask
 
-from ..config import METRICS_PATH
+from good_ai.good_ai.context.get_context import get_context
+
 from ..helper import snake_case_to_text
 
 
 def create_dash_app(function_name: str) -> Flask:
-    app = Dash(function_name, requests_pathname_prefix=METRICS_PATH + "/")
+    app = Dash(function_name, requests_pathname_prefix=get_context().metrics_path + "/")
 
     markdown_text = f"""
     # {snake_case_to_text(function_name)} - metrics
diff --git a/good_ai/src/good_ai/good_ai/models/load_model.py b/good_ai/src/good_ai/good_ai/models/load_model.py
index b30a611..147edb0 100644
--- a/good_ai/src/good_ai/good_ai/models/load_model.py
+++ b/good_ai/src/good_ai/good_ai/models/load_model.py
@@ -5,7 +5,7 @@ from joblib import load
 
 from good_ai.open_s3 import LargeFile
 
-from ..set_default_config import set_default_config_if_uninitialized
+from ..context import get_context
 
 logger = logging.getLogger("models")
 
@@ -13,7 +13,7 @@ logger = logging.getLogger("models")
 def load_model(
     key: str, version: Optional[int] = None, return_path: bool = False
 ) -> Tuple[Any, int]:
-    set_default_config_if_uninitialized()
+    get_context()  # will setup LargeFile if there was no config set
 
     file = LargeFile(name=key, mode="rb", version=version)
 
diff --git a/good_ai/src/good_ai/good_ai/models/save_model.py b/good_ai/src/good_ai/good_ai/models/save_model.py
index 1702b42..8053145 100644
--- a/good_ai/src/good_ai/good_ai/models/save_model.py
+++ b/good_ai/src/good_ai/good_ai/models/save_model.py
@@ -4,17 +4,16 @@ from typing import Optional, Union
 
 from joblib import dump
 
+from good_ai.good_ai.context.get_context import get_context
 from good_ai.open_s3 import LargeFile
 
-from ..set_default_config import set_default_config_if_uninitialized
-
 logger = logging.getLogger("models")
 
 
 def save_model(
     model: Union[Path, str, object], key: str, keep_last_n: Optional[int] = None
 ) -> int:
-    set_default_config_if_uninitialized()
+    get_context()  # will setup LargeFile if there was no config set
 
     file = LargeFile(name=key, mode="wb", keep_last_n=keep_last_n)
 
diff --git a/good_ai/src/good_ai/good_ai/tracing/persistence/__init__.py b/good_ai/src/good_ai/good_ai/persistence/__init__.py
similarity index 100%
rename from good_ai/src/good_ai/good_ai/tracing/persistence/__init__.py
rename to good_ai/src/good_ai/good_ai/persistence/__init__.py
diff --git a/good_ai/src/good_ai/good_ai/tracing/persistence/mongodb_driver.py b/good_ai/src/good_ai/good_ai/persistence/mongodb_driver.py
similarity index 100%
rename from good_ai/src/good_ai/good_ai/tracing/persistence/mongodb_driver.py
rename to good_ai/src/good_ai/good_ai/persistence/mongodb_driver.py
diff --git a/good_ai/src/good_ai/good_ai/tracing/persistence/persistence_driver.py b/good_ai/src/good_ai/good_ai/persistence/persistence_driver.py
similarity index 100%
rename from good_ai/src/good_ai/good_ai/tracing/persistence/persistence_driver.py
rename to good_ai/src/good_ai/good_ai/persistence/persistence_driver.py
diff --git a/good_ai/src/good_ai/good_ai/tracing/persistence/tinydb_driver.py b/good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py
similarity index 100%
rename from good_ai/src/good_ai/good_ai/tracing/persistence/tinydb_driver.py
rename to good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py
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
deleted file mode 100644
index e99ad38..0000000
--- a/good_ai/src/good_ai/good_ai/set_default_config.py
+++ /dev/null
@@ -1,55 +0,0 @@
-import logging
-import random
-from pathlib import Path
-
-from good_ai.good_ai.tracing.tracing_context import TracingContext
-from good_ai.open_s3 import LargeFile
-
-from .tracing import PersistenceDriver, TinyDbDriver
-
-logger = logging.getLogger("good_ai")
-
-_initialized = False
-
-
-def set_default_config_if_uninitialized() -> None:
-    if not _initialized:
-        set_default_config()
-
-
-def set_default_config(
-    log_level: int = logging.INFO,
-    s3_config: Path = Path("s3.ini"),
-    seed: int = 42,
-    tracing_db_driver: PersistenceDriver = TinyDbDriver(Path("tracing_database.json")),
-) -> None:
-    global _initialized
-    logging.basicConfig(level=log_level)
-
-    _initialize_large_file(s3_config)
-    _set_seed(seed)
-
-    TracingContext.persistence_driver = tracing_db_driver
-
-    _initialized = True
-
-    logger.info("Defaults: configured ✅")
-
-
-def _initialize_large_file(s3_config: Path) -> None:
-    if s3_config.exists():
-        LargeFile.configure_credentials_from_file(s3_config)
-    else:
-        logger.info(
-            f"Provided S3 config ({s3_config.resolve()}) not found, skipping LargeFile initialisation"
-        )
-
-
-def _set_seed(seed: int) -> None:
-    random.seed(seed)
-    try:
-        import numpy
-
-        numpy.random.seed(seed + 1)
-    except ImportError:
-        pass
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 12773c1..47ec644 100644
--- a/good_ai/src/good_ai/good_ai/tracing/__init__.py
+++ b/good_ai/src/good_ai/good_ai/tracing/__init__.py
@@ -1,2 +1 @@
-from .persistence import MongoDbDriver, PersistenceDriver, TinyDbDriver
 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 f710ee3..550c768 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,14 +5,13 @@ from datetime import datetime
 from types import TracebackType
 from typing import Any, DefaultDict, List, Optional, Type
 
+from ..context import get_context
 from ..views import Model, Trace
-from .persistence import PersistenceDriver
 
 logger = logging.getLogger("good_ai")
 
 
 class TracingContext:
-    persistence_driver: PersistenceDriver
     _contexts: DefaultDict[int, List["TracingContext"]] = defaultdict(lambda: [])
 
     def __init__(self) -> None:
@@ -62,7 +61,7 @@ class TracingContext:
 
         if type is None:
             assert self._trace is not None
-            self.persistence_driver.save_document(self._trace.dict())
+            get_context().persistence.save_document(self._trace.dict())
         else:
             logger.exception(f"Could not finish operation: {exception}")