async def relay_handler(self, request):
data = await request.read()
url = request.url.with_host(self.outbound_host).with_port(
self.outbound_port)
with get_tracer().async_span(
service_name=self.__class__.__name__,
span_name=f"[2]{url.path} relay",
request_headers=request.headers,
):
try:
client = self.get_client()
async with client.request(request.method,
url,
data=data,
headers=request.headers) as resp:
body = await resp.read()
except aiohttp.client_exceptions.ClientConnectionError:
return aiohttp.web.Response(status=503,
body=b"Service Unavailable")
return aiohttp.web.Response(
status=resp.status,
body=body,
headers=resp.headers,
)
def api_func_with_tracing():
with get_tracer().span(
service_name=f"BentoService.{self.bento_service.name}",
span_name=f"InferenceAPI {api.name} HTTP route",
request_headers=request.headers,
):
return api_func()
async def request_dispatcher(self, request):
with get_tracer().async_span(
service_name=self.__class__.__name__,
span_name="[1]http request",
is_root=True,
standalone=True,
sample_rate=0.001,
):
api_route = request.match_info.get("path")
if api_route in self.batch_handlers:
req = HTTPRequest(
tuple((k.decode(), v.decode())
for k, v in request.raw_headers),
await request.read(),
)
try:
resp = await self.batch_handlers[api_route](req)
except RemoteException as e:
# known remote exception
logger.error(traceback.format_exc())
resp = aiohttp.web.Response(
status=e.payload.status,
headers=e.payload.headers,
body=e.payload.body,
)
except Exception: # pylint: disable=broad-except
logger.error(traceback.format_exc())
resp = aiohttp.web.HTTPInternalServerError()
else:
resp = await self.relay_handler(request)
return resp
def wrapped_func(*args, **kwargs):
with get_tracer().span(
service_name=f"BentoService.{self.service.name}",
span_name=f"InferenceAPI {self.name} user defined callback function",
):
if append_arg and append_arg in kwargs:
tasks = kwargs.pop(append_arg)
elif append_arg in kwargs:
tasks = kwargs[append_arg]
else:
tasks = []
try:
return self._user_func(*args, **kwargs)
except Exception as e: # pylint: disable=broad-except
logger.error("Error caught in API function:", exc_info=1)
if self.batch:
for task in tasks:
if not task.is_discarded:
task.discard(
http_status=500,
err_msg=f"Exception happened in API function: {e}",
)
return [None] * sum(
1 if t.batch is None else t.batch for t in tasks
)
else:
task = tasks
if not task.is_discarded:
task.discard(
http_status=500,
err_msg=f"Exception happened in API function: {e}",
)
return [None] * (1 if task.batch is None else task.batch)
def handle_batch_request(self, requests: Sequence[HTTPRequest]):
with get_tracer().span(
service_name=f"BentoService.{self.service.name}",
span_name=f"InferenceAPI {self.name} handle batch requests",
):
inf_tasks = map(self.input_adapter.from_http_request, requests)
results = self.infer(inf_tasks)
responses = tuple(map(self.output_adapter.to_http_response, results))
for inf_task, response in zip(inf_tasks, responses):
response.headers['X-Request-Id'] = inf_task.task_id
return responses
def wrapped_func(*args, **kwargs):
with get_tracer().span(
service_name=f"BentoService.{self.service.name}",
span_name=
f"InferenceAPI {self.name} user defined callback function",
):
try:
results = self._user_func(*args, **kwargs)
# TODO: add checks for datatype validation so that FastAPI doesn't trip
return results
except Exception as e:
return e
async def _batch_handler_template(self, requests, api_route):
'''
batch request handler
params:
* requests: list of aiohttp request
* api_route: called API name
raise:
* RemoteException: known exceptions from model server
* Exception: other exceptions
'''
from aiohttp.client_exceptions import ClientConnectionError
from aiohttp.web import Response
headers = {MARSHAL_REQUEST_HEADER: "true"}
api_url = f"http://{self.outbound_host}:{self.outbound_port}/{api_route}"
with get_tracer().async_span(
service_name=self.__class__.__name__,
span_name=f"[2]merged {api_route}",
request_headers=headers,
):
reqs_s = DataLoader.merge_requests(requests)
try:
client = self.get_client()
async with client.post(api_url, data=reqs_s,
headers=headers) as resp:
raw = await resp.read()
except ClientConnectionError as e:
raise RemoteException(e,
payload=HTTPResponse(
status=503,
body=b"Service Unavailable"))
if resp.status != 200:
raise RemoteException(
f"Bad response status from model server:\n{resp.status}\n{raw}",
payload=HTTPResponse(
status=resp.status,
headers=tuple(resp.headers.items()),
body=raw,
),
)
merged = DataLoader.split_responses(raw)
return tuple(
Response(
body=i.body, headers=i.headers, status=i.status or 500)
for i in merged)
