266 lines
8.1 KiB
Python
266 lines
8.1 KiB
Python
import inspect
|
|
from functools import lru_cache, partial, wraps
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
Generic,
|
|
Iterable,
|
|
List,
|
|
Optional,
|
|
Type,
|
|
TypeVar,
|
|
cast,
|
|
overload,
|
|
)
|
|
|
|
from fastapi import APIRouter, FastAPI, status
|
|
from pydantic import BaseModel, create_model
|
|
|
|
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 ..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 .routes.bootstrap_dashboard import bootstrap_dashboard
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class GreatAI(Generic[T]):
|
|
def __init__(self, func: Callable[..., Any], version: str, return_raw_result: bool):
|
|
is_asynchronous = inspect.iscoroutinefunction(func)
|
|
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
|
|
|
|
wraps(func)(self)
|
|
|
|
self._version = version
|
|
|
|
self.app = FastAPI(
|
|
title=self.name,
|
|
version=self.version,
|
|
description=self.documentation
|
|
+ f"\n\nFind out more in the [dashboard]({DASHBOARD_PATH}).",
|
|
docs_url=None,
|
|
redoc_url=None,
|
|
)
|
|
|
|
@overload
|
|
@staticmethod
|
|
def create(
|
|
func: Optional[Callable[..., T]] = None,
|
|
) -> "GreatAI[T]":
|
|
...
|
|
|
|
@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]"]:
|
|
...
|
|
|
|
@staticmethod
|
|
def create(
|
|
func: Optional[Callable[..., T]] = None,
|
|
*,
|
|
version: str = "0.0.1",
|
|
return_raw_result: bool = False,
|
|
disable_rest_api: bool = False,
|
|
disable_docs: bool = False,
|
|
disable_dashboard: bool = False,
|
|
):
|
|
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
|
|
)
|
|
|
|
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 process_batch(
|
|
self,
|
|
batch: Iterable[Any],
|
|
concurrency: Optional[int] = None,
|
|
do_not_persist_traces: bool = False,
|
|
) -> List[Trace[T]]:
|
|
return list(
|
|
parallel_map(
|
|
freeze_arguments(
|
|
partial(
|
|
self._cached_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__)
|
|
|
|
@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}"
|
|
|
|
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()
|
|
)
|
|
)
|
|
)
|
|
|
|
def _bootstrap_rest_api(self, disable_docs: bool, disable_dashboard: bool) -> None:
|
|
self._bootstrap_prediction_endpoint()
|
|
|
|
if not disable_docs:
|
|
bootstrap_docs_endpoints(self.app)
|
|
|
|
if not disable_dashboard:
|
|
bootstrap_dashboard(
|
|
self.app,
|
|
function_name=self._func.__name__,
|
|
documentation=self.documentation,
|
|
)
|
|
bootstrap_trace_endpoints(self.app)
|
|
|
|
bootstrap_feedback_endpoints(self.app)
|
|
self._bootstrap_meta_endpoints()
|
|
|
|
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 ...,
|
|
)
|
|
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)
|