From 85a06096bfa34640c184077c2f132963fcf38a5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= Date: Sat, 9 Apr 2022 22:02:00 +0200 Subject: [PATCH] Add log_argument MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: András Schmelczer --- .vscode/settings.json | 2 + example/dash-test.py | 0 example/{main.py => main_batch.py} | 6 +- example/main_service.py | 10 ++ example/predict_domain.py | 5 +- good_ai/src/good_ai/good_ai/__init__.py | 1 + .../src/good_ai/good_ai/context/context.py | 1 + .../good_ai/good_ai/context/get_context.py | 5 + .../good_ai/good_ai/deploy/process_batch.py | 9 +- .../good_ai/good_ai/deploy/process_single.py | 1 - good_ai/src/good_ai/good_ai/deploy/serve.py | 1 - .../good_ai/good_ai/exceptions/__init__.py | 1 + .../exceptions/argument_validation_error.py | 2 + .../src/good_ai/good_ai/helper/__init__.py | 2 + .../src/good_ai/good_ai/helper/filter_args.py | 17 ++ .../src/good_ai/good_ai/helper/get_args.py | 14 ++ .../src/good_ai/good_ai/metrics/__init__.py | 2 + .../good_ai/metrics/create_dash_app.py | 165 ++++++++++++++---- .../good_ai/metrics/get_description.py | 20 +++ .../good_ai/good_ai/metrics/log_argument.py | 37 ++++ .../src/good_ai/good_ai/metrics/log_metric.py | 26 +++ .../src/good_ai/good_ai/models/use_model.py | 4 +- .../good_ai/persistence/persistence_driver.py | 11 +- .../good_ai/persistence/tinydb_driver.py | 12 +- .../good_ai/tracing/tracing_context.py | 15 +- good_ai/src/good_ai/good_ai/views/trace.py | 19 +- good_ai/src/good_ai/utilities/parallel_map.py | 2 +- 27 files changed, 332 insertions(+), 58 deletions(-) delete mode 100644 example/dash-test.py rename example/{main.py => main_batch.py} (82%) create mode 100644 example/main_service.py create mode 100644 good_ai/src/good_ai/good_ai/exceptions/__init__.py create mode 100644 good_ai/src/good_ai/good_ai/exceptions/argument_validation_error.py create mode 100644 good_ai/src/good_ai/good_ai/helper/filter_args.py create mode 100644 good_ai/src/good_ai/good_ai/helper/get_args.py create mode 100644 good_ai/src/good_ai/good_ai/metrics/get_description.py create mode 100644 good_ai/src/good_ai/good_ai/metrics/log_argument.py create mode 100644 good_ai/src/good_ai/good_ai/metrics/log_metric.py diff --git a/.vscode/settings.json b/.vscode/settings.json index 4e5b3c7..61ce14b 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -2,6 +2,8 @@ "cSpell.words": [ "boto", "botocore", + "iloc", + "inplace", "plotly", "pydantic", "pyplot", diff --git a/example/dash-test.py b/example/dash-test.py deleted file mode 100644 index e69de29..0000000 diff --git a/example/main.py b/example/main_batch.py similarity index 82% rename from example/main.py rename to example/main_batch.py index 2d1def9..f9148f8 100644 --- a/example/main.py +++ b/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/ss-data-0/s2-corpus-1583.json") as f: + with open(".cache/data-1/s2-corpus-323.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()) diff --git a/example/main_service.py b/example/main_service.py new file mode 100644 index 0000000..f770fed --- /dev/null +++ b/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) diff --git a/example/predict_domain.py b/example/predict_domain.py index e6d6938..a91d23a 100644 --- a/example/predict_domain.py +++ b/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 diff --git a/good_ai/src/good_ai/good_ai/__init__.py b/good_ai/src/good_ai/good_ai/__init__.py index bc171fa..c2f67f4 100644 --- a/good_ai/src/good_ai/good_ai/__init__.py +++ b/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 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 2358759..53e040c 100644 --- a/good_ai/src/good_ai/good_ai/context/context.py +++ b/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 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 b44a54b..f8ae061 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 @@ -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 ✅") 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 9dc0f3a..92885a4 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 @@ -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) 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 index 19d8346..6d5855c 100644 --- a/good_ai/src/good_ai/good_ai/deploy/process_single.py +++ b/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 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 3ffde5c..c23f7e4 100644 --- a/good_ai/src/good_ai/good_ai/deploy/serve.py +++ b/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 diff --git a/good_ai/src/good_ai/good_ai/exceptions/__init__.py b/good_ai/src/good_ai/good_ai/exceptions/__init__.py new file mode 100644 index 0000000..ed0a91a --- /dev/null +++ b/good_ai/src/good_ai/good_ai/exceptions/__init__.py @@ -0,0 +1 @@ +from .argument_validation_error import ArgumentValidationError diff --git a/good_ai/src/good_ai/good_ai/exceptions/argument_validation_error.py b/good_ai/src/good_ai/good_ai/exceptions/argument_validation_error.py new file mode 100644 index 0000000..122057b --- /dev/null +++ b/good_ai/src/good_ai/good_ai/exceptions/argument_validation_error.py @@ -0,0 +1,2 @@ +class ArgumentValidationError(Exception): + pass 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 4fdd18c..c0a687a 100644 --- a/good_ai/src/good_ai/good_ai/helper/__init__.py +++ b/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 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 new file mode 100644 index 0000000..a089cea --- /dev/null +++ b/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 diff --git a/good_ai/src/good_ai/good_ai/helper/get_args.py b/good_ai/src/good_ai/good_ai/helper/get_args.py new file mode 100644 index 0000000..f31ffe1 --- /dev/null +++ b/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} diff --git a/good_ai/src/good_ai/good_ai/metrics/__init__.py b/good_ai/src/good_ai/good_ai/metrics/__init__.py index dfa9e7d..d1e5adf 100644 --- a/good_ai/src/good_ai/good_ai/metrics/__init__.py +++ b/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 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 9c116a3..d3ddf24 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 @@ -1,49 +1,152 @@ import pandas as pd -import plotly.express as px -from dash import Dash, dcc, html +from dash import Dash, dash_table, 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""" - # {snake_case_to_text(function_name)} - metrics - > A human-friendly framework for robust end-to-end AI deployments + documents = get_context().persistence.get_documents() - ## Using the API - - You can find the available endpoints at [/docs](/docs). - - ## Metrics - - The recent traces and aggregated metrics are presented below. - """ - - df = pd.read_csv( - "https://gist.githubusercontent.com/chriddyp/5d1ea79569ed194d432e56108a04d188/raw/a9f9e8076b837d541398e999dcbac2b2826a81f8/gdp-life-exp-2007.csv" - ) - - fig = px.scatter( - df, - x="gdp per capita", - y="life expectancy", - size="population", - color="continent", - hover_name="country", - log_x=True, - size_max=60, + df = pd.DataFrame( + [ + { + "id": d.evaluation_id, + "created": d.created, + "execution_time_ms": d.execution_time_ms, + "models": ", ".join(f"{m.key}:{m.version}" for m in d.models), + "evaluation": d.evaluation, + **d.logged_values, + } + for d in documents + ] ) + print(df) app.layout = html.Div( - [ - dcc.Markdown(children=markdown_text), - dcc.Graph(id="life-exp-vs-gdp", figure=fig), + children=[ + get_description(function_name), + 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 diff --git a/good_ai/src/good_ai/good_ai/metrics/get_description.py b/good_ai/src/good_ai/good_ai/metrics/get_description.py new file mode 100644 index 0000000..8976963 --- /dev/null +++ b/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) 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 new file mode 100644 index 0000000..713c308 --- /dev/null +++ b/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 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 new file mode 100644 index 0000000..f308b1b --- /dev/null +++ b/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 diff --git a/good_ai/src/good_ai/good_ai/models/use_model.py b/good_ai/src/good_ai/good_ai/models/use_model.py index 189ba65..eef7ef5 100644 --- a/good_ai/src/good_ai/good_ai/models/use_model.py +++ b/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)) 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 dd37537..5ea4400 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 @@ -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 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 f5208d0..2b7fc81 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 @@ -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: - return self._db.insert(document) + def save_document(self, trace: Trace) -> str: + 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()] 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 550c768..bf8840e 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 @@ -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: - self._input = input + def log_value(self, name: str, value: Any) -> None: + 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}") diff --git a/good_ai/src/good_ai/good_ai/views/trace.py b/good_ai/src/good_ai/good_ai/views/trace.py index 2d0c988..054f5ae 100644 --- a/good_ai/src/good_ai/good_ai/views/trace.py +++ b/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())) diff --git a/good_ai/src/good_ai/utilities/parallel_map.py b/good_ai/src/good_ai/utilities/parallel_map.py index 12a0c30..0132850 100644 --- a/good_ai/src/good_ai/utilities/parallel_map.py +++ b/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]