[:fontawesome-solid-chart-simple: Train it](train.ipynb){ .md-button .md-button--primary }
@@ -32,7 +32,7 @@ We load and preprocess the dataset while relying on [great_ai.utilities.clean][g
After training and evaluating a model, it is exported using [great_ai.save_model][].
??? tip "Remote storage"
- To store your model remotely, you need to set your credentials before calling `save_model`.
+ To store your model remotely, you must set your credentials before calling `save_model`.
For example, to use [AWS S3](https://aws.amazon.com/s3){ target=_blank }:
```python
@@ -68,9 +68,9 @@ def predict_domain(sentence, model):
```
1. [@use_model][great_ai.use_model] loads and injects your model into the `predict_domain` function's `model` argument.
- You can freely reference it knowing that the function is always provided with it.
+ You can freely reference it, knowing that the function is always provided with it.
-Finally, we test the model's inference function through the GreatAI dashboard. [The only thing left is to deploy the hardened-service properly.](/how-to-guides/use-service)
+Finally, we test the model's inference function through the GreatAI dashboard. [The only thing left is to deploy the hardened service properly.](/how-to-guides/use-service)
[:material-book: Learn about all the features](/how-to-guides/create-service){ .md-button .md-button--primary }
diff --git a/docs/tutorial/train.ipynb b/docs/tutorial/train.ipynb
index fd4e12c..61b8e1e 100644
--- a/docs/tutorial/train.ipynb
+++ b/docs/tutorial/train.ipynb
@@ -30,9 +30,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "First, we have to get some data. After downloading it [from here](https://github.com/allenai/scibert/tree/master/data/text_classification/mag), we might notice that the dataset is in [JSON Lines](https://jsonlines.org/) format (each line is a seperate JSON document). \n",
+ "First, we have to get some data. After downloading it [from here](https://github.com/allenai/scibert/tree/master/data/text_classification/mag), we might notice that the dataset is in [JSON Lines](https://jsonlines.org/) format (each line is a separate JSON document). \n",
"\n",
- "Let's write a function which takes a single line, and returns the sentence and the corresponding label from it. Before returning, the sentence is also [cleaned](/reference/utilities/#great_ai.utilities.clean.clean) to remove any LaTeX, XML, unicode, PDF-extraction artifacts."
+ "Let's write a function which takes a single line and returns the sentence and the corresponding label from it. Before returning, the sentence is also [cleaned](/reference/utilities/#great_ai.utilities.clean.clean) to remove any LaTeX, XML, unicode, PDF-extraction artifacts."
]
},
{
@@ -238,7 +238,7 @@
"\n",
"You might wonder that *\"this is great, but besides some utility functions (`clean`, `simple_parallel_map`, ...) what more value does GreatAI add?\"*. This would be a valid argument because the scope of GreatAI actually only starts here.\n",
"\n",
- "> Not coincidentally, this is the point where the scope of Data Science ends but it's still a grey-zone for software engineering.\n",
+ "> Not coincidentally, this is the point where the scope of Data Science ends but it's still a grey zone for software engineering.\n",
"\n",
"In order to use this model in production, we have to make it available on some possibly shared infrastructure."
]
diff --git a/great_ai/__init__.py b/great_ai/__init__.py
index a80ed2d..da7dba8 100644
--- a/great_ai/__init__.py
+++ b/great_ai/__init__.py
@@ -1,5 +1,6 @@
-"""GreatAI."""
-__version__ = "0.1.5"
+"""Transform your prototype AI code into production-ready software."""
+
+__version__ = "0.1.11"
from .context import configure
diff --git a/great_ai/__main__.py b/great_ai/__main__.py
index 3cc12be..ec4915d 100644
--- a/great_ai/__main__.py
+++ b/great_ai/__main__.py
@@ -76,9 +76,11 @@ def serve() -> None:
finally:
if args.file_name.endswith(".ipynb"):
- Path(get_script_name_of_notebook(args.file_name)).unlink(
- missing_ok=True
- )
+ try:
+ Path(get_script_name_of_notebook(args.file_name)).unlink()
+ except FileNotFoundError:
+ # missing_ok only exists >= Python 3.8
+ pass
else:
class EventHandler(PatternMatchingEventHandler):
@@ -124,9 +126,11 @@ def serve() -> None:
observer.stop()
restart_handler.stop_server()
if args.file_name.endswith(".ipynb"):
- Path(get_script_name_of_notebook(args.file_name)).unlink(
- missing_ok=True
- )
+ try:
+ Path(get_script_name_of_notebook(args.file_name)).unlink()
+ except FileNotFoundError:
+ # missing_ok only exists >= Python 3.8
+ pass
observer.join()
@@ -171,6 +175,7 @@ def find_app(file_name: str) -> Optional[str]:
finally:
logging.disable(logging.NOTSET)
+ app_name = None
for name, value in module.__dict__.items():
if isinstance(value, GreatAI):
app_name = name
@@ -179,7 +184,7 @@ def find_app(file_name: str) -> Optional[str]:
logger.info(f"Found `{app_name}` to be the GreatAI app ")
else:
raise MissingArgumentError(
- "GreatAI app could not be found, make sure to use `@GreatAI.deploy` on your prediction function"
+ "GreatAI app could not be found, make sure to use `@GreatAI.create` on your prediction function"
)
return f"{file_name}:{app_name}.app"
diff --git a/great_ai/context.py b/great_ai/context.py
index 29327dc..24412a5 100644
--- a/great_ai/context.py
+++ b/great_ai/context.py
@@ -4,7 +4,7 @@ from logging import DEBUG, Logger
from pathlib import Path
from typing import Any, Dict, Optional, Type, Union, cast
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
from great_ai import __version__
@@ -34,8 +34,7 @@ class Context(BaseModel):
dashboard_table_size: int
route_config: RouteConfig
- class Config:
- arbitrary_types_allowed = True
+ model_config = ConfigDict(arbitrary_types_allowed=True)
def to_flat_dict(self) -> Dict[str, Any]:
return {
@@ -137,9 +136,11 @@ def configure(
),
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,
+ 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,
route_config=route_config,
diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py
index c269021..ac14590 100644
--- a/great_ai/deploy/great_ai.py
+++ b/great_ai/deploy/great_ai.py
@@ -16,12 +16,12 @@ from typing import (
overload,
)
-from async_lru import alru_cache
from fastapi import FastAPI
-from tqdm.cli import tqdm
+from tqdm import tqdm
from ..constants import DASHBOARD_PATH
from ..context import get_context
+from ..external.async_lru import alru_cache
from ..helper import freeze_arguments, get_function_metadata_store, snake_case_to_text
from ..models.use_model import model_versions
from ..parameters.automatically_decorate_parameters import (
@@ -42,11 +42,15 @@ V = TypeVar("V")
class GreatAI(Generic[T, V]):
- """Wrapper for a prediction function providing the implementation of SE4ML best-practices.
+ """Wrapper for a prediction function providing the implementation of SE4ML best practices.
Provides caching (with argument freezing), a TracingContext during execution, the
scaffolding of HTTP endpoints using FastAPI and a dashboard using Dash.
+ IMPORTANT: when a request is served from cache, no new trace is created. Thus, the
+ same trace can be returned multiple times. If this is undesirable turn off caching
+ using `configure(prediction_cache_size=0)`.
+
Supports wrapping async and synchronous functions while also maintaining correct
typing.
@@ -99,15 +103,13 @@ class GreatAI(Generic[T, V]):
# Overloaded function signatures 1 and 2 overlap with incompatible return types
# https://github.com/python/mypy/issues/12759
func: Callable[..., Awaitable[V]],
- ) -> "GreatAI[Awaitable[Trace[V]], V]":
- ...
+ ) -> "GreatAI[Awaitable[Trace[V]], V]": ...
@overload
@staticmethod
def create(
func: Callable[..., V],
- ) -> "GreatAI[Trace[V], V]":
- ...
+ ) -> "GreatAI[Trace[V], V]": ...
@staticmethod
def create(
@@ -142,7 +144,7 @@ class GreatAI(Generic[T, V]):
>>> my_function('3').output
Traceback (most recent call last):
...
- TypeError: type of a must be int; got str instead
+ TypeError: argument a is not of the expected type: ...
Args:
func: The prediction function that needs to be decorated.
@@ -166,8 +168,7 @@ class GreatAI(Generic[T, V]):
concurrency: Optional[int] = None,
unpack_arguments: Literal[True],
do_not_persist_traces: bool = ...,
- ) -> List[Trace[V]]:
- ...
+ ) -> List[Trace[V]]: ...
@overload
def process_batch(
@@ -177,8 +178,7 @@ class GreatAI(Generic[T, V]):
concurrency: Optional[int] = None,
unpack_arguments: Literal[False] = ...,
do_not_persist_traces: bool = ...,
- ) -> List[Trace[V]]:
- ...
+ ) -> List[Trace[V]]: ...
def process_batch(
self,
@@ -230,22 +230,24 @@ class GreatAI(Generic[T, V]):
),
)
+ # inner/inner_async return the class' T (bound to Trace | Awaitable[Trace]);
+ # in this method T resolves to Trace[V], but that is not provable to mypy, so
+ # cast to the concrete shape parallel_map expects.
+ map_function = cast(
+ Callable[[Any], Union[Trace[V], Awaitable[Trace[V]]]],
+ inner_async if get_function_metadata_store(self).is_asynchronous else inner,
+ )
+
return list(
tqdm(
- parallel_map(
- inner_async
- if get_function_metadata_store(self).is_asynchronous
- else inner,
- batch,
- concurrency=concurrency,
- ),
+ parallel_map(map_function, batch, concurrency=concurrency),
total=len(batch),
)
)
@staticmethod
def _get_cached_traced_function(
- func: Callable[..., Union[V, Awaitable[V]]]
+ func: Callable[..., Union[V, Awaitable[V]]],
) -> Callable[..., T]:
@lru_cache(maxsize=get_context().prediction_cache_size)
def func_in_tracing_context_sync(
diff --git a/great_ai/deploy/routes/bootstrap_dashboard.py b/great_ai/deploy/routes/bootstrap_dashboard.py
index 2989c95..6bafc5f 100644
--- a/great_ai/deploy/routes/bootstrap_dashboard.py
+++ b/great_ai/deploy/routes/bootstrap_dashboard.py
@@ -1,7 +1,7 @@
from pathlib import Path
+from a2wsgi import WSGIMiddleware
from fastapi import FastAPI
-from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.responses import FileResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
@@ -19,7 +19,10 @@ def bootstrap_dashboard(app: FastAPI, function_name: str, documentation: str) ->
def get_favicon() -> FileResponse:
return FileResponse(PATH / "dashboard/assets/favicon.ico")
- app.mount(DASHBOARD_PATH, WSGIMiddleware(dash_app))
+ # a2wsgi types its WSGI input (Flask) and ASGI output more strictly than
+ # Starlette's loose WSGIApp/ASGIApp aliases, so mypy flags the mount despite
+ # this being the canonical, runtime-correct way to bridge the Dash app.
+ app.mount(DASHBOARD_PATH, WSGIMiddleware(dash_app)) # type: ignore[arg-type]
@app.get("/", include_in_schema=False)
def redirect_to_entrypoint() -> RedirectResponse:
diff --git a/great_ai/deploy/routes/bootstrap_feedback_endpoints.py b/great_ai/deploy/routes/bootstrap_feedback_endpoints.py
index 8b219ee..b14df7e 100644
--- a/great_ai/deploy/routes/bootstrap_feedback_endpoints.py
+++ b/great_ai/deploy/routes/bootstrap_feedback_endpoints.py
@@ -31,7 +31,7 @@ def bootstrap_feedback_endpoints(app: FastAPI) -> None:
return trace.feedback
@router.delete("/", status_code=status.HTTP_204_NO_CONTENT)
- def delete_feedback(trace_id: str) -> Any:
+ def delete_feedback(trace_id: str) -> Response:
trace = get_context().tracing_database.get(trace_id)
if trace is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
diff --git a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py
index 0ed6f5b..b3787ec 100644
--- a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py
+++ b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py
@@ -22,10 +22,10 @@ def bootstrap_prediction_endpoint(
try:
if inspect.iscoroutinefunction(func):
return await cast(Callable[..., Awaitable[Trace]], func)(
- **cast(BaseModel, input_value).dict()
+ **cast(BaseModel, input_value).model_dump()
)
return cast(Callable[..., Trace], func)(
- **cast(BaseModel, input_value).dict()
+ **cast(BaseModel, input_value).model_dump()
)
except Exception as e:
raise HTTPException(
diff --git a/great_ai/deploy/routes/dashboard/create_dash_app.py b/great_ai/deploy/routes/dashboard/create_dash_app.py
index 6ad1ade..e76ab98 100644
--- a/great_ai/deploy/routes/dashboard/create_dash_app.py
+++ b/great_ai/deploy/routes/dashboard/create_dash_app.py
@@ -26,14 +26,49 @@ def create_dash_app(function_name: str, version: str, function_docs: str) -> Fla
app = Dash(
function_name,
requests_pathname_prefix=DASHBOARD_PATH + "/",
- server=Flask(__name__),
+ server=Flask(__name__), # type: ignore[arg-type]
title=function_name,
- update_title=None,
+ update_title=None, # type: ignore[arg-type]
external_stylesheets=[
"/assets/index.css",
],
)
+ execution_time_histogram_container = html.Div()
+ configuration_container = html.Div(
+ className="configuration-container",
+ )
+ table = get_traces_table()
+ traces_table_container = html.Div(
+ [
+ html.Header(
+ [
+ html.H2("Latest traces"),
+ html.P(
+ "Recent traces and aggregated metrics are presented below. Try filtering the table."
+ ),
+ html.A(
+ "Filtering syntax.",
+ href="https://dash.plotly.com/datatable/filtering",
+ target="_blank",
+ ),
+ ]
+ ),
+ table,
+ ],
+ className="traces-table-container",
+ )
+ parallel_coordinates = dcc.Graph(
+ className="parallel-coordinates", config={"displaylogo": False}
+ )
+ interval = dcc.Interval(
+ interval=2 * 1000, # in milliseconds
+ n_intervals=0,
+ max_intervals=(
+ 1 if get_context().is_production else -1
+ ), # will be incremented in production upon each successful request
+ )
+
app.layout = html.Main(
[
html.Div(
@@ -49,43 +84,15 @@ def create_dash_app(function_name: str, version: str, function_docs: str) -> Fla
function_docs=function_docs,
accent_color=accent_color,
),
- execution_time_histogram_container := html.Div(),
+ execution_time_histogram_container,
],
),
- configuration_container := html.Div(
- className="configuration-container",
- ),
- traces_table_container := html.Div(
- [
- html.Header(
- [
- html.H2("Latest traces"),
- html.P(
- "Recent traces and aggregated metrics are presented below. Try filtering the table."
- ),
- html.A(
- "Filtering syntax.",
- href="https://dash.plotly.com/datatable/filtering",
- target="_blank",
- ),
- ]
- ),
- table := get_traces_table(),
- ],
- className="traces-table-container",
- ),
- parallel_coordinates := dcc.Graph(
- className="parallel-coordinates", config={"displaylogo": False}
- ),
+ configuration_container,
+ traces_table_container,
+ parallel_coordinates,
html.Div(className="space-filler"),
get_footer(),
- interval := dcc.Interval(
- interval=2 * 1000, # in milliseconds
- n_intervals=0,
- max_intervals=1
- if get_context().is_production
- else -1, # will be incremented in production upon each successful request
- ),
+ interval,
]
)
@@ -152,7 +159,7 @@ def create_dash_app(function_name: str, version: str, function_docs: str) -> Fla
take=page_size,
conjunctive_filters=non_null_conjunctive_filters,
conjunctive_tags=[ONLINE_TAG_NAME],
- sort_by=[SortBy.parse_obj(s) for s in sort_by],
+ sort_by=[SortBy.model_validate(s) for s in sort_by],
)
if non_null_conjunctive_filters:
@@ -237,7 +244,7 @@ def update_charts(
fig.update_layout(
margin=dict(l=0, r=0, b=0, t=0, pad=0),
)
- execution_time_histogram.figure = fig
+ execution_time_histogram.figure = fig # type: ignore[attr-defined]
parallel_coords_fig = go.Figure(
go.Parcoords(
diff --git a/great_ai/deploy/routes/dashboard/get_traces_table.py b/great_ai/deploy/routes/dashboard/get_traces_table.py
index bcf621e..1550f48 100644
--- a/great_ai/deploy/routes/dashboard/get_traces_table.py
+++ b/great_ai/deploy/routes/dashboard/get_traces_table.py
@@ -1,10 +1,12 @@
+from typing import Any
+
from dash import dash_table
from ....context import get_context
-def get_traces_table() -> dash_table.DataTable:
- return dash_table.DataTable(
+def get_traces_table() -> Any:
+ return dash_table.DataTable( # type: ignore[attr-defined]
page_current=0,
page_size=get_context().dashboard_table_size,
page_action="custom",
diff --git a/great_ai/external/__init__.py b/great_ai/external/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/great_ai/external/async_lru.py b/great_ai/external/async_lru.py
new file mode 100644
index 0000000..c209d5a
--- /dev/null
+++ b/great_ai/external/async_lru.py
@@ -0,0 +1,192 @@
+import asyncio
+from collections import OrderedDict
+from functools import _CacheInfo, _make_key, partial, wraps
+
+__version__ = "1.0.3"
+
+__all__ = ("alru_cache",)
+
+
+def unpartial(fn):
+ while hasattr(fn, "func"):
+ fn = fn.func
+
+ return fn
+
+
+def _done_callback(fut, task):
+ if task.cancelled():
+ fut.cancel()
+ return
+
+ exc = task.exception()
+ if exc is not None:
+ fut.set_exception(exc)
+ return
+
+ fut.set_result(task.result())
+
+
+def _cache_invalidate(wrapped, typed, *args, **kwargs):
+ key = _make_key(args, kwargs, typed)
+
+ exists = key in wrapped._cache
+
+ if exists:
+ wrapped._cache.pop(key)
+
+ return exists
+
+
+def _cache_clear(wrapped):
+ wrapped.hits = wrapped.misses = 0
+ wrapped._cache = OrderedDict()
+ wrapped.tasks = set()
+
+
+def _open(wrapped):
+ if not wrapped.closed:
+ raise RuntimeError("alru_cache is not closed")
+
+ was_closed = (
+ wrapped.hits == wrapped.misses == len(wrapped.tasks) == len(wrapped._cache) == 0
+ )
+
+ if not was_closed:
+ raise RuntimeError("alru_cache was not closed correctly")
+
+ wrapped.closed = False
+
+
+def _close(wrapped, *, cancel=False, return_exceptions=True):
+ if wrapped.closed:
+ raise RuntimeError("alru_cache is closed")
+
+ wrapped.closed = True
+
+ if cancel:
+ for task in wrapped.tasks:
+ if not task.done(): # not sure is it possible
+ task.cancel()
+
+ return _wait_closed(wrapped, return_exceptions=return_exceptions)
+
+
+async def _wait_closed(wrapped, *, return_exceptions):
+ wait_closed = asyncio.gather(*wrapped.tasks, return_exceptions=return_exceptions)
+
+ wait_closed.add_done_callback(partial(_close_waited, wrapped))
+
+ ret = await wait_closed
+
+ # hack to get _close_waited callback to be executed
+ await asyncio.sleep(0)
+
+ return ret
+
+
+def _close_waited(wrapped, _):
+ wrapped.cache_clear()
+
+
+def _cache_info(wrapped, maxsize):
+ return _CacheInfo(
+ wrapped.hits,
+ wrapped.misses,
+ maxsize,
+ len(wrapped._cache),
+ )
+
+
+def __cache_touch(wrapped, key):
+ try:
+ wrapped._cache.move_to_end(key)
+ except KeyError: # not sure is it possible
+ pass
+
+
+def _cache_hit(wrapped, key):
+ wrapped.hits += 1
+ __cache_touch(wrapped, key)
+
+
+def _cache_miss(wrapped, key):
+ wrapped.misses += 1
+ __cache_touch(wrapped, key)
+
+
+def alru_cache(
+ fn=None,
+ maxsize=128,
+ typed=False,
+ *,
+ cache_exceptions=True,
+):
+ def wrapper(fn):
+ _origin = unpartial(fn)
+
+ if not asyncio.iscoroutinefunction(_origin):
+ raise RuntimeError("Coroutine function is required, got {}".format(fn))
+
+ # functools.partialmethod support
+ if hasattr(fn, "_make_unbound_method"):
+ fn = fn._make_unbound_method()
+
+ @wraps(fn)
+ async def wrapped(*fn_args, **fn_kwargs):
+ if wrapped.closed:
+ raise RuntimeError("alru_cache is closed for {}".format(wrapped))
+
+ loop = asyncio.get_event_loop()
+
+ key = _make_key(fn_args, fn_kwargs, typed)
+
+ fut = wrapped._cache.get(key)
+
+ if fut is not None:
+ if not fut.done():
+ _cache_hit(wrapped, key)
+ return await asyncio.shield(fut)
+
+ exc = fut._exception
+
+ if exc is None or cache_exceptions:
+ _cache_hit(wrapped, key)
+ return fut.result()
+
+ # exception here and cache_exceptions == False
+ wrapped._cache.pop(key)
+
+ fut = loop.create_future()
+ task = loop.create_task(fn(*fn_args, **fn_kwargs))
+ task.add_done_callback(partial(_done_callback, fut))
+
+ wrapped.tasks.add(task)
+ task.add_done_callback(wrapped.tasks.remove)
+
+ wrapped._cache[key] = fut
+
+ if maxsize is not None and len(wrapped._cache) > maxsize:
+ wrapped._cache.popitem(last=False)
+
+ _cache_miss(wrapped, key)
+ return await asyncio.shield(fut)
+
+ _cache_clear(wrapped)
+ wrapped._origin = _origin
+ wrapped.closed = False
+ wrapped.cache_info = partial(_cache_info, wrapped, maxsize)
+ wrapped.cache_clear = partial(_cache_clear, wrapped)
+ wrapped.invalidate = partial(_cache_invalidate, wrapped, typed)
+ wrapped.close = partial(_close, wrapped)
+ wrapped.open = partial(_open, wrapped)
+
+ return wrapped
+
+ if fn is None:
+ return wrapper
+
+ if callable(fn) or hasattr(fn, "_make_unbound_method"):
+ return wrapper(fn)
+
+ raise NotImplementedError("{} decorating is not supported".format(fn))
diff --git a/great_ai/helper/assert_function_is_not_finalised.py b/great_ai/helper/assert_function_is_not_finalised.py
index 08da383..0d6f950 100644
--- a/great_ai/helper/assert_function_is_not_finalised.py
+++ b/great_ai/helper/assert_function_is_not_finalised.py
@@ -6,8 +6,8 @@ from .get_function_metadata_store import get_function_metadata_store
def assert_function_is_not_finalised(func: Callable) -> None:
error_message = (
- "The outer-most (first) decorator has to be `@GreatAI.deploy`. "
- + f"In the case of `{func.__name__}`, it is not: fix this by moving `@GreatAI.deploy` to the top."
+ "The outer-most (first) decorator has to be `@GreatAI.create`. "
+ + f"In the case of `{func.__name__}`, it is not: fix this by moving `@GreatAI.create` to the top."
)
if get_function_metadata_store(func).is_finalised:
diff --git a/great_ai/helper/freeze_arguments.py b/great_ai/helper/freeze_arguments.py
index c50b13c..e6ff56f 100644
--- a/great_ai/helper/freeze_arguments.py
+++ b/great_ai/helper/freeze_arguments.py
@@ -1,12 +1,24 @@
import inspect
-from functools import wraps
-from typing import Any, Callable, Mapping, Sequence, Set, TypeVar, Union, cast
+from functools import lru_cache, wraps
+from typing import Any, Callable, Mapping, Sequence, Set, Type, TypeVar, Union, cast
from pydantic import BaseModel
F = TypeVar("F", bound=Callable)
+@lru_cache(maxsize=None)
+def _hashable_model_type(model_type: Type[BaseModel]) -> Type[BaseModel]:
+ # The subclass is cached per model type: Pydantic v2's __eq__ requires both
+ # operands to share a type, so creating a fresh subclass on every freeze()
+ # call would make even identical models compare unequal.
+ class HashableValue(model_type): # type: ignore
+ def __hash__(self) -> int:
+ return hash(frozenset((k, freeze(v)) for k, v in self.model_dump().items()))
+
+ return HashableValue
+
+
class FrozenDict(dict):
def __hash__(self) -> int: # type: ignore
return hash(frozenset((k, freeze(v)) for k, v in self.items()))
@@ -55,11 +67,6 @@ def freeze(value: Union[Sequence[Any], Mapping[str, Any], Set[Any], BaseModel])
return FrozenSet(value)
if isinstance(value, BaseModel):
-
- class HashableValue(type(value)): # type: ignore
- def __hash__(self) -> int:
- return hash(frozenset((k, freeze(v)) for k, v in self.dict().items()))
-
- return HashableValue(**value.dict())
+ return _hashable_model_type(type(value))(**value.model_dump())
return value
diff --git a/great_ai/large_file/helper/__init__.py b/great_ai/large_file/helper/__init__.py
index 9639296..5eb33e8 100644
--- a/great_ai/large_file/helper/__init__.py
+++ b/great_ai/large_file/helper/__init__.py
@@ -1,2 +1,3 @@
+from .cached_property import cached_property
from .human_readable_to_byte import human_readable_to_byte
from .progress_bar import DownloadProgressBar, UploadProgressBar
diff --git a/great_ai/large_file/helper/cached_property.py b/great_ai/large_file/helper/cached_property.py
new file mode 100644
index 0000000..d141d3c
--- /dev/null
+++ b/great_ai/large_file/helper/cached_property.py
@@ -0,0 +1,54 @@
+from threading import RLock
+
+_NOT_FOUND = object()
+
+
+class cached_property:
+ def __init__(self, func): # type: ignore
+ self.func = func
+ self.attrname = None
+ self.__doc__ = func.__doc__
+ self.lock = RLock()
+
+ def __set_name__(self, owner, name): # type: ignore
+ if self.attrname is None:
+ self.attrname = name
+ elif name != self.attrname:
+ raise TypeError(
+ "Cannot assign the same cached_property to two different names "
+ f"({self.attrname!r} and {name!r})."
+ )
+
+ def __get__(self, instance, owner=None): # type: ignore
+ if instance is None:
+ return self
+ if self.attrname is None:
+ raise TypeError(
+ "Cannot use cached_property instance without calling __set_name__ on it."
+ )
+ try:
+ cache = instance.__dict__
+ except (
+ AttributeError
+ ): # not all objects have __dict__ (e.g. class defines slots)
+ msg = (
+ f"No '__dict__' attribute on {type(instance).__name__!r} "
+ f"instance to cache {self.attrname!r} property."
+ )
+ raise TypeError(msg) from None
+ val = cache.get(self.attrname, _NOT_FOUND)
+ if val is _NOT_FOUND:
+ with self.lock:
+ # check if another thread filled cache while we awaited lock
+ val = cache.get(self.attrname, _NOT_FOUND)
+ if val is _NOT_FOUND:
+ val = self.func(instance)
+ try:
+ cache[self.attrname] = val
+ except TypeError:
+ msg = (
+ f"The '__dict__' attribute on {type(instance).__name__!r} instance "
+ f"does not support item assignment for caching {self.attrname!r} property."
+ )
+ raise TypeError(msg) from None
+ return val
diff --git a/great_ai/large_file/large_file/large_file_base.py b/great_ai/large_file/large_file/large_file_base.py
index 4f412e1..92fa5d9 100644
--- a/great_ai/large_file/large_file/large_file_base.py
+++ b/great_ai/large_file/large_file/large_file_base.py
@@ -1,5 +1,7 @@
import os
+import re
import shutil
+import sys
import tempfile
from abc import ABC, abstractmethod
from functools import lru_cache
@@ -21,31 +23,26 @@ ARCHIVE_EXTENSION = ".tar.gz"
class LargeFileBase(ABC):
- """
- Store large files remotely. Use local cache for speed up.
+ """Base for LargeFile implementations with different backends.
+
+ Store large files remotely using the familiar API of `open()`. With built-in
+ versioning, pruning and local cache.
+
+ By default, files are stored in the ".cache" folder and the least recently used is
+ deleted after the overall size reaches 30 GBs.
Examples:
+ >>> LargeFileBase.cache_path = Path(".cache")
- ```
- with LargeFile("test.txt", "w", keep_last_n=3) as f:
- for i in range(1000000):
- f.write('test\n')
+ >>> LargeFileBase.max_cache_size = "30GB"
- with LargeFile("test.txt", "r") as f:
- print(f.readlines()[0])
-
- path_to_cached_text_file = LargeFile("test.txt", version=0).get()
- ```
-
- By default, files are stored in the ".cache" folder and the
- least recently use is deleted after the overall size reaches 30 GBs.
-
- Change it with the following properties.
-
- ```
- LargeFile.cache_path = Path(".cache")
- LargeFile.max_cache_size = "30GB"
- ```
+ Attributes:
+ initialized: Tell whether `configure_credentials` or
+ `configure_credentials_from_file` has been already called.
+ cache_path: Storage location for cached files.
+ max_cache_size: Delete files until the folder at `cache_path` is smaller than
+ this value. Examples: "5 GB", "10MB", "0.3 TB". Set to `None` for no
+ automatic cache-pruning.
"""
initialized = False
@@ -65,6 +62,12 @@ class LargeFileBase(ABC):
keep_last_n: Optional[int] = None,
cache_only_mode: bool = False,
):
+ clean_name = re.sub(r"[^!a-zA-Z0-9._-]+", "", name)
+ if clean_name != name:
+ raise ValueError(
+ f"Given name contains illegal characters, consider changing it to: `{clean_name}`"
+ )
+
self._name = name
self._version = version
self._mode = mode
@@ -73,7 +76,13 @@ class LargeFileBase(ABC):
self._buffering = buffering
self._encoding = encoding
+
+ if errors is not None and sys.version_info[1] < 8:
+ raise RuntimeError(
+ "The `errors` kwarg is only supported in 3.8 <= Python versions."
+ )
self._errors = errors
+
self._newline = newline
LargeFileBase.cache_path.mkdir(parents=True, exist_ok=True)
@@ -86,25 +95,36 @@ class LargeFileBase(ABC):
cls,
secrets: Union[Path, str, ConfigFile],
) -> None:
+ """Load file and feed its content to `configure_credentials`.
+
+ Extra keys are ignored.
+ """
+
if not isinstance(secrets, ConfigFile):
secrets = ConfigFile(secrets)
cls.configure_credentials(**{k.lower(): v for k, v in secrets.items()})
@classmethod
def configure_credentials(cls, **kwargs: str) -> None:
+ """Configure required credentials for the LargeFile backend."""
+
cls.initialized = True
def __enter__(self) -> IO:
+ params = dict(
+ mode=self._mode,
+ buffering=self._buffering,
+ encoding=self._encoding,
+ newline=self._newline,
+ delete=False,
+ prefix="large_file-",
+ )
+
+ if sys.version_info[1] >= 8:
+ params["errors"] = self._errors
+
self._file: IO[Any] = (
- tempfile.NamedTemporaryFile(
- mode=self._mode,
- buffering=self._buffering,
- encoding=self._encoding,
- newline=self._newline,
- errors=self._errors,
- delete=False,
- prefix="large_file-",
- )
+ tempfile.NamedTemporaryFile(**params) # type: ignore
if "w" in self._mode
else open(
self.get(),
@@ -135,16 +155,22 @@ class LargeFileBase(ABC):
return False
- @property
- def _local_name(self) -> str:
- return f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}{self.version}"
-
@property
def version(self) -> int:
+ """Numeric version of the file proxied by this LargeFile instance."""
+
return cast(int, self._version)
@lru_cache(1)
def get(self, hide_progress: bool = False) -> Path:
+ """Return path to the proxy of a file (or directory).
+
+ If not available in the local cache, an attempt is made to download it.
+
+ Args:
+ hide_progress: Do not show a progress update after each 10% of progress.
+ """
+
remote_path = next(
i.remote_path for i in self._instances if i.version == self._version
)
@@ -171,6 +197,14 @@ class LargeFileBase(ABC):
return destination
def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
+ """Upload a file (or directory) as a new version of `key`.
+
+ The file/directory is compressed before upload.
+
+ Args:
+ hide_progress: Do not show a progress update after each 10% of progress.
+ """
+
if isinstance(path, str):
path = Path(path)
@@ -208,9 +242,26 @@ class LargeFileBase(ABC):
self.clean_up()
def delete(self) -> None:
+ """Delete all versions of the files under this `key`."""
+
self._keep_last_n = 0
self._delete_old_remote_versions()
+ @property
+ def versions_pretty(self) -> str:
+ """Formatted string of all available versions."""
+ return ", ".join((str(i.version) for i in self._instances))
+
+ def clean_up(self) -> None:
+ """Delete local and remote versions according to currently set cache and retention policy."""
+
+ self._delete_old_remote_versions()
+ self._prune_cache()
+
+ @property
+ def _local_name(self) -> str:
+ return f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}{self.version}"
+
def _find_instances(self) -> None:
if self._cache_only_mode:
self._instances = self._find_instances_from_cache()
@@ -269,14 +320,6 @@ class LargeFileBase(ABC):
else:
raise ValueError("Unsupported file mode.")
- @property
- def versions_pretty(self) -> str:
- return ", ".join((str(i.version) for i in self._instances))
-
- def clean_up(self) -> None:
- self._delete_old_remote_versions()
- self._prune_cache()
-
def _prune_cache(self) -> None:
self.cache_path.mkdir(parents=True, exist_ok=True)
diff --git a/great_ai/large_file/large_file/large_file_local.py b/great_ai/large_file/large_file/large_file_local.py
index 6b31762..c1c5a12 100644
--- a/great_ai/large_file/large_file/large_file_local.py
+++ b/great_ai/large_file/large_file/large_file_local.py
@@ -9,6 +9,30 @@ logger = get_logger("large_file")
class LargeFileLocal(LargeFileBase):
+ """LargeFile implementation using local filesystem as a backend.
+
+ Store large files remotely using the familiar API of `open()`. With built-in
+ versioning, pruning and local cache.
+
+ IMPORTANT: If LargeFileLocal.max_cache_size is too small, it won't be enough to
+ store all your files and they can end up deleted.
+
+ See parent for more details.
+
+ Examples:
+ >>> LargeFileLocal.cache_path = Path(".cache")
+
+ >>> LargeFileLocal.max_cache_size = "30GB"
+
+ >>> with LargeFileLocal("my_test.txt", "w", keep_last_n=2) as f:
+ ... f.write('test')
+ 4
+
+ >>> with LargeFileLocal("my_test.txt") as f:
+ ... print(f.read())
+ test
+ """
+
def __init__(
self,
name: str,
@@ -40,6 +64,7 @@ class LargeFileLocal(LargeFileBase):
def _download(
self, remote_path: Any, local_path: Path, hide_progress: bool
) -> None:
+ # This will never be called because the file must be in the cache
raise NotImplementedError()
def _upload(self, local_path: Path, hide_progress: bool) -> None:
diff --git a/great_ai/large_file/large_file/large_file_mongo.py b/great_ai/large_file/large_file/large_file_mongo.py
index 673ac5c..0376ce9 100644
--- a/great_ai/large_file/large_file/large_file_mongo.py
+++ b/great_ai/large_file/large_file/large_file_mongo.py
@@ -1,13 +1,13 @@
import re
-from functools import cached_property
from pathlib import Path
from typing import Any, List
-from gridfs import DEFAULT_CHUNK_SIZE, Database, GridFSBucket
+from gridfs import DEFAULT_CHUNK_SIZE, GridFSBucket
from pymongo import MongoClient
+from pymongo.database import Database
from ...utilities import get_logger
-from ..helper import DownloadProgressBar, UploadProgressBar
+from ..helper import DownloadProgressBar, UploadProgressBar, cached_property
from ..models import DataInstance
from .large_file_base import LargeFileBase
@@ -18,6 +18,14 @@ MONGO_NAME_VERSION_SEPARATOR = "_"
class LargeFileMongo(LargeFileBase):
+ """LargeFile implementation using GridFS (MongoDB) as a backend.
+
+ Store large files remotely using the familiar API of `open()`. With built-in
+ versioning, pruning and local cache.
+
+ See parent for more details.
+ """
+
mongo_connection_string = None
mongo_database = None
@@ -37,7 +45,7 @@ class LargeFileMongo(LargeFileBase):
def _client(self) -> GridFSBucket:
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."
+ "Please configure the MongoDB access options by calling LargeFileMongo.configure_credentials or set cache_only_mode=True in the constructor."
)
db: Database = MongoClient(self.mongo_connection_string)[self.mongo_database]
@@ -53,7 +61,6 @@ class LargeFileMongo(LargeFileBase):
),
version=int(f.name.split(MONGO_NAME_VERSION_SEPARATOR)[-1]),
remote_path=(f._id, f.length),
- origin="mongodb",
)
for f in self._client.find(
{
diff --git a/great_ai/large_file/large_file/large_file_s3.py b/great_ai/large_file/large_file/large_file_s3.py
index 0104eba..3e32a59 100644
--- a/great_ai/large_file/large_file/large_file_s3.py
+++ b/great_ai/large_file/large_file/large_file_s3.py
@@ -1,11 +1,10 @@
-from functools import cached_property
from pathlib import Path
from typing import Any, List, Optional
import boto3
from ...utilities import get_logger
-from ..helper import DownloadProgressBar, UploadProgressBar
+from ..helper import DownloadProgressBar, UploadProgressBar, cached_property
from ..models import DataInstance
from .large_file_base import LargeFileBase
@@ -16,31 +15,12 @@ S3_NAME_VERSION_SEPARATOR = "/"
class LargeFileS3(LargeFileBase):
- """
- Store large files in S3. Use local cache for speed up.
+ """LargeFile implementation using S3-compatible storage as a backend.
- Examples:
+ Store large files remotely using the familiar API of `open()`. With built-in
+ versioning, pruning and local cache.
- ```
- with LargeFile("test.txt", "w", keep_last_n=3) as f:
- for i in range(1000000):
- f.write('test\n')
-
- with LargeFile("test.txt", "r") as f:
- print(f.readlines()[0])
-
- path_to_cached_text_file = LargeFile("test.txt", version=0).get()
- ```
-
- By default, files are stored in the ".cache" folder and the
- least recently use is deleted after the overall size reaches 30 GBs.
-
- Change it with the following properties.
-
- ```
- LargeFile.cache_path = Path(".cache")
- LargeFile.max_cache_size = "30GB"
- ```
+ See parent for more details.
"""
region_name = None
@@ -76,7 +56,7 @@ class LargeFileS3(LargeFileBase):
or self.bucket_name is None
):
raise ValueError(
- "Please configure the S3 access options by calling LargeFileS3.configure_credentials or set offline_mode=True in the constructor."
+ "Please configure the S3 access options by calling LargeFileS3.configure_credentials or set cache_only_mode=True in the constructor."
)
return boto3.client(
@@ -120,9 +100,13 @@ class LargeFileS3(LargeFileBase):
Bucket=self.bucket_name,
Key=remote_path,
Filename=str(local_path),
- Callback=None
- if hide_progress
- else DownloadProgressBar(name=str(remote_path), size=size, logger=logger),
+ Callback=(
+ None
+ if hide_progress
+ else DownloadProgressBar(
+ name=str(remote_path), size=size, logger=logger
+ )
+ ),
)
def _upload(self, local_path: Path, hide_progress: bool) -> None:
@@ -133,9 +117,11 @@ class LargeFileS3(LargeFileBase):
Filename=str(local_path),
Bucket=self.bucket_name,
Key=key,
- Callback=None
- if hide_progress
- else UploadProgressBar(path=local_path, logger=logger),
+ Callback=(
+ None
+ if hide_progress
+ else UploadProgressBar(path=local_path, logger=logger)
+ ),
)
def _delete_old_remote_versions(self) -> None:
diff --git a/great_ai/models/use_model.py b/great_ai/models/use_model.py
index f5e90c1..e18eff8 100644
--- a/great_ai/models/use_model.py
+++ b/great_ai/models/use_model.py
@@ -50,12 +50,12 @@ def use_model(
>>> my_function(4)
7
- Args:
- key: The model's name as stored by the LargeFile implementation.
- version: The model's version as stored by the LargeFile implementation.
- model_kwarg_name: the parameter to use for injecting the loaded model
- Returns:
- A decorator for model injection.
+ Args:
+ key: The model's name as stored by the LargeFile implementation.
+ version: The model's version as stored by the LargeFile implementation.
+ model_kwarg_name: the parameter to use for injecting the loaded model
+ Returns:
+ A decorator for model injection.
"""
assert (
diff --git a/great_ai/parameters/log_metric.py b/great_ai/parameters/log_metric.py
index 30a1f4c..f307573 100644
--- a/great_ai/parameters/log_metric.py
+++ b/great_ai/parameters/log_metric.py
@@ -5,11 +5,27 @@ from ..context import get_context
from ..tracing import TracingContext
-def log_metric(argument_name: str, value: Any) -> None:
+def log_metric(argument_name: str, value: Any, disable_logging: bool = False) -> None:
+ """Log a key (argument_name)-value pair that is persisted inside the trace.
+
+ The name of the function from where this is called is also stored.
+
+ Args:
+ argument_name: The key for storing the value.
+ value: Value to log. Must be JSON-serialisable.
+ disable_logging: If True, only persist in trace but don't show in console
+ """
+
tracing_context = TracingContext.get_current_tracing_context()
- caller = inspect.stack()[1].function
- actual_name = f"metric:{caller}:{argument_name}"
+ try:
+ caller = inspect.stack()[1].function
+ actual_name = f"metric:{caller}:{argument_name}"
+ except:
+ # inspect might not work in notebooks
+ actual_name = f"metric:{argument_name}"
+
if tracing_context:
tracing_context.log_value(name=actual_name, value=value)
- get_context().logger.info(f"{actual_name}={value}")
+ if not disable_logging:
+ get_context().logger.info(f"{actual_name}={value}")
diff --git a/great_ai/parameters/parameter.py b/great_ai/parameters/parameter.py
index 43696d3..df3a7d4 100644
--- a/great_ai/parameters/parameter.py
+++ b/great_ai/parameters/parameter.py
@@ -1,7 +1,7 @@
from functools import wraps
from typing import Any, Callable, Dict, TypeVar, cast
-from typeguard import check_type
+from typeguard import TypeCheckError, check_type
from ..errors import ArgumentValidationError
from ..helper import get_arguments, get_function_metadata_store
@@ -19,6 +19,10 @@ def parameter(
) -> Callable[[F], F]:
"""Control the validation and logging of function parameters.
+ Basically, a parameter decorator. Unfortunately, Python does not have that concept,
+ thus, it's a method decorator that expects the name of the to-be-decorated
+ parameter.
+
Examples:
>>> @parameter('a')
... def my_function(a: int):
@@ -28,7 +32,7 @@ def parameter(
>>> my_function('3')
Traceback (most recent call last):
...
- TypeError: type of a must be int; got str instead
+ TypeError: argument a is not of the expected type: ...
>>> @parameter('positive_a', validate=lambda v: v > 0)
... def my_function(positive_a: int):
@@ -39,7 +43,7 @@ def parameter(
great_ai.errors.argument_validation_error.ArgumentValidationError: ...
Args:
- parameter_name: Name of parameter to consider
+ parameter_name: Name of parameter to consider.
validate: Optional validate to run against the concrete argument.
ArgumentValidationError is thrown when the return value is False.
disable_logging: Do not save the value in any active TracingContext.
@@ -61,13 +65,16 @@ def parameter(
expected_type = func.__annotations__.get(parameter_name)
if expected_type is not None:
- check_type(parameter_name, argument, expected_type)
+ try:
+ check_type(argument, expected_type)
+ except TypeCheckError as error:
+ raise TypeError(
+ f"argument {parameter_name} is not of the expected type: {error}"
+ ) from error
if not validate(argument):
raise ArgumentValidationError(
- f"""Argument {parameter_name} in {
- func.__name__
- } did not pass validation"""
+ f"Argument {parameter_name} in {func.__name__} did not pass validation"
)
context = TracingContext.get_current_tracing_context()
diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py
index f00a6ec..506d74d 100644
--- a/great_ai/persistence/mongodb_driver.py
+++ b/great_ai/persistence/mongodb_driver.py
@@ -90,7 +90,7 @@ class MongoDbDriver(TracingDatabaseDriver):
value = client[self.mongo_database].traces.find_one(id)
if value:
- value = Trace.parse_obj(value)
+ value = Trace.model_validate(value)
return value
@@ -150,7 +150,7 @@ class MongoDbDriver(TracingDatabaseDriver):
]
with client[self.mongo_database].traces.find(**query) as cursor:
- documents = [Trace[Any].parse_obj(t) for t in cursor]
+ documents = [Trace[Any].model_validate(t) for t in cursor]
return documents, count
def update(self, id: str, new_version: Trace) -> None:
diff --git a/great_ai/persistence/parallel_tinydb_driver.py b/great_ai/persistence/parallel_tinydb_driver.py
index 92b5fc8..6e278a1 100644
--- a/great_ai/persistence/parallel_tinydb_driver.py
+++ b/great_ai/persistence/parallel_tinydb_driver.py
@@ -29,16 +29,16 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
path_to_db = Path(DEFAULT_TRACING_DB_FILENAME)
def save(self, trace: Trace) -> str:
- return self._safe_execute(lambda db: db.insert(trace.dict()))
+ return self._safe_execute(lambda db: db.insert(trace.model_dump()))
def save_batch(self, documents: List[Trace]) -> List[str]:
- traces = [d.dict() for d in documents]
+ traces = [d.model_dump() for d in documents]
return self._safe_execute(lambda db: db.insert_multiple(traces))
def get(self, id: str) -> Optional[Trace]:
value = self._safe_execute(lambda db: db.get(lambda d: d["trace_id"] == id))
if value:
- value = Trace.parse_obj(value)
+ value = Trace.model_validate(value)
return value
def query(
@@ -51,7 +51,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
since: Optional[datetime] = None,
until: Optional[datetime] = None,
has_feedback: Optional[bool] = None,
- sort_by: Sequence[SortBy] = []
+ sort_by: Sequence[SortBy] = [],
) -> Tuple[List[Trace], int]:
def does_match(d: Dict[str, Any]) -> bool:
return (
@@ -63,15 +63,15 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
)
)
- documents: List[Trace] = [
- Trace.parse_obj(t)
- for t in self._safe_execute(lambda db: db.search(does_match))
- ]
-
+ documents = self._safe_execute(lambda db: db.search(does_match))
if not documents:
return [], 0
- df = pd.DataFrame([d.to_flat_dict() for d in documents])
+ traces: List[Trace] = [Trace.model_validate(d) for d in documents]
+ # The DataFrame keeps its default 0..n-1 index, so filtering/sorting below
+ # preserve the row labels that index back into `traces`. This avoids
+ # reconstructing Traces from the lossy flattened rows.
+ df = pd.DataFrame([t.to_flat_dict() for t in traces])
for f in conjunctive_filters:
operator = f.operator.lower()
@@ -80,7 +80,12 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
getattr(df[f.property], operator_mapping[f.operator])(f.value)
]
elif operator == "contains":
- df = df.loc[df[f.property].str.contains(f.value, case=False)]
+ df = df.loc[
+ df[f.property].str.contains(
+ str(int(f.value)) if isinstance(f.value, float) else f.value,
+ case=False,
+ )
+ ]
if sort_by:
df.sort_values(
@@ -91,22 +96,23 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
count = len(df)
result = df.iloc[skip:] if take is None else df.iloc[skip : skip + take]
- return [
- next(d for d in documents if d.trace_id == trace_id)
- for trace_id in result["trace_id"]
- ], count
+ return [traces[i] for i in result.index], count
def update(self, id: str, new_version: Trace) -> None:
self._safe_execute(
- lambda db: db.update(new_version.dict(), lambda d: d["trace_id"] == id)
+ lambda db: db.update(
+ new_version.model_dump(), lambda d: d["trace_id"] == id
+ )
)
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]) -> None:
- for i in ids:
- self.delete(i)
+ with lock:
+ with TinyDB(self.path_to_db) as db:
+ for id in ids:
+ db.remove(lambda d: d["trace_id"] == id)
def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any:
with lock:
diff --git a/great_ai/persistence/tracing_database_driver.py b/great_ai/persistence/tracing_database_driver.py
index 810aea4..3491237 100644
--- a/great_ai/persistence/tracing_database_driver.py
+++ b/great_ai/persistence/tracing_database_driver.py
@@ -54,7 +54,7 @@ class TracingDatabaseDriver(ABC):
until: Optional[datetime] = None,
since: Optional[datetime] = None,
has_feedback: Optional[bool] = None,
- sort_by: Sequence[SortBy] = []
+ sort_by: Sequence[SortBy] = [],
) -> Tuple[List[Trace], int]:
pass
diff --git a/great_ai/remote/call_remote_great_ai_async.py b/great_ai/remote/call_remote_great_ai_async.py
index 751ddf3..eac83ee 100644
--- a/great_ai/remote/call_remote_great_ai_async.py
+++ b/great_ai/remote/call_remote_great_ai_async.py
@@ -28,7 +28,7 @@ async def call_remote_great_ai_async(
base_uri: Address of the remote instance, example: 'http://localhost:6060'
data: The input sent as a json to the remote instance.
retry_count: Retry on any HTTP communication failure.
- timeout_in_seconds: Overall permissable max length of the request. `None` means
+ timeout_in_seconds: Overall permissible max length of the request. `None` means
no timeout.
model_class: A subtype of BaseModel to be used for deserialising the `.output`
of the trace.
@@ -64,7 +64,7 @@ async def call_remote_great_ai_async(
)
try:
if model_class is not None:
- trace["output"] = model_class.parse_obj(trace["output"])
- return Trace.parse_obj(trace)
+ trace["output"] = model_class.model_validate(trace["output"])
+ return Trace.model_validate(trace)
except Exception:
raise RemoteCallError("Could not parse Trace")
diff --git a/great_ai/tracing/add_ground_truth.py b/great_ai/tracing/add_ground_truth.py
index 4080db2..15d87a9 100644
--- a/great_ai/tracing/add_ground_truth.py
+++ b/great_ai/tracing/add_ground_truth.py
@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timezone
from math import ceil
from random import shuffle
from typing import Any, Iterable, List, TypeVar, Union, cast
@@ -23,7 +23,7 @@ def add_ground_truth(
tags: Union[List[str], str] = [],
train_split_ratio: float = 1,
test_split_ratio: float = 0,
- validation_split_ratio: float = 0
+ validation_split_ratio: float = 0,
) -> None:
"""Add training data (with optional train-test splitting).
@@ -94,7 +94,7 @@ def add_ground_truth(
)
shuffle(split_tags)
- created = datetime.utcnow().isoformat()
+ created = datetime.now(timezone.utc).replace(tzinfo=None).isoformat()
traces = [
cast(
Trace[T],
diff --git a/great_ai/tracing/query_ground_truth.py b/great_ai/tracing/query_ground_truth.py
index bdbf9df..d30739f 100644
--- a/great_ai/tracing/query_ground_truth.py
+++ b/great_ai/tracing/query_ground_truth.py
@@ -10,13 +10,13 @@ def query_ground_truth(
*,
since: Optional[datetime] = None,
until: Optional[datetime] = None,
- return_max_count: Optional[int] = None
+ return_max_count: Optional[int] = None,
) -> List[Trace]:
"""Return training samples.
Combines, filters, and returns data-points that have been either added by
`add_ground_truth` or were the result of a prediction after which their trace got
- a feedback through the RESP API-s `/traces/{trace_id}/feedback` endpoint
+ feedback through the RESP API-s `/traces/{trace_id}/feedback` endpoint
(end-to-end feedback).
Filtering can be used to only return points matching all given tags (or the single
diff --git a/great_ai/tracing/tracing_context.py b/great_ai/tracing/tracing_context.py
index 28dc52b..9180061 100644
--- a/great_ai/tracing/tracing_context.py
+++ b/great_ai/tracing/tracing_context.py
@@ -1,5 +1,5 @@
from contextvars import ContextVar
-from datetime import datetime
+from datetime import datetime, timezone
from time import perf_counter
from types import TracebackType
from typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar, cast
@@ -23,7 +23,7 @@ class TracingContext(Generic[T]):
self._models: List[Model] = []
self._values: Dict[str, Any] = {}
self._trace: Optional[Trace[T]] = None
- self._start_datetime = datetime.utcnow()
+ self._start_datetime = datetime.now(timezone.utc).replace(tzinfo=None)
self._start_time = perf_counter()
self._name = function_name
@@ -49,15 +49,19 @@ class TracingContext(Generic[T]):
logged_values=self._values,
models=self._models,
output=output,
- exception=None
- if exception is None
- else f"{type(exception).__name__}: {exception}",
+ exception=(
+ None
+ if exception is None
+ else f"{type(exception).__name__}: {exception}"
+ ),
tags=[
self._name,
ONLINE_TAG_NAME,
- PRODUCTION_TAG_NAME
- if get_context().is_production
- else DEVELOPMENT_TAG_NAME,
+ (
+ PRODUCTION_TAG_NAME
+ if get_context().is_production
+ else DEVELOPMENT_TAG_NAME
+ ),
],
),
)
@@ -98,5 +102,5 @@ class TracingContext(Generic[T]):
_current_tracing_context: ContextVar[Optional[TracingContext]] = ContextVar(
- "_current_tracing_context"
+ "_current_tracing_context", default=None
)
diff --git a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py b/great_ai/utilities/evaluate_ranking/evaluate_ranking.py
index 38b3a92..a613c93 100644
--- a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py
+++ b/great_ai/utilities/evaluate_ranking/evaluate_ranking.py
@@ -4,7 +4,7 @@ from typing import Dict, List, Optional, TypeVar
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
-from sklearn.metrics import average_precision_score, precision_recall_curve
+from sklearn.metrics import average_precision_score
from ..unique import unique
from .draw_f1_iso_lines import draw_f1_iso_lines
@@ -34,7 +34,7 @@ def evaluate_ranking(
`expected`)
title: Title of the plot.
disable_interpolation: Do not interpolate.
- axes: Matplotlib axes for ploting inside a subplot.
+ axes: Matplotlib axes for plotting inside a subplot.
output_svg: If specified, save the chart as an svg to the given Path.
reverse_order: Reverse the ranking specified by `expected`.
plot: Display a plot on the screen.
@@ -46,11 +46,11 @@ def evaluate_ranking(
assert 0 <= target_recall <= 1
if plot and axes is None:
- fig = plt.figure(figsize=(20, 10))
+ fig = plt.figure(figsize=(10, 10))
fig.patch.set_facecolor("white")
ax = plt.axes()
else:
- ax = axes
+ ax = axes # type: ignore[assignment]
classes = sorted(unique(expected), reverse=reverse_order)
str_classes = [str(c) for c in classes]
@@ -67,15 +67,23 @@ def evaluate_ranking(
(v < classes[i]) if reverse_order else (v > classes[i])
for v in expected
]
- precision, recall, _ = precision_recall_curve(
- binarized_expected, actual_scores
+
+ sorted_expected_actual = sorted(
+ zip(binarized_expected, actual_scores), key=lambda v: v[1], reverse=True
)
+ precision = []
+ recall = []
+ correct = 0
+ for all, (e, score) in enumerate(sorted_expected_actual, start=1):
+ correct += int(e)
+ precision.append(correct / all)
+ recall.append(all / len(sorted_expected_actual))
if not disable_interpolation:
- for j in range(1, len(precision)):
- precision[j] = max(precision[j - 1], precision[j])
+ for j in range(len(precision) - 2, -1, -1):
+ precision[j] = max(precision[j], precision[j + 1])
- closest_recall_index = np.argmin(np.abs(recall - target_recall))
+ closest_recall_index = np.argmin(np.abs(np.array(recall) - target_recall))
precision_at_closest_recall = precision[closest_recall_index]
average_precision = average_precision_score(
binarized_expected, actual_scores
diff --git a/great_ai/utilities/external/pylatexenc/latex2text/__init__.py b/great_ai/utilities/external/pylatexenc/latex2text/__init__.py
index d8ae865..41b23b0 100644
--- a/great_ai/utilities/external/pylatexenc/latex2text/__init__.py
+++ b/great_ai/utilities/external/pylatexenc/latex2text/__init__.py
@@ -1549,7 +1549,7 @@ def latex2text(
"in favor of the `pylatexenc.latex2text.LatexNodes2Text` class.",
)
- (nodelist, tpos, tlen) = latexwalker.get_latex_nodes(
+ nodelist, tpos, tlen = latexwalker.get_latex_nodes(
content, keep_inline_math=keep_inline_math, tolerant_parsing=tolerant_parsing
)
diff --git a/great_ai/utilities/external/pylatexenc/latex2text/__main__.py b/great_ai/utilities/external/pylatexenc/latex2text/__main__.py
index a19f9b6..4f9808f 100644
--- a/great_ai/utilities/external/pylatexenc/latex2text/__main__.py
+++ b/great_ai/utilities/external/pylatexenc/latex2text/__main__.py
@@ -289,7 +289,7 @@ def main(argv=None):
latex, tolerant_parsing=args.tolerant_parsing, strict_braces=args.strict_braces
)
- (nodelist, pos, len_) = lw.get_latex_nodes()
+ nodelist, pos, len_ = lw.get_latex_nodes()
ln2t = LatexNodes2Text(
math_mode=args.math_mode,
diff --git a/great_ai/utilities/external/pylatexenc/latex2text/_defaultspecs.py b/great_ai/utilities/external/pylatexenc/latex2text/_defaultspecs.py
index e4a3e76..86e3ef0 100644
--- a/great_ai/utilities/external/pylatexenc/latex2text/_defaultspecs.py
+++ b/great_ai/utilities/external/pylatexenc/latex2text/_defaultspecs.py
@@ -1774,7 +1774,7 @@ def make_accented_char(node, combining, l2tobj):
for u in unicode_accents_list:
- (mname, mcombining) = u
+ mname, mcombining = u
_latex_specs_base["macros"].append(
MacroTextSpec(
mname, lambda x, l2tobj, c=mcombining: make_accented_char(x, c, l2tobj)
diff --git a/great_ai/utilities/external/pylatexenc/latexencode/_partial_latex_encoder.py b/great_ai/utilities/external/pylatexenc/latexencode/_partial_latex_encoder.py
index 393551d..6162ad6 100644
--- a/great_ai/utilities/external/pylatexenc/latexencode/_partial_latex_encoder.py
+++ b/great_ai/utilities/external/pylatexenc/latexencode/_partial_latex_encoder.py
@@ -71,7 +71,7 @@ class PartialLatexToLatexEncoder(UnicodeToLatexEncoder):
# keyword arguments:
keep_latex_chars=r"\${}^_",
conversion_rules=None,
- **kwargs
+ **kwargs,
):
base_conversion_rules = conversion_rules
@@ -89,7 +89,7 @@ class PartialLatexToLatexEncoder(UnicodeToLatexEncoder):
)
]
+ base_conversion_rules,
- **kwargs
+ **kwargs,
)
self.keep_latex_chars = keep_latex_chars
diff --git a/great_ai/utilities/external/pylatexenc/latexencode/_unicode_to_latex_encoder.py b/great_ai/utilities/external/pylatexenc/latexencode/_unicode_to_latex_encoder.py
index fbc5a47..af54052 100644
--- a/great_ai/utilities/external/pylatexenc/latexencode/_unicode_to_latex_encoder.py
+++ b/great_ai/utilities/external/pylatexenc/latexencode/_unicode_to_latex_encoder.py
@@ -617,7 +617,7 @@ class UnicodeToLatexEncoder(object):
res = rulecallable(s, p.pos)
if res is None:
return None
- (consumed, repl) = res
+ consumed, repl = res
self._apply_replacement(p, repl, consumed, rule)
return True
diff --git a/great_ai/utilities/external/pylatexenc/latexwalker/__init__.py b/great_ai/utilities/external/pylatexenc/latexwalker/__init__.py
index b60bd69..60685e9 100644
--- a/great_ai/utilities/external/pylatexenc/latexwalker/__init__.py
+++ b/great_ai/utilities/external/pylatexenc/latexwalker/__init__.py
@@ -477,7 +477,7 @@ class LatexNode(object):
parsing_state=None,
pos=None,
len=None,
- **kwargs
+ **kwargs,
):
# Important: subclasses must specify a list of fields they set in the
@@ -607,7 +607,7 @@ class LatexGroupNode(LatexNode):
"nodelist",
"delimiters",
),
- **kwargs
+ **kwargs,
)
self.nodelist = nodelist
self.delimiters = delimiters
@@ -639,7 +639,7 @@ class LatexCommentNode(LatexNode):
"comment",
"comment_post_space",
),
- **kwargs
+ **kwargs,
)
self.comment = comment
@@ -715,7 +715,7 @@ class LatexMacroNode(LatexNode):
super(LatexMacroNode, self).__init__(
_fields=("macroname", "nodeargd", "macro_post_space"),
_redundant_fields=("nodeoptarg", "nodeargs"),
- **kwargs
+ **kwargs,
)
self.macroname = macroname
@@ -801,7 +801,7 @@ class LatexEnvironmentNode(LatexNode):
"optargs",
"args",
),
- **kwargs
+ **kwargs,
)
self.environmentname = environmentname
@@ -1264,7 +1264,7 @@ class LatexWalker(object):
environments=True,
keep_inline_math=None,
parsing_state=None,
- **kwargs
+ **kwargs,
):
r"""
Parses the latex content given to the constructor (and stored in `self.s`),
@@ -1628,7 +1628,7 @@ class LatexWalker(object):
r"Expected expression, got \end",
self.s,
pos,
- **self.pos_to_lineno_colno(pos, as_dict=True)
+ **self.pos_to_lineno_colno(pos, as_dict=True),
)
else:
return self._mknodeposlen(
@@ -1674,7 +1674,7 @@ class LatexWalker(object):
"Expected expression, got closing brace '{}'".format(tok.arg),
self.s,
pos,
- **self.pos_to_lineno_colno(pos, as_dict=True)
+ **self.pos_to_lineno_colno(pos, as_dict=True),
)
return self._mknodeposlen(
LatexCharsNode,
@@ -1717,7 +1717,7 @@ class LatexWalker(object):
"Unknown token type: {}".format(tok.tok),
self.s,
pos,
- **self.pos_to_lineno_colno(pos, as_dict=True)
+ **self.pos_to_lineno_colno(pos, as_dict=True),
)
def get_latex_maybe_optional_arg(self, pos, parsing_state=None):
@@ -1821,10 +1821,10 @@ class LatexWalker(object):
pos=pos,
msg="get_latex_braced_group: not an opening brace/bracket: %s"
% (self.s[pos]),
- **self.pos_to_lineno_colno(pos, as_dict=True)
+ **self.pos_to_lineno_colno(pos, as_dict=True),
)
- (nodelist, npos, nlen) = self.get_latex_nodes(
+ nodelist, npos, nlen = self.get_latex_nodes(
firsttok.pos + firsttok.len,
stop_upon_closing_brace=(brace_type, closing_brace),
parsing_state=parsing_state,
@@ -1883,12 +1883,14 @@ class LatexWalker(object):
pos=pos,
msg=r"get_latex_environment: expected \begin{%s}: %s"
% (
- environmentname
- if environmentname is not None
- else "",
+ (
+ environmentname
+ if environmentname is not None
+ else ""
+ ),
firsttok.arg,
),
- **self.pos_to_lineno_colno(pos, as_dict=True)
+ **self.pos_to_lineno_colno(pos, as_dict=True),
)
if environmentname is None:
environmentname = firsttok.arg
@@ -1915,9 +1917,9 @@ class LatexWalker(object):
argsresult = (None, pos, 0, {})
if len(argsresult) == 4:
- (argd, apos, alen, adic) = argsresult
+ argd, apos, alen, adic = argsresult
else:
- (argd, apos, alen) = argsresult
+ argd, apos, alen = argsresult
adic = {}
pos = apos + alen
@@ -1930,7 +1932,7 @@ class LatexWalker(object):
math_mode_delimiter="{" + environmentname + "}",
)
- (nodelist, npos, nlen) = self.get_latex_nodes(
+ nodelist, npos, nlen = self.get_latex_nodes(
pos,
stop_upon_end_environment=environmentname,
parsing_state=parsing_state_inner,
@@ -1975,7 +1977,7 @@ class LatexWalker(object):
s=self.s,
pos=tok.pos,
msg="End of input while parsing {}".format(what),
- **self.pos_to_lineno_colno(tok.pos, as_dict=True)
+ **self.pos_to_lineno_colno(tok.pos, as_dict=True),
)
if getattr(e, "pos", None) is not None and e.lineno is None and e.colno is None:
@@ -2227,7 +2229,7 @@ class LatexWalker(object):
s=self.s,
pos=tok.pos,
msg="Unexpected mismatching closing brace: '%s'" % (tok.arg),
- **self.pos_to_lineno_colno(tok.pos, as_dict=True)
+ **self.pos_to_lineno_colno(tok.pos, as_dict=True),
)
return True
@@ -2239,7 +2241,7 @@ class LatexWalker(object):
s=self.s,
pos=tok.pos,
msg=("Unexpected closing environment: '{}'".format(tok.arg)),
- **self.pos_to_lineno_colno(tok.pos, as_dict=True)
+ **self.pos_to_lineno_colno(tok.pos, as_dict=True),
)
elif tok.arg != stop_upon_end_environment:
# p.push_lastchars(tok_to_pos_and_chars_from_ppos(tok))
@@ -2252,7 +2254,7 @@ class LatexWalker(object):
tok.arg, stop_upon_end_environment
)
),
- **self.pos_to_lineno_colno(tok.pos, as_dict=True)
+ **self.pos_to_lineno_colno(tok.pos, as_dict=True),
)
return True
@@ -2276,7 +2278,7 @@ class LatexWalker(object):
tok.arg,
stop_upon_closing_mathmode,
),
- **self.pos_to_lineno_colno(tok.pos, as_dict=True)
+ **self.pos_to_lineno_colno(tok.pos, as_dict=True),
)
# all ok, this is a new math mode opening. Keep an assert
# in case we forget to include some math-mode delimiters in
@@ -2291,7 +2293,7 @@ class LatexWalker(object):
s=self.s,
pos=tok.pos,
msg="Unexpected closing math mode: '{}'".format(tok.arg),
- **self.pos_to_lineno_colno(tok.pos, as_dict=True)
+ **self.pos_to_lineno_colno(tok.pos, as_dict=True),
)
# we have encountered a new math inline, parse the math expression
@@ -2306,7 +2308,7 @@ class LatexWalker(object):
)
try:
- (mathinline_nodelist, mpos, mlen) = self.get_latex_nodes(
+ mathinline_nodelist, mpos, mlen = self.get_latex_nodes(
p.pos,
stop_upon_closing_mathmode=corresponding_closing_mathmode,
parsing_state=parsing_state_inner,
@@ -2316,7 +2318,7 @@ class LatexWalker(object):
_maketuple(
'math mode "{}"'.format(tok.arg),
tok.pos,
- *self.pos_to_lineno_colno(tok.pos)
+ *self.pos_to_lineno_colno(tok.pos),
)
)
raise
@@ -2354,7 +2356,7 @@ class LatexWalker(object):
if tok.tok == "brace_open":
# another braced group to read.
try:
- (groupnode, bpos, blen) = self.get_latex_braced_group(
+ groupnode, bpos, blen = self.get_latex_braced_group(
tok.pos, brace_type=tok.arg, parsing_state=p.parsing_state
)
# except LatexWalkerEndOfStream as e:
@@ -2376,7 +2378,7 @@ class LatexWalker(object):
if tok.tok == "begin_environment":
# an environment to read.
try:
- (envnode, epos, elen) = self.get_latex_environment(
+ envnode, epos, elen = self.get_latex_environment(
tok.pos, environmentname=tok.arg, parsing_state=p.parsing_state
)
except LatexWalkerParseError as e:
@@ -2384,7 +2386,7 @@ class LatexWalker(object):
_maketuple(
'begin environment "{}"'.format(tok.arg),
tok.pos,
- *self.pos_to_lineno_colno(tok.pos)
+ *self.pos_to_lineno_colno(tok.pos),
)
)
raise
@@ -2415,9 +2417,9 @@ class LatexWalker(object):
margsresult = (None, tok.pos + tok.len, 0, {})
if len(margsresult) == 4:
- (nodeargd, mapos, malen, mdic) = margsresult
+ nodeargd, mapos, malen, mdic = margsresult
else:
- (nodeargd, mapos, malen) = margsresult
+ nodeargd, mapos, malen = margsresult
mdic = {}
p.pos = mapos + malen
@@ -2473,9 +2475,9 @@ class LatexWalker(object):
if res is not None:
# specials expects arguments, read them
if len(res) == 4:
- (nodeargd, mapos, malen, spdic) = res
+ nodeargd, mapos, malen, spdic = res
else:
- (nodeargd, mapos, malen) = res
+ nodeargd, mapos, malen = res
spdic = {}
p.pos = mapos + malen
@@ -2505,7 +2507,7 @@ class LatexWalker(object):
s=self.s,
pos=p.pos,
msg="Unknown token: {!r}".format(tok),
- **self.pos_to_lineno_colno(p.pos, as_dict=True)
+ **self.pos_to_lineno_colno(p.pos, as_dict=True),
)
while True:
@@ -2531,7 +2533,7 @@ class LatexWalker(object):
msg="Unexpected end of stream, was expecting {}".format(
expecting
),
- **self.pos_to_lineno_colno(len(self.s), as_dict=True)
+ **self.pos_to_lineno_colno(len(self.s), as_dict=True),
)
if self.tolerant_parsing:
r_endnow = True
@@ -2651,7 +2653,7 @@ def get_latex_nodes(
stop_upon_closing_brace=None,
stop_upon_end_environment=None,
stop_upon_closing_mathmode=None,
- **parse_flags
+ **parse_flags,
):
"""
Parses latex content `s`.
diff --git a/great_ai/utilities/external/pylatexenc/latexwalker/__main__.py b/great_ai/utilities/external/pylatexenc/latexwalker/__main__.py
index fe95f0b..46eccff 100644
--- a/great_ai/utilities/external/pylatexenc/latexwalker/__main__.py
+++ b/great_ai/utilities/external/pylatexenc/latexwalker/__main__.py
@@ -184,7 +184,7 @@ def main(argv=None):
latex, tolerant_parsing=args.tolerant_parsing, strict_braces=args.strict_braces
)
- (nodelist, pos, len_) = latexwalker.get_latex_nodes()
+ nodelist, pos, len_ = latexwalker.get_latex_nodes()
if args.output_format == "human":
print("\n--- NODES ---\n")
diff --git a/great_ai/utilities/external/pylatexenc/macrospec/__init__.py b/great_ai/utilities/external/pylatexenc/macrospec/__init__.py
index 9975837..66ab0de 100644
--- a/great_ai/utilities/external/pylatexenc/macrospec/__init__.py
+++ b/great_ai/utilities/external/pylatexenc/macrospec/__init__.py
@@ -34,7 +34,6 @@ macros and environments, specifying how they should be parsed by
`pylatexenc 2.0`.
"""
-
import sys
if sys.version_info.major > 2:
@@ -154,7 +153,7 @@ class EnvironmentSpec(object):
environmentname,
args_parser=MacroStandardArgsParser(),
is_math_mode=False,
- **kwargs
+ **kwargs,
):
super(EnvironmentSpec, self).__init__(**kwargs)
self.environmentname = environmentname
@@ -776,9 +775,9 @@ class LatexContextDb(object):
new_context.add_context_category(
cat,
macros=self.d[cat]["macros"].values() if keep_macros else [],
- environments=self.d[cat]["environments"].values()
- if keep_environments
- else [],
+ environments=(
+ self.d[cat]["environments"].values() if keep_environments else []
+ ),
specials=self.d[cat]["specials"].values() if keep_specials else [],
)
diff --git a/great_ai/utilities/external/pylatexenc/macrospec/_argparsers.py b/great_ai/utilities/external/pylatexenc/macrospec/_argparsers.py
index 7fb20e8..abadc40 100644
--- a/great_ai/utilities/external/pylatexenc/macrospec/_argparsers.py
+++ b/great_ai/utilities/external/pylatexenc/macrospec/_argparsers.py
@@ -296,7 +296,7 @@ class MacroStandardArgsParser(object):
for j, argt in enumerate(self.argspec):
if argt == "{":
- (node, np, nl) = w.get_latex_expression(
+ node, np, nl = w.get_latex_expression(
p, strict_braces=False, parsing_state=get_inner_parsing_state(j)
)
p = np + nl
@@ -315,7 +315,7 @@ class MacroStandardArgsParser(object):
if optarginfotuple is None:
argnlist.append(None)
continue
- (node, np, nl) = optarginfotuple
+ node, np, nl = optarginfotuple
p = np + nl
argnlist.append(node)
diff --git a/great_ai/utilities/parallel_map/manage_communication.py b/great_ai/utilities/parallel_map/manage_communication.py
index d6c0740..4d53d2b 100644
--- a/great_ai/utilities/parallel_map/manage_communication.py
+++ b/great_ai/utilities/parallel_map/manage_communication.py
@@ -40,7 +40,9 @@ def manage_communication(
except Exception as e:
if ignore_exceptions:
logger.error(
- f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
+ f"""Exception {e} encountered in input, traceback:\n{
+ traceback.format_exc()
+ }"""
)
else:
raise
@@ -51,10 +53,12 @@ def manage_communication(
if r.exception is not None:
if ignore_exceptions:
logger.error(
- f"Exception {r.exception} encountered in worker, traceback:\n{r.worker_traceback}"
+ f"""Exception {
+ r.exception
+ } encountered in worker, traceback:\n{r.worker_traceback}"""
)
else:
- raise WorkerException from r.exception
+ raise WorkerException(r.exception)
if unordered:
diff --git a/great_ai/utilities/parallel_map/map_result.py b/great_ai/utilities/parallel_map/map_result.py
index 81d3dbb..481ce5d 100644
--- a/great_ai/utilities/parallel_map/map_result.py
+++ b/great_ai/utilities/parallel_map/map_result.py
@@ -4,5 +4,5 @@ from typing import Any, NamedTuple, Optional
class MapResult(NamedTuple):
order: int
value: Any
- exception: Optional[Exception] = None
+ exception: Optional[str] = None
worker_traceback: Optional[str] = None
diff --git a/great_ai/utilities/parallel_map/mapper_function.py b/great_ai/utilities/parallel_map/mapper_function.py
index 921065a..1aa4721 100644
--- a/great_ai/utilities/parallel_map/mapper_function.py
+++ b/great_ai/utilities/parallel_map/mapper_function.py
@@ -43,7 +43,11 @@ def mapper_function(
),
)
except Exception as e:
- return MapResult(i, None, e, traceback.format_exc())
+ # `e` has to be stringified in order to avoid any
+ # surprising serialization error when returning it
+ return MapResult(
+ i, None, str(e), traceback.format_exc()
+ )
async def main() -> List[MapResult]:
return await asyncio.gather(
@@ -57,7 +61,9 @@ def mapper_function(
last_chunk.append(MapResult(i, func(value)))
except Exception as e:
last_chunk.append(
- MapResult(i, None, e, traceback.format_exc())
+ # `e` has to be stringified in order to avoid any
+ # surprising serialization error when returning it
+ MapResult(i, None, str(e), traceback.format_exc())
)
except queue.Empty:
pass
diff --git a/great_ai/utilities/parallel_map/parallel_map.py b/great_ai/utilities/parallel_map/parallel_map.py
index 81cd0bd..f8c8e1f 100644
--- a/great_ai/utilities/parallel_map/parallel_map.py
+++ b/great_ai/utilities/parallel_map/parallel_map.py
@@ -31,8 +31,7 @@ def parallel_map(
chunk_size: Optional[int] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[Optional[V]]:
- ...
+) -> Iterable[Optional[V]]: ...
@overload
@@ -44,8 +43,7 @@ def parallel_map(
ignore_exceptions: Literal[True],
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[Optional[V]]:
- ...
+) -> Iterable[Optional[V]]: ...
@overload
@@ -57,8 +55,7 @@ def parallel_map(
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[V]:
- ...
+) -> Iterable[V]: ...
@overload
@@ -70,8 +67,7 @@ def parallel_map(
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[V]:
- ...
+) -> Iterable[V]: ...
def parallel_map(
@@ -91,7 +87,7 @@ def parallel_map(
or ignored.
The new processes are forked if the OS allows it, otherwise, new Python processes
- are bootstrapped which can incur some startup cost. Each process processes a single
+ are bootstrapped which can incur some start-up cost. Each process processes a single
chunk at once.
Examples:
diff --git a/great_ai/utilities/parallel_map/simple_parallel_map.py b/great_ai/utilities/parallel_map/simple_parallel_map.py
index 0981a57..7a506b3 100644
--- a/great_ai/utilities/parallel_map/simple_parallel_map.py
+++ b/great_ai/utilities/parallel_map/simple_parallel_map.py
@@ -1,6 +1,6 @@
from typing import Awaitable, Callable, List, Optional, Sequence, TypeVar, Union
-from tqdm.cli import tqdm
+from tqdm import tqdm
from .parallel_map import parallel_map
diff --git a/great_ai/utilities/parallel_map/threaded_parallel_map.py b/great_ai/utilities/parallel_map/threaded_parallel_map.py
index faf0b48..21d80e5 100644
--- a/great_ai/utilities/parallel_map/threaded_parallel_map.py
+++ b/great_ai/utilities/parallel_map/threaded_parallel_map.py
@@ -29,8 +29,7 @@ def threaded_parallel_map(
chunk_size: Optional[int] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[Optional[V]]:
- ...
+) -> Iterable[Optional[V]]: ...
@overload
@@ -42,8 +41,7 @@ def threaded_parallel_map(
ignore_exceptions: Literal[True],
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[Optional[V]]:
- ...
+) -> Iterable[Optional[V]]: ...
@overload
@@ -55,8 +53,7 @@ def threaded_parallel_map(
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[V]:
- ...
+) -> Iterable[V]: ...
@overload
@@ -68,8 +65,7 @@ def threaded_parallel_map(
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
-) -> Iterable[V]:
- ...
+) -> Iterable[V]: ...
def threaded_parallel_map(
diff --git a/great_ai/views/outputs/classification_output.py b/great_ai/views/outputs/classification_output.py
index f955cc2..7e3cccc 100644
--- a/great_ai/views/outputs/classification_output.py
+++ b/great_ai/views/outputs/classification_output.py
@@ -6,4 +6,4 @@ from ..hashable_base_model import HashableBaseModel
class ClassificationOutput(HashableBaseModel):
label: Union[str, int]
confidence: float
- explanation: Optional[Any]
+ explanation: Optional[Any] = None
diff --git a/great_ai/views/outputs/regression_output.py b/great_ai/views/outputs/regression_output.py
index 96c1daf..997c92e 100644
--- a/great_ai/views/outputs/regression_output.py
+++ b/great_ai/views/outputs/regression_output.py
@@ -5,4 +5,4 @@ from ..hashable_base_model import HashableBaseModel
class RegressionOutput(HashableBaseModel):
value: Union[int, float]
- explanation: Optional[Any]
+ explanation: Optional[Any] = None
diff --git a/great_ai/views/outputs/sequence_labeling_output.py b/great_ai/views/outputs/sequence_labeling_output.py
index 131d8f2..7dceb8c 100644
--- a/great_ai/views/outputs/sequence_labeling_output.py
+++ b/great_ai/views/outputs/sequence_labeling_output.py
@@ -7,9 +7,9 @@ class LabeledToken(HashableBaseModel):
token: str
tag: Literal["B", "I", "O", "E", "S"]
confidence: float
- explanation: Optional[Any]
+ explanation: Optional[Any] = None
class SequenceLabelingOutput(HashableBaseModel):
labeled_tokens: List[LabeledToken]
- explanation: Optional[Any]
+ explanation: Optional[Any] = None
diff --git a/great_ai/views/query.py b/great_ai/views/query.py
index 7257a7b..ece1d21 100644
--- a/great_ai/views/query.py
+++ b/great_ai/views/query.py
@@ -1,7 +1,7 @@
from datetime import datetime
from typing import List, Optional, Sequence
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
from .filter import Filter
from .sort_by import SortBy
@@ -15,8 +15,8 @@ class Query(BaseModel):
until: Optional[datetime] = None
has_feedback: Optional[bool] = None
- class Config:
- schema_extra = {
+ model_config = ConfigDict(
+ json_schema_extra={
"example": {
"filter": [
{
@@ -33,3 +33,4 @@ class Query(BaseModel):
"has_feedback": False,
}
}
+ )
diff --git a/great_ai/views/trace.py b/great_ai/views/trace.py
index 0fa36cc..f7c6328 100644
--- a/great_ai/views/trace.py
+++ b/great_ai/views/trace.py
@@ -1,7 +1,7 @@
from pprint import pformat
from typing import Any, Dict, Generic, List, Optional, TypeVar
-from pydantic import Extra
+from pydantic import ConfigDict
from .hashable_base_model import HashableBaseModel
from .model import Model
@@ -9,19 +9,35 @@ from .model import Model
T = TypeVar("T")
-class Trace(Generic[T], HashableBaseModel):
+class Trace(HashableBaseModel, Generic[T]):
+ """Universal structure for storing prediction traces and training data.
+
+ Attributes:
+ trace_id: UUID4 identifier for uniquely referring to a trace.
+ created: Timestamp of its (original) construction.
+ original_execution_time_ms: Wall-time elapsed while its generating
+ TracingContext was alive.
+ logged_values: Values persisted through using `@parameter` or `log_metric()`.
+ models: Marks left by each encountered `@use_model` decorated function.
+ exception: Exception description if any was encountered.
+ output: Return value of the function wrapped by GreatAI.
+ feedback: Feedback obtained using the REST API of `add_ground_truth`.
+ tags: Tags used for filtering traces. Contains the name of the original
+ function, value of `ENVIRONMENT`, its split if has any, and either
+ `ground_truth` or `online` depending on the origin of the Trace.
+ """
+
trace_id: str
created: str
original_execution_time_ms: float
logged_values: Dict[str, Any]
models: List[Model]
- exception: Optional[str]
- output: Optional[T]
+ exception: Optional[str] = None
+ output: Optional[T] = None
feedback: Any = None
tags: List[str]
- class Config:
- extra = Extra.ignore
+ model_config = ConfigDict(extra="ignore")
@property
def input(self) -> Any:
@@ -62,7 +78,7 @@ class Trace(Generic[T], HashableBaseModel):
def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
return {
**(
- self.dict()
+ self.model_dump()
if include_original
else {
"trace_id": self.trace_id,
@@ -71,7 +87,11 @@ class Trace(Generic[T], HashableBaseModel):
}
),
**{
- k: pformat(v, indent=2, compact=True)
+ k: (
+ v
+ if (isinstance(v, float) or isinstance(v, int))
+ else pformat(v, indent=2, compact=True)
+ )
for k, v in self.logged_values.items()
},
"models_flat": self.models_flat,
@@ -82,6 +102,7 @@ class Trace(Generic[T], HashableBaseModel):
}
def __repr__(self) -> str:
- return f"""Trace[{type(self.output).__name__}]({
- pformat(self.dict(), indent=2, compact=True).replace('{ ', '{', 1)
- })"""
+ formatted = pformat(self.model_dump(), indent=2, compact=True).replace(
+ "{ ", "{", 1
+ )
+ return f"Trace[{type(self.output).__name__}]({formatted})"
diff --git a/mkdocs.yaml b/mkdocs.yaml
index 2ef77cf..1014fc7 100644
--- a/mkdocs.yaml
+++ b/mkdocs.yaml
@@ -6,7 +6,7 @@ repo_url: https://github.com/schmelczer/great-ai
repo_name: schmelczer/great-ai
edit_uri: edit/main/docs/
-copyright: GNU General Public License v3
+copyright: GNU General Public License v3
extra:
generator: false
@@ -22,7 +22,7 @@ extra:
name: great-ai on DockerHub
- icon: fontawesome/brands/python
link: https://pypi.org/project/great-ai
- name: great-ai on PyPI
+ name: great-ai on PyPI
theme:
name: "material"
@@ -48,8 +48,8 @@ theme:
name: Switch to light mode
watch:
-- docs
-- mkdocs.yaml
+ - docs
+ - mkdocs.yaml
plugins:
- git-revision-date-localized:
@@ -63,7 +63,6 @@ plugins:
- search
- mkdocs-jupyter:
include_source: true
- ignore: ["docs/examples/scibert/*.ipynb"]
markdown_extensions:
- pymdownx.highlight:
@@ -89,27 +88,32 @@ markdown_extensions:
nav:
- Overview: index.md
- - Tutorial:
- - Tutorial overview: tutorial/index.md
- - tutorial/train.ipynb
- - tutorial/deploy.ipynb
+ - Tutorial:
+ - Tutorial overview: tutorial/index.md
+ - tutorial/train.ipynb
+ - tutorial/deploy.ipynb
- User Guides:
- - how-to-guides/install.md
- - how-to-guides/create-service.md
- - how-to-guides/configure-service.md
- - how-to-guides/use-service.md
- - how-to-guides/handle-training-data.md
- - how-to-guides/large_file.md
- - how-to-guides/call-remote.md
+ - how-to-guides/install.md
+ - how-to-guides/create-service.md
+ - how-to-guides/configure-service.md
+ - how-to-guides/use-service.md
+ - how-to-guides/handle-training-data.md
+ - how-to-guides/large-file.md
+ - how-to-guides/call-remote.md
- Reference:
- - reference/index.md
- - reference/utilities.md
- - reference/large-file.md
- - reference/views.md
- - Examples:
- - Explainable Naive Bayes:
- - examples/simple/data.ipynb
- - examples/simple/train.ipynb
- - examples/simple/deploy.ipynb
- - Explainable SciBERT: examples/scibert.md
+ - reference/index.md
+ - reference/utilities.md
+ - reference/large-file.md
+ - reference/views.md
+ - Examples:
+ - Explainable Naive Bayes:
+ - examples/simple/data.ipynb
+ - examples/simple/train.ipynb
+ - examples/simple/deploy.ipynb
+ - Explainable SciBERT:
+ - examples/scibert/index.md
+ - examples/scibert/data.ipynb
+ - examples/scibert/train.ipynb
+ - examples/scibert/deploy.ipynb
+ - examples/scibert/additional-files.md
- explanation.md
diff --git a/pyproject.toml b/pyproject.toml
index 1e9b0cb..a02460d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,52 +21,58 @@ classifiers = [
"Natural Language :: English",
]
keywords = ["SE4ML", "MLOps", "AI engineering", "general", "robust", "end-to-end", "automated", "trustworthy", "ai", "deployment"]
-requires-python = ">= 3.8"
+requires-python = ">= 3.10"
dependencies = [
- "scikit-learn",
- "matplotlib",
- "numpy",
- "nbconvert",
- "ipython",
- "unidecode >= 1.3.0",
- "syntok >= 1.4.0",
- "langcodes[data] >= 3.3.0",
- "langdetect >= 1.0.9",
- "tinydb >= 4.7.0",
- "boto3 >= 1.23.0",
- "plotly >= 5.8.0",
- "pandas",
- "dash >= 2.4.0",
- "fastapi >= 0.70.0",
- "uvicorn[standard] >= 0.18.0",
- "watchdog >= 2.1.0",
- "typeguard >= 2.10.0",
- "pymongo >= 4.0.0",
- "dill >= 0.3.5.0",
- 'tqdm',
- "async_lru >= 1.0.0",
- "httpx >= 0.20.0",
+ "scikit-learn >= 1.3, < 2",
+ "matplotlib >= 3.7, < 4",
+ "numpy >= 1.24, < 3",
+ "nbconvert >= 7.0, < 8",
+ "ipython >= 8.0, < 10",
+ "unidecode >= 1.3.0, < 2",
+ "syntok >= 1.4.0, < 2",
+ "langcodes[data] >= 3.3.0, < 4",
+ "langdetect >= 1.0.9, < 2",
+ "tinydb >= 4.7.0, < 5",
+ "boto3 >= 1.34, < 2",
+ "plotly >= 5.20, < 7",
+ "pandas >= 2.0, < 4",
+ "dash >= 2.16, < 5",
+ "fastapi >= 0.110, < 1",
+ "uvicorn[standard] >= 0.27, < 1",
+ "a2wsgi >= 1.9, < 2",
+ "watchdog >= 3.0, < 7",
+ "typeguard >= 4.0, < 5",
+ "pydantic >= 2.5, < 3",
+ "pymongo >= 4.6, < 5",
+ "dill >= 0.3.6, < 1",
+ "tqdm >= 4.64, < 5",
+ "httpx >= 0.24, < 1",
]
[project.optional-dependencies]
dev = [
- "flit",
- "mkdocs",
- "mkdocstrings[python]",
- "mkdocs-material",
- "mkdocs-jupyter",
- "mkdocs-git-revision-date-localized-plugin",
- "autoflake",
- "isort",
- "black[jupyter]",
- "mypy",
- "flake8",
- "pytest",
- "pytest-cov",
- "pytest-subtests",
- "pytest-asyncio",
+ "flit >= 3.9, < 4",
+ "mkdocs >= 1.5, < 2",
+ "mkdocstrings[python] >= 0.24, < 1",
+ "mkdocs-material >= 9.5, < 10",
+ "mkdocs-jupyter >= 0.24, < 1",
+ "mkdocs-git-revision-date-localized-plugin >= 1.2, < 2",
+ "autoflake >= 2.0, < 3",
+ "isort >= 5.12, < 7",
+ "black[jupyter] >= 25.1, < 26",
+ "mypy >= 1.8, < 2",
+ "flake8 >= 7.0, < 8",
+ "tox >= 4.0, < 5",
+ "pytest >= 7.4, < 9",
+ "pytest-cov >= 4.1, < 8",
+ "pytest-asyncio >= 0.23, < 2",
]
+[tool.pytest.ini_options]
+# `...` in doctests (e.g. third-party exception messages that drift between
+# dependency versions) is treated as a wildcard rather than a literal.
+doctest_optionflags = ["ELLIPSIS"]
+
[project.urls]
Documentation = "https://great-ai.scoutinscience.com"
GitHub = "https://github.com/schmelczer/great-ai"
diff --git a/scripts/check-python.sh b/scripts/check-python.sh
index 668308c..1b36198 100755
--- a/scripts/check-python.sh
+++ b/scripts/check-python.sh
@@ -3,7 +3,7 @@
set -e
echo "Installing dependencies if necessary"
-python3 -m pip install --upgrade autoflake isort black[jupyter] mypy flake8
+python3 -m pip install autoflake isort black[jupyter] mypy flake8
for dir in "$@"; do
echo "Checking $dir"
@@ -16,10 +16,10 @@ for dir in "$@"; do
if find . -name "*.py" 2>/dev/null | grep -q .; then
yes | python3 -m mypy . --install-types > /dev/null || true
- python3 -m mypy --namespace-packages --ignore-missing-imports --install-types --non-interactive --disallow-untyped-defs --disallow-incomplete-defs --follow-imports=silent --exclude=external/ --exclude=/build/ --pretty .
+ python3 -m mypy --namespace-packages --ignore-missing-imports --install-types --non-interactive --disallow-untyped-defs --disallow-incomplete-defs --follow-imports=silent --exclude=external/ --exclude=/build/ --exclude='(^|/)__main__\.py$' --pretty .
fi
- python3 -m flake8 . --count --show-source --statistics --exclude=__init__.py,.env,external --ignore=E501,E402,F821,W503,E722,E203
+ python3 -m flake8 . --count --show-source --statistics --exclude=__init__.py,.env,external --extend-ignore=E501,E402,F821,W503,E722,E203,E704
cd -
done
diff --git a/tests/external/__init__.py b/tests/external/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/external/conftest.py b/tests/external/conftest.py
new file mode 100644
index 0000000..1fb0b17
--- /dev/null
+++ b/tests/external/conftest.py
@@ -0,0 +1,47 @@
+import asyncio
+import gc
+import os
+
+import pytest
+from great_ai.external.async_lru import _CacheInfo
+
+asyncio.set_event_loop(None)
+
+
+@pytest.fixture
+def event_loop(request):
+ loop = asyncio.new_event_loop()
+ loop.set_debug(bool(os.environ.get("PYTHONASYNCIODEBUG")))
+
+ yield loop
+
+ loop.call_soon(loop.stop)
+ loop.run_forever()
+ loop.close()
+
+ gc.collect()
+
+
+@pytest.fixture
+def loop(event_loop, request):
+ asyncio.set_event_loop(None)
+ request.addfinalizer(lambda: asyncio.set_event_loop(None))
+
+ return event_loop
+
+
+@pytest.fixture
+def check_lru(request):
+ def _check_lru(wrapped, *, hits, misses, cache, tasks, maxsize=128):
+ assert wrapped.hits == hits
+ assert wrapped.misses == misses
+ assert len(wrapped._cache) == cache
+ assert len(wrapped.tasks) == tasks
+ assert wrapped.cache_info() == _CacheInfo(
+ hits=hits,
+ misses=misses,
+ maxsize=maxsize,
+ currsize=cache,
+ )
+
+ return _check_lru
diff --git a/tests/external/test_basic.py b/tests/external/test_basic.py
new file mode 100644
index 0000000..604662d
--- /dev/null
+++ b/tests/external/test_basic.py
@@ -0,0 +1,187 @@
+import asyncio
+from functools import partial
+
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+alru_cache_attrs = [
+ "hits",
+ "misses",
+ "tasks",
+ "closed",
+ "cache_info",
+ "cache_clear",
+ "invalidate",
+ "close",
+ "open",
+]
+
+alru_cache_calable_attrs = alru_cache_attrs.copy()
+for attr in ["hits", "misses", "tasks", "closed"]:
+ alru_cache_calable_attrs.remove(attr)
+
+
+def test_alru_cache_not_callable(loop):
+ with pytest.raises(NotImplementedError):
+ alru_cache("foo")
+
+
+def test_alru_cache_not_coroutine(loop):
+ with pytest.raises(RuntimeError):
+
+ @alru_cache
+ def not_coro(val):
+ return val
+
+
+def test_alru_cache_deco(loop, check_lru):
+ asyncio.set_event_loop(loop)
+
+ @alru_cache
+ async def coro():
+ pass
+
+ assert asyncio.iscoroutinefunction(coro)
+
+ for attr in alru_cache_attrs:
+ assert hasattr(coro, attr)
+ for attr in alru_cache_calable_attrs:
+ assert callable(getattr(coro, attr))
+
+ assert isinstance(coro._cache, dict)
+ assert isinstance(coro.tasks, set)
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ awaitable = coro()
+ assert asyncio.iscoroutine(awaitable)
+ loop.run_until_complete(awaitable)
+
+
+def test_alru_cache_deco_called(check_lru, loop):
+ asyncio.set_event_loop(loop)
+
+ @alru_cache()
+ async def coro():
+ pass
+
+ assert asyncio.iscoroutinefunction(coro)
+
+ for attr in alru_cache_attrs:
+ assert hasattr(coro, attr)
+ for attr in alru_cache_calable_attrs:
+ assert callable(getattr(coro, attr))
+
+ assert isinstance(coro._cache, dict)
+ assert isinstance(coro.tasks, set)
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ awaitable = coro()
+ assert asyncio.iscoroutine(awaitable)
+ loop.run_until_complete(awaitable)
+
+
+def test_alru_cache_fn_called(check_lru, loop):
+ asyncio.set_event_loop(loop)
+
+ async def coro():
+ pass
+
+ coro_wrapped = alru_cache(coro)
+
+ assert asyncio.iscoroutinefunction(coro_wrapped)
+
+ for attr in alru_cache_attrs:
+ assert hasattr(coro_wrapped, attr)
+ for attr in alru_cache_calable_attrs:
+ assert callable(getattr(coro_wrapped, attr))
+
+ assert isinstance(coro_wrapped._cache, dict)
+ assert isinstance(coro_wrapped.tasks, set)
+ check_lru(coro_wrapped, hits=0, misses=0, cache=0, tasks=0)
+
+ awaitable = coro_wrapped()
+ assert asyncio.iscoroutine(awaitable)
+ loop.run_until_complete(awaitable)
+
+
+def test_alru_cache_origin(loop):
+ asyncio.set_event_loop(loop)
+
+ async def coro():
+ pass
+
+ coro_wrapped = alru_cache(coro)
+
+ assert coro_wrapped._origin is coro
+
+ coro_wrapped = alru_cache(partial(coro))
+
+ assert coro_wrapped._origin is coro
+
+
+@pytest.mark.asyncio
+async def test_alru_cache_await_same_result_async(check_lru, loop):
+ calls = 0
+ val = object()
+
+ @alru_cache()
+ async def coro():
+ nonlocal calls
+ calls += 1
+
+ return val
+
+ coros = [coro() for _ in range(100)]
+ ret = await asyncio.gather(*coros)
+ expected = [val] * 100
+ assert ret == expected
+ check_lru(coro, hits=99, misses=1, cache=1, tasks=0)
+
+ assert calls == 1
+ assert await coro() is val
+ check_lru(coro, hits=100, misses=1, cache=1, tasks=0)
+
+
+@pytest.mark.asyncio
+async def test_alru_cache_await_same_result_coroutine(check_lru, loop):
+ calls = 0
+ val = object()
+
+ @alru_cache()
+ async def coro():
+ nonlocal calls
+ calls += 1
+
+ return val
+
+ coros = [coro() for _ in range(100)]
+ ret = await asyncio.gather(*coros)
+ expected = [val] * 100
+ assert ret == expected
+ check_lru(coro, hits=99, misses=1, cache=1, tasks=0)
+
+ assert calls == 1
+ assert await coro() is val
+ check_lru(coro, hits=100, misses=1, cache=1, tasks=0)
+
+
+@pytest.mark.asyncio
+async def test_alru_cache_dict_not_shared(check_lru, loop):
+ async def coro(val):
+ return val
+
+ coro1 = alru_cache()(coro)
+ coro2 = alru_cache()(coro)
+
+ ret1 = await coro1(1)
+ check_lru(coro1, hits=0, misses=1, cache=1, tasks=0)
+
+ ret2 = await coro2(1)
+ check_lru(coro2, hits=0, misses=1, cache=1, tasks=0)
+
+ assert ret1 == ret2
+
+ assert coro1._cache[1].result() == coro2._cache[1].result()
+ assert coro1._cache != coro2._cache
+ assert coro1._cache.keys() == coro2._cache.keys()
+ assert coro1._cache is not coro2._cache
diff --git a/tests/external/test_cache_clear.py b/tests/external/test_cache_clear.py
new file mode 100644
index 0000000..4e59ea6
--- /dev/null
+++ b/tests/external/test_cache_clear.py
@@ -0,0 +1,22 @@
+import asyncio
+
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_cache_clear(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ return val
+
+ inputs = [1, 2, 3]
+ coros = [coro(v) for v in inputs]
+ ret = await asyncio.gather(*coros)
+ assert ret == inputs
+ check_lru(coro, hits=0, misses=3, cache=3, tasks=0)
+
+ coro.cache_clear()
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
diff --git a/tests/external/test_cache_info.py b/tests/external/test_cache_info.py
new file mode 100644
index 0000000..1a39b96
--- /dev/null
+++ b/tests/external/test_cache_info.py
@@ -0,0 +1,38 @@
+import asyncio
+
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_cache_info(check_lru, loop):
+ @alru_cache(maxsize=4)
+ async def coro(val):
+ return val
+
+ inputs = [1, 2, 3]
+ coros = [coro(v) for v in inputs]
+ ret = await asyncio.gather(*coros)
+ assert ret == inputs
+ check_lru(coro, hits=0, misses=3, cache=3, tasks=0, maxsize=4)
+
+ coro.cache_clear()
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0, maxsize=4)
+
+ inputs = [1, 1, 1]
+ coros = [coro(v) for v in inputs]
+ ret = await asyncio.gather(*coros)
+ assert ret == inputs
+ check_lru(coro, hits=2, misses=1, cache=1, tasks=0, maxsize=4)
+
+ coro.cache_clear()
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0, maxsize=4)
+
+ inputs = [1, 2, 3, 4] * 2
+ coros = [coro(v) for v in inputs]
+ ret = await asyncio.gather(*coros)
+ assert ret == inputs
+ check_lru(coro, hits=4, misses=4, cache=4, tasks=0, maxsize=4)
diff --git a/tests/external/test_cache_invalidate.py b/tests/external/test_cache_invalidate.py
new file mode 100644
index 0000000..48dfd8c
--- /dev/null
+++ b/tests/external/test_cache_invalidate.py
@@ -0,0 +1,85 @@
+import asyncio
+
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_cache_invalidate(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ return val
+
+ inputs = [1, 2, 3]
+
+ coro.invalidate(1)
+ coro.invalidate(2)
+ coro.invalidate(3)
+
+ coros = [coro(v) for v in inputs]
+
+ ret = await asyncio.gather(*coros)
+
+ assert ret == inputs
+ check_lru(coro, hits=0, misses=3, cache=3, tasks=0)
+
+ coro.invalidate(1)
+ check_lru(coro, hits=0, misses=3, cache=2, tasks=0)
+ coro.invalidate(2)
+ check_lru(coro, hits=0, misses=3, cache=1, tasks=0)
+ coro.invalidate(3)
+ check_lru(coro, hits=0, misses=3, cache=0, tasks=0)
+
+ inputs = [1, 2, 3]
+ coros = [coro(v) for v in inputs]
+
+ ret = await asyncio.gather(*coros)
+
+ assert ret == inputs
+ check_lru(coro, hits=0, misses=6, cache=3, tasks=0)
+
+
+async def test_cache_invalidate_multiple_args(check_lru, loop):
+ @alru_cache()
+ async def coro(*args):
+ return len(args)
+
+ for i, size in enumerate(range(10)):
+ args = tuple(range(size))
+ ret = await coro(*args)
+ assert ret == size
+ check_lru(coro, hits=0, misses=i + 1, cache=1, tasks=0)
+ coro.invalidate(*args)
+ check_lru(coro, hits=0, misses=i + 1, cache=0, tasks=0)
+
+ for size in range(10):
+ args = tuple(range(size))
+ ret = await coro(*args)
+ assert ret == size
+ check_lru(coro, hits=0, misses=20, cache=10, tasks=0)
+
+
+async def test_cache_invalidate_multiple_args_different_order(check_lru, loop):
+ @alru_cache()
+ async def coro(*args):
+ return len(args)
+
+ for i, size in enumerate(range(2, 10)):
+ args = tuple(range(size))
+ rev_args = tuple(reversed(args))
+ ret = await coro(*args)
+ assert ret == size
+ check_lru(coro, hits=0, misses=2 * i + 1, cache=i + 1, tasks=0)
+ ret = await coro(*rev_args)
+ # The reversed args should be a miss
+ check_lru(coro, hits=0, misses=2 * i + 2, cache=i + 2, tasks=0)
+ coro.invalidate(*rev_args)
+ # The reversed args should be invalidated
+ check_lru(coro, hits=0, misses=2 * i + 2, cache=i + 1, tasks=0)
+
+ for i, size in enumerate(range(2, 10)):
+ args = tuple(range(size))
+ ret = await coro(*args)
+ assert ret == size
+ check_lru(coro, hits=i + 1, misses=16, cache=8, tasks=0)
diff --git a/tests/external/test_close.py b/tests/external/test_close.py
new file mode 100644
index 0000000..3008899
--- /dev/null
+++ b/tests/external/test_close.py
@@ -0,0 +1,170 @@
+import asyncio
+
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_cache_close(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ await asyncio.sleep(0.2)
+
+ return val
+
+ assert not coro.closed
+
+ inputs = [1, 2, 3, 4, 5]
+
+ coros = [coro(v) for v in inputs]
+
+ gather = asyncio.gather(*coros)
+
+ await asyncio.sleep(0.1)
+
+ check_lru(coro, hits=0, misses=5, cache=5, tasks=5)
+
+ close = coro.close()
+
+ with pytest.raises(RuntimeError):
+ await coro(1)
+
+ check_lru(coro, hits=0, misses=5, cache=5, tasks=5)
+
+ ret_close = await close
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ ret_gather = await gather
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ assert set(ret_close) == set(ret_gather) == set(inputs)
+
+ with pytest.raises(RuntimeError):
+ coro.close()
+
+
+async def test_cache_close_cancel_return_exceptions(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ await asyncio.sleep(0.2)
+
+ return val
+
+ inputs = [1, 2, 3, 4, 5]
+
+ coros = [coro(v) for v in inputs]
+
+ gather = asyncio.gather(*coros)
+
+ await asyncio.sleep(0.1)
+
+ close = coro.close(cancel=True)
+
+ check_lru(coro, hits=0, misses=5, cache=5, tasks=5)
+
+ ret_close = await close
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ for err in ret_close:
+ assert isinstance(err, asyncio.CancelledError)
+
+ with pytest.raises(asyncio.CancelledError):
+ await gather
+
+
+async def test_cache_close_cancel_not_return_exceptions(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ await asyncio.sleep(0.2)
+
+ return val
+
+ inputs = [1, 2, 3, 4, 5]
+
+ coros = [coro(v) for v in inputs]
+
+ gather = asyncio.gather(*coros)
+
+ await asyncio.sleep(0.1)
+
+ close = coro.close(cancel=True, return_exceptions=False)
+
+ check_lru(coro, hits=0, misses=5, cache=5, tasks=5)
+
+ with pytest.raises(asyncio.CancelledError):
+ await close
+
+ with pytest.raises(asyncio.CancelledError):
+ await gather
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+
+async def test_cache_close_return_exceptions(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ await asyncio.sleep(0.2)
+
+ raise ZeroDivisionError
+
+ inputs = [1, 2, 3, 4, 5]
+
+ coros = [coro(v) for v in inputs]
+
+ gather = asyncio.gather(*coros)
+
+ await asyncio.sleep(0.1)
+
+ close = coro.close()
+
+ check_lru(coro, hits=0, misses=5, cache=5, tasks=5)
+
+ with pytest.raises(ZeroDivisionError):
+ await gather
+
+ check_lru(coro, hits=0, misses=5, cache=5, tasks=0)
+
+ ret_close = await close
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ for err in ret_close:
+ assert isinstance(err, ZeroDivisionError)
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+
+async def test_cache_close_not_return_exceptions(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ await asyncio.sleep(0.2)
+
+ raise ZeroDivisionError
+
+ inputs = [1, 2, 3, 4, 5]
+
+ coros = [coro(v) for v in inputs]
+
+ gather = asyncio.gather(*coros, return_exceptions=True)
+
+ await asyncio.sleep(0.1)
+
+ close = coro.close(return_exceptions=False)
+
+ check_lru(coro, hits=0, misses=5, cache=5, tasks=5)
+
+ with pytest.raises(ZeroDivisionError):
+ await close
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ ret_gather = await gather
+
+ for err in ret_gather:
+ assert isinstance(err, ZeroDivisionError)
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
diff --git a/tests/external/test_exception.py b/tests/external/test_exception.py
new file mode 100644
index 0000000..3d4aa25
--- /dev/null
+++ b/tests/external/test_exception.py
@@ -0,0 +1,48 @@
+import asyncio
+
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_alru_cache_exception(check_lru, loop):
+ @alru_cache(cache_exceptions=True)
+ async def coro(val):
+ 1 / 0
+
+ inputs = [1, 1, 1]
+ coros = [coro(v) for v in inputs]
+
+ ret = await asyncio.gather(*coros, return_exceptions=True)
+
+ check_lru(coro, hits=2, misses=1, cache=1, tasks=0)
+
+ for item in ret:
+ assert isinstance(item, ZeroDivisionError)
+
+ with pytest.raises(ZeroDivisionError):
+ await coro(1)
+
+ check_lru(coro, hits=3, misses=1, cache=1, tasks=0)
+
+
+async def test_alru_not_cache_exception(check_lru, loop):
+ @alru_cache(cache_exceptions=False)
+ async def coro(val):
+ 1 / 0
+
+ inputs = [1, 1, 1]
+ coros = [coro(v) for v in inputs]
+
+ ret = await asyncio.gather(*coros, return_exceptions=True)
+
+ check_lru(coro, hits=2, misses=1, cache=1, tasks=0)
+
+ for item in ret:
+ assert isinstance(item, ZeroDivisionError)
+
+ with pytest.raises(ZeroDivisionError):
+ await coro(1)
+
+ check_lru(coro, hits=2, misses=2, cache=1, tasks=0)
diff --git a/tests/external/test_internals.py b/tests/external/test_internals.py
new file mode 100644
index 0000000..142c1fa
--- /dev/null
+++ b/tests/external/test_internals.py
@@ -0,0 +1,374 @@
+import asyncio
+from collections import OrderedDict
+from functools import _make_key, partial
+from unittest import mock
+
+import pytest
+from great_ai.external.async_lru import (
+ __cache_touch,
+ _cache_clear,
+ _cache_hit,
+ _cache_info,
+ _cache_invalidate,
+ _cache_miss,
+ _close,
+ _close_waited,
+ _done_callback,
+ _open,
+ _wait_closed,
+)
+
+
+class Wrapped:
+ pass
+
+
+@pytest.mark.asyncio
+async def test_done_callback_cancelled():
+ loop = asyncio.get_running_loop()
+ task = loop.create_future()
+ fut = loop.create_future()
+
+ task.add_done_callback(partial(_done_callback, fut))
+
+ task.cancel()
+
+ await asyncio.sleep(0)
+
+ assert fut.cancelled()
+
+
+@pytest.mark.asyncio
+async def test_done_callback_exception():
+ loop = asyncio.get_running_loop()
+ task = loop.create_future()
+ fut = loop.create_future()
+
+ task.add_done_callback(partial(_done_callback, fut))
+
+ exc = ZeroDivisionError()
+
+ task.set_exception(exc)
+
+ await asyncio.sleep(0)
+
+ with pytest.raises(ZeroDivisionError):
+ await fut
+
+ with pytest.raises(ZeroDivisionError):
+ fut.result()
+
+ assert fut.exception() is exc
+
+
+@pytest.mark.asyncio
+async def test_done_callback():
+ loop = asyncio.get_running_loop()
+ task = loop.create_future()
+ fut = loop.create_future()
+
+ task.add_done_callback(partial(_done_callback, fut))
+
+ task.set_result(1)
+
+ await asyncio.sleep(0)
+
+ assert fut.result() == 1
+
+
+def test_cache_invalidate_typed():
+ wrapped = Wrapped()
+ wrapped._cache = {}
+
+ args = (1,)
+ kwargs = {"1": 1}
+
+ from_cache = _cache_invalidate(wrapped, True, *args, **kwargs)
+
+ assert not from_cache
+
+ key = _make_key(args, kwargs, True)
+
+ wrapped._cache[key] = 0
+
+ from_cache = _cache_invalidate(wrapped, True, *args, **kwargs)
+
+ assert from_cache
+
+ assert len(wrapped._cache) == 0
+
+ wrapped._cache[key] = 0
+
+ args = (1.0,)
+
+ from_cache = _cache_invalidate(wrapped, True, *args, **kwargs)
+
+ assert not from_cache
+
+ wrapped._cache[key] = 1
+
+
+def test_cache_invalidate_not_typed():
+ wrapped = Wrapped()
+ wrapped._cache = {}
+
+ args = (1,)
+ kwargs = {"1": 1}
+
+ from_cache = _cache_invalidate(wrapped, False, *args, **kwargs)
+
+ assert not from_cache
+
+ key = _make_key(args, kwargs, False)
+
+ wrapped._cache[key] = 0
+
+ from_cache = _cache_invalidate(wrapped, False, *args, **kwargs)
+
+ assert from_cache
+
+ assert len(wrapped._cache) == 0
+
+ wrapped._cache[key] = 0
+
+ args = (1.0,)
+
+ from_cache = _cache_invalidate(wrapped, False, *args, **kwargs)
+
+ assert from_cache
+
+ assert len(wrapped._cache) == 0
+
+
+def test_cache_clear():
+ wrapped = Wrapped()
+
+ attrs = ["hits", "_cache", "tasks"]
+ for attr in attrs:
+ assert not hasattr(wrapped, attr)
+
+ _cache_clear(wrapped)
+
+ for attr in attrs:
+ assert hasattr(wrapped, attr)
+
+ assert wrapped.hits == wrapped.misses == 0
+ assert isinstance(wrapped._cache, dict)
+ assert len(wrapped._cache) == 0
+ assert isinstance(wrapped.tasks, set)
+ assert len(wrapped.tasks) == 0
+
+ _cache = wrapped._cache
+ tasks = wrapped.tasks
+
+ _cache_clear(wrapped)
+
+ assert wrapped._cache is not _cache
+ assert wrapped.tasks is not tasks
+
+
+def test_open():
+ wrapped = Wrapped()
+ wrapped.hits = wrapped.misses = 1
+ wrapped._cache = {}
+ wrapped.tasks = set()
+ wrapped.closed = True
+
+ with pytest.raises(RuntimeError):
+ _open(wrapped)
+
+ wrapped.hits = wrapped.misses = 0
+
+ _open(wrapped)
+
+ assert not wrapped.closed
+
+ with pytest.raises(RuntimeError):
+ _open(wrapped)
+
+
+def test_close(loop):
+ wrapped = Wrapped()
+ wrapped.closed = False
+ wrapped.tasks = set()
+
+ awaitable = _close(wrapped, cancel=False, return_exceptions=True)
+ loop.run_until_complete(awaitable)
+
+ assert wrapped.closed
+
+ with pytest.raises(RuntimeError):
+ _close(wrapped, cancel=False, return_exceptions=True)
+
+ fut = loop.create_future()
+ wrapped.closed = False
+ wrapped.tasks = {fut}
+
+ awaitable = _close(wrapped, cancel=True, return_exceptions=True)
+ loop.run_until_complete(awaitable)
+
+ assert fut.cancelled()
+
+ fut = loop.create_future()
+ fut.set_result(None)
+ wrapped.closed = False
+ wrapped.tasks = {fut}
+
+ awaitable = _close(wrapped, cancel=True, return_exceptions=True)
+ loop.run_until_complete(awaitable)
+
+ assert not fut.cancelled()
+
+ fut = loop.create_future()
+ fut.set_exception(ZeroDivisionError)
+ wrapped.closed = False
+ wrapped.tasks = {fut}
+
+ awaitable = _close(wrapped, cancel=True, return_exceptions=True)
+ loop.run_until_complete(awaitable)
+
+ assert not fut.cancelled()
+
+
+@pytest.mark.asyncio
+async def test_wait_closed():
+ loop = asyncio.get_running_loop()
+ wrapped = Wrapped()
+ wrapped.tasks = set()
+
+ with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
+ ret = await _wait_closed(
+ wrapped,
+ return_exceptions=True,
+ )
+ assert ret == []
+ mocked.assert_called_once()
+
+ asyncio.set_event_loop(loop)
+ with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
+ ret = await _wait_closed(
+ wrapped,
+ return_exceptions=True,
+ )
+ assert ret == []
+ mocked.assert_called_once()
+ asyncio.set_event_loop(None)
+
+ fut = loop.create_future()
+ fut.set_result(None)
+ wrapped.tasks = {fut}
+ with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
+ ret = await _wait_closed(
+ wrapped,
+ return_exceptions=True,
+ )
+ assert ret == [None]
+ mocked.assert_called_once()
+
+ exc = ZeroDivisionError()
+ fut = loop.create_future()
+ fut.set_exception(exc)
+ wrapped.tasks = {fut}
+ with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
+ ret = await _wait_closed(
+ wrapped,
+ return_exceptions=True,
+ )
+ assert ret == [exc]
+ mocked.assert_called_once()
+
+ fut = loop.create_future()
+ fut.set_exception(ZeroDivisionError)
+ wrapped.tasks = {fut}
+ with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
+ with pytest.raises(ZeroDivisionError):
+ await _wait_closed(
+ wrapped,
+ return_exceptions=False,
+ )
+ # the awaited result raised before _wait_closed could flush its done
+ # callback, so yield once to let the scheduled _close_waited run
+ await asyncio.sleep(0)
+ mocked.assert_called_once()
+
+
+def test_close_waited():
+ wrapped = Wrapped()
+ # _close_waited calls wrapped.cache_clear(); mock that directly (patching the
+ # module-level _cache_clear would not affect the already-bound partial).
+ wrapped.cache_clear = mock.Mock()
+
+ _close_waited(wrapped, None)
+
+ wrapped.cache_clear.assert_called_once()
+
+
+def test_cache_info():
+ wrapped = Wrapped()
+ wrapped._cache = {}
+ wrapped.hits = wrapped.misses = 0
+
+ assert (0, 0, 3, 0) == _cache_info(wrapped, 3)
+
+ wrapped._cache[1] = 1
+
+ assert (0, 0, 1, 1) == _cache_info(wrapped, 1)
+
+ wrapped.hits = 2
+ wrapped.misses = 3
+ wrapped._cache[2] = 2
+
+ assert (2, 3, 5, 2) == _cache_info(wrapped, 5)
+
+
+def test__cache_touch():
+ wrapped = Wrapped()
+
+ wrapped._cache = OrderedDict()
+ wrapped._cache[1] = 1
+ wrapped._cache[2] = 2
+
+ __cache_touch(wrapped, 1)
+ assert list(wrapped._cache) == [2, 1]
+
+ __cache_touch(wrapped, 2)
+ assert list(wrapped._cache) == [1, 2]
+
+ # test KeyError
+ __cache_touch(wrapped, 100)
+
+
+def test_cache_hit():
+ wrapped = Wrapped()
+ wrapped.hits = 1
+ wrapped._cache = OrderedDict()
+ wrapped._cache[1] = 1
+
+ with mock.patch("great_ai.external.async_lru.__cache_touch") as mocked:
+ _cache_hit(wrapped, 1)
+
+ mocked.assert_called_once()
+
+ assert wrapped.hits == 2
+
+ _cache_hit(wrapped, 1)
+
+ assert wrapped.hits == 3
+
+
+def test_cache_miss():
+ wrapped = Wrapped()
+ wrapped.misses = 1
+ wrapped._cache = OrderedDict()
+ wrapped._cache[1] = 1
+
+ with mock.patch("great_ai.external.async_lru.__cache_touch") as mocked:
+ _cache_miss(wrapped, 1)
+
+ mocked.assert_called_once()
+
+ assert wrapped.misses == 2
+
+ _cache_miss(wrapped, 1)
+
+ assert wrapped.misses == 3
diff --git a/tests/external/test_open.py b/tests/external/test_open.py
new file mode 100644
index 0000000..b749df4
--- /dev/null
+++ b/tests/external/test_open.py
@@ -0,0 +1,39 @@
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_alru_cache_open(check_lru, loop):
+ @alru_cache()
+ async def coro(val):
+ return val
+
+ await coro(1)
+
+ check_lru(coro, hits=0, misses=1, cache=1, tasks=0)
+
+ with pytest.raises(RuntimeError):
+ coro.open()
+
+ close = coro.close()
+
+ assert coro.closed
+
+ with pytest.raises(RuntimeError):
+ await coro()
+
+ with pytest.raises(RuntimeError):
+ coro.open()
+
+ await close
+
+ check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+ coro.open()
+
+ ret = await coro(1)
+
+ assert ret == 1
+
+ check_lru(coro, hits=0, misses=1, cache=1, tasks=0)
diff --git a/tests/external/test_partialmethod.py b/tests/external/test_partialmethod.py
new file mode 100644
index 0000000..dfa888f
--- /dev/null
+++ b/tests/external/test_partialmethod.py
@@ -0,0 +1,50 @@
+import asyncio
+from functools import partial, partialmethod
+
+import pytest
+from great_ai.external.async_lru import alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_partialmethod_basic(check_lru, loop):
+ class Obj:
+ async def _coro(self, val):
+ return val
+
+ coro = alru_cache(partialmethod(_coro, 2))
+
+ obj = Obj()
+
+ coros = [obj.coro() for _ in range(5)]
+
+ check_lru(obj.coro, hits=0, misses=0, cache=0, tasks=0)
+
+ ret = await asyncio.gather(*coros)
+
+ check_lru(obj.coro, hits=4, misses=1, cache=1, tasks=0)
+
+ assert ret == [2, 2, 2, 2, 2]
+
+
+async def test_partialmethod_partial(check_lru, loop):
+ class Obj:
+ def __init__(self):
+ self.coro = alru_cache(partial(self._coro, 2))
+
+ async def __coro(self, val1, val2):
+ return val1 + val2
+
+ _coro = partialmethod(__coro, 1)
+
+ obj = Obj()
+
+ coros = [obj.coro() for _ in range(5)]
+
+ check_lru(obj.coro, hits=0, misses=0, cache=0, tasks=0)
+
+ ret = await asyncio.gather(*coros)
+
+ check_lru(obj.coro, hits=4, misses=1, cache=1, tasks=0)
+
+ assert ret == [3, 3, 3, 3, 3]
diff --git a/tests/external/test_size.py b/tests/external/test_size.py
new file mode 100644
index 0000000..c683c8e
--- /dev/null
+++ b/tests/external/test_size.py
@@ -0,0 +1,58 @@
+import asyncio
+
+import pytest
+from great_ai.external.async_lru import _make_key, alru_cache
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_alru_cache_removing_lru_keys(check_lru, loop):
+ @alru_cache(maxsize=3)
+ async def coro(val):
+ return val
+
+ key5 = _make_key((5,), {}, False)
+ key4 = _make_key((4,), {}, False)
+ key3 = _make_key((3,), {}, False)
+ key2 = _make_key((2,), {}, False)
+ key1 = _make_key((1,), {}, False)
+
+ for i, v in enumerate([3, 4, 5]):
+ await coro(v)
+ check_lru(coro, hits=0, misses=i + 1, cache=i + 1, tasks=0, maxsize=3)
+
+ check_lru(coro, hits=0, misses=3, cache=3, tasks=0, maxsize=3)
+ assert list(coro._cache) == [key3, key4, key5]
+
+ for v in [3, 2, 1]:
+ await coro(v)
+ check_lru(coro, hits=1, misses=5, cache=3, tasks=0, maxsize=3)
+ assert list(coro._cache) == [key3, key2, key1]
+
+
+async def test_alru_cache_none_max_size(check_lru, loop):
+ @alru_cache(maxsize=None)
+ async def coro(val):
+ return val
+
+ inputs = [1, 2, 3, 4] * 2
+ coros = [coro(v) for v in inputs]
+
+ ret = await asyncio.gather(*coros)
+
+ check_lru(coro, hits=4, misses=4, cache=4, tasks=0, maxsize=None)
+ assert ret == inputs
+
+
+async def test_alru_cache_zero_max_size(check_lru, loop):
+ @alru_cache(maxsize=0)
+ async def coro(val):
+ return val
+
+ inputs = [1, 2, 3, 4] * 2
+ coros = [coro(v) for v in inputs]
+
+ ret = await asyncio.gather(*coros)
+
+ check_lru(coro, hits=0, misses=8, cache=0, tasks=0, maxsize=0)
+ assert ret == inputs
diff --git a/tests/external/test_utils.py b/tests/external/test_utils.py
new file mode 100644
index 0000000..aa52d5d
--- /dev/null
+++ b/tests/external/test_utils.py
@@ -0,0 +1,19 @@
+from functools import partial
+
+from great_ai.external.async_lru import unpartial
+
+
+def test_unpartial():
+ def foo():
+ pass
+
+ assert unpartial(foo) is foo
+
+ bar = partial(foo)
+
+ assert unpartial(bar) is foo
+
+ for _ in range(10):
+ bar = partial(bar)
+
+ assert unpartial(bar) is foo
diff --git a/tests/large_file/test_large_file.py b/tests/large_file/test_large_file.py
index 87571fa..9afc615 100644
--- a/tests/large_file/test_large_file.py
+++ b/tests/large_file/test_large_file.py
@@ -24,6 +24,20 @@ def test_uninitialized() -> None:
LargeFileS3("test-file")
+def test_bad_file_name() -> None:
+ with pytest.raises(ValueError):
+ LargeFileS3("../no-no", cache_only_mode=True)
+
+ with pytest.raises(ValueError):
+ LargeFileS3("is this wrong?", cache_only_mode=True)
+
+ with pytest.raises(FileNotFoundError):
+ LargeFileS3("!378378_fer-ef", cache_only_mode=True) # this should work
+
+ with pytest.raises(FileNotFoundError):
+ LargeFileS3("3", cache_only_mode=True) # this should work
+
+
def test_bad_file_modes() -> None:
with pytest.raises(ValueError):
LargeFileS3("test-file", "w", version=3)
diff --git a/tests/test_basic_starters.py b/tests/test_basic_starters.py
index d7a8052..1303b22 100644
--- a/tests/test_basic_starters.py
+++ b/tests/test_basic_starters.py
@@ -25,13 +25,13 @@ def test_create_trivial_cases() -> None:
def test_create_with_other_decorator() -> None:
@GreatAI.create
- @lru_cache
+ @lru_cache()
def hello_world_1(name: str) -> str:
return f"Hello {name}!"
assert hello_world_1("andras").output == "Hello andras!"
- @lru_cache
+ @lru_cache()
@GreatAI.create
def hello_world_2(name: str) -> str:
return f"Hello {name}!"
diff --git a/tests/utilities/test_evaluate_ranking.py b/tests/utilities/test_evaluate_ranking.py
index b0ef4ad..8111b6a 100644
--- a/tests/utilities/test_evaluate_ranking.py
+++ b/tests/utilities/test_evaluate_ranking.py
@@ -16,7 +16,7 @@ def test_default() -> None:
reverse_order=True,
)
- assert results == {"d": 1.0, "c": 0.6666666666666666, "b": 0.4}
+ assert results == {"d": 5 / 6, "c": 2 / 3, "b": 2 / 5}
results = evaluate_ranking(
["a", "a", "b", "b", "c", "d", "d", "d"],
@@ -26,11 +26,7 @@ def test_default() -> None:
disable_interpolation=True,
)
- assert results == {
- "a": 0.6666666666666666,
- "b": 0.3333333333333333,
- "c": 0.2857142857142857,
- }
+ assert results == {"a": 0.6, "b": 0.4, "c": 0.2}
def test_mismatching_lengths() -> None:
@@ -81,7 +77,11 @@ def test_empty() -> None:
def test_save() -> None:
path = Path("test.svg")
- path.unlink(missing_ok=True)
+ try:
+ path.unlink()
+ except FileNotFoundError:
+ # missing_ok only exists >= Python 3.8
+ pass
evaluate_ranking(
["a", "a", "b", "b", "c", "d", "d", "d"],
diff --git a/tests/utilities/test_parallel_map.py b/tests/utilities/test_parallel_map.py
index 9fe2b8f..4774f7b 100644
--- a/tests/utilities/test_parallel_map.py
+++ b/tests/utilities/test_parallel_map.py
@@ -23,9 +23,7 @@ def test_with_iterable() -> None:
expected = [v**3 for v in range(10)]
- assert (
- list(parallel_map(lambda x: x**3, my_generator(), chunk_size=1)) == expected
- )
+ assert list(parallel_map(lambda x: x**3, my_generator(), chunk_size=1)) == expected
def test_simple_case_invalid_values() -> None:
@@ -44,14 +42,10 @@ def test_this_process_exception() -> None:
assert False
with pytest.raises(AssertionError):
- list(
- parallel_map(lambda v: v**2, my_generator(), concurrency=2, chunk_size=2)
- )
+ list(parallel_map(lambda v: v**2, my_generator(), concurrency=2, chunk_size=2))
with pytest.raises(AssertionError):
- list(
- parallel_map(lambda v: v**2, my_generator(), concurrency=1, chunk_size=2)
- )
+ list(parallel_map(lambda v: v**2, my_generator(), concurrency=1, chunk_size=2))
def test_ignore_this_process_exception() -> None:
diff --git a/tox.ini b/tox.ini
new file mode 100644
index 0000000..f853c1f
--- /dev/null
+++ b/tox.ini
@@ -0,0 +1,13 @@
+[tox]
+envlist = py37,py38,py39,py310
+isolated_build = True
+
+[testenv]
+alwayscopy = True
+setenv =
+ PY_IGNORE_IMPORTMISMATCH = 1
+deps =
+ pytest
+ pytest-cov
+ pytest-asyncio
+commands = pytest --doctest-modules --asyncio-mode=strict