Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: add response headers for error responses #4622

Merged
merged 1 commit into from
Apr 1, 2024
Merged

fix: add response headers for error responses #4622

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 22 additions & 13 deletions src/_bentoml_impl/server/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,10 +171,12 @@ async def handle_uncaught_exception(self, req: Request, exc: Exception) -> Respo
from starlette.responses import JSONResponse

log_exception(req, sys.exc_info())
return JSONResponse(
resp = JSONResponse(
{"error": "An unexpected error has occurred, please check the server log."},
status_code=500,
)
self._add_response_headers(resp)
return resp

async def handle_validation_error(self, req: Request, exc: Exception) -> Response:
from starlette.responses import JSONResponse
Expand All @@ -185,7 +187,9 @@ async def handle_validation_error(self, req: Request, exc: Exception) -> Respons
"error": f"{exc.error_count()} validation error for {exc.title}",
"detail": exc.errors(include_context=False),
}
return JSONResponse(data, status_code=400)
resp = JSONResponse(data, status_code=400)
self._add_response_headers(resp)
return resp

async def handle_bentoml_exception(self, req: Request, exc: Exception) -> Response:
from starlette.responses import JSONResponse
Expand All @@ -194,12 +198,14 @@ async def handle_bentoml_exception(self, req: Request, exc: Exception) -> Respon
assert isinstance(exc, BentoMLException)
status = exc.error_code.value
if 400 <= status < 500 and status not in (401, 403):
return JSONResponse(
resp = JSONResponse(
{"error": f"BentoService error handling API request: {exc}"},
status_code=status,
)
else:
return JSONResponse("", status_code=status)
resp = JSONResponse("", status_code=status)
self._add_response_headers(resp)
return resp

def __call__(self) -> Starlette:
app = super().__call__()
Expand Down Expand Up @@ -307,6 +313,17 @@ def create_instance(self) -> None:
self._service_instance = self.service()
set_current_service(self._service_instance)

def _add_response_headers(self, resp: Response) -> None:
from bentoml._internal.context import trace_context

if trace_context.request_id is not None:
resp.headers["X-BentoML-Request-ID"] = str(trace_context.request_id)
if (
BentoMLContainer.http.response.trace_id.get()
and trace_context.trace_id is not None
):
resp.headers["X-BentoML-Trace-ID"] = str(trace_context.trace_id)

async def destroy_instance(self) -> None:
from _bentoml_sdk.service.dependency import cleanup

Expand Down Expand Up @@ -435,8 +452,6 @@ async def api_endpoint(self, name: str, request: Request) -> Response:

from _bentoml_sdk.io_models import ARGS
from _bentoml_sdk.io_models import KWARGS
from bentoml._internal.container import BentoMLContainer
from bentoml._internal.context import trace_context
from bentoml._internal.utils import get_original_func
from bentoml._internal.utils.http import set_cookies

Expand Down Expand Up @@ -501,13 +516,7 @@ async def inner() -> t.AsyncGenerator[t.Any, None]:
response.status_code = ctx.response.status_code
response.headers.update(ctx.response.metadata)
set_cookies(response, ctx.response.cookies)
if trace_context.request_id is not None:
response.headers["X-BentoML-Request-ID"] = str(trace_context.request_id)
if (
BentoMLContainer.http.response.trace_id.get()
and trace_context.trace_id is not None
):
response.headers["X-BentoML-Trace-ID"] = str(trace_context.trace_id)
self._add_response_headers(response)
# clean the request resources after the response is consumed.
response.background = BackgroundTask(request.close)
return response
Loading