Add log_argument

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

.vscode/settings.json

@@ -2,6 +2,8 @@
   "cSpell.words": [
     "boto",
     "botocore",
+    "iloc",
+    "inplace",
     "plotly",
     "pydantic",
     "pyplot",

example/main_batch.py

@@ -7,13 +7,11 @@ from predict_domain import predict_domain
 from good_ai import process_batch, serve
 
 if __name__ == "__main__":
-    serve(predict_domain)
+    with open(".cache/data-1/s2-corpus-323.json") as f:
-    
-    with open(".cache/ss-data-0/s2-corpus-1583.json") as f:
         raw = json.load(f)
 
     shuffle(raw)
-    data = {f'{r["title"]} {r["abstract"]}': r["domain"] for r in raw[:5]}
+    data = {f'{r["title"]} {r["abstract"]}': r["domain"] for r in raw[:10]}
 
     results = process_batch(predict_domain, data.keys())
 

example/main_service.py

@@ -0,0 +1,10 @@
+import json
+from random import shuffle
+
+from devtools import debug
+from predict_domain import predict_domain
+
+from good_ai import process_batch, serve
+
+if __name__ == "__main__":
+    serve(predict_domain)

example/predict_domain.py

@@ -1,15 +1,18 @@
 import re
 from typing import Dict, Iterable, List
 
+
 from config import model_key
 from models import DomainPrediction
 from preprocess import preprocess
 from sklearn.pipeline import Pipeline
 
-from good_ai import use_model
+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")
 def predict_domain(
     text: str, model: Pipeline, cut_off_probability: float = 0.2

good_ai/src/good_ai/good_ai/__init__.py

@@ -1,3 +1,4 @@
 from .context import set_default_config
 from .deploy import process_batch, process_single, serve
+from .metrics import log_argument, log_metric
 from .models import save_model, use_model

good_ai/src/good_ai/good_ai/context/context.py

@@ -7,6 +7,7 @@ class Context(BaseModel):
     metrics_path: str
     persistence: PersistenceDriver
     is_production: bool
+    is_threadsafe: bool
 
     class Config:
         arbitrary_types_allowed = True

good_ai/src/good_ai/good_ai/context/get_context.py

@@ -11,6 +11,7 @@ from .context import Context
 
 logger = logging.getLogger("good_ai")
 
+
 _context: Optional[Context] = None
 PRODUCTION_KEY = "production"
 
@@ -36,10 +37,14 @@ def set_default_config(
     _initialize_large_file(s3_config)
     _set_seed(seed)
 
+    is_threadsafe = not isinstance(persistence_driver, TinyDbDriver)
+    if not is_threadsafe:
+        logger.warn("The selected persistence driver (TinyDbDriver) is not threadsafe")
     _context = Context(
         metrics_path="/metrics",
         persistence=persistence_driver,
         is_production=is_production,
+        is_threadsafe=is_threadsafe,
     )
 
     logger.info("Defaults: configured ✅")

good_ai/src/good_ai/good_ai/deploy/process_batch.py

@@ -1,10 +1,14 @@
+import logging
 from typing import Any, Callable, Iterable, Optional, Sequence
 
 from good_ai.utilities.parallel_map import parallel_map
 
+from ..context import get_context
 from ..tracing import TracingContext
 from ..views import Trace
 
+logger = logging.getLogger("good_ai")
+
 
 def process_batch(
     function: Callable[..., Any],
@@ -13,9 +17,12 @@ def process_batch(
 ) -> Sequence[Trace]:
     def inner(input: Any) -> Trace:
         with TracingContext() as t:
-            t.log_input(input)
             result = function(input)
             output = t.log_output(result)
         return output
 
+    if not get_context().is_threadsafe:
+        concurrency = 1
+        logger.warn("Concurrency is ignored")
+
     return parallel_map(inner, batch, concurrency=concurrency)

good_ai/src/good_ai/good_ai/deploy/process_single.py

@@ -6,7 +6,6 @@ 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

@@ -32,7 +32,6 @@ def serve(
     @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

good_ai/src/good_ai/good_ai/exceptions/__init__.py

@@ -0,0 +1 @@
+from .argument_validation_error import ArgumentValidationError

good_ai/src/good_ai/good_ai/exceptions/argument_validation_error.py

@@ -0,0 +1,2 @@
+class ArgumentValidationError(Exception):
+    pass

good_ai/src/good_ai/good_ai/helper/__init__.py

@@ -1 +1,3 @@
+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

@@ -0,0 +1,17 @@
+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/helper/get_args.py

@@ -0,0 +1,14 @@
+import inspect
+from typing import Any, Callable, Dict, List
+
+
+def get_args(
+    func: Callable[..., Any], args: List[Any], kwargs: Dict[str, 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
+    ]
+    return {**dict(zip(filter_keys, args)), **kwargs}

good_ai/src/good_ai/good_ai/metrics/__init__.py

@@ -1 +1,3 @@
 from .create_dash_app import create_dash_app
+from .log_argument import log_argument
+from .log_metric import log_metric

good_ai/src/good_ai/good_ai/metrics/create_dash_app.py

@@ -1,49 +1,152 @@
 import pandas as pd
-import plotly.express as px
+from dash import Dash, dash_table, html
-from dash import Dash, dcc, html
+from dash.dependencies import Input, Output
 from flask import Flask
 
 from good_ai.good_ai.context.get_context import get_context
 
-from ..helper import snake_case_to_text
+from .get_description import get_description
 
 
 def create_dash_app(function_name: str) -> Flask:
     app = Dash(function_name, requests_pathname_prefix=get_context().metrics_path + "/")
 
-    markdown_text = f"""
+    documents = get_context().persistence.get_documents()
-    # {snake_case_to_text(function_name)} - metrics
-    > A human-friendly framework for robust end-to-end AI deployments
 
-    ## Using the API
+    df = pd.DataFrame(
-
+        [
-    You can find the available endpoints at [/docs](/docs).
+            {
-
+                "id": d.evaluation_id,
-    ## Metrics
+                "created": d.created,
-
+                "execution_time_ms": d.execution_time_ms,
-    The recent traces and aggregated metrics are presented below.
+                "models": ", ".join(f"{m.key}:{m.version}" for m in d.models),
-    """
+                "evaluation": d.evaluation,
-
+                **d.logged_values,
-    df = pd.read_csv(
+            }
-        "https://gist.githubusercontent.com/chriddyp/5d1ea79569ed194d432e56108a04d188/raw/a9f9e8076b837d541398e999dcbac2b2826a81f8/gdp-life-exp-2007.csv"
+            for d in documents
-    )
+        ]
-
-    fig = px.scatter(
-        df,
-        x="gdp per capita",
-        y="life expectancy",
-        size="population",
-        color="continent",
-        hover_name="country",
-        log_x=True,
-        size_max=60,
     )
+    print(df)
 
     app.layout = html.Div(
-        [
+        children=[
-            dcc.Markdown(children=markdown_text),
+            get_description(function_name),
-            dcc.Graph(id="life-exp-vs-gdp", figure=fig),
+            html.Div(
+                dash_table.DataTable(
+                    id="table-paging-with-graph",
+                    columns=[{"name": i, "id": i} for i in df.columns],
+                    page_current=0,
+                    page_size=20,
+                    page_action="custom",
+                    filter_action="custom",
+                    filter_query="",
+                    sort_action="custom",
+                    sort_mode="multi",
+                    sort_by=[],
+                ),
+                style={"height": 750, "overflowY": "scroll"},
+                className="six columns",
+            ),
+            html.Div(id="table-paging-with-graph-container", className="five columns"),
         ]
     )
 
+    operators = [
+        ["ge ", ">="],
+        ["le ", "<="],
+        ["lt ", "<"],
+        ["gt ", ">"],
+        ["ne ", "!="],
+        ["eq ", "="],
+        ["contains "],
+        ["datestartswith "],
+    ]
+
+    def split_filter_part(filter_part):
+        for operator_type in operators:
+            for operator in operator_type:
+                if operator in filter_part:
+                    name_part, value_part = filter_part.split(operator, 1)
+                    name = name_part[name_part.find("{") + 1 : name_part.rfind("}")]
+
+                    value_part = value_part.strip()
+                    v0 = value_part[0]
+                    if v0 == value_part[-1] and v0 in ("'", '"', "`"):
+                        value = value_part[1:-1].replace("\\" + v0, v0)
+                    else:
+                        try:
+                            value = float(value_part)
+                        except ValueError:
+                            value = value_part
+
+                    # word operators need spaces after them in the filter string,
+                    # but we don't want these later
+                    return name, operator_type[0].strip(), value
+
+        return [None] * 3
+
+    @app.callback(
+        Output("table-paging-with-graph", "data"),
+        Input("table-paging-with-graph", "page_current"),
+        Input("table-paging-with-graph", "page_size"),
+        Input("table-paging-with-graph", "sort_by"),
+        Input("table-paging-with-graph", "filter_query"),
+    )
+    def update_table(page_current, page_size, sort_by, filter):
+        filtering_expressions = filter.split(" && ")
+        dff = df
+        for filter_part in filtering_expressions:
+            col_name, operator, filter_value = split_filter_part(filter_part)
+
+            if operator in ("eq", "ne", "lt", "le", "gt", "ge"):
+                # these operators match pandas series operator method names
+                dff = dff.loc[getattr(dff[col_name], operator)(filter_value)]
+            elif operator == "contains":
+                dff = dff.loc[dff[col_name].str.contains(filter_value)]
+            elif operator == "datestartswith":
+                # this is a simplification of the front-end filtering logic,
+                # only works with complete fields in standard format
+                dff = dff.loc[dff[col_name].str.startswith(filter_value)]
+
+        if len(sort_by):
+            dff = dff.sort_values(
+                [col["column_id"] for col in sort_by],
+                ascending=[col["direction"] == "asc" for col in sort_by],
+                inplace=False,
+            )
+
+        return dff.iloc[
+            page_current * page_size : (page_current + 1) * page_size
+        ].to_dict("records")
+
+    # @app.callback(
+    #     Output('table-paging-with-graph-container', "children"),
+    #     Input('table-paging-with-graph', "data"))
+    # def update_graph(rows):
+    #     dff = pd.DataFrame(rows)
+    #     return html.Div(
+    #         [
+    #             dcc.Graph(
+    #                 id=column,
+    #                 figure={
+    #                     "data": [
+    #                         {
+    #                             "x": dff["country"],
+    #                             "y": dff[column] if column in dff else [],
+    #                             "type": "bar",
+    #                             "marker": {"color": "#0074D9"},
+    #                         }
+    #                     ],
+    #                     "layout": {
+    #                         "xaxis": {"automargin": True},
+    #                         "yaxis": {"automargin": True},
+    #                         "height": 250,
+    #                         "margin": {"t": 10, "l": 10, "r": 10},
+    #                     },
+    #                 },
+    #             )
+    #             for column in ["pop", "lifeExp", "gdpPercap"]
+    #         ]
+    #     )
+
     return app.server

good_ai/src/good_ai/good_ai/metrics/get_description.py

@@ -0,0 +1,20 @@
+from dash.dcc import Markdown
+
+from ..helper import snake_case_to_text
+
+
+def get_description(function_name: str) -> Markdown:
+    markdown_text = f"""
+    # {snake_case_to_text(function_name)} - metrics
+    > A human-friendly framework for robust end-to-end AI deployments
+
+    ## Using the API
+
+    You can find the available endpoints at [/docs](/docs).
+
+    ## Metrics
+
+    The recent traces and aggregated metrics are presented below.
+    """
+
+    return Markdown(markdown_text)

good_ai/src/good_ai/good_ai/metrics/log_argument.py

@@ -0,0 +1,37 @@
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, Type
+
+from ..exceptions import ArgumentValidationError
+from ..helper import get_args
+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}"
+
+        @wraps(func)
+        def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+            arguments = get_args(func, args, kwargs)
+            argument = arguments[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__} did not pass validation"
+                )
+
+            context = TracingContext.get_current_context()
+            if context:
+                context.log_value(name=actual_name, value=argument)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

good_ai/src/good_ai/good_ai/metrics/log_metric.py

@@ -0,0 +1,26 @@
+from functools import wraps
+from typing import Any, Callable, Dict, List
+
+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

good_ai/src/good_ai/good_ai/models/use_model.py

@@ -1,5 +1,5 @@
 from functools import wraps
-from typing import Any, Callable, Literal, Union
+from typing import Any, Callable, Dict, List, Literal, Union
 
 from ..tracing import TracingContext
 from ..views import Model
@@ -23,7 +23,7 @@ def use_model(
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         @wraps(func)
-        def wrapper(*args: Any, **kwargs: Any) -> Any:
+        def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
             context = TracingContext.get_current_context()
             if context:
                 context.log_model(Model(key=key, version=actual_version))

good_ai/src/good_ai/good_ai/persistence/persistence_driver.py

@@ -1,8 +1,15 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict
+
+from black import List
+
+from good_ai.good_ai.views.trace import Trace
 
 
 class PersistenceDriver(ABC):
     @abstractmethod
-    def save_document(self, document: Dict[str, Any]) -> str:
+    def save_document(self, document: Trace) -> str:
+        pass
+
+    @abstractmethod
+    def get_documents(self) -> List[Trace]:
         pass

good_ai/src/good_ai/good_ai/persistence/tinydb_driver.py

@@ -1,8 +1,11 @@
 from pathlib import Path
-from typing import Any, Dict
+from uuid import uuid4
 
+from black import List
 from tinydb import TinyDB
+from tinydb.table import Document
 
+from ..views import Trace
 from .persistence_driver import PersistenceDriver
 
 
@@ -11,5 +14,8 @@ class TinyDbDriver(PersistenceDriver):
         super().__init__()
         self._db = TinyDB(path_to_db)
 
-    def save_document(self, document: Dict[str, Any]) -> str:
+    def save_document(self, trace: Trace) -> str:
-        return self._db.insert(document)
+        return self._db.insert(Document(trace.dict(), doc_id=uuid4().int))
+
+    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

@@ -3,7 +3,7 @@ import threading
 from collections import defaultdict
 from datetime import datetime
 from types import TracebackType
-from typing import Any, DefaultDict, List, Optional, Type
+from typing import Any, DefaultDict, Dict, List, Optional, Type
 
 from ..context import get_context
 from ..views import Model, Trace
@@ -16,25 +16,26 @@ class TracingContext:
 
     def __init__(self) -> None:
         self._models: List[Model] = []
-        self._input: Any = None
+        self._values: Dict[str, Any] = {}
         self._output: Any = None
         self._trace: Optional[Trace] = None
         self._start_time = datetime.utcnow()
 
-    def log_input(self, input: Any) -> None:
+    def log_value(self, name: str, value: Any) -> None:
-        self._input = input
+        self._values[name] = value
 
     def log_model(self, model: Model) -> None:
         self._models.append(model)
 
-    def log_output(self, output: Any) -> Trace:
+    def log_output(self, output: Any, evaluation_id: Optional[str] = None) -> Trace:
         self._output = output
 
         delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
         self._trace = Trace(
+            evaluation_id=evaluation_id,
             created=self._start_time.isoformat(),
             execution_time_ms=delta_time,
-            input=self._input,
+            logged_values=self._values,
             models=self._models,
             output=self._output,
         )
@@ -61,7 +62,7 @@ class TracingContext:
 
         if type is None:
             assert self._trace is not None
-            get_context().persistence.save_document(self._trace.dict())
+            get_context().persistence.save_document(self._trace)
         else:
             logger.exception(f"Could not finish operation: {exception}")
 

good_ai/src/good_ai/good_ai/views/trace.py

@@ -1,16 +1,27 @@
-from typing import Any, List
+from typing import Any, Dict, List, Optional
 from uuid import uuid4
 
-from pydantic import BaseModel
+from pydantic import BaseModel, validator
 
 from .model import Model
 
 
 class Trace(BaseModel):
-    id = str(uuid4())
+    evaluation_id: Optional[str]
     created: str
     execution_time_ms: float
-    input: Any
+    logged_values: Dict[str, Any]
     models: List[Model]
     output: Any
     evaluation: Any = None
+
+    @validator("evaluation_id", always=True)
+    def validate_single_set(
+        cls, v: Optional[str], values: Dict[str, Any]
+    ) -> Optional[str]:
+        if not v:
+            return str(uuid4())
+        return v
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple(self.__dict__.values()))

good_ai/src/good_ai/utilities/parallel_map.py

@@ -23,7 +23,7 @@ def parallel_map(
     if not chunk_size:
         chunk_size = max(1, ceil(len(values) / concurrency / 10))
 
-    if concurrency == 1:
+    if concurrency == 1 or len(values) <= chunk_size:
         iterable = values if disable_progress else tqdm(values)
         return [function(v) for v in iterable]