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,
)
