async def _batch_handler_template(self, requests, api_name): ''' batch request handler params: * requests: list of aiohttp request * api_name: called API name raise: * RemoteException: known exceptions from model server * Exception: other exceptions ''' headers = {self._MARSHAL_FLAG: "true"} api_url = f"http://{self.outbound_host}:{self.outbound_port}/{api_name}" with async_trace( ZIPKIN_API_URL, service_name=self.__class__.__name__, span_name=f"[2]merged {api_name}", ) as trace_ctx: headers.update(make_http_headers(trace_ctx)) reqs_s = DataLoader.merge_requests(requests) async with aiohttp.ClientSession() as client: async with client.post(api_url, data=reqs_s, headers=headers) as resp: raw = await resp.read() if resp.status != 200: raise RemoteException( f"Bad response status from model server:\n{resp.status}\n{raw}", payload=SimpleResponse(resp.status, resp.headers, raw), ) merged = DataLoader.split_responses(raw) return tuple( aiohttp.web.Response( body=i.data, headers=i.headers, status=i.status) for i in merged)
async def request_dispatcher(self, request): with async_trace( ZIPKIN_API_URL, service_name=self.__class__.__name__, span_name="[1]http request", is_root=True, standalone=True, sample_rate=0.001, ): api_name = request.match_info.get("name") if api_name 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_name](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
async def relay_handler(self, request): data = await request.read() headers = dict(request.headers) url = request.url.with_host(self.outbound_host).with_port(self.outbound_port) with async_trace( ZIPKIN_API_URL, service_name=self.__class__.__name__, span_name=f"[2]{url.path} relay", ) as trace_ctx: headers.update(make_http_headers(trace_ctx)) async with aiohttp.ClientSession() as client: async with client.request( request.method, url, data=data, headers=request.headers ) as resp: body = await resp.read() return aiohttp.web.Response( status=resp.status, body=body, headers=resp.headers, )