PaddleX/paddlex/inference/serving/basic_serving/_app.py

# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import contextlib
import json
from queue import Queue
from threading import Thread
from typing import (
    Any,
    AsyncGenerator,
    Callable,
    Dict,
    Generic,
    List,
    Optional,
    Tuple,
    TypedDict,
    TypeVar,
)

from typing_extensions import ParamSpec, TypeGuard

from ....utils import logging
from ....utils.deps import class_requires_deps, function_requires_deps, is_dep_available
from ...pipelines import BasePipeline
from ..infra.config import AppConfig
from ..infra.models import AIStudioNoResultResponse
from ..infra.utils import call_async, generate_log_id

if is_dep_available("aiohttp"):
    import aiohttp
if is_dep_available("fastapi"):
    import fastapi
    from fastapi.encoders import jsonable_encoder
    from fastapi.exceptions import RequestValidationError
    from fastapi.responses import JSONResponse
if is_dep_available("starlette"):
    from starlette.exceptions import HTTPException

PipelineT = TypeVar("PipelineT", bound=BasePipeline)
P = ParamSpec("P")
R = TypeVar("R")


class _Error(TypedDict):
    error: str


def _is_error(obj: object) -> TypeGuard[_Error]:
    return (
        isinstance(obj, dict)
        and obj.keys() == {"error"}
        and isinstance(obj["error"], str)
    )


# XXX: Since typing info (e.g., the pipeline class) cannot be easily obtained
# without abstraction leaks, generic classes do not offer additional benefits
# for type hinting. However, I would stick with the current design, as it does
# not introduce runtime overhead at the moment and may prove useful in the
# future.
@class_requires_deps("fastapi")
class PipelineWrapper(Generic[PipelineT]):
    def __init__(self, pipeline: PipelineT) -> None:
        super().__init__()
        self._pipeline = pipeline
        # HACK: We work around a bug in Paddle Inference by performing all
        # inference in the same thread.
        self._queue = Queue()
        self._closed = False
        self._loop = asyncio.get_running_loop()
        self._thread = Thread(target=self._worker, daemon=False)
        self._thread.start()

    @property
    def pipeline(self) -> PipelineT:
        return self._pipeline

    async def infer(self, *args: Any, **kwargs: Any) -> List[Any]:
        def _infer(*args, **kwargs) -> List[Any]:
            output: list = []
            with contextlib.closing(self._pipeline.predict(*args, **kwargs)) as it:
                for item in it:
                    if _is_error(item):
                        raise fastapi.HTTPException(
                            status_code=500, detail=item["error"]
                        )
                    output.append(item)

            return output

        return await self.call(_infer, *args, **kwargs)

    async def call(self, func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
        if self._closed:
            raise RuntimeError("`PipelineWrapper` has already been closed")
        fut = self._loop.create_future()
        self._queue.put((func, args, kwargs, fut))
        return await fut

    async def close(self):
        if not self._closed:
            self._queue.put(None)
            await call_async(self._thread.join)
            self._closed = True

    def _worker(self):
        while not self._closed:
            item = self._queue.get()
            if item is None:
                break
            func, args, kwargs, fut = item
            try:
                result = func(*args, **kwargs)
                self._loop.call_soon_threadsafe(fut.set_result, result)
            except Exception as e:
                self._loop.call_soon_threadsafe(fut.set_exception, e)
            finally:
                self._queue.task_done()


@class_requires_deps("aiohttp")
class AppContext(Generic[PipelineT]):
    def __init__(self, *, config: AppConfig) -> None:
        super().__init__()
        self._config = config
        self.extra: Dict[str, Any] = {}
        self._pipeline: Optional[PipelineWrapper[PipelineT]] = None
        self._aiohttp_session: Optional[aiohttp.ClientSession] = None

    @property
    def config(self) -> AppConfig:
        return self._config

    @property
    def pipeline(self) -> PipelineWrapper[PipelineT]:
        if not self._pipeline:
            raise AttributeError("`pipeline` has not been set.")
        return self._pipeline

    @pipeline.setter
    def pipeline(self, val: PipelineWrapper[PipelineT]) -> None:
        self._pipeline = val

    @property
    def aiohttp_session(self) -> "aiohttp.ClientSession":
        if not self._aiohttp_session:
            raise AttributeError("`aiohttp_session` has not been set.")
        return self._aiohttp_session

    @aiohttp_session.setter
    def aiohttp_session(self, val: "aiohttp.ClientSession") -> None:
        self._aiohttp_session = val


@function_requires_deps("fastapi", "aiohttp", "starlette")
def create_app(
    *, pipeline: PipelineT, app_config: AppConfig, app_aiohttp_session: bool = True
) -> Tuple["fastapi.FastAPI", AppContext[PipelineT]]:
    @contextlib.asynccontextmanager
    async def _app_lifespan(app: "fastapi.FastAPI") -> AsyncGenerator[None, None]:
        ctx.pipeline = PipelineWrapper[PipelineT](pipeline)
        try:
            if app_aiohttp_session:
                async with aiohttp.ClientSession(
                    cookie_jar=aiohttp.DummyCookieJar()
                ) as aiohttp_session:
                    ctx.aiohttp_session = aiohttp_session
                    yield
            else:
                yield
        finally:
            await ctx.pipeline.close()

    # Should we control API versions?
    app = fastapi.FastAPI(lifespan=_app_lifespan)
    ctx = AppContext[PipelineT](config=app_config)
    app.state.context = ctx

    @app.get("/health", operation_id="checkHealth")
    async def _check_health() -> AIStudioNoResultResponse:
        return AIStudioNoResultResponse(
            logId=generate_log_id(), errorCode=0, errorMsg="Healthy"
        )

    @app.exception_handler(RequestValidationError)
    async def _validation_exception_handler(
        request: fastapi.Request, exc: RequestValidationError
    ) -> JSONResponse:
        json_compatible_data = jsonable_encoder(
            AIStudioNoResultResponse(
                logId=generate_log_id(),
                errorCode=422,
                errorMsg=json.dumps(exc.errors()),
            )
        )
        return JSONResponse(content=json_compatible_data, status_code=422)

    @app.exception_handler(HTTPException)
    async def _http_exception_handler(
        request: fastapi.Request, exc: HTTPException
    ) -> JSONResponse:
        json_compatible_data = jsonable_encoder(
            AIStudioNoResultResponse(
                logId=generate_log_id(), errorCode=exc.status_code, errorMsg=exc.detail
            )
        )
        return JSONResponse(content=json_compatible_data, status_code=exc.status_code)

    @app.exception_handler(Exception)
    async def _unexpected_exception_handler(
        request: fastapi.Request, exc: Exception
    ) -> JSONResponse:
        # XXX: The default server will duplicate the error message. Is it
        # necessary to log the exception info here?
        logging.exception("Unhandled exception")
        json_compatible_data = jsonable_encoder(
            AIStudioNoResultResponse(
                logId=generate_log_id(),
                errorCode=500,
                errorMsg="Internal server error",
            )
        )
        return JSONResponse(content=json_compatible_data, status_code=500)

    return app, ctx


# TODO: Precise type hints
@function_requires_deps("fastapi")
def primary_operation(
    app: "fastapi.FastAPI", path: str, operation_id: str, **kwargs: Any
) -> Callable:
    return app.post(
        path,
        operation_id=operation_id,
        responses={
            422: {"model": AIStudioNoResultResponse},
            500: {"model": AIStudioNoResultResponse},
        },
        response_model_exclude_none=True,
        **kwargs,
    )