andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Fix typing and minor issues

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-07 14:12:44 +02:00
parent 2db2253578
commit 72ab627a34
54 changed files with 635 additions and 589 deletions
Download patch file Download diff file Expand all files Collapse all files

2
great_ai/__init__.py
View file

@ -3,7 +3,7 @@ __version__ = "0.0.12"
from .context import configure
from .deploy import GreatAI
from .deploy import GreatAI, RouteConfig
from .exceptions import (
ArgumentValidationError,
MissingArgumentError,

7
great_ai/__main__.py
View file

@ -8,7 +8,6 @@ from threading import Event
from typing import Optional
import uvicorn
from parse_arguments import parse_arguments
from uvicorn._subprocess import get_subprocess
from uvicorn.config import LOGGING_CONFIG, Config
from uvicorn.supervisors.basereload import BaseReload
@ -21,6 +20,8 @@ from great_ai.deploy import GreatAI
from great_ai.exceptions import ArgumentValidationError, MissingArgumentError
from great_ai.utilities import get_logger
from .parse_arguments import parse_arguments
logger = get_logger(SERVER_NAME)
@ -28,11 +29,11 @@ GREAT_AI_LOGGING_CONFIG = {
**LOGGING_CONFIG,
"formatters": {
"default": {
"()": "great_ai.logger.CustomFormatter",
"()": "great_ai.utilities.logger.CustomFormatter",
"fmt": "%(asctime)s | %(levelname)8s | %(message)s",
},
"access": {
"()": "great_ai.logger.CustomFormatter",
"()": "great_ai.utilities.logger.CustomFormatter",
"fmt": "%(asctime)s | %(levelname)8s | %(message)s", # noqa: E501
},
},

18
great_ai/context.py
View file

@ -56,12 +56,12 @@ def configure(
*,
log_level: int = DEBUG,
seed: int = 42,
tracing_database: Optional[Type[TracingDatabaseDriver]] = None,
tracing_database_factory: Optional[Type[TracingDatabaseDriver]] = None,
large_file_implementation: Optional[Type[LargeFileBase]] = None,
should_log_exception_stack: Optional[bool] = None,
prediction_cache_size: int = 512,
disable_se4ml_banner: bool = False,
dashboard_table_size: int = 50,
dashboard_table_size: int = 20,
) -> None:
global _context
logger = get_logger("great_ai", level=log_level)
@ -76,17 +76,17 @@ def configure(
_set_seed(seed)
tracing_database = _initialize_tracing_database(tracing_database, logger=logger)()
tracing_database_factory = _initialize_tracing_database(
tracing_database_factory, logger=logger
)
tracing_database = tracing_database_factory()
if not tracing_database.is_production_ready:
message = f"The selected tracing database ({tracing_database_factory.__name__}) is not recommended for production"
if is_production:
logger.error(
f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production"
)
logger.error(message)
else:
logger.warning(
f"The selected tracing database ({type(tracing_database).__name__}) is not recommended for production"
)
logger.warning(message)
_context = Context(
tracing_database=tracing_database,

1
great_ai/deploy/__init__.py
View file

@ -1 +1,2 @@
from .great_ai import GreatAI
from .routes import RouteConfig

320
great_ai/deploy/great_ai.py
View file

@ -1,266 +1,216 @@
import inspect
from functools import lru_cache, partial, wraps
from textwrap import dedent
from typing import (
Any,
Awaitable,
Callable,
Generic,
Iterable,
List,
Optional,
Type,
Sequence,
TypeVar,
Union,
cast,
overload,
)
from fastapi import APIRouter, FastAPI, status
from pydantic import BaseModel, create_model
from async_lru import alru_cache
from fastapi import FastAPI
from ..constants import DASHBOARD_PATH
from ..context import get_context
from ..helper import (
freeze_arguments,
get_function_metadata_store,
snake_case_to_text,
use_http_exceptions,
)
from ..helper import freeze_arguments, get_function_metadata_store, snake_case_to_text
from ..models import model_versions
from ..parameters import automatically_decorate_parameters
from ..tracing.tracing_context import TracingContext
from ..utilities import parallel_map
from ..views import ApiMetadata, CacheStatistics, HealthCheckResponse, Trace
from .routes import (
bootstrap_docs_endpoints,
bootstrap_feedback_endpoints,
bootstrap_trace_endpoints,
)
from ..views import ApiMetadata, Trace
from .routes.bootstrap_dashboard import bootstrap_dashboard
from .routes.bootstrap_docs_endpoints import bootstrap_docs_endpoints
from .routes.bootstrap_feedback_endpoints import bootstrap_feedback_endpoints
from .routes.bootstrap_meta_endpoints import bootstrap_meta_endpoints
from .routes.bootstrap_prediction_endpoint import bootstrap_prediction_endpoint
from .routes.bootstrap_trace_endpoints import bootstrap_trace_endpoints
from .routes.route_config import RouteConfig
T = TypeVar("T")
T = TypeVar("T", bound=Union[Trace, Awaitable[Trace]])
V = TypeVar("V")
class GreatAI(Generic[T]):
def __init__(self, func: Callable[..., Any], version: str, return_raw_result: bool):
is_asynchronous = inspect.iscoroutinefunction(func)
class GreatAI(Generic[T, V]):
__name__: str
__doc__: str
def __init__(
self,
func: Callable[..., Union[V, Awaitable[V]]],
version: str,
route_config: RouteConfig,
):
func = automatically_decorate_parameters(func)
get_function_metadata_store(func).is_finalised = True
self._func = func
def func_in_tracing_context_sync(
*args: Any, do_not_persist_traces: bool = False, **kwargs: Any
) -> Trace[T]:
with TracingContext[T](
func.__name__, do_not_persist_traces=do_not_persist_traces
) as t:
result = func(*args, **kwargs)
output = t.finalise(output=result)
return result if return_raw_result else output
async def func_in_tracing_context_async(
*args: Any, do_not_persist_traces: bool = False, **kwargs: Any
) -> Trace[T]:
with TracingContext[T](
func.__name__, do_not_persist_traces=do_not_persist_traces
) as t:
result = await func(*args, **kwargs)
output = t.finalise(output=result)
return result if return_raw_result else output
func_in_tracing_context = (
func_in_tracing_context_async
if is_asynchronous
else func_in_tracing_context_sync
)
self._cached_func = lru_cache(get_context().prediction_cache_size)(
func_in_tracing_context
) # cannot put decorator on method, because it require the context to be setup
self._cached_func = self._get_cached_traced_function(func)
self._wrapped_func = wraps(func)(freeze_arguments(self._cached_func))
wraps(func)(self)
self.__doc__ = f"GreatAI wrapper for interacting with the `{self.__name__}` function.\n\n{dedent(self.__doc__ or '')}"
self._version = version
self.version = version
flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions)
if flat_model_versions:
self.version += f"+{flat_model_versions}"
self.app = FastAPI(
title=self.name,
title=snake_case_to_text(self.__name__),
version=self.version,
description=self.documentation
description=self.__doc__
+ f"\n\nFind out more in the [dashboard]({DASHBOARD_PATH}).",
docs_url=None,
redoc_url=None,
)
self._bootstrap_rest_api(route_config)
@overload
@staticmethod
def create(
func: Optional[Callable[..., T]] = None,
) -> "GreatAI[T]":
def create( # type: ignore
# "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]":
...
@overload
@staticmethod
def create(
version: str,
return_raw_result: bool,
disable_rest_api: bool,
disable_docs: bool,
disable_dashboard: bool,
) -> Callable[[Callable[..., T]], "GreatAI[T]"]:
func: Callable[..., V],
) -> "GreatAI[Trace[V], V]":
...
@overload
@staticmethod
def create(
func: Optional[Union[Callable[..., V], Callable[..., Awaitable[V]]]] = ...,
*,
version: str = ...,
route_config: RouteConfig = ...,
) -> Callable:
...
@staticmethod
def create(
func: Optional[Callable[..., T]] = None,
func: Optional[Callable] = None,
*,
version: str = "0.0.1",
return_raw_result: bool = False,
disable_rest_api: bool = False,
disable_docs: bool = False,
disable_dashboard: bool = False,
):
route_config: RouteConfig = RouteConfig(),
) -> Union[Callable, "GreatAI"]:
if func is None:
return cast(
Callable[[Callable[..., T]], GreatAI[T]],
partial(
GreatAI.create,
version=version,
return_raw_result=return_raw_result,
disable_rest_api=disable_rest_api,
disable_docs=disable_docs,
disable_dashboard=disable_dashboard,
),
)
instance = GreatAI[T](
func, version=version, return_raw_result=return_raw_result
@overload
def inner(func: Awaitable[V]) -> GreatAI[Awaitable[Trace[V]], V]:
...
@overload
def inner(func: Callable[..., V]) -> GreatAI[Trace[V], V]:
...
def inner(func): # type: ignore
return GreatAI.create(
func,
version=version,
route_config=route_config,
)
return inner
return GreatAI[T, V](
func,
version=version,
route_config=route_config,
)
if not disable_rest_api:
instance._bootstrap_rest_api(
disable_docs=disable_docs, disable_dashboard=disable_dashboard
)
return instance
@freeze_arguments
def __call__(self, *args: Any, **kwargs: Any) -> Trace[T]:
return self._cached_func(*args, **kwargs)
def __call__(self, *args: Any, **kwargs: Any) -> T:
return self._wrapped_func(*args, **kwargs)
def process_batch(
self,
batch: Iterable[Any],
batch: Sequence,
concurrency: Optional[int] = None,
do_not_persist_traces: bool = False,
) -> List[Trace[T]]:
do_not_persist_traces: Optional[bool] = False,
) -> List[Trace[V]]:
return list(
parallel_map(
freeze_arguments(
partial(
self._cached_func, do_not_persist_traces=do_not_persist_traces
)
partial(
self._wrapped_func, do_not_persist_traces=do_not_persist_traces
),
batch,
concurrency=concurrency,
)
)
@property
def name(self) -> str:
return snake_case_to_text(self._func.__name__)
@staticmethod
def _get_cached_traced_function(
func: Callable[..., Union[V, Awaitable[V]]]
) -> Callable[..., T]:
@lru_cache(maxsize=get_context().prediction_cache_size)
def func_in_tracing_context_sync(
*args: Any,
do_not_persist_traces: bool = False,
**kwargs: Any,
) -> Trace[V]:
with TracingContext[V](
func.__name__, do_not_persist_traces=do_not_persist_traces
) as t:
result = func(*args, **kwargs)
return t.finalise(output=result)
@property
def version(self) -> str:
flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions)
if flat_model_versions:
flat_model_versions = f"+{flat_model_versions}"
@alru_cache(maxsize=get_context().prediction_cache_size)
async def func_in_tracing_context_async(
*args: Any,
do_not_persist_traces: bool = False,
**kwargs: Any,
) -> Trace[V]:
with TracingContext[V](
func.__name__, do_not_persist_traces=do_not_persist_traces
) as t:
result = await cast(Callable[..., Awaitable], func)(*args, **kwargs)
return t.finalise(output=result)
return f"{self._version}{flat_model_versions}"
@property
def documentation(self) -> str:
return (
f"GreatAI wrapper for interacting with the `{self._func.__name__}` function.\n\n"
+ (
"\n".join(
line.strip()
for line in (self._func.__doc__ or "").split("\n")
if line.strip()
)
)
func_in_tracing_context_async
if get_function_metadata_store(func).is_asynchronous
else func_in_tracing_context_sync
)
def _bootstrap_rest_api(self, disable_docs: bool, disable_dashboard: bool) -> None:
self._bootstrap_prediction_endpoint()
def _bootstrap_rest_api(self, route_config: RouteConfig) -> None:
if route_config.prediction_endpoint_enabled:
bootstrap_prediction_endpoint(self.app, self._wrapped_func)
if not disable_docs:
if route_config.docs_endpoints_enabled:
bootstrap_docs_endpoints(self.app)
if not disable_dashboard:
if route_config.dashboard_enabled:
bootstrap_dashboard(
self.app,
function_name=self._func.__name__,
documentation=self.documentation,
function_name=self.__name__,
documentation=self.__doc__,
)
if route_config.trace_endpoints_enabled:
bootstrap_trace_endpoints(self.app)
bootstrap_feedback_endpoints(self.app)
self._bootstrap_meta_endpoints()
if route_config.feedback_endpoints_enabled:
bootstrap_feedback_endpoints(self.app)
def _bootstrap_prediction_endpoint(self) -> None:
router = APIRouter(
tags=["predictions"],
)
schema = self._get_schema()
@router.post(
"/predict", status_code=status.HTTP_200_OK, response_model=Trace[T]
)
@use_http_exceptions
def predict(input_value: schema) -> Trace[T]: # type: ignore
return self(**cast(BaseModel, input_value).dict())
self.app.include_router(router)
def _get_schema(self) -> Type[BaseModel]:
signature = inspect.signature(self._func)
parameters = {
p.name: (
p.annotation if p.annotation != inspect._empty else Any,
p.default if p.default != inspect._empty else ...,
if route_config.meta_endpoints_enabled:
bootstrap_meta_endpoints(
self.app,
self._cached_func,
ApiMetadata(
name=self.__name__,
version=self.version,
documentation=self.__doc__,
configuration=get_context().to_flat_dict(),
),
)
for p in signature.parameters.values()
if p.name in get_function_metadata_store(self._func).input_parameter_names
}
schema: Type[BaseModel] = create_model("InputModel", **parameters) # type: ignore
return schema
def _bootstrap_meta_endpoints(self) -> None:
router = APIRouter(
tags=["meta"],
)
@router.get("/health", status_code=status.HTTP_200_OK)
def check_health() -> HealthCheckResponse:
hits, misses, maxsize, cache_size = self._cached_func.cache_info()
cache_statistics = CacheStatistics(
hits=hits, misses=misses, size=cache_size, max_size=maxsize
)
return HealthCheckResponse(
is_healthy=True, cache_statistics=cache_statistics
)
@router.get(
"/version", response_model=ApiMetadata, status_code=status.HTTP_200_OK
)
def get_version() -> ApiMetadata:
return ApiMetadata(
name=self.name,
version=self.version,
documentation=self.documentation,
configuration=get_context().to_flat_dict(),
)
self.app.include_router(router)

3
great_ai/deploy/routes/__init__.py
View file

@ -1,4 +1,7 @@
from .bootstrap_dashboard import bootstrap_dashboard
from .bootstrap_docs_endpoints import bootstrap_docs_endpoints
from .bootstrap_feedback_endpoints import bootstrap_feedback_endpoints
from .bootstrap_meta_endpoints import bootstrap_meta_endpoints
from .bootstrap_prediction_endpoint import bootstrap_prediction_endpoint
from .bootstrap_trace_endpoints import bootstrap_trace_endpoints
from .route_config import RouteConfig

2
great_ai/deploy/routes/bootstrap_dashboard.py
View file

@ -6,7 +6,7 @@ from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from ...constants import DASHBOARD_PATH
from .dashboard import create_dash_app
from .dashboard.create_dash_app import create_dash_app
PATH = Path(__file__).parent.resolve()

26
great_ai/deploy/routes/bootstrap_meta_endpoints.py Normal file
View file

@ -0,0 +1,26 @@
from typing import Any
from fastapi import APIRouter, FastAPI, status
from ...views import ApiMetadata, CacheStatistics, HealthCheckResponse
def bootstrap_meta_endpoints(app: FastAPI, func: Any, metadata: ApiMetadata) -> None:
router = APIRouter(
tags=["meta"],
)
@router.get("/health", status_code=status.HTTP_200_OK)
def check_health() -> HealthCheckResponse:
hits, misses, maxsize, cache_size = func.cache_info()
cache_statistics = CacheStatistics(
hits=hits, misses=misses, size=cache_size, max_size=maxsize
)
return HealthCheckResponse(is_healthy=True, cache_statistics=cache_statistics)
@router.get("/version", response_model=ApiMetadata, status_code=status.HTTP_200_OK)
def get_version() -> ApiMetadata:
return metadata
app.include_router(router)

51
great_ai/deploy/routes/bootstrap_prediction_endpoint.py Normal file
View file

@ -0,0 +1,51 @@
import inspect
from typing import Any, Awaitable, Callable, Type, Union, cast
from fastapi import APIRouter, FastAPI, HTTPException, status
from pydantic import BaseModel, create_model
from ...helper import get_function_metadata_store
from ...views import Trace
def bootstrap_prediction_endpoint(
app: FastAPI, func: Callable[..., Union[Trace, Awaitable[Trace]]]
) -> None:
router = APIRouter(
tags=["predictions"],
)
schema = _get_schema(func)
@router.post("/predict", status_code=status.HTTP_200_OK, response_model=Trace)
async def predict(input_value: schema) -> Trace: # type: ignore
try:
if inspect.iscoroutinefunction(func):
return await cast(Callable[..., Awaitable[Trace]], func)(
**cast(BaseModel, input_value).dict()
)
return cast(Callable[..., Trace], func)(
**cast(BaseModel, input_value).dict()
)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"The following exception has occurred: {type(e).__name__}: {e}",
)
app.include_router(router)
def _get_schema(func: Callable) -> Type[BaseModel]:
signature = inspect.signature(func)
parameters = {
p.name: (
p.annotation if p.annotation != inspect._empty else Any,
p.default if p.default != inspect._empty else ...,
)
for p in signature.parameters.values()
if p.name in get_function_metadata_store(func).input_parameter_names
}
schema: Type[BaseModel] = create_model("InputModel", **parameters) # type: ignore
return schema

0
great_ai/deploy/routes/dashboard/assets/__init__.py
View file

8
great_ai/deploy/routes/dashboard/assets/index.css
View file

@ -129,13 +129,17 @@ main > header > div > h1 {
.version-tag {
border-radius: var(--border-radius);
background: #ddd;
display: inline-block;
font-size: 1rem;
padding: 3px 6px;
padding: 3px 8px;
margin-left: var(--small-padding)
}
main > header .version-tag {
background: var(--background-color);
vertical-align: 4px;
}
main > header > *:nth-child(2) {
min-width: 250px;
max-width: 550px;

3
great_ai/deploy/routes/dashboard/create_dash_app.py
View file

@ -21,12 +21,13 @@ from .get_traces_table import get_traces_table
def create_dash_app(function_name: str, version: str, function_docs: str) -> Flask:
accent_color = text_to_hex_color(function_name)
function_name = snake_case_to_text(function_name)
app = Dash(
function_name,
requests_pathname_prefix=DASHBOARD_PATH + "/",
server=Flask(__name__),
title=snake_case_to_text(function_name),
title=function_name,
update_title=None,
external_stylesheets=[
"/assets/index.css",

6
great_ai/deploy/routes/dashboard/get_footer.py
View file

@ -1,5 +1,7 @@
from dash import html
from great_ai import __version__
from ....constants import GITHUB_LINK
@ -8,7 +10,9 @@ def get_footer() -> html.Footer:
[
html.Div(
[
html.H6("GreatAI"),
html.H6(
["GreatAI", html.Span(__version__, className="version-tag")]
),
html.P(
"A human-friendly framework for robust end-to-end AI deployments."
),

10
great_ai/deploy/routes/route_config.py Normal file
View file

@ -0,0 +1,10 @@
from pydantic import BaseModel
class RouteConfig(BaseModel):
prediction_endpoint_enabled: bool = True
docs_endpoints_enabled: bool = True
dashboard_enabled: bool = True
feedback_endpoints_enabled: bool = True
trace_endpoints_enabled: bool = True
meta_endpoints_enabled: bool = True

1
great_ai/helper/__init__.py
View file

@ -5,4 +5,3 @@ from .hashable_base_model import HashableBaseModel
from .snake_case_to_text import snake_case_to_text
from .strip_lines import strip_lines
from .text_to_hex_color import text_to_hex_color
from .use_http_exceptions import use_http_exceptions

4
great_ai/helper/assert_function_is_not_finalised.py
View file

@ -1,10 +1,10 @@
from typing import Any, Callable
from typing import Callable
from ..exceptions import WrongDecoratorOrderError
from .get_function_metadata_store import get_function_metadata_store
def assert_function_is_not_finalised(func: Callable[..., Any]) -> None:
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."

39
great_ai/helper/freeze_arguments.py
View file

@ -1,41 +1,50 @@
import inspect
from functools import wraps
from typing import Any, Callable, Dict, List, Set, Union
from typing import Any, Callable, Mapping, Sequence, Set, TypeVar, Union, cast
from pydantic import BaseModel
F = TypeVar("F", bound=Callable)
class FrozenDict(dict):
def __hash__(self) -> int:
def __hash__(self) -> int: # type: ignore
return hash(frozenset((k, freeze(v)) for k, v in self.items()))
class FrozenList(list):
def __hash__(self) -> int:
def __hash__(self) -> int: # type: ignore
return hash(tuple(freeze(i) for i in self))
class FrozenSet(set):
def __hash__(self) -> int:
def __hash__(self) -> int: # type: ignore
return hash(frozenset(freeze(i) for i in self))
def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
"""Transform mutable dictionary
Into immutable
Useful to be compatible with cache
source: https://stackoverflow.com/questions/6358481/using-functools-lru-cache-with-dictionary-arguments
"""
def freeze_arguments(func: F) -> F:
@wraps(func)
def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
def wrapper(*args: Any, **kwargs: Any) -> Any:
args = tuple(freeze(arg) for arg in args)
kwargs = {k: freeze(v) for k, v in kwargs.items()}
return func(*args, **kwargs)
return wrapper
@wraps(func)
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
return await wrapper(*args, **kwargs)
return cast(F, async_wrapper if inspect.iscoroutinefunction(func) else wrapper)
def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
def freeze(value: Union[Sequence[Any], Mapping[str, Any], Set[Any], BaseModel]) -> Any:
"""
>>> class MyClass(BaseModel):
... a: int
>>> my_object = MyClass(a=3)
>>> my_other_object = MyClass(a=4)
>>> freeze(my_object) == freeze(my_other_object), freeze(my_object) == freeze(my_object)
(False, True)
"""
if isinstance(value, dict):
return FrozenDict(value)
@ -47,7 +56,7 @@ def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
if isinstance(value, BaseModel):
class HashableValue(type(value)):
class HashableValue(type(value)): # type: ignore
def __hash__(self) -> int:
return hash(frozenset((k, freeze(v)) for k, v in self.dict().items()))

2
great_ai/helper/get_arguments.py
View file

@ -3,7 +3,7 @@ from typing import Any, Callable, Dict, Mapping, Sequence
def get_arguments(
func: Callable[..., Any], args: Sequence[Any], kwargs: Mapping[str, Any]
func: Callable, args: Sequence[Any], kwargs: Mapping[str, Any]
) -> Dict[str, Any]:
"""Return mapping from parameter names to actual argument values"""

6
great_ai/helper/get_function_metadata_store.py
View file

@ -1,12 +1,14 @@
import inspect
from typing import Any, Callable, cast
from ..views.function_metadata import FunctionMetadata
def get_function_metadata_store(func: Callable[..., Any]) -> FunctionMetadata:
def get_function_metadata_store(func: Callable) -> FunctionMetadata:
any_func = cast(Any, func)
if not hasattr(any_func, "_great_ai_metadata"):
any_func._great_ai_metadata = FunctionMetadata()
is_asynchronous = inspect.iscoroutinefunction(func)
any_func._great_ai_metadata = FunctionMetadata(is_asynchronous=is_asynchronous)
return any_func._great_ai_metadata

6
great_ai/helper/text_to_hex_color.py
View file

@ -4,10 +4,14 @@ from hashlib import md5
def text_to_hex_color(text: str) -> str:
ascii_bytes = text.encode("ascii")
digest = md5(
ascii_bytes
).hexdigest() # the built-in hash function is salted differently in each process
integer = int(digest, 16)
hue = integer % 6311 / 6311.0
rgb = colorsys.hsv_to_rgb(hue, 0.75, 0.6)
rgb = colorsys.hsv_to_rgb(hue, 0.8, 0.6)
return "#" + "".join("%02X" % round(i * 255) for i in rgb)

20
great_ai/helper/use_http_exceptions.py
View file

@ -1,20 +0,0 @@
from functools import wraps
from typing import Any, Callable, Dict, List, TypeVar, cast
from fastapi import HTTPException, status
F = TypeVar("F", bound=Callable[..., Any])
def use_http_exceptions(func: F) -> F:
@wraps(func)
def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
try:
return func(*args, **kwargs)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"The following exception has occurred: {type(e).__name__}: {e}",
)
return cast(F, wrapper)

4
great_ai/large_file/large_file/large_file_base.py
View file

@ -89,9 +89,7 @@ class LargeFileBase(ABC):
cls.configure_credentials(**ConfigFile(secrets_path))
@classmethod
def configure_credentials(
cls,
) -> None:
def configure_credentials(cls, **kwargs: str) -> None:
cls.initialized = True
def __enter__(self) -> IO:

4
great_ai/large_file/large_file/large_file_mongo.py
View file

@ -1,7 +1,7 @@
import re
from functools import cached_property
from pathlib import Path
from typing import Any, List, Mapping
from typing import Any, List
from gridfs import DEFAULT_CHUNK_SIZE, Database, GridFSBucket
from pymongo import MongoClient
@ -27,7 +27,7 @@ class LargeFileMongo(LargeFileBase):
*,
mongo_connection_string: str,
mongo_database: str,
**_: Mapping[str, Any],
**_: Any,
) -> None:
cls.mongo_connection_string = mongo_connection_string
cls.mongo_database = mongo_database

4
great_ai/large_file/large_file/large_file_s3.py
View file

@ -1,6 +1,6 @@
from functools import cached_property
from pathlib import Path
from typing import Any, List, Mapping, Optional
from typing import Any, List, Optional
import boto3
@ -58,7 +58,7 @@ class LargeFileS3(LargeFileBase):
aws_secret_access_key: str,
large_files_bucket_name: str,
aws_endpoint_url: Optional[str] = None,
**_: Mapping[str, Any],
**_: Any,
) -> None:
cls.region_name = aws_region_name
cls.access_key_id = aws_access_key_id

6
great_ai/models/use_model.py
View file

@ -21,7 +21,7 @@ from ..helper.assert_function_is_not_finalised import assert_function_is_not_fin
from ..tracing.tracing_context import TracingContext
from ..views import Model
F = TypeVar("F", bound=Callable[..., Any])
F = TypeVar("F", bound=Callable)
def use_model(
@ -70,4 +70,6 @@ def _load_model(key: str, version: Optional[int] = None) -> Tuple[Any, int]:
return path, file.version
with file as f:
return load(f), file.version
loaded = load(f)
return loaded, file.version

4
great_ai/parameters/automatically_decorate_parameters.py
View file

@ -1,10 +1,10 @@
import inspect
from typing import Any, Callable, TypeVar
from typing import Callable, TypeVar
from ..helper.get_function_metadata_store import get_function_metadata_store
from .parameter import parameter
F = TypeVar("F", bound=Callable[..., Any])
F = TypeVar("F", bound=Callable)
def automatically_decorate_parameters(func: F) -> F:

5
great_ai/parameters/parameter.py
View file

@ -8,14 +8,13 @@ from ..helper import get_arguments, get_function_metadata_store
from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
from ..tracing.tracing_context import TracingContext
T = TypeVar("T")
F = TypeVar("F", bound=Callable[..., Any])
F = TypeVar("F", bound=Callable)
def parameter(
parameter_name: str,
*,
validator: Callable[[T], bool] = lambda _: True,
validator: Callable[[Any], bool] = lambda _: True,
disable_logging: bool = False,
) -> Callable[[F], F]:
def decorator(func: F) -> F:

64
great_ai/persistence/mongodb_driver.py
View file

@ -1,5 +1,5 @@
from datetime import datetime
from typing import Any, List, Mapping, Optional, Sequence, Tuple
from typing import Any, Dict, List, Optional, Sequence, Tuple
from pymongo import MongoClient
@ -20,6 +20,9 @@ operator_mapping = {
class MongodbDriver(TracingDatabaseDriver):
is_production_ready = True
mongo_connection_string: str
mongo_database: str
def __init__(self) -> None:
super().__init__()
if self.mongo_connection_string is None or self.mongo_database is None:
@ -33,7 +36,7 @@ class MongodbDriver(TracingDatabaseDriver):
*,
mongo_connection_string: str,
mongo_database: str,
**_: Mapping[str, Any],
**_: Any,
) -> None:
cls.mongo_connection_string = mongo_connection_string
cls.mongo_database = mongo_database
@ -43,21 +46,23 @@ class MongodbDriver(TracingDatabaseDriver):
serialized = trace.to_flat_dict()
serialized["_id"] = trace.trace_id
with MongoClient(self.mongo_connection_string) as client:
return client[self.mongo_database].traces.insert_one(serialized)
with MongoClient[Any](self.mongo_connection_string) as client:
return client[self.mongo_database].traces.insert_one(serialized).inserted_id
def save_batch(self, documents: List[Trace]) -> List[str]:
serialized = [d.to_flat_dict() for d in documents]
for s in serialized:
s["_id"] = s["trace_id"]
with MongoClient(self.mongo_connection_string) as client:
return client[self.mongo_database].traces.insert_many(
serialized, ordered=False
with MongoClient[Any](self.mongo_connection_string) as client:
return (
client[self.mongo_database]
.traces.insert_many(serialized, ordered=False)
.inserted_ids
)
def get(self, id: str) -> Optional[Trace]:
with MongoClient(self.mongo_connection_string) as client:
with MongoClient[Any](self.mongo_connection_string) as client:
value = client[self.mongo_database].traces.find_one(id)
if value:
@ -83,15 +88,8 @@ class MongodbDriver(TracingDatabaseDriver):
sort_by: Sequence[SortBy] = [],
) -> Tuple[List[Trace], int]:
query = {
"filter": {
"$and": [{"tags": tag} for tag in conjunctive_tags]
+ [
{f.property: {self._get_operator(f): f.value}}
for f in conjunctive_filters
]
+ [{}]
},
query: Dict[str, Any] = {
"filter": {},
"sort": [
(col.column_id, 1 if col.direction == "asc" else -1) for col in sort_by
],
@ -103,35 +101,43 @@ class MongodbDriver(TracingDatabaseDriver):
if take:
query["limit"] = take
and_query: List[Dict[str, Any]] = [{}]
and_query.extend({"tags": tag} for tag in conjunctive_tags)
and_query.extend(
{f.property: {self._get_operator(f): f.value}} for f in conjunctive_filters
)
if since:
query["filter"]["$and"].append({"created": {"$gte": since}})
and_query.append({"created": {"$gte": since}})
if until:
query["filter"]["$and"].append({"created": {"$lte": until}})
and_query.append({"created": {"$lte": until}})
if has_feedback is not None:
query["filter"]["$and"].append(
and_query.append(
{"feedback": {"$ne": None}} if has_feedback else {"feedback": None}
)
query["filter"]["$and"] = and_query
with MongoClient(self.mongo_connection_string) as client:
with MongoClient[Any](self.mongo_connection_string) as client:
values = client[self.mongo_database].traces.find(**query)
documents = [Trace.parse_obj(t) for t in values]
documents = [Trace[Any].parse_obj(t) for t in values]
return documents, len(documents)
def update(self, id: str, new_version: Trace) -> None:
serialized = new_version.to_flat_dict()
serialized["_id"] = new_version.trace_id
with MongoClient(self.mongo_connection_string) as client:
client[self.mongo_database].traces.update_one(id, new_version)
with MongoClient[Any](self.mongo_connection_string) as client:
client[self.mongo_database].traces.update_one({"_id": id}, serialized)
def delete(self, id: str) -> None:
with MongoClient(self.mongo_connection_string) as client:
client[self.mongo_database].traces.delete_one(id)
with MongoClient[Any](self.mongo_connection_string) as client:
client[self.mongo_database].traces.delete_one({"_id": id})
def delete_batch(self, ids: List[str]) -> List[str]:
def delete_batch(self, ids: List[str]) -> None:
delete_filter = {"_id": {"$in": ids}}
with MongoClient(self.mongo_connection_string) as client:
return client[self.mongo_database].traces.delete_many(delete_filter)
with MongoClient[Any](self.mongo_connection_string) as client:
client[self.mongo_database].traces.delete_many(delete_filter)

14
great_ai/persistence/parallel_tinydb_driver.py
View file

@ -1,7 +1,7 @@
from datetime import datetime
from multiprocessing import Lock
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, cast
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
import pandas as pd
from tinydb import TinyDB
@ -48,14 +48,8 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
def does_match(d: Dict[str, Any]) -> bool:
return (
not set(conjunctive_tags) - set(d["tags"])
and (
since is None
or cast(datetime, datetime.fromisoformat(d["created"])) >= since
)
and (
until is None
or cast(datetime, datetime.fromisoformat(d["created"])) <= until
)
and (since is None or datetime.fromisoformat(d["created"]) >= since)
and (until is None or datetime.fromisoformat(d["created"]) <= until)
and (
has_feedback is None or has_feedback == (d["feedback"] is not None)
)
@ -102,7 +96,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
def delete(self, id: str) -> None:
self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id))
def delete_batch(self, ids: List[str]) -> List[str]:
def delete_batch(self, ids: List[str]) -> None:
for i in ids:
self.delete(i)

2
great_ai/tracing/add_ground_truth.py
View file

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

4
great_ai/tracing/delete_ground_truth.py
View file

@ -1,5 +1,5 @@
from datetime import datetime
from typing import List, Optional, Union
from typing import List, Optional, Union, cast
from ..context import get_context
@ -19,4 +19,4 @@ def delete_ground_truth(
conjunctive_tags=tags, until=until, since=since, has_feedback=True
)
db.delete_batch([i.trace_id for i in items])
db.delete_batch([cast(str, i.trace_id) for i in items])

2
great_ai/utilities/chunk.py
View file

@ -3,7 +3,7 @@ from typing import Iterable, List, TypeVar
T = TypeVar("T")
def chunk(values: Iterable[T], chunk_size: int) -> Iterable[T]:
def chunk(values: Iterable[T], chunk_size: int) -> Iterable[List[T]]:
assert chunk_size >= 1
result: List[T] = []

16
great_ai/utilities/config_file/config_file.py
View file

@ -1,6 +1,6 @@
import os
from pathlib import Path
from typing import Dict, Iterable, Tuple, Union
from typing import Dict, ItemsView, Iterator, KeysView, Mapping, Union, ValuesView
from ..logger import get_logger
from .parse_error import ParseError
@ -11,7 +11,7 @@ ENVIRONMENT_VARIABLE_KEY_PREFIX = "ENV"
logger = get_logger("ConfigFile")
class ConfigFile:
class ConfigFile(Mapping[str, str]):
def __init__(self, path: Union[Path, str]) -> None:
if not isinstance(path, Path):
path = Path(path)
@ -24,7 +24,7 @@ class ConfigFile:
self._parse()
def _parse(self):
def _parse(self) -> None:
with open(self._path, encoding="utf-8") as f:
lines: str = f.read()
@ -72,20 +72,20 @@ class ConfigFile:
__getitem__ = __getattr__
def __iter__(self) -> Iterable[Tuple[str, str]]:
def __iter__(self) -> Iterator[str]:
return iter(self._key_values)
def __len__(self) -> int:
return len(self._key_values)
def keys(self):
def keys(self) -> KeysView[str]:
return self._key_values.keys()
def values(self):
def values(self) -> ValuesView[str]:
return self._key_values.values()
def items(self):
def items(self) -> ItemsView[str, str]:
return self._key_values.items()
def __repr__(self):
def __repr__(self) -> str:
return f"{type(self).__name__}(path={self._path}) {self._key_values}"

10
great_ai/utilities/evaluate_ranking/evaluate_ranking.py
View file

@ -1,5 +1,5 @@
from pathlib import Path
from typing import Dict, List, Optional, Union
from typing import Dict, List, Optional, TypeVar
import matplotlib
import matplotlib.pyplot as plt
@ -9,9 +9,11 @@ from sklearn.metrics import average_precision_score, precision_recall_curve
from ..unique import unique
from .draw_f1_iso_lines import draw_f1_iso_lines
T = TypeVar("T", str, float)
def evaluate_ranking(
expected: List[Union[str, float]],
expected: List[T],
actual_scores: List[float],
target_recall: float,
title: Optional[str] = "",
@ -20,7 +22,7 @@ def evaluate_ranking(
output_svg: Optional[Path] = None,
reverse_order: bool = False,
plot: bool = True,
) -> Dict[Union[str, float], float]:
) -> Dict[T, float]:
assert 0 <= target_recall <= 1
if plot and axes is None:
@ -39,7 +41,7 @@ def evaluate_ranking(
draw_f1_iso_lines(axes=ax)
results: Dict[Union[str, float], float] = {}
results: Dict[T, float] = {}
for i in range(len(classes) - 1):
binarized_expected = [
(v < classes[i]) if reverse_order else (v > classes[i])

14
great_ai/utilities/parallel_map/get_config.py
View file

@ -17,14 +17,15 @@ def get_config(
) -> ParallelMapConfiguration:
is_input_sequence = hasattr(input_values, "__len__")
input_length = len(input_values) if is_input_sequence else None # type: ignore
if concurrency is None:
concurrency = os.cpu_count() or 1
assert concurrency >= 1, "At least one mapper process has to be created"
if chunk_size is None:
if is_input_sequence:
chunk_size = max(1, ceil(len(input_values) / concurrency / 10))
if input_length is not None:
chunk_size = max(1, ceil(input_length / concurrency / 10))
else:
raise ValueError(
"The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_size`"
@ -32,19 +33,14 @@ def get_config(
assert chunk_size >= 1, "Chunks have to contain at least one element"
chunk_count: Optional[int] = None
if is_input_sequence:
chunk_count = ceil(len(input_values) / chunk_size)
if input_length is not None:
chunk_count = ceil(input_length / chunk_size)
if chunk_count < concurrency:
logger.warning(
f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks"
)
concurrency = chunk_count
if concurrency == 1:
logger.warning("Running in series, there is no reason for parallelism")
input_length = len(input_values) if is_input_sequence else None
config = ParallelMapConfiguration(
concurrency=concurrency,
chunk_count=chunk_count,

32
great_ai/utilities/parallel_map/manage_communication.py
View file

@ -1,10 +1,11 @@
import multiprocessing as mp
import queue
import traceback
from typing import Dict, Iterable, TypeVar, Union
from typing import Dict, Iterable, List, TypeVar, Union
from ..chunk import chunk
from ..logger import get_logger
from .map_result import MapResult
from .worker_exception import WorkerException
logger = get_logger("parallel_map")
@ -27,6 +28,7 @@ def manage_communication(
next_output_index = 0
read_input_length = 0
is_iteration_over = False
while not is_iteration_over or next_output_index < read_input_length:
if not is_iteration_over:
try:
@ -36,30 +38,30 @@ def manage_communication(
except StopIteration:
is_iteration_over = True
except Exception as e:
if not ignore_exceptions:
raise
else:
if ignore_exceptions:
logger.error(
f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
)
else:
raise
try:
result_chunk = output_queue.get_nowait()
for index, value, exception in result_chunk:
if exception is not None:
e, tb = exception
if not ignore_exceptions:
raise WorkerException from e
else:
result_chunk: List[MapResult] = output_queue.get_nowait()
for r in result_chunk:
if r.exception is not None:
if ignore_exceptions:
logger.error(
f"Exception {e} encountered in worker, traceback:\n{tb}"
f"Exception {r.exception} encountered in worker, traceback:\n{r.worker_traceback}"
)
else:
raise WorkerException from r.exception
if unordered:
yield value
yield r.value
next_output_index += 1
else:
indexed_results[index] = value
indexed_results[r.order] = r.value
if not unordered:
while next_output_index in indexed_results:

36
great_ai/utilities/parallel_map/manage_serial.py
View file

@ -1,36 +0,0 @@
import traceback
from typing import Callable, Iterable, TypeVar
from ..logger import get_logger
from .worker_exception import WorkerException
logger = get_logger("parallel_map")
T = TypeVar("T")
V = TypeVar("V")
def manage_serial(
*,
function: Callable[[T], V],
input_values: Iterable[T],
ignore_exceptions: bool,
) -> Iterable[V]:
try:
for v in input_values:
try:
yield function(v)
except Exception as e:
if not ignore_exceptions:
raise WorkerException from e
else:
logger.error(
f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}"
)
except Exception as e:
if not ignore_exceptions:
raise
else:
logger.error(
f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
)

8
great_ai/utilities/parallel_map/map_result.py Normal file
View file

@ -0,0 +1,8 @@
from typing import Any, NamedTuple, Optional
class MapResult(NamedTuple):
order: int
value: Any
exception: Optional[Exception] = None
worker_traceback: Optional[str] = None

63
great_ai/utilities/parallel_map/mapper_function.py
View file

@ -1,42 +1,71 @@
import asyncio
import inspect
import multiprocessing as mp
import queue
import threading
import traceback
from typing import Callable, Union
from multiprocessing.synchronize import Event
from typing import Any, Awaitable, Callable, List, TypeVar, Union, cast
import dill
from .map_result import MapResult
T = TypeVar("T")
V = TypeVar("V")
def mapper_function(
input_queue: Union[mp.Queue, queue.Queue],
output_queue: Union[mp.Queue, queue.Queue],
should_stop: Union[mp.Event, threading.Event],
map_function: Union[bytes, Callable],
should_stop: Union[Event, threading.Event],
func: Union[bytes, Callable[[T], V], Callable[[T], Awaitable[V]]],
) -> None:
try:
if isinstance(map_function, bytes):
map_function = dill.loads(map_function)
if isinstance(func, bytes):
func = cast(Callable[[T], V], dill.loads(func))
last_chunk = None
is_asynchronous = inspect.iscoroutinefunction(func)
last_chunk: List[MapResult] = []
while not should_stop.wait(0.1):
if last_chunk is None:
if not last_chunk:
try:
input_chunk = input_queue.get_nowait()
last_chunk = []
for i, v in input_chunk:
result, exception = None, None
try:
result = map_function(v)
except Exception as e:
exception = e, traceback.format_exc()
last_chunk.append((i, result, exception))
if is_asynchronous:
async def safe(i: int, value: T) -> Any:
try:
return MapResult(
i,
await cast(Callable[[T], Awaitable[V]], func)(
value
),
)
except Exception as e:
return MapResult(i, None, e, traceback.format_exc())
async def main() -> List[MapResult]:
return await asyncio.gather(
*[safe(i, v) for i, v in input_chunk]
)
last_chunk = asyncio.run(main())
else:
for i, value in input_chunk:
try:
last_chunk.append(MapResult(i, func(value)))
except Exception as e:
last_chunk.append(
MapResult(i, None, e, traceback.format_exc())
)
except queue.Empty:
pass
if last_chunk is not None:
if last_chunk:
try:
output_queue.put_nowait(last_chunk)
last_chunk = None
last_chunk = []
except queue.Full:
pass
except (KeyboardInterrupt, BrokenPipeError):

112
great_ai/utilities/parallel_map/parallel_map.py
View file

@ -1,11 +1,20 @@
import multiprocessing as mp
from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
from typing import (
Awaitable,
Callable,
Iterable,
Literal,
Optional,
Sequence,
TypeVar,
Union,
overload,
)
import dill
from .get_config import get_config
from .manage_communication import manage_communication
from .manage_serial import manage_serial
from .mapper_function import mapper_function
from .worker_exception import WorkerException
@ -15,79 +24,72 @@ V = TypeVar("V")
@overload
def parallel_map(
function: Callable[[T], V],
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Sequence[T],
*,
chunk_size: Optional[int],
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: Optional[Literal[False]],
) -> Iterable[V]:
...
@overload
def parallel_map(
function: Callable[[T], V],
input_values: Iterable[T],
*,
chunk_size: int,
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: Optional[Literal[False]],
) -> Iterable[V]:
...
@overload
def parallel_map(
function: Callable[[T], V],
input_values: Sequence[T],
*,
chunk_size: Optional[int],
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: True,
ignore_exceptions: Literal[True],
chunk_size: Optional[int] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[Optional[V]]:
...
@overload
def parallel_map(
function: Callable[[T], V],
input_values: Iterable[T],
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: int,
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: True,
ignore_exceptions: Literal[True],
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[Optional[V]]:
...
@overload
def parallel_map(
function,
input_values,
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Sequence[T],
*,
chunk_size=None,
concurrency=None,
unordered=False,
ignore_exceptions=False,
):
chunk_size: Optional[int] = ...,
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[V]:
...
@overload
def parallel_map(
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: int,
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[V]:
...
def parallel_map(
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: Optional[int] = None,
ignore_exceptions: bool = False,
concurrency: Optional[int] = None,
unordered: bool = False,
) -> Iterable[Optional[V]]:
config = get_config(
function=function,
function=func,
input_values=input_values,
chunk_size=chunk_size,
concurrency=concurrency,
)
if config.concurrency == 1:
yield from manage_serial(
function=function,
input_values=input_values,
ignore_exceptions=ignore_exceptions,
)
ctx = (
mp.get_context("fork")
if "fork" in mp.get_all_start_methods()
@ -99,7 +101,7 @@ def parallel_map(
output_queue = manager.Queue(config.concurrency * 2)
should_stop = ctx.Event()
serialized_map_function = dill.dumps(function, byref=True, recurse=False)
serialized_map_function = dill.dumps(func, byref=True, recurse=False)
processes = [
ctx.Process(
@ -110,7 +112,7 @@ def parallel_map(
input_queue=input_queue,
output_queue=output_queue,
should_stop=should_stop,
map_function=serialized_map_function,
func=serialized_map_function,
),
)
for i in range(config.concurrency)

114
great_ai/utilities/parallel_map/threaded_parallel_map.py
View file

@ -1,10 +1,19 @@
import queue
import threading
from typing import Callable, Iterable, Literal, Optional, Sequence, TypeVar, overload
from typing import (
Awaitable,
Callable,
Iterable,
Literal,
Optional,
Sequence,
TypeVar,
Union,
overload,
)
from .get_config import get_config
from .manage_communication import manage_communication
from .manage_serial import manage_serial
from .mapper_function import mapper_function
T = TypeVar("T")
@ -13,81 +22,74 @@ V = TypeVar("V")
@overload
def threaded_parallel_map(
function: Callable[[T], V],
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Sequence[T],
*,
chunk_size: Optional[int],
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: Optional[Literal[False]],
) -> Iterable[V]:
...
@overload
def threaded_parallel_map(
function: Callable[[T], V],
input_values: Iterable[T],
*,
chunk_size: int,
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: Optional[Literal[False]],
) -> Iterable[V]:
...
@overload
def threaded_parallel_map(
function: Callable[[T], V],
input_values: Sequence[T],
*,
chunk_size: Optional[int],
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: True,
ignore_exceptions: Literal[True],
chunk_size: Optional[int] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[Optional[V]]:
...
@overload
def threaded_parallel_map(
function: Callable[[T], V],
input_values: Iterable[T],
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: int,
concurrency: Optional[int],
unordered: Optional[bool],
ignore_exceptions: True,
ignore_exceptions: Literal[True],
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[Optional[V]]:
...
@overload
def threaded_parallel_map(
function,
input_values,
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Sequence[T],
*,
chunk_size=None,
concurrency=None,
unordered=False,
ignore_exceptions=False,
):
chunk_size: Optional[int] = ...,
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[V]:
...
@overload
def threaded_parallel_map(
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: int,
ignore_exceptions: Literal[False] = ...,
concurrency: Optional[int] = ...,
unordered: bool = ...,
) -> Iterable[V]:
...
def threaded_parallel_map(
func: Callable[[T], Union[V, Awaitable[V]]],
input_values: Union[Iterable[T], Sequence[T]],
*,
chunk_size: Optional[int] = None,
ignore_exceptions: bool = False,
concurrency: Optional[int] = None,
unordered: bool = False,
) -> Iterable[Optional[V]]:
config = get_config(
function=function,
function=func,
input_values=input_values,
chunk_size=chunk_size,
concurrency=concurrency,
)
if config.concurrency == 1:
yield from manage_serial(
function=function,
input_values=input_values,
ignore_exceptions=ignore_exceptions,
)
input_queue = queue.Queue(config.concurrency * 2)
output_queue = queue.Queue(config.concurrency * 2)
input_queue: queue.Queue = queue.Queue(config.concurrency * 2)
output_queue: queue.Queue = queue.Queue(config.concurrency * 2)
should_stop = threading.Event()
threads = [
@ -99,7 +101,7 @@ def threaded_parallel_map(
input_queue=input_queue,
output_queue=output_queue,
should_stop=should_stop,
map_function=function,
func=func,
),
)
for i in range(config.concurrency)

1
great_ai/views/function_metadata.py
View file

@ -4,6 +4,7 @@ from pydantic import BaseModel
class FunctionMetadata(BaseModel):
is_asynchronous: bool
input_parameter_names: List[str] = []
model_parameter_names: List[str] = []
is_finalised: bool = False

8
great_ai/views/trace.py
View file

@ -25,10 +25,10 @@ class Trace(Generic[T], HashableBaseModel):
extra = Extra.ignore
@validator("trace_id", always=True)
def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
if not v:
return str(uuid4())
return v
def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> str:
if v:
return v
return str(uuid4())
@property
def input(self) -> Any: