Move files
This commit is contained in:
parent
3cf28379e8
commit
00cc8225c5
159 changed files with 31 additions and 49 deletions
|
|
@ -1,267 +0,0 @@
|
|||
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 great_ai.utilities import parallel_map
|
||||
|
||||
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 ..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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue