_app.py 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import asyncio
  15. import contextlib
  16. import json
  17. from typing import (
  18. Any,
  19. AsyncGenerator,
  20. Callable,
  21. Dict,
  22. Generic,
  23. List,
  24. Optional,
  25. Tuple,
  26. TypedDict,
  27. TypeVar,
  28. )
  29. import aiohttp
  30. import fastapi
  31. from fastapi.encoders import jsonable_encoder
  32. from fastapi.exceptions import RequestValidationError
  33. from fastapi.responses import JSONResponse
  34. from starlette.exceptions import HTTPException
  35. from typing_extensions import ParamSpec, TypeGuard
  36. from ....utils import logging
  37. from ...pipelines import BasePipeline
  38. from ..infra.config import AppConfig
  39. from ..infra.models import NoResultResponse
  40. from ..infra.utils import call_async, generate_log_id
  41. _PipelineT = TypeVar("_PipelineT", bound=BasePipeline)
  42. _P = ParamSpec("_P")
  43. _R = TypeVar("_R")
  44. class _Error(TypedDict):
  45. error: str
  46. def _is_error(obj: object) -> TypeGuard[_Error]:
  47. return (
  48. isinstance(obj, dict)
  49. and obj.keys() == {"error"}
  50. and isinstance(obj["error"], str)
  51. )
  52. # XXX: Since typing info (e.g., the pipeline class) cannot be easily obtained
  53. # without abstraction leaks, generic classes do not offer additional benefits
  54. # for type hinting. However, I would stick with the current design, as it does
  55. # not introduce runtime overhead at the moment and may prove useful in the
  56. # future.
  57. class PipelineWrapper(Generic[_PipelineT]):
  58. def __init__(self, pipeline: _PipelineT) -> None:
  59. super().__init__()
  60. self._pipeline = pipeline
  61. self._lock = asyncio.Lock()
  62. @property
  63. def pipeline(self) -> _PipelineT:
  64. return self._pipeline
  65. async def infer(self, *args: Any, **kwargs: Any) -> List[Any]:
  66. def _infer() -> List[Any]:
  67. output: list = []
  68. with contextlib.closing(self._pipeline(*args, **kwargs)) as it:
  69. for item in it:
  70. if _is_error(item):
  71. raise fastapi.HTTPException(
  72. status_code=500, detail=item["error"]
  73. )
  74. output.append(item)
  75. return output
  76. return await self.call(_infer)
  77. async def call(
  78. self, func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs
  79. ) -> _R:
  80. async with self._lock:
  81. return await call_async(func, *args, **kwargs)
  82. class AppContext(Generic[_PipelineT]):
  83. def __init__(self, *, config: AppConfig) -> None:
  84. super().__init__()
  85. self._config = config
  86. self.extra: Dict[str, Any] = {}
  87. self._pipeline: Optional[PipelineWrapper[_PipelineT]] = None
  88. self._aiohttp_session: Optional[aiohttp.ClientSession] = None
  89. @property
  90. def config(self) -> AppConfig:
  91. return self._config
  92. @property
  93. def pipeline(self) -> PipelineWrapper[_PipelineT]:
  94. if not self._pipeline:
  95. raise AttributeError("`pipeline` has not been set.")
  96. return self._pipeline
  97. @pipeline.setter
  98. def pipeline(self, val: PipelineWrapper[_PipelineT]) -> None:
  99. self._pipeline = val
  100. @property
  101. def aiohttp_session(self) -> aiohttp.ClientSession:
  102. if not self._aiohttp_session:
  103. raise AttributeError("`aiohttp_session` has not been set.")
  104. return self._aiohttp_session
  105. @aiohttp_session.setter
  106. def aiohttp_session(self, val: aiohttp.ClientSession) -> None:
  107. self._aiohttp_session = val
  108. def create_app(
  109. *, pipeline: _PipelineT, app_config: AppConfig, app_aiohttp_session: bool = True
  110. ) -> Tuple[fastapi.FastAPI, AppContext[_PipelineT]]:
  111. @contextlib.asynccontextmanager
  112. async def _app_lifespan(app: fastapi.FastAPI) -> AsyncGenerator[None, None]:
  113. ctx.pipeline = PipelineWrapper[_PipelineT](pipeline)
  114. if app_aiohttp_session:
  115. async with aiohttp.ClientSession(
  116. cookie_jar=aiohttp.DummyCookieJar()
  117. ) as aiohttp_session:
  118. ctx.aiohttp_session = aiohttp_session
  119. yield
  120. else:
  121. yield
  122. # Should we control API versions?
  123. app = fastapi.FastAPI(lifespan=_app_lifespan)
  124. ctx = AppContext[_PipelineT](config=app_config)
  125. app.state.context = ctx
  126. @app.get("/health", operation_id="checkHealth")
  127. async def _check_health() -> NoResultResponse:
  128. return NoResultResponse(
  129. logId=generate_log_id(), errorCode=0, errorMsg="Healthy"
  130. )
  131. @app.exception_handler(RequestValidationError)
  132. async def _validation_exception_handler(
  133. request: fastapi.Request, exc: RequestValidationError
  134. ) -> JSONResponse:
  135. json_compatible_data = jsonable_encoder(
  136. NoResultResponse(
  137. logId=generate_log_id(),
  138. errorCode=422,
  139. errorMsg=json.dumps(exc.errors()),
  140. )
  141. )
  142. return JSONResponse(content=json_compatible_data, status_code=422)
  143. @app.exception_handler(HTTPException)
  144. async def _http_exception_handler(
  145. request: fastapi.Request, exc: HTTPException
  146. ) -> JSONResponse:
  147. json_compatible_data = jsonable_encoder(
  148. NoResultResponse(
  149. logId=generate_log_id(), errorCode=exc.status_code, errorMsg=exc.detail
  150. )
  151. )
  152. return JSONResponse(content=json_compatible_data, status_code=exc.status_code)
  153. @app.exception_handler(Exception)
  154. async def _unexpected_exception_handler(
  155. request: fastapi.Request, exc: Exception
  156. ) -> JSONResponse:
  157. # XXX: The default server will duplicate the error message. Is it
  158. # necessary to log the exception info here?
  159. logging.exception("Unhandled exception")
  160. json_compatible_data = jsonable_encoder(
  161. NoResultResponse(
  162. logId=generate_log_id(),
  163. errorCode=500,
  164. errorMsg="Internal server error",
  165. )
  166. )
  167. return JSONResponse(content=json_compatible_data, status_code=500)
  168. return app, ctx
  169. # TODO: Precise type hints
  170. def primary_operation(
  171. app: fastapi.FastAPI, path: str, operation_id: str, **kwargs: Any
  172. ) -> Callable:
  173. return app.post(
  174. path,
  175. operation_id=operation_id,
  176. responses={422: {"model": NoResultResponse}, 500: {"model": NoResultResponse}},
  177. response_model_exclude_none=True,
  178. **kwargs,
  179. )