async def _handle_data_request(
self, requests: List[DataRequest],
endpoint: Optional[str]) -> Tuple[DataRequest, Dict]:
self.logger.debug(f'recv {len(requests)} DataRequest(s)')
DataRequestHandler.merge_routes(requests)
uses_before_metadata = None
if self.uses_before_address:
(
response,
uses_before_metadata,
) = await self.connection_pool.send_requests_once(
requests, deployment='uses_before', timeout=self.timeout_send)
requests = [response]
worker_send_tasks = self.connection_pool.send_requests(
requests=requests,
deployment=self._deployment_name,
polling_type=self._polling[endpoint],
timeout=self.timeout_send,
)
worker_results = await asyncio.gather(*worker_send_tasks)
if len(worker_results) == 0:
raise RuntimeError(
f'Head {self.name} did not receive a response when sending message to worker pods'
)
worker_results, metadata = zip(*worker_results)
response_request = worker_results[0]
uses_after_metadata = None
if self.uses_after_address:
(
response_request,
uses_after_metadata,
) = await self.connection_pool.send_requests_once(
worker_results,
deployment='uses_after',
timeout=self.timeout_send)
elif len(worker_results) > 1 and self._reduce:
DataRequestHandler.reduce_requests(worker_results)
elif len(worker_results) > 1 and not self._reduce:
# worker returned multiple responsed, but the head is configured to skip reduction
# just concatenate the docs in this case
response_request.data.docs = DataRequestHandler.get_docs_from_request(
requests, field='docs')
merged_metadata = self._merge_metadata(metadata, uses_after_metadata,
uses_before_metadata)
return response_request, merged_metadata
async def _handle_data_request(
self, requests: List[DataRequest],
endpoint: Optional[str]) -> Tuple[DataRequest, Dict]:
self.logger.debug(f'recv {len(requests)} DataRequest(s)')
DataRequestHandler.merge_routes(requests)
uses_before_metadata = None
if self.uses_before_address:
(
response,
uses_before_metadata,
) = await self.connection_pool.send_requests_once(
requests, deployment='uses_before')
requests = [response]
elif len(requests) > 1 and not self._has_uses:
requests = [DataRequestHandler.reduce_requests(requests)]
worker_send_tasks = self.connection_pool.send_requests(
requests=requests,
deployment=self._deployment_name,
polling_type=self._polling[endpoint],
)
worker_results = await asyncio.gather(*worker_send_tasks)
if len(worker_results) == 0:
raise RuntimeError(
f'Head {self.name} did not receive a response when sending message to worker pods'
)
worker_results, metadata = zip(*worker_results)
response_request = worker_results[0]
uses_after_metadata = None
if self.uses_after_address:
(
response_request,
uses_after_metadata,
) = await self.connection_pool.send_requests_once(
worker_results, deployment='uses_after')
elif len(worker_results) > 1:
DataRequestHandler.reduce_requests(worker_results)
merged_metadata = self._merge_metadata(metadata, uses_after_metadata,
uses_before_metadata)
return response_request, merged_metadata
async def _wait_previous_and_send(
self,
request: DataRequest,
previous_task: Optional[asyncio.Task],
connection_pool: GrpcConnectionPool,
endpoint: Optional[str],
executor_endpoint_mapping: Optional[Dict] = None,
target_executor_pattern: Optional[str] = None,
):
# Check my condition and send request with the condition
metadata = {}
if previous_task is not None:
result = await previous_task
request, metadata = result[0], result[1]
if metadata and 'is-error' in metadata:
return request, metadata
elif request is not None:
self.parts_to_send.append(request)
# this is a specific needs
if len(self.parts_to_send) == self.number_of_parts:
self.start_time = datetime.utcnow()
if self._filter_condition is not None:
self._update_requests()
if self._reduce and len(self.parts_to_send) > 1:
self.parts_to_send = [
DataRequestHandler.reduce_requests(
self.parts_to_send)
]
# avoid sending to executor which does not bind to this endpoint
if endpoint is not None and executor_endpoint_mapping is not None:
if (endpoint
not in executor_endpoint_mapping[self.name]
and __default_endpoint__
not in executor_endpoint_mapping[self.name]):
return request, metadata
if target_executor_pattern is not None and not re.match(
target_executor_pattern, self.name):
return request, metadata
# otherwise, send to executor and get response
try:
resp, metadata = await connection_pool.send_requests_once(
requests=self.parts_to_send,
deployment=self.name,
head=True,
endpoint=endpoint,
timeout=self._timeout_send,
retries=self._retries,
)
except InternalNetworkError as err:
self._handle_internalnetworkerror(err)
self.end_time = datetime.utcnow()
if metadata and 'is-error' in metadata:
self.status = resp.header.status
return resp, metadata
return None, {}