async def start_wait(self, interval, call): with async_trace( ZIPKIN_API_URL, service_name=self.__class__.__name__, span_name="[1]parade task", sample_rate=1, is_root=True, ): try: # await asyncio.sleep(interval) await asyncio.sleep(0.02) with async_trace( ZIPKIN_API_URL, service_name=self.__class__.__name__, span_name="[2]call", ): async with self.outbound_sema: self.status = self.STATUS_CLOSED self.batch_output = await call( self.batch_input[:self.cur]) self.status = self.STATUS_RETURNED async with self.returned: self.returned.notify_all() except Exception as e: # noqa TODO raise e finally: # make sure parade is closed if self.status != self.STATUS_OPEN: self.status = self.STATUS_CLOSED
async def _batch_handler_template(self, requests, api_name): 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) try: async with aiohttp.ClientSession() as client: async with client.post(api_url, data=reqs_s, headers=headers) as resp: raw = await resp.read() merged = DataLoader.split_responses(raw) except (aiohttp.ClientConnectorError, aiohttp.ServerDisconnectedError): return (aiohttp.web.HTTPServiceUnavailable, ) * len(requests) if merged is None: return (aiohttp.web.HTTPInternalServerError, ) * len(requests) return tuple( aiohttp.web.Response( body=i.data, headers=i.headers, status=i.status) for i in merged)
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=f"[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 = SimpleRequest(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.data, ) except Exception: # pylint: disable=broad-except logger.error(traceback.format_exc()) resp = aiohttp.web.InternalServerError() else: resp = await self.relay_handler(request) return resp
async def request_dispatcher(self, request): with async_trace( ZIPKIN_API_URL, service_name=self.__class__.__name__, span_name=f"[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 = SimpleRequest(request.raw_headers, await request.read()) resp = await self.batch_handlers[api_name](req) else: resp = await self.relay_handler(request) return resp
async def _relay_handler(self, request, api_name): data = await request.read() headers = dict(request.headers) api_url = f"http://{self.target_host}:{self.target_port}/{api_name}" with async_trace( ZIPKIN_API_URL, service_name=self.__class__.__name__, span_name=f"[2]{api_name} relay", ) as trace_ctx: headers.update(make_http_headers(trace_ctx)) async with aiohttp.ClientSession() as client: async with client.request(request.method, api_url, data=data, headers=request.headers) as resp: body = await resp.read() return aiohttp.web.Response( status=resp.status, body=body, headers=resp.headers, )
async def relay_handler(self, request: aiohttp.web.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, )