Minor fixes

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

.vscode/settings.json

@@ -6,6 +6,7 @@
     "inplace",
     "levelno",
     "plotly",
+    "psutil",
     "pydantic",
     "pyplot",
     "sklearn",

.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": {

example/predict_domain.py

@@ -11,13 +11,13 @@ 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)

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:

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,
     )
 

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

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

good_ai/src/good_ai/good_ai/helper/filter_args.py

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

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

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))
-
+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=metric)
-            return func(*args, **kwargs)
-
-        return wrapper
-
-    return decorator
+        context.log_value(name=actual_name, value=value)

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

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)

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

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()]

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

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)

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]