Revert "Disable streaming on send task results for consistency"

This reverts commit f6ac53e.

Signed-off-by: Shah, Karan <[email protected]>

Author: MasterSkepticista

openfl/protocols/aggregator.proto

@@ -12,7 +12,7 @@ service Aggregator {
   rpc Ping(PingRequest) returns (PingResponse) {}
   rpc GetTasks(GetTasksRequest) returns (GetTasksResponse) {}
   rpc GetAggregatedTensors(GetAggregatedTensorsRequest) returns (GetAggregatedTensorsResponse) {}
-  rpc SendLocalTaskResults(TaskResults) returns (SendLocalTaskResultsResponse) {}
+  rpc SendLocalTaskResults(stream DataStream) returns (SendLocalTaskResultsResponse) {}
   rpc InteropRelay(InteropMessage) returns (InteropMessage) {}
 }
 

openfl/transport/grpc/aggregator_client.py

@@ -10,7 +10,7 @@
 
 import grpc
 
-from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc
+from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc, utils
 from openfl.transport.grpc.common import create_header, create_insecure_channel, create_tls_channel
 
 logger = logging.getLogger(__name__)
@@ -416,7 +416,8 @@ def send_local_task_results(
             tensors=named_tensors,
         )
 
-        response = self.stub.SendLocalTaskResults(request)
+        # convert (potentially) long list of tensors into stream
+        response = self.stub.SendLocalTaskResults(utils.proto_to_datastream(request))
         self.validate_response(response)
 
     @_atomic_connection

openfl/transport/grpc/aggregator_server.py

@@ -10,7 +10,7 @@
 
 import grpc
 
-from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc
+from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc, utils
 from openfl.transport.grpc.common import create_grpc_server, create_header, synchronized
 
 logger = logging.getLogger(__name__)
@@ -256,7 +256,7 @@ def GetAggregatedTensors(self, request, context):
 
     @synchronized
     def SendLocalTaskResults(self, request, context):  # NOQA:N802
-        """Store collaborator task results on the aggregator.
+        """Request a model download from aggregator.
 
         This method handles a request from a collaborator to send the results
         of a local task.
@@ -270,14 +270,23 @@ def SendLocalTaskResults(self, request, context):  # NOQA:N802
             aggregator_pb2.SendLocalTaskResultsResponse: The response to the
                 request.
         """
-        self.validate_collaborator(request, context)
-        self.check_request(request)
+        try:
+            proto = aggregator_pb2.TaskResults()
+            proto = utils.datastream_to_proto(proto, request)
+        except RuntimeError:
+            raise RuntimeError(
+                "Empty stream message, reestablishing connection from client to resume training..."
+            )
 
-        collaborator_name = request.header.sender
-        task_name = request.task_name
-        round_number = request.round_number
-        data_size = request.data_size
-        named_tensors = request.tensors
+        self.validate_collaborator(proto, context)
+        # all messages get sanity checked
+        self.check_request(proto)
+
+        collaborator_name = proto.header.sender
+        task_name = proto.task_name
+        round_number = proto.round_number
+        data_size = proto.data_size
+        named_tensors = proto.tensors
         self.aggregator.send_local_task_results(
             collaborator_name, round_number, task_name, data_size, named_tensors
         )