_app.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
  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. TypeVar,
  27. )
  28. import aiohttp
  29. import fastapi
  30. from fastapi.encoders import jsonable_encoder
  31. from fastapi.exceptions import RequestValidationError
  32. from fastapi.responses import JSONResponse
  33. from starlette.exceptions import HTTPException
  34. from typing_extensions import ParamSpec
  35. from ....utils import logging
  36. from ...pipelines_new import BasePipeline
  37. from ..infra.config import AppConfig
  38. from ..infra.models import NoResultResponse
  39. from ..infra.utils import call_async, generate_log_id
  40. _PipelineT = TypeVar("_PipelineT", bound=BasePipeline)
  41. _P = ParamSpec("_P")
  42. _R = TypeVar("_R")
  43. # XXX: Since typing info (e.g., the pipeline class) cannot be easily obtained
  44. # without abstraction leaks, generic classes do not offer additional benefits
  45. # for type hinting. However, I would stick with the current design, as it does
  46. # not introduce runtime overhead at the moment and may prove useful in the
  47. # future.
  48. class PipelineWrapper(Generic[_PipelineT]):
  49. def __init__(self, pipeline: _PipelineT) -> None:
  50. super().__init__()
  51. self._pipeline = pipeline
  52. self._lock = asyncio.Lock()
  53. @property
  54. def pipeline(self) -> _PipelineT:
  55. return self._pipeline
  56. async def infer(self, *args: Any, **kwargs: Any) -> List[Any]:
  57. def _infer() -> List[Any]:
  58. output = list(self._pipeline(*args, **kwargs))
  59. if (
  60. len(output) == 1
  61. and isinstance(output[0], dict)
  62. and output[0].keys() == {"error"}
  63. ):
  64. raise fastapi.HTTPException(status_code=500, detail=output[0]["error"])
  65. return output
  66. return await self.call(_infer)
  67. async def call(
  68. self, func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs
  69. ) -> _R:
  70. async with self._lock:
  71. return await call_async(func, *args, **kwargs)
  72. class AppContext(Generic[_PipelineT]):
  73. def __init__(self, *, config: AppConfig) -> None:
  74. super().__init__()
  75. self._config = config
  76. self.extra: Dict[str, Any] = {}
  77. self._pipeline: Optional[PipelineWrapper[_PipelineT]] = None
  78. self._aiohttp_session: Optional[aiohttp.ClientSession] = None
  79. @property
  80. def config(self) -> AppConfig:
  81. return self._config
  82. @property
  83. def pipeline(self) -> PipelineWrapper[_PipelineT]:
  84. if not self._pipeline:
  85. raise AttributeError("`pipeline` has not been set.")
  86. return self._pipeline
  87. @pipeline.setter
  88. def pipeline(self, val: PipelineWrapper[_PipelineT]) -> None:
  89. self._pipeline = val
  90. @property
  91. def aiohttp_session(self) -> aiohttp.ClientSession:
  92. if not self._aiohttp_session:
  93. raise AttributeError("`aiohttp_session` has not been set.")
  94. return self._aiohttp_session
  95. @aiohttp_session.setter
  96. def aiohttp_session(self, val: aiohttp.ClientSession) -> None:
  97. self._aiohttp_session = val
  98. def create_app(
  99. *, pipeline: _PipelineT, app_config: AppConfig, app_aiohttp_session: bool = True
  100. ) -> Tuple[fastapi.FastAPI, AppContext[_PipelineT]]:
  101. @contextlib.asynccontextmanager
  102. async def _app_lifespan(app: fastapi.FastAPI) -> AsyncGenerator[None, None]:
  103. ctx.pipeline = PipelineWrapper[_PipelineT](pipeline)
  104. if app_aiohttp_session:
  105. async with aiohttp.ClientSession(
  106. cookie_jar=aiohttp.DummyCookieJar()
  107. ) as aiohttp_session:
  108. ctx.aiohttp_session = aiohttp_session
  109. yield
  110. else:
  111. yield
  112. # Should we control API versions?
  113. app = fastapi.FastAPI(lifespan=_app_lifespan)
  114. ctx = AppContext[_PipelineT](config=app_config)
  115. app.state.context = ctx
  116. @app.get("/health", operation_id="checkHealth")
  117. async def _check_health() -> NoResultResponse:
  118. return NoResultResponse(
  119. logId=generate_log_id(), errorCode=0, errorMsg="Healthy"
  120. )
  121. @app.exception_handler(RequestValidationError)
  122. async def _validation_exception_handler(
  123. request: fastapi.Request, exc: RequestValidationError
  124. ) -> JSONResponse:
  125. json_compatible_data = jsonable_encoder(
  126. NoResultResponse(
  127. logId=generate_log_id(),
  128. errorCode=422,
  129. errorMsg=json.dumps(exc.errors()),
  130. )
  131. )
  132. return JSONResponse(content=json_compatible_data, status_code=422)
  133. @app.exception_handler(HTTPException)
  134. async def _http_exception_handler(
  135. request: fastapi.Request, exc: HTTPException
  136. ) -> JSONResponse:
  137. json_compatible_data = jsonable_encoder(
  138. NoResultResponse(
  139. logId=generate_log_id(), errorCode=exc.status_code, errorMsg=exc.detail
  140. )
  141. )
  142. return JSONResponse(content=json_compatible_data, status_code=exc.status_code)
  143. @app.exception_handler(Exception)
  144. async def _unexpected_exception_handler(
  145. request: fastapi.Request, exc: Exception
  146. ) -> JSONResponse:
  147. # XXX: The default server will duplicate the error message. Is it
  148. # necessary to log the exception info here?
  149. logging.exception("Unhandled exception")
  150. json_compatible_data = jsonable_encoder(
  151. NoResultResponse(
  152. logId=generate_log_id(),
  153. errorCode=500,
  154. errorMsg="Internal server error",
  155. )
  156. )
  157. return JSONResponse(content=json_compatible_data, status_code=500)
  158. return app, ctx
  159. # TODO: Precise type hints
  160. def primary_operation(
  161. app: fastapi.FastAPI, path: str, operation_id: str, **kwargs: Any
  162. ) -> Callable:
  163. return app.post(
  164. path,
  165. operation_id=operation_id,
  166. responses={422: {"model": NoResultResponse}, 500: {"model": NoResultResponse}},
  167. response_model_exclude_none=True,
  168. **kwargs,
  169. )