Fix typing and minor issues
This commit is contained in:
parent
2db2253578
commit
72ab627a34
54 changed files with 635 additions and 589 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue