Repository: jina-ai/jina
File: jina/serve/runtimes/head/__init__.py (Python)

Code before:
                raise uses_before_result
            else:
                (response, uses_before_metadata) = uses_before_result
                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,
            retries=self._retries,
        )
        total_shards = len(worker_send_tasks)

        all_worker_results = await asyncio.gather(*worker_send_tasks)
        worker_results = list(
            filter(lambda x: isinstance(x, Tuple), all_worker_results)
        )
        exceptions = list(
            filter(
                lambda x: isinstance(x, (AioRpcError, InternalNetworkError)),
                all_worker_results,
            )
        )
        failed_shards = len(exceptions)
        if failed_shards:
            self.logger.warning(f'{failed_shards} shards out of {total_shards} failed.')

        if len(worker_results) == 0:
            if exceptions:
                # raise the underlying error first
                raise exceptions[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:
            uses_after_result = await self.connection_pool.send_requests_once(
                worker_results,
                deployment='uses_after',
                timeout=self.timeout_send,
                retries=self._retries,
            )
            if isinstance(uses_after_result, (AioRpcError, InternalNetworkError)):
                raise uses_after_result
            else:
                (response_request, uses_after_metadata) = uses_after_result

Review comment (suggestion):
Just a suggestion, bu personally I would pull this into a separate method, seeing that `_handle_data_request()` is already quite long. 

Code after:
                all_worker_results,
            )
        )
        total_shards = len(worker_send_tasks)
        failed_shards = len(exceptions)
        if failed_shards:
            self.logger.warning(f'{failed_shards} shards out of {total_shards} failed.')

        return worker_results, exceptions, total_shards, failed_shards

    async def _handle_data_request(
        self, requests: List[DataRequest], endpoint: Optional[str]
    ) -> Tuple[DataRequest, Dict]:
        self.logger.debug(f'recv {len(requests)} DataRequest(s)')

        WorkerRequestHandler.merge_routes(requests)

        uses_before_metadata = None
        if self.uses_before_address:
            result = await self.connection_pool.send_requests_once(
                requests,
                deployment='uses_before',
                timeout=self.timeout_send,
                retries=self._retries,
            )
            if issubclass(type(result), BaseException):
                raise result
            else:
                response, uses_before_metadata = result
                requests = [response]

        (
            worker_results,
            exceptions,
            total_shards,
            failed_shards,
        ) = await self._gather_worker_tasks(requests, endpoint)

        if len(worker_results) == 0:
            if exceptions:
                # raise the underlying error first
                raise exceptions[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: