From 4b72bddeb940acc337ff65c6aba9df583e0683c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= <andras@schmelczer.dev>
Date: Sun, 10 Apr 2022 12:59:05 +0200
Subject: [PATCH] Minor fixes
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                         |  1 +
 .vscode/tasks.json                            |  4 +--
 example/predict_domain.py                     |  6 ++--
 .../src/good_ai/good_ai/context/context.py    |  1 -
 .../good_ai/good_ai/context/get_context.py    | 10 +++----
 .../good_ai/good_ai/deploy/process_batch.py   |  2 +-
 .../src/good_ai/good_ai/helper/__init__.py    |  1 -
 .../src/good_ai/good_ai/helper/filter_args.py | 17 -----------
 .../good_ai/good_ai/metrics/log_argument.py   | 14 ++++++---
 .../src/good_ai/good_ai/metrics/log_metric.py | 30 +++++--------------
 .../good_ai/good_ai/persistence/__init__.py   |  1 +
 .../persistence/parallel_tinydb_driver.py     | 30 +++++++++++++++++++
 .../good_ai/persistence/persistence_driver.py |  2 ++
 .../good_ai/persistence/tinydb_driver.py      |  4 ++-
 .../good_ai/tracing/tracing_context.py        |  2 +-
 .../utilities/logger/custom_formatter.py      |  2 +-
 good_ai/src/good_ai/utilities/parallel_map.py |  6 ++++
 17 files changed, 73 insertions(+), 60 deletions(-)
 delete mode 100644 good_ai/src/good_ai/good_ai/helper/filter_args.py
 create mode 100644 good_ai/src/good_ai/good_ai/persistence/parallel_tinydb_driver.py

diff --git a/.vscode/settings.json b/.vscode/settings.json
index c9af87b..097fc0d 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -6,6 +6,7 @@
     "inplace",
     "levelno",
     "plotly",
+    "psutil",
     "pydantic",
     "pyplot",
     "sklearn",
diff --git a/.vscode/tasks.json b/.vscode/tasks.json
index b26b1a5..6f1581d 100644
--- a/.vscode/tasks.json
+++ b/.vscode/tasks.json
@@ -4,9 +4,9 @@
     {
       "label": "Format and lint Python",
       "type": "shell",
-      "command": "source .env/bin/activate && scripts/format-python.sh good_ai && scripts/check-python.sh good_ai",
+      "command": "source .env/bin/activate && scripts/format-python.sh good_ai && scripts/format-python.sh example",
       "windows": {
-        "command": ".env/bin/activate.bat; scripts/format-python.sh good_ai; scripts\\check-python.sh good_ai"
+        "command": ".env/bin/activate.bat; scripts/format-python.sh good_ai; scripts\\format-python.sh example"
       },
       "group": "test",
       "presentation": {
diff --git a/example/predict_domain.py b/example/predict_domain.py
index a91d23a..f940cb7 100644
--- a/example/predict_domain.py
+++ b/example/predict_domain.py
@@ -11,14 +11,14 @@ from good_ai import use_model, log_argument, log_metric
 from good_ai.utilities.clean import clean
 
 
-@log_metric('text_length', calculate=lambda text: len(text))
-@log_argument('text', expected_type=str, validator=lambda t: len(t) > 0)
 @use_model(model_key, version="latest")
+@log_argument('text', validator=lambda t: len(t) > 0)
 def predict_domain(
     text: str, model: Pipeline, cut_off_probability: float = 0.2
 ) -> List[DomainPrediction]:
     assert 0 <= cut_off_probability <= 1
-    
+    log_metric('text_length', len(text))
+
     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/context/context.py b/good_ai/src/good_ai/good_ai/context/context.py
index 078a7b6..efd55c6 100644
--- a/good_ai/src/good_ai/good_ai/context/context.py
+++ b/good_ai/src/good_ai/good_ai/context/context.py
@@ -9,7 +9,6 @@ class Context(BaseModel):
     metrics_path: str
     persistence: PersistenceDriver
     is_production: bool
-    is_threadsafe: bool
     logger: Logger
 
     class Config:
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
index 1a6d4f7..cb177f3 100644
--- a/good_ai/src/good_ai/good_ai/context/get_context.py
+++ b/good_ai/src/good_ai/good_ai/context/get_context.py
@@ -7,7 +7,7 @@ from typing import Optional, cast
 from good_ai.open_s3 import LargeFile
 from good_ai.utilities.logger import create_logger
 
-from ..persistence import PersistenceDriver, TinyDbDriver
+from ..persistence import PersistenceDriver, TinyDbDriver, ParallelTinyDbDriver
 from .context import Context
 
 _context: Optional[Context] = None
@@ -25,7 +25,7 @@ def set_default_config(
     log_level: int = INFO,
     s3_config: Path = Path("s3.ini"),
     seed: int = 42,
-    persistence_driver: PersistenceDriver = TinyDbDriver(Path("tracing_database.json")),
+    persistence_driver: PersistenceDriver = ParallelTinyDbDriver(Path("tracing_database.json")),
     is_production_mode_override: Optional[bool] = None,
 ) -> None:
     global _context
@@ -38,14 +38,12 @@ def set_default_config(
     _initialize_large_file(s3_config, logger=logger)
     _set_seed(seed)
 
-    is_threadsafe = not isinstance(persistence_driver, TinyDbDriver)
-    if not is_threadsafe:
-        logger.warn("The selected persistence driver (TinyDbDriver) is not threadsafe")
+    if not persistence_driver.is_threadsafe:
+        logger.warn(f"The selected persistence driver ({type(persistence_driver).__name__}) is not threadsafe")
     _context = Context(
         metrics_path="/metrics",
         persistence=persistence_driver,
         is_production=is_production,
-        is_threadsafe=is_threadsafe,
         logger=logger,
     )
 
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 07bc0bb..79076f6 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
@@ -18,7 +18,7 @@ def process_batch(
             output = t.log_output(result)
         return output
 
-    if not get_context().is_threadsafe:
+    if not get_context().persistence.is_threadsafe:
         concurrency = 1
         get_context().logger.warn("Concurrency is ignored")
 
diff --git a/good_ai/src/good_ai/good_ai/helper/__init__.py b/good_ai/src/good_ai/good_ai/helper/__init__.py
index c0a687a..e84cb5d 100644
--- a/good_ai/src/good_ai/good_ai/helper/__init__.py
+++ b/good_ai/src/good_ai/good_ai/helper/__init__.py
@@ -1,3 +1,2 @@
-from .filter_args import filter_args
 from .get_args import get_args
 from .snake_case_to_text import snake_case_to_text
diff --git a/good_ai/src/good_ai/good_ai/helper/filter_args.py b/good_ai/src/good_ai/good_ai/helper/filter_args.py
deleted file mode 100644
index a089cea..0000000
--- a/good_ai/src/good_ai/good_ai/helper/filter_args.py
+++ /dev/null
@@ -1,17 +0,0 @@
-import inspect
-from typing import Any, Callable, Dict
-
-
-def filter_args(
-    dict_to_filter: Dict[str, Any], func: Callable[..., Any]
-) -> Dict[str, Any]:
-    signature = inspect.signature(func)
-    filter_keys = [
-        param.name
-        for param in signature.parameters.values()
-        if param.kind == param.POSITIONAL_OR_KEYWORD
-    ]
-    filtered_dict = {
-        filter_key: dict_to_filter[filter_key] for filter_key in filter_keys
-    }
-    return filtered_dict
diff --git a/good_ai/src/good_ai/good_ai/metrics/log_argument.py b/good_ai/src/good_ai/good_ai/metrics/log_argument.py
index 713c308..cc6c427 100644
--- a/good_ai/src/good_ai/good_ai/metrics/log_argument.py
+++ b/good_ai/src/good_ai/good_ai/metrics/log_argument.py
@@ -1,5 +1,5 @@
 from functools import wraps
-from typing import Any, Callable, Dict, List, Optional, Type
+from typing import Any, Callable, Dict, List
 
 from ..exceptions import ArgumentValidationError
 from ..helper import get_args
@@ -9,20 +9,26 @@ from ..tracing import TracingContext
 def log_argument(
     argument_name: str,
     *,
-    expected_type: Optional[Type] = None,
     validator: Callable[[Any], bool] = lambda _: True,
 ) -> Callable[..., Any]:
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-        actual_name = f"{func.__name__}:{argument_name}"
+        actual_name = f"arg:{func.__name__}:{argument_name}"
 
         @wraps(func)
         def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
             arguments = get_args(func, args, kwargs)
             argument = arguments[argument_name]
 
+            expected_type = func.__annotations__.get(argument_name)
+
             if (
                 expected_type is not None and not isinstance(argument, expected_type)
-            ) or not validator(argument):
+            ):
+                raise ArgumentValidationError(
+                    f"Argument {argument_name} in {func.__name__} has the wrong type, expected: {expected_type.__name__}, got: {type(argument).__name__}"
+                )
+
+            if not validator(argument):
                 raise ArgumentValidationError(
                     f"Argument {argument_name} in {func.__name__} did not pass validation"
                 )
diff --git a/good_ai/src/good_ai/good_ai/metrics/log_metric.py b/good_ai/src/good_ai/good_ai/metrics/log_metric.py
index f308b1b..4a8263d 100644
--- a/good_ai/src/good_ai/good_ai/metrics/log_metric.py
+++ b/good_ai/src/good_ai/good_ai/metrics/log_metric.py
@@ -1,26 +1,12 @@
-from functools import wraps
-from typing import Any, Callable, Dict, List
+import inspect
+from typing import Any
 
-from ..helper import filter_args, get_args
 from ..tracing import TracingContext
 
 
-def log_metric(
-    argument_name: str, *, calculate: Callable[[Any], bool] = lambda _: True
-) -> Callable[..., Any]:
-    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-        actual_name = f"{func.__name__}:{argument_name}"
-
-        @wraps(func)
-        def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
-            arguments = get_args(func, args, kwargs)
-            metric = calculate(**filter_args(arguments, calculate))
-
-            context = TracingContext.get_current_context()
-            if context:
-                context.log_value(name=actual_name, value=metric)
-            return func(*args, **kwargs)
-
-        return wrapper
-
-    return decorator
+def log_metric(argument_name: str, value: Any) -> None:
+    context = TracingContext.get_current_context()
+    caller = inspect.stack()[1].function
+    actual_name = f"metric:{caller}:{argument_name}"
+    if context:
+        context.log_value(name=actual_name, value=value)
diff --git a/good_ai/src/good_ai/good_ai/persistence/__init__.py b/good_ai/src/good_ai/good_ai/persistence/__init__.py
index 7bddbe7..d4aeaaf 100644
--- a/good_ai/src/good_ai/good_ai/persistence/__init__.py
+++ b/good_ai/src/good_ai/good_ai/persistence/__init__.py
@@ -1,3 +1,4 @@
 from .mongodb_driver import MongoDbDriver
 from .persistence_driver import PersistenceDriver
 from .tinydb_driver import TinyDbDriver
+from .parallel_tinydb_driver import ParallelTinyDbDriver
\ No newline at end of file
diff --git a/good_ai/src/good_ai/good_ai/persistence/parallel_tinydb_driver.py b/good_ai/src/good_ai/good_ai/persistence/parallel_tinydb_driver.py
new file mode 100644
index 0000000..4802a12
--- /dev/null
+++ b/good_ai/src/good_ai/good_ai/persistence/parallel_tinydb_driver.py
@@ -0,0 +1,30 @@
+from pathlib import Path
+from typing import Any, Callable
+from black import List
+from tinydb import TinyDB
+from multiprocessing import Process, Lock
+
+from ..views import Trace
+from .persistence_driver import PersistenceDriver
+
+
+lock = Lock()
+
+
+class ParallelTinyDbDriver(PersistenceDriver):
+    is_threadsafe = True
+    
+    def __init__(self, path_to_db: Path) -> None:
+        super().__init__()
+        self._path_to_db = path_to_db
+
+    def save_document(self, trace: Trace) -> str:
+        return self._safe_execute(lambda db: db.insert(trace.dict()))
+
+    def get_documents(self) -> List[Trace]:
+        return self._safe_execute(lambda db: [Trace.parse_obj(t) for t in db.all()])
+
+    def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any:
+        with lock:
+            with TinyDB(self._path_to_db) as db:
+                return func(db)
diff --git a/good_ai/src/good_ai/good_ai/persistence/persistence_driver.py b/good_ai/src/good_ai/good_ai/persistence/persistence_driver.py
index 5ea4400..f74d723 100644
--- a/good_ai/src/good_ai/good_ai/persistence/persistence_driver.py
+++ b/good_ai/src/good_ai/good_ai/persistence/persistence_driver.py
@@ -6,6 +6,8 @@ from good_ai.good_ai.views.trace import Trace
 
 
 class PersistenceDriver(ABC):
+    is_threadsafe: bool
+
     @abstractmethod
     def save_document(self, document: Trace) -> str:
         pass
diff --git a/good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py b/good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py
index 2b7fc81..b7f21ec 100644
--- a/good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py
+++ b/good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py
@@ -10,12 +10,14 @@ from .persistence_driver import PersistenceDriver
 
 
 class TinyDbDriver(PersistenceDriver):
+    is_threadsafe = False
+    
     def __init__(self, path_to_db: Path) -> None:
         super().__init__()
         self._db = TinyDB(path_to_db)
 
     def save_document(self, trace: Trace) -> str:
-        return self._db.insert(Document(trace.dict(), doc_id=uuid4().int))
+        return self._db.insert(trace.dict())
 
     def get_documents(self) -> List[Trace]:
         return [Trace.parse_obj(t) for t in self._db.all()]
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 ec990f2..c9fcbbf 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
@@ -63,4 +63,4 @@ class TracingContext:
         else:
             get_context().logger.exception(f"Could not finish operation: {exception}")
 
-        return True
+        return False
diff --git a/good_ai/src/good_ai/utilities/logger/custom_formatter.py b/good_ai/src/good_ai/utilities/logger/custom_formatter.py
index a1bcfd0..0760ce3 100644
--- a/good_ai/src/good_ai/utilities/logger/custom_formatter.py
+++ b/good_ai/src/good_ai/utilities/logger/custom_formatter.py
@@ -15,7 +15,7 @@ class CustomFormatter(logging.Formatter):
             logging.CRITICAL: BOLD_RED + self.fmt + RESET,
         }
 
-    def format(self, record: logging.LogRecord) -> None:
+    def format(self, record: logging.LogRecord) -> str:
         log_fmt = self.FORMATS.get(record.levelno)
         formatter = logging.Formatter(log_fmt)
         return formatter.format(record)
diff --git a/good_ai/src/good_ai/utilities/parallel_map.py b/good_ai/src/good_ai/utilities/parallel_map.py
index 0132850..3999303 100644
--- a/good_ai/src/good_ai/utilities/parallel_map.py
+++ b/good_ai/src/good_ai/utilities/parallel_map.py
@@ -4,6 +4,9 @@ from typing import Any, Callable, Iterable, List, Optional
 import multiprocess as mp
 import psutil
 from tqdm.auto import tqdm
+from .logger import create_logger
+
+logger = create_logger('parallel_map')
 
 
 def parallel_map(
@@ -23,7 +26,10 @@ def parallel_map(
     if not chunk_size:
         chunk_size = max(1, ceil(len(values) / concurrency / 10))
 
+    logger.info(f"Starting parallel map, concurrency: {concurrency}, chunk size: {chunk_size}")
+
     if concurrency == 1 or len(values) <= chunk_size:
+        logger.warn(f"Running in series, there is no reason for parallelism")
         iterable = values if disable_progress else tqdm(values)
         return [function(v) for v in iterable]