Integrate MongoDB

This commit is contained in:
Andras Schmelczer 2022-06-25 13:01:44 +02:00
parent 750ba7b0d4
commit cf0ac4b161
30 changed files with 1024 additions and 769 deletions

View file

@ -13,6 +13,7 @@
"iloc",
"initialisation",
"initialised",
"initialising",
"inplace",
"ipynb",
"joblib",

View file

@ -11,15 +11,6 @@
"> The blue boxes show the steps implemented in this notebook."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"MAX_CHUNK_COUNT = 4"
]
},
{
"cell_type": "markdown",
"metadata": {},
@ -31,6 +22,15 @@
"In this case, we download the semantic scholar dataset from a public S3 bucket."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"MAX_CHUNK_COUNT = 1"
]
},
{
"cell_type": "code",
"execution_count": 2,
@ -39,7 +39,7 @@
{
"data": {
"text/plain": [
"'Processing 4 out of the 6002 available chunks'"
"'Processing 1 out of the 6002 available chunks'"
]
},
"execution_count": 2,
@ -49,6 +49,7 @@
],
"source": [
"import urllib.request\n",
"from random import shuffle\n",
"\n",
"manifest = (\n",
" urllib.request.urlopen(\n",
@ -58,7 +59,9 @@
" .decode()\n",
") # a list of available chunks separated by '\\n' characters\n",
"\n",
"chunks = manifest.split()[:MAX_CHUNK_COUNT]\n",
"lines = manifest.split()\n",
"shuffle(lines)\n",
"chunks = lines[:MAX_CHUNK_COUNT]\n",
"\n",
"f\"Processing {len(chunks)} out of the {len(manifest.split())} available chunks\""
]
@ -83,23 +86,11 @@
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[38;5;226m2022-06-19 14:59:12,562 | WARNING | Limiting concurrency to 4 because there are only 4 chunks\u001b[0m\n",
"\u001b[38;5;39m2022-06-19 14:59:12,563 | INFO | Starting parallel map (concurrency: 4, chunk size: 1)\u001b[0m\n"
"\u001b[38;5;226m2022-06-25 11:20:01,955 | WARNING | Limiting concurrency to 1 because there are only 1 chunks\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:20:01,956 | INFO | Starting parallel map (concurrency: 1, chunk size: 1)\u001b[0m\n",
"\u001b[38;5;226m2022-06-25 11:20:01,956 | WARNING | Running in series, there is no reason for parallelism\u001b[0m\n",
"100%|██████████| 1/1 [04:02<00:00, 242.61s/it]\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ff8fc113515944cfa75127f4aba3d491",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/4 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
@ -121,15 +112,13 @@
"\n",
" # Transform\n",
" return [\n",
" # Create pairs of `(text, [...domains])`\n",
" # The text is cleaned to remove PDF extraction, web scraping, and other common artifacts\n",
" (\n",
" clean(\n",
" f'{c[\"title\"]} {c[\"paperAbstract\"]} {c[\"journalName\"]} {c[\"venue\"]}',\n",
" convert_to_ascii=True,\n",
" ),\n",
" ), # The text is cleaned to remove PDF extraction, web scraping, and other common artifacts\n",
" c[\"fieldsOfStudy\"],\n",
" )\n",
" ) # Create pairs of `(text, [...domains])`\n",
" for c in chunk\n",
" if c[\"fieldsOfStudy\"] and is_english(predict_language(c[\"paperAbstract\"]))\n",
" ]\n",
@ -160,16 +149,15 @@
"\n",
"Upload the dataset (or a part of it) to a central repository using `great_ai.add_ground_truth`. This step automatically tags each datapoint with a split label according to the ratios we set. Additional tags can be also given.\n",
"\n",
"#### Use a different repository\n",
"#### Production-ready backend\n",
"\n",
"For the sake of simplicity, the tutorial uses the local hard drive (`great_ai.ParallelTinyDbDriver`) as the central repository.\n",
"This can be simply changed, for example, by the following snippet:\n",
"The MongoDB driver is automatically configured if `mongo.ini` exists with the following scheme:\n",
"\n",
"```python\n",
"from great_ai import configure, MongoDbDriver\n",
"\n",
"configure(tracing_database=MongoDbDriver('mongodb://localhost:27017_or_something_like_that'))\n",
"```"
"```ini\n",
"mongo_connection_string=mongodb://localhost:27017/\n",
"mongo_database=my_great_ai_db\n",
"```\n",
"> You can install MongoDB from [here](https://www.mongodb.com/docs/manual/installation) or [use it as a service](https://www.mongodb.com/cloud/atlas/register)\n"
]
},
{
@ -181,9 +169,17 @@
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[38;5;226m2022-06-19 15:03:30,300 | WARNING | Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
"\u001b[38;5;226m2022-06-19 15:03:30,301 | WARNING | The selected persistence driver (ParallelTinyDbDriver) is not recommended for production\u001b[0m\n",
"\u001b[38;5;39m2022-06-19 15:03:30,301 | INFO | Options: configured ✅\u001b[0m\n"
"\u001b[38;5;226m2022-06-25 11:24:04,668 | WARNING | Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,669 | INFO | Found credentials file (/data/projects/great-ai/examples/simple/mongo.ini), initialising MongodbDriver\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,670 | INFO | Found credentials file (/data/projects/great-ai/examples/simple/mongo.ini), initialising LargeFileMongo\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,671 | INFO | Settings: configured ✅\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,672 | INFO | 🔩 tracing_database: MongodbDriver\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,672 | INFO | 🔩 large_file_implementation: LargeFileMongo\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,673 | INFO | 🔩 is_production: False\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,673 | INFO | 🔩 should_log_exception_stack: True\u001b[0m\n",
"\u001b[38;5;39m2022-06-25 11:24:04,674 | INFO | 🔩 prediction_cache_size: 512\u001b[0m\n",
"\u001b[38;5;226m2022-06-25 11:24:04,674 | WARNING | You still need to check whether you follow all best practices before trusting your deployment.\u001b[0m\n",
"\u001b[38;5;226m2022-06-25 11:24:04,674 | WARNING | > Find out more at https://se-ml.github.io/practices/\u001b[0m\n"
]
}
],

File diff suppressed because one or more lines are too long

Binary file not shown.

After

Width:  |  Height:  |  Size: 160 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 64 KiB

View file

@ -1,3 +1,2 @@
[DEFAULT]
connection_string=mongodb://localhost:27017/
database=large_file_tests
mongo_connection_string=mongodb://localhost:27017/ # change this
mongo_database=great_ai_db2 # this will be automatically created

File diff suppressed because it is too large Load diff

View file

@ -1,4 +1,3 @@
[DEFAULT]
aws_region_name = your_region_like_eu-west-2
aws_access_key_id = YOUR_ACCESS_KEY_ID
aws_secret_access_key = YOUR_VERY_SECRET_ACCESS_KEY

View file

@ -78,7 +78,6 @@ The package can be used as a module from the command-line to give you more flexi
Create an .ini file (or use _~/.aws/credentials_). It may look like this:
```ini
[DEFAULT]
aws_region_name = your_region_like_eu-west-2
aws_access_key_id = YOUR_ACCESS_KEY_ID
aws_secret_access_key = YOUR_VERY_SECRET_ACCESS_KEY

View file

@ -20,7 +20,6 @@ packages = find:
include_package_data = True
python_requires = >=3.8
install_requires =
click < 8.1.0
unidecode >= 1.3.0
multiprocess >= 0.70.0.0
tqdm >= 4.0.0

View file

@ -19,7 +19,7 @@ from .great_ai.context import _is_in_production_mode
from .great_ai.deploy import GreatAI
from .great_ai.exceptions import ArgumentValidationError, MissingArgumentError
from .parse_arguments import parse_arguments
from .utilities.logger import get_logger
from .utilities import get_logger
logger = get_logger(SERVER_NAME)

View file

@ -7,5 +7,5 @@ from .output_models import (
RegressionOutput,
)
from .parameters import log_metric, parameter
from .persistence import MongoDbDriver, ParallelTinyDbDriver, TracingDatabaseDriver
from .tracing import add_ground_truth, query_ground_truth
from .persistence import MongodbDriver, ParallelTinyDbDriver, TracingDatabaseDriver
from .tracing import add_ground_truth, delete_ground_truth, query_ground_truth

View file

@ -1,16 +1,18 @@
from pathlib import Path
from great_ai.large_file import LargeFileLocal, LargeFileMongo, LargeFileS3
from ..large_file import LargeFileMongo, LargeFileS3
from .persistence.mongodb_driver import MongodbDriver
ENV_VAR_KEY = "ENVIRONMENT"
PRODUCTION_KEY = "production"
DEFAULT_TRACING_DB_FILENAME = "tracing_database.json"
DASHBOARD_PATH = "/dashboard"
MONGO_CONFIG_PATHS = ["mongodb.ini", "mongo.ini", "mongo_db.ini", "mongo-db.ini"]
DEFAULT_TRACING_DATABASE_CONFIG_PATHS = {
MongodbDriver: MONGO_CONFIG_PATHS,
}
DEFAULT_LARGE_FILE_CONFIG_PATHS = {
LargeFileLocal: None,
LargeFileMongo: Path("mongodb.ini"),
LargeFileS3: Path("s3.ini"),
LargeFileS3: ["s3.ini", "b2.ini"],
LargeFileMongo: MONGO_CONFIG_PATHS,
}
GITHUB_LINK = "https://github.com/ScoutinScience/great-ai"
@ -25,4 +27,4 @@ ONLINE_TAG_NAME = "online"
SERVER_NAME = "GreatAI-Server"
SE4ML_WEBSITE = 'https://se-ml.github.io/practices/'
SE4ML_WEBSITE = "https://se-ml.github.io/practices/"

View file

@ -7,11 +7,11 @@ from typing import Any, Dict, Optional, Type, cast
from pydantic import BaseModel
from great_ai.large_file import LargeFile, LargeFileLocal
from great_ai.utilities.logger import get_logger
import yaml
from great_ai.utilities import get_logger
from .constants import (
DEFAULT_LARGE_FILE_CONFIG_PATHS,
DEFAULT_TRACING_DB_FILENAME,
DEFAULT_TRACING_DATABASE_CONFIG_PATHS,
ENV_VAR_KEY,
PRODUCTION_KEY,
SE4ML_WEBSITE,
@ -26,6 +26,7 @@ class Context(BaseModel):
logger: Logger
should_log_exception_stack: bool
prediction_cache_size: int
dashboard_table_size: int
class Config:
arbitrary_types_allowed = True
@ -37,6 +38,7 @@ class Context(BaseModel):
"is_production": self.is_production,
"should_log_exception_stack": self.should_log_exception_stack,
"prediction_cache_size": self.prediction_cache_size,
"dashboard_table_size": self.dashboard_table_size,
}
@ -54,13 +56,12 @@ def configure(
*,
log_level: int = DEBUG,
seed: int = 42,
tracing_database: TracingDatabaseDriver = ParallelTinyDbDriver(
Path(DEFAULT_TRACING_DB_FILENAME)
),
large_file_implementation: Type[LargeFile] = LargeFileLocal,
tracing_database: Optional[Type[TracingDatabaseDriver]] = None,
large_file_implementation: Optional[Type[LargeFile]] = None,
should_log_exception_stack: Optional[bool] = None,
prediction_cache_size: int = 512,
disable_se4ml_banner: bool=False
disable_se4ml_banner: bool = False,
dashboard_table_size: int = 50,
) -> None:
global _context
logger = get_logger("great_ai", level=log_level)
@ -72,9 +73,11 @@ def configure(
)
is_production = _is_in_production_mode(logger=logger)
_initialize_large_file(large_file_implementation, logger=logger)
_set_seed(seed)
tracing_database = _initialize_tracing_database(tracing_database, logger=logger)()
if not tracing_database.is_production_ready:
if is_production:
logger.error(
@ -87,23 +90,27 @@ def configure(
_context = Context(
tracing_database=tracing_database,
large_file_implementation=large_file_implementation,
large_file_implementation=_initialize_large_file(
large_file_implementation, logger=logger
),
is_production=is_production,
logger=logger,
should_log_exception_stack=not is_production
if should_log_exception_stack is None
else should_log_exception_stack,
prediction_cache_size=prediction_cache_size,
dashboard_table_size=dashboard_table_size,
)
logger.info("Setting: configured ✅")
logger.info("Settings: configured ✅")
for k, v in get_context().to_flat_dict().items():
logger.info(f'🔩 {k}: {v}')
logger.info(f"🔩 {k}: {v}")
if not is_production and not disable_se4ml_banner:
logger.warning(f'You still need to check whether you follow all best practices so that you and others can trust your deployment.')
logger.warning(f'> Find out more at {SE4ML_WEBSITE}')
logger.warning(
"You still need to check whether you follow all best practices before trusting your deployment."
)
logger.warning(f"> Find out more at {SE4ML_WEBSITE}")
def _is_in_production_mode(logger: Optional[Logger]) -> bool:
@ -129,24 +136,48 @@ def _is_in_production_mode(logger: Optional[Logger]) -> bool:
return is_production
def _initialize_large_file(large_file: Type[LargeFile], logger: Logger) -> None:
path = DEFAULT_LARGE_FILE_CONFIG_PATHS[large_file]
if path is None:
return
def _initialize_tracing_database(
selected: Optional[Type[TracingDatabaseDriver]], logger: Logger
) -> Type[TracingDatabaseDriver]:
for tracing_driver, paths in DEFAULT_TRACING_DATABASE_CONFIG_PATHS.items():
if selected is None or selected == tracing_driver:
if tracing_driver.initialized:
logger.warning(
f"{tracing_driver.__name__} has been already configured: skipping initialisation"
)
return tracing_driver
for p in paths:
if Path(p).exists():
logger.info(
f"Found credentials file ({Path(p).absolute()}), initialising {tracing_driver.__name__}"
)
tracing_driver.configure_credentials_from_file(p)
return tracing_driver
logger.warning(
"Cannot find credentials files, defaulting to using ParallelTinyDbDriver"
)
return ParallelTinyDbDriver
def _initialize_large_file(
selected: Optional[Type[LargeFile]], logger: Logger
) -> Type[LargeFile]:
for large_file, paths in DEFAULT_LARGE_FILE_CONFIG_PATHS.items():
if selected is None or selected == large_file:
if large_file.initialized:
logger.warning(
f"{large_file.__name__} has been already configured: skipping initialisation"
)
return
if path.exists():
large_file.configure_credentials_from_file(path)
logger.info(f"{large_file.__name__} initialised with config ({path.resolve()})")
else:
logger.warning(
f"Default {large_file.__name__} config ({path.resolve()}) not found, skipping {large_file.__name__} initialisation"
return large_file
for p in paths:
if Path(p).exists():
logger.info(
f"Found credentials file ({Path(p).absolute()}), initialising {large_file.__name__}"
)
large_file.configure_credentials_from_file(p)
return large_file
logger.warning("Cannot find credentials files, defaulting to using LargeFileLocal")
return LargeFileLocal
def _set_seed(seed: int) -> None:

View file

@ -13,13 +13,12 @@ from typing import (
cast,
)
import yaml
from fastapi import APIRouter, FastAPI, status
from pydantic import BaseModel, create_model
from great_ai.great_ai.deploy.routes.bootstrap_dashboard import bootstrap_dashboard
from great_ai.great_ai.views.cache_statistics import CacheStatistics
from great_ai.utilities.parallel_map import parallel_map
from great_ai.utilities import parallel_map
from ..constants import DASHBOARD_PATH
from ..context import get_context
@ -40,6 +39,7 @@ from .routes import (
T = TypeVar("T")
class GreatAI(Generic[T]):
def __init__(self, func: Callable[..., Any], version: str):
func = automatically_decorate_parameters(func)
@ -70,7 +70,6 @@ class GreatAI(Generic[T]):
redoc_url=None,
)
@staticmethod
def deploy(
func: Optional[Callable[..., T]] = None,

View file

@ -20,6 +20,10 @@ def bootstrap_trace_endpoints(app: FastAPI) -> None:
) -> List[Trace]:
return get_context().tracing_database.query(
conjunctive_filters=query.filter,
conjunctive_tags=query.conjunctive_tags,
since=query.since,
until=query.until,
has_feedback=query.has_feedback,
sort_by=query.sort,
skip=skip,
take=take,

View file

@ -1,5 +1,5 @@
from math import ceil
from typing import Any, Dict, List, Sequence, Tuple
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
import pandas as pd
import plotly.express as px
@ -8,12 +8,12 @@ from dash import Dash, dcc, html
from dash.dependencies import Input, Output
from flask import Flask
from great_ai.utilities.unique import unique
from great_ai.utilities import unique
from ....constants import DASHBOARD_PATH, ONLINE_TAG_NAME
from ....context import get_context
from ....helper import snake_case_to_text, text_to_hex_color
from ....views import SortBy
from ....views import SortBy, Trace
from .get_description import get_description
from .get_filter_from_datatable import get_filter_from_datatable
from .get_footer import get_footer
@ -105,19 +105,32 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
@app.callback(
Output(table, "data"),
Output(table, "page_count"),
Output(table, "columns"),
Output(traces_table_container, "style"),
Output(execution_time_histogram_container, "children"),
Output(parallel_coordinates, "figure"),
Output(parallel_coordinates, "style"),
Input(table, "page_current"),
Input(table, "page_size"),
Input(table, "sort_by"),
Input(table, "filter_query"),
Input(interval, "n_intervals"),
)
def update_table(
def update_page(
page_current: int,
page_size: int,
sort_by: List[SortBy],
sort_by: List[Dict[str, Union[str, int]]],
filter_query: str,
n_intervals: int,
) -> Tuple[List[Dict[str, Any]], int]:
) -> Tuple[
List[Dict[str, Any]],
int,
List[Dict[str, Sequence[str]]],
Dict[str, Any],
Any,
go.Figure,
Dict[str, Any],
]:
conjunctive_filters = (
[get_filter_from_datatable(f) for f in filter_query.split(" && ")]
if filter_query
@ -128,30 +141,35 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
elements, count = get_context().tracing_database.query(
skip=page_current * page_size,
take=page_size,
conjunctive_tags=[ONLINE_TAG_NAME],
conjunctive_filters=non_null_conjunctive_filters,
sort_by=sort_by,
conjunctive_tags=[ONLINE_TAG_NAME],
sort_by=[SortBy.parse_obj(s) for s in sort_by],
)
columns, style = update_layout(elements[0] if elements else None)
execution_time_histogram, parallel_coords_fig, style = update_charts(
elements=elements, function_name=function_name, accent_color=accent_color
)
return (
[e.to_flat_dict() for e in elements],
[e.to_flat_dict(include_original=False) for e in elements],
max(1, ceil(count / page_size)),
columns,
style,
execution_time_histogram,
parallel_coords_fig,
style,
)
@app.callback(
Output(table, "columns"),
Output(traces_table_container, "style"),
Input(interval, "n_intervals"),
)
return app.server
def update_layout(
n_intervals: int,
first_element: Optional[Trace],
) -> Tuple[List[Dict[str, Sequence[str]]], Dict[str, Any]]:
elements, count = get_context().tracing_database.query(
take=1, conjunctive_tags=[ONLINE_TAG_NAME]
)
if elements:
keys = list(elements[0].to_flat_dict().keys())
if first_element:
keys = list(first_element.to_flat_dict(include_original=False).keys())
header_height = max(len(i.split(":")) for i in keys)
columns = [
{
@ -166,31 +184,13 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
return (
columns,
{"display": "block" if count > 0 else "none"},
{"display": "none" if first_element is None else "block"},
)
@app.callback(
Output(execution_time_histogram_container, "children"),
Output(parallel_coordinates, "figure"),
Output(parallel_coordinates, "style"),
Input(table, "filter_query"),
Input(interval, "n_intervals"),
)
def update_charts(
filter_query: str, n_intervals: int
elements: List[Trace], function_name: str, accent_color: str
) -> Tuple[Any, go.Figure, Dict[str, Any]]:
conjunctive_filters = (
[get_filter_from_datatable(f) for f in filter_query.split(" && ")]
if filter_query
else []
)
non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
elements, count = get_context().tracing_database.query(
conjunctive_tags=[ONLINE_TAG_NAME],
conjunctive_filters=non_null_conjunctive_filters,
)
if not elements:
return (
html.Span(
@ -201,7 +201,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
{"display": "none"},
)
flat_elements = [e.to_flat_dict() for e in elements]
flat_elements = [e.to_flat_dict(include_original=False) for e in elements]
execution_time_histogram = dcc.Graph(config={"displaylogo": False})
df = pd.DataFrame(flat_elements)
@ -224,8 +224,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
dimensions=[
get_dimension_descriptor(df, c)
for c in df.columns
if c
not in {"trace_id", "created", "output", "exception", "feedback"}
if c not in {"trace_id", "created", "output", "exception", "feedback"}
and "_flat" not in c
],
line_color=accent_color,
@ -233,8 +232,6 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
)
return execution_time_histogram, parallel_coords_fig, {}
return app.server
def get_dimension_descriptor(df: pd.DataFrame, column: str) -> Dict[str, Any]:
dimension: Dict[str, Any] = {

View file

@ -1,10 +1,12 @@
from dash import dash_table
from great_ai.great_ai.context import get_context
def get_traces_table() -> dash_table.DataTable:
return dash_table.DataTable(
page_current=0,
page_size=20,
page_size=get_context().dashboard_table_size,
page_action="custom",
filter_action="custom",
sort_action="custom",
@ -24,7 +26,6 @@ def get_traces_table() -> dash_table.DataTable:
"background-color": "white",
"font-weight": "bold",
},
style_table={"max-height": "70vh", "overflow": "auto"},
merge_duplicate_headers=True,
style_cell_conditional=[
{"if": {"column_id": "output"}, "width": 1500},

View file

@ -1,3 +1,3 @@
from .mongodb_driver import MongoDbDriver
from .mongodb_driver import MongodbDriver
from .parallel_tinydb_driver import ParallelTinyDbDriver
from .tracing_database_driver import TracingDatabaseDriver

View file

@ -1,5 +1,137 @@
from datetime import datetime
from typing import Any, List, Mapping, Optional, Sequence, Tuple
from pymongo import MongoClient
from ..views import Filter, SortBy, Trace
from .tracing_database_driver import TracingDatabaseDriver
operator_mapping = {
"=": "$eq",
"!=": "$ne",
"<": "$lt",
"<=": "$lte",
">": "$gt",
">=": "$gte",
"contains": "$regex",
}
class MongoDbDriver(TracingDatabaseDriver):
class MongodbDriver(TracingDatabaseDriver):
is_production_ready = True
def __init__(self) -> None:
super().__init__()
if self.mongo_connection_string is None or self.mongo_database is None:
raise ValueError(
"Please configure the MongoDB access options by calling MongodbDriver.configure_credentials"
)
@classmethod
def configure_credentials( # type: ignore
cls,
*,
mongo_connection_string: str,
mongo_database: str,
**_: Mapping[str, Any],
) -> None:
cls.mongo_connection_string = mongo_connection_string
cls.mongo_database = mongo_database
super().configure_credentials()
def save(self, trace: Trace) -> str:
serialized = trace.to_flat_dict()
serialized["_id"] = trace.trace_id
with MongoClient(self.mongo_connection_string) as client:
return client[self.mongo_database].traces.insert_one(serialized)
def save_batch(self, documents: List[Trace]) -> List[str]:
serialized = [d.to_flat_dict() for d in documents]
for s in serialized:
s["_id"] = s["trace_id"]
with MongoClient(self.mongo_connection_string) as client:
return client[self.mongo_database].traces.insert_many(
serialized, ordered=False
)
def get(self, id: str) -> Optional[Trace]:
with MongoClient(self.mongo_connection_string) as client:
value = client[self.mongo_database].traces.find_one(id)
if value:
value = Trace.parse_obj(value)
return value
def _get_operator(self, filter: Filter) -> str:
if filter.operator == "contains" and not isinstance(filter.value, str):
return operator_mapping["="]
return operator_mapping[filter.operator]
def query(
self,
*,
skip: int = 0,
take: Optional[int] = None,
conjunctive_filters: Sequence[Filter] = [],
conjunctive_tags: Sequence[str] = [],
since: Optional[datetime] = None,
until: Optional[datetime] = None,
has_feedback: Optional[bool] = None,
sort_by: Sequence[SortBy] = [],
) -> Tuple[List[Trace], int]:
query = {
"filter": {
"$and": [{"tags": tag} for tag in conjunctive_tags]
+ [
{f.property: {self._get_operator(f): f.value}}
for f in conjunctive_filters
]
+ [{}]
},
"sort": [
(col.column_id, 1 if col.direction == "asc" else -1) for col in sort_by
],
}
if skip:
query["skip"] = skip
if take:
query["limit"] = take
if since:
query["filter"]["$and"].append({"created": {"$gte": since}})
if until:
query["filter"]["$and"].append({"created": {"$lte": until}})
if has_feedback is not None:
query["filter"]["$and"].append(
{"feedback": {"$ne": None}} if has_feedback else {"feedback": None}
)
with MongoClient(self.mongo_connection_string) as client:
values = client[self.mongo_database].traces.find(**query)
documents = [Trace.parse_obj(t) for t in values]
return documents, len(documents)
def update(self, id: str, new_version: Trace) -> None:
serialized = new_version.dict()
serialized["_id"] = new_version.trace_id
with MongoClient(self.mongo_connection_string) as client:
client[self.mongo_database].traces.update_one(id, new_version)
def delete(self, id: str) -> None:
with MongoClient(self.mongo_connection_string) as client:
client[self.mongo_database].traces.delete_one(id)
def delete_batch(self, ids: List[str]) -> List[str]:
delete_filter = {"_id": {"$in": ids}}
with MongoClient(self.mongo_connection_string) as client:
return client[self.mongo_database].traces.delete_many(delete_filter)

View file

@ -9,6 +9,7 @@ from tinydb import TinyDB
from ..views import Filter, SortBy, Trace
from .tracing_database_driver import TracingDatabaseDriver
DEFAULT_TRACING_DB_FILENAME = "tracing_database.json"
lock = Lock()
@ -17,10 +18,7 @@ operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=
class ParallelTinyDbDriver(TracingDatabaseDriver):
is_production_ready = False
def __init__(self, path_to_db: Path) -> None:
super().__init__()
self._path_to_db = path_to_db
path_to_db = Path(DEFAULT_TRACING_DB_FILENAME)
def save(self, trace: Trace) -> str:
return self._safe_execute(lambda db: db.insert(trace.dict()))
@ -43,8 +41,9 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
conjunctive_filters: Sequence[Filter] = [],
conjunctive_tags: Sequence[str] = [],
since: Optional[datetime] = None,
sort_by: Sequence[SortBy] = [],
has_feedback: Optional[bool] = None
until: Optional[datetime] = None,
has_feedback: Optional[bool] = None,
sort_by: Sequence[SortBy] = []
) -> Tuple[List[Trace], int]:
def does_match(d: Dict[str, Any]) -> bool:
return (
@ -53,6 +52,10 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
since is None
or cast(datetime, datetime.fromisoformat(d["created"])) >= since
)
and (
until is None
or cast(datetime, datetime.fromisoformat(d["created"])) <= until
)
and (
has_feedback is None or has_feedback == (d["feedback"] is not None)
)
@ -99,7 +102,11 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
def delete(self, id: str) -> None:
self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id))
def delete_batch(self, ids: List[str]) -> List[str]:
for i in ids:
self.delete(i)
def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any:
with lock:
with TinyDB(self._path_to_db) as db:
with TinyDB(self.path_to_db) as db:
return func(db)

View file

@ -1,12 +1,29 @@
from abc import ABC, abstractmethod
from abc import ABC, abstractclassmethod, abstractmethod
from datetime import datetime
from typing import List, Optional, Sequence, Tuple
from pathlib import Path
from typing import List, Optional, Sequence, Tuple, Union
from great_ai.utilities import ConfigFile
from ..views import Filter, SortBy, Trace
class TracingDatabaseDriver(ABC):
is_production_ready: bool
initialized: bool = False
@classmethod
def configure_credentials_from_file(
cls,
secrets_path: Union[Path, str],
) -> None:
cls.configure_credentials(**ConfigFile(secrets_path))
@abstractclassmethod
def configure_credentials(
cls,
) -> None:
cls.initialized = True
@abstractmethod
def save(self, document: Trace) -> str:
@ -31,9 +48,10 @@ class TracingDatabaseDriver(ABC):
take: Optional[int] = None,
conjunctive_filters: Sequence[Filter] = [],
conjunctive_tags: Sequence[str] = [],
until: Optional[datetime] = None,
since: Optional[datetime] = None,
sort_by: Sequence[SortBy] = [],
has_feedback: Optional[bool] = None
has_feedback: Optional[bool] = None,
sort_by: Sequence[SortBy] = []
) -> Tuple[List[Trace], int]:
pass
@ -44,3 +62,10 @@ class TracingDatabaseDriver(ABC):
@abstractmethod
def delete(self, id: str) -> None:
pass
@abstractmethod
def delete_batch(
self,
ids: List[str],
) -> None:
pass

View file

@ -1,3 +1,4 @@
from .add_ground_truth import add_ground_truth
from .delete_ground_truth import delete_ground_truth
from .query_ground_truth import query_ground_truth
from .tracing_context import TracingContext

View file

@ -51,7 +51,7 @@ def add_ground_truth(
created = datetime.utcnow().isoformat()
traces = [
Trace[T](
Trace(
created=created,
original_execution_time_ms=0,
logged_values=X if isinstance(X, dict) else {"input": X},

View file

@ -0,0 +1,22 @@
from datetime import datetime
from typing import List, Optional, Union
from ..context import get_context
def delete_ground_truth(
conjunctive_tags: Union[List[str], str] = [],
*,
until: Optional[datetime] = None,
since: Optional[datetime] = None,
) -> None:
tags = (
conjunctive_tags if isinstance(conjunctive_tags, list) else [conjunctive_tags]
)
db = get_context().tracing_database
items, length = db.query(
conjunctive_tags=tags, until=until, since=since, has_feedback=True
)
db.delete_batch([i.trace_id for i in items])

View file

@ -1,10 +1,8 @@
from datetime import datetime
from typing import List, Optional, TypeVar, Union
from typing import List, Optional, Union
from ..context import get_context
from ..views.trace import Trace
T = TypeVar("T")
from ..views import Trace
def query_ground_truth(
@ -12,7 +10,7 @@ def query_ground_truth(
*,
since: Optional[datetime] = None,
return_max_count: Optional[int] = None
) -> List[Trace[T]]:
) -> List[Trace]:
tags = (
conjunctive_tags if isinstance(conjunctive_tags, list) else [conjunctive_tags]
)

View file

@ -1,4 +1,5 @@
from typing import List
from datetime import datetime
from typing import List, Optional, Sequence
from pydantic import BaseModel
@ -9,6 +10,10 @@ from .sort_by import SortBy
class Query(BaseModel):
filter: List[Filter] = []
sort: List[SortBy] = []
conjunctive_tags: Sequence[str] = []
since: Optional[datetime] = None
until: Optional[datetime] = None
has_feedback: Optional[bool] = None
class Config:
schema_extra = {
@ -24,5 +29,7 @@ class Query(BaseModel):
{"column_id": "original_execution_time_ms", "direction": "asc"},
{"column_id": "id", "direction": "desc"},
],
"conjunctive_tags": ["online"],
"has_feedback": False,
}
}

View file

@ -1,8 +1,8 @@
from typing import Literal
from typing_extensions import TypedDict
from pydantic import BaseModel
class SortBy(TypedDict):
class SortBy(BaseModel):
column_id: str
direction: Literal["asc", "desc"]

View file

@ -1,8 +1,8 @@
from pprint import pformat
from typing import Any, Dict, Generic, List, Optional, TypeVar
from uuid import uuid4
import yaml
from pydantic import validator
from pydantic import Extra, validator
from ..helper import HashableBaseModel
from .model import Model
@ -10,7 +10,7 @@ from .model import Model
T = TypeVar("T")
class Trace(HashableBaseModel, Generic[T]):
class Trace(Generic[T], HashableBaseModel):
trace_id: Optional[str]
created: str
original_execution_time_ms: float
@ -21,6 +21,9 @@ class Trace(HashableBaseModel, Generic[T]):
feedback: Any = None
tags: List[str]
class Config:
extra = Extra.ignore
@validator("trace_id", always=True)
def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
if not v:
@ -41,26 +44,42 @@ class Trace(HashableBaseModel, Generic[T]):
@property
def output_flat(self) -> str:
return yaml.dump(self.output, stream=None)
return pformat(self.output, indent=2, compact=True)
@property
def exception_flat(self) -> str:
return (
"null"
if self.exception is None
else pformat(self.exception, indent=2, compact=True)
)
@property
def feedback_flat(self) -> str:
return (
"null" if self.feedback is None else yaml.dump(self.feedback, stream=None)
"null"
if self.feedback is None
else pformat(self.feedback, indent=2, compact=True)
)
@property
def tags_flat(self) -> str:
return ",\n".join(self.tags)
def to_flat_dict(self) -> Dict[str, Any]:
def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
return {
**(
self.dict()
if include_original
else {
"trace_id": self.trace_id,
"created": self.created,
"original_execution_time_ms": self.original_execution_time_ms,
}
),
**self.logged_values,
"models_flat": self.models_flat,
"exception": "null" if self.exception is None else self.exception,
"exception_flat": self.exception_flat,
"output_flat": self.output_flat,
"feedback_flat": self.feedback_flat,
"tags_flat": self.tags_flat,

View file

@ -19,29 +19,29 @@ MONGO_NAME_VERSION_SEPARATOR = "_"
class LargeFileMongo(LargeFile):
connection_string = None
database = None
mongo_connection_string = None
mongo_database = None
@classmethod
def configure_credentials( # type: ignore
cls,
*,
connection_string: str,
database: str,
mongo_connection_string: str,
mongo_database: str,
**_: Mapping[str, Any],
) -> None:
cls.connection_string = connection_string
cls.database = database
cls.mongo_connection_string = mongo_connection_string
cls.mongo_database = mongo_database
super().configure_credentials()
@cached_property
def _client(self) -> GridFSBucket:
if self.connection_string is None or self.database is None:
if self.mongo_connection_string is None or self.mongo_database is None:
raise ValueError(
"Please configure the MongoDB access options by calling LargeFileMongo.configure_credentials or set offline_mode=True in the constructor."
)
db: Database = MongoClient(self.connection_string)[self.database]
db: Database = MongoClient(self.mongo_connection_string)[self.mongo_database]
return GridFSBucket(db)
def _find_remote_instances(self) -> List[DataInstance]: