fix: check for result types after the task has completed its execution

jina/serve/networking.py

@@ -11,6 +11,7 @@
 from grpc_health.v1 import health_pb2, health_pb2_grpc
 from grpc_reflection.v1alpha.reflection_pb2 import ServerReflectionRequest
 from grpc_reflection.v1alpha.reflection_pb2_grpc import ServerReflectionStub
+
 from jina import __default_endpoint__
 from jina.enums import PollingType
 from jina.excepts import EstablishGrpcConnectionError, InternalNetworkError
@@ -852,10 +853,7 @@ def send_requests_once(
                 timeout=timeout,
                 retries=retries,
             )
-            if isinstance(result, (AioRpcError, InternalNetworkError)):
-                raise result
-            else:
-                return result
+            return result
         else:
             self._logger.debug(
                 f'no available connections for deployment {deployment} and shard {shard_id}'
@@ -1015,6 +1013,7 @@ async def task_wrapper():
                         connection_list=connections,
                     )
                     if error:
+                        print(f'--->returning error: {type(error)}')
                         return error
 
         return asyncio.create_task(task_wrapper())

jina/serve/runtimes/gateway/graph/topology_graph.py

@@ -7,6 +7,7 @@
 
 import grpc.aio
 from grpc.aio import AioRpcError
+
 from jina import __default_endpoint__
 from jina.excepts import InternalNetworkError
 from jina.serve.networking import GrpcConnectionPool
@@ -159,7 +160,7 @@ async def _wait_previous_and_send(
                         return request, metadata
                     # otherwise, send to executor and get response
                     try:
-                        resp, metadata = await connection_pool.send_requests_once(
+                        result = await connection_pool.send_requests_once(
                             requests=self.parts_to_send,
                             deployment=self.name,
                             metadata=self._metadata,
@@ -168,6 +169,10 @@ async def _wait_previous_and_send(
                             timeout=self._timeout_send,
                             retries=self._retries,
                         )
+                        if issubclass(type(result), BaseException):
+                            raise result
+                        else:
+                            resp, metadata = result
                         if WorkerRequestHandler._KEY_RESULT in resp.parameters:
                             # Accumulate results from each Node and then add them to the original
                             self.result_in_params_returned = resp.parameters[

jina/serve/runtimes/head/__init__.py

@@ -10,6 +10,7 @@
 from grpc.aio import AioRpcError
 from grpc_health.v1 import health, health_pb2, health_pb2_grpc
 from grpc_reflection.v1alpha import reflection
+
 from jina.enums import PollingType
 from jina.excepts import InternalNetworkError
 from jina.helper import get_full_version
@@ -318,7 +319,7 @@ async def _gather_worker_tasks(self, requests, endpoint):
         )
         exceptions = list(
             filter(
-                lambda x: isinstance(x, (AioRpcError, InternalNetworkError)),
+                lambda x: issubclass(type(x), BaseException),
                 all_worker_results,
             )
         )
@@ -338,16 +339,17 @@ async def _handle_data_request(
 
         uses_before_metadata = None
         if self.uses_before_address:
-            (
-                response,
-                uses_before_metadata,
-            ) = await self.connection_pool.send_requests_once(
+            result = await self.connection_pool.send_requests_once(
                 requests,
                 deployment='uses_before',
                 timeout=self.timeout_send,
                 retries=self._retries,
             )
-            requests = [response]
+            if issubclass(type(result), BaseException):
+                raise result
+            else:
+                response, uses_before_metadata = result
+                requests = [response]
 
         (
             worker_results,
@@ -369,15 +371,16 @@ async def _handle_data_request(
         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(
+            result = await self.connection_pool.send_requests_once(
                 worker_results,
                 deployment='uses_after',
                 timeout=self.timeout_send,
                 retries=self._retries,
             )
+            if issubclass(type(result), BaseException):
+                raise result
+            else:
+                response_request, uses_after_metadata = result
         elif len(worker_results) > 1 and self._reduce:
             response_request = WorkerRequestHandler.reduce_requests(worker_results)
         elif len(worker_results) > 1 and not self._reduce: