Refactor by adding context

Signed-off-by: András Schmelczer <andras@schmelczer.dev>

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)
 

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

good_ai/src/good_ai/good_ai/config.py

@@ -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 ‼️")

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

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

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

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

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:

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)

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

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:

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

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)
 

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)
 

good_ai/src/good_ai/good_ai/set_default_config.py

@@ -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

good_ai/src/good_ai/good_ai/tracing/__init__.py

@@ -1,2 +1 @@
-from .persistence import MongoDbDriver, PersistenceDriver, TinyDbDriver
 from .tracing_context import TracingContext

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}")