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

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