Fix typing and minor issues

This commit is contained in:
Andras Schmelczer 2022-07-07 14:12:44 +02:00
commit 72ab627a34
54 changed files with 609 additions and 563 deletions

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,

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
},
},

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,

View file

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

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)

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

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()

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)

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

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;

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",

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."
),

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

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

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."

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()))

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"""

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

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)

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)

Some files were not shown because too many files have changed in this diff Show more