143 lines
4.8 KiB
Python
143 lines
4.8 KiB
Python
import inspect
|
|
from types import UnionType
|
|
from typing import (
|
|
Annotated,
|
|
Any,
|
|
Awaitable,
|
|
Callable,
|
|
Type,
|
|
Union,
|
|
cast,
|
|
get_args,
|
|
get_origin,
|
|
)
|
|
|
|
import numpy as np
|
|
from fastapi import APIRouter, FastAPI, HTTPException, status
|
|
from numpy.typing import NDArray
|
|
from pydantic import (
|
|
AfterValidator,
|
|
BaseModel,
|
|
PlainSerializer,
|
|
PlainValidator,
|
|
WithJsonSchema,
|
|
create_model,
|
|
)
|
|
|
|
from ...errors import ArgumentValidationError
|
|
from ...helper import get_function_metadata_store
|
|
from ...helper.get_function_type_hints import get_function_type_hints
|
|
from ...helper.numpy_to_python import numpy_to_python
|
|
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)
|
|
return_type = _get_api_type(get_function_type_hints(func).get("return", Any))
|
|
response_model = Trace[return_type] # type: ignore[valid-type]
|
|
positional_names = [
|
|
p.name
|
|
for p in inspect.signature(func).parameters.values()
|
|
if p.kind == p.POSITIONAL_ONLY and p.name in schema.model_fields
|
|
]
|
|
|
|
@router.post(
|
|
"/predict", status_code=status.HTTP_200_OK, response_model=response_model
|
|
)
|
|
async def predict(input_value: schema) -> Trace: # type: ignore
|
|
try:
|
|
# model_dump() recursively converts nested models to dictionaries,
|
|
# losing the validated argument types expected by the function.
|
|
arguments = {
|
|
name: getattr(input_value, name) for name in schema.model_fields
|
|
}
|
|
positional = [arguments.pop(name) for name in positional_names]
|
|
if inspect.iscoroutinefunction(func):
|
|
return await cast(Callable[..., Awaitable[Trace]], func)(
|
|
*positional, **arguments
|
|
)
|
|
return cast(Callable[..., Trace], func)(*positional, **arguments)
|
|
except ArgumentValidationError as e:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
|
detail=str(e),
|
|
) from e
|
|
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)
|
|
type_hints = get_function_type_hints(func)
|
|
parameters = {
|
|
p.name: (
|
|
_get_api_type(type_hints.get(p.name, 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
|
|
|
|
|
|
def _validate_numpy_array(value: Any) -> np.ndarray:
|
|
if not isinstance(value, (list, np.ndarray)):
|
|
raise ValueError("Expected an array (a JSON list or NumPy ndarray)")
|
|
return np.asarray(value)
|
|
|
|
|
|
_NumpyArray = Annotated[
|
|
np.ndarray,
|
|
PlainValidator(_validate_numpy_array),
|
|
PlainSerializer(lambda value: value.tolist(), when_used="json"),
|
|
WithJsonSchema({"type": "array", "items": {}}),
|
|
]
|
|
|
|
|
|
def _get_api_type(annotation: Any) -> Any:
|
|
"""Represent NumPy arrays as JSON arrays without losing other annotations."""
|
|
origin = get_origin(annotation)
|
|
if annotation in (np.ndarray, NDArray) or origin in (np.ndarray, NDArray):
|
|
return _NumpyArray
|
|
if isinstance(annotation, type) and issubclass(annotation, np.generic):
|
|
scalar_types = (
|
|
(np.bool_, bool),
|
|
(np.integer, int),
|
|
(np.floating, float),
|
|
(np.complexfloating, complex),
|
|
(np.str_, str),
|
|
(np.bytes_, bytes),
|
|
)
|
|
for numpy_type, python_type in scalar_types:
|
|
if issubclass(annotation, numpy_type):
|
|
return Annotated[
|
|
python_type,
|
|
AfterValidator(annotation),
|
|
PlainSerializer(
|
|
numpy_to_python, return_type=python_type, when_used="json"
|
|
),
|
|
]
|
|
arguments = get_args(annotation)
|
|
if origin is Annotated:
|
|
return Annotated[_get_api_type(arguments[0]), *arguments[1:]]
|
|
converted = tuple(_get_api_type(argument) for argument in arguments)
|
|
if all(old is new for old, new in zip(arguments, converted)):
|
|
return annotation
|
|
if origin in (Union, UnionType):
|
|
return Union[converted]
|
|
if hasattr(annotation, "copy_with"):
|
|
return annotation.copy_with(converted)
|
|
return origin[converted]
|