Additional error handling

securefederatedai · Jan 24, 2023 · 954d675 · 954d675
1 parent af3d3fc
commit 954d675
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/openfl/transport/grpc/aggregator_client.py b/openfl/transport/grpc/aggregator_client.py
@@ -137,16 +137,19 @@ class AsyncRetryOnRpcErrorUnaryStreamClientInterceptor(
     async def _intercept_call(self, continuation, client_call_details, request_or_iterator):
         """Intercept the async call to the gRPC server."""
         started_time = time.time()
+        is_connected = False
         while True:
             try:
                 async_iter = await continuation(client_call_details, request_or_iterator)
                 async for response in async_iter:
                     yield response
+                    is_connected = True
             except grpc.aio.AioRpcError as error:
                 if (
                         error.code() == grpc.StatusCode.UNAVAILABLE
                         and error.details() == 'Socket closed'
-                ):
+                ) or is_connected:
+                    logger.info(f'Rpc error: {error.code()}, details: {error.details()}.')
                     return
                 await self.handle_error_for_retry(error, started_time)