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, {}