[WIP][Workflow Interface] Resolving "Response code: StatusCode.RESOURCE_EXHAUSTED" (Issue #1565) in FederatedRuntime (#1572)

* grpc resource issue partial fix

Signed-off-by: Ishant Thakare <[email protected]>

* update director.proto

Signed-off-by: Ishant Thakare <[email protected]>

* fix rpc get_flow_state

Signed-off-by: Ishant Thakare <[email protected]>

* code cleanup

Signed-off-by: Ishant Thakare <[email protected]>

* code cleanup

Signed-off-by: Ishant Thakare <[email protected]>

* copyright changes

Signed-off-by: Ishant Thakare <[email protected]>

* update checkpoint rpc

Signed-off-by: Ishant Thakare <[email protected]>

* incorporated review comments

Signed-off-by: Ishant Thakare <[email protected]>

---------

Signed-off-by: Ishant Thakare <[email protected]>
Co-authored-by: Noopur <[email protected]>
Co-authored-by: Payal Chaurasiya <[email protected]>

Author: 3 people

openfl/experimental/workflow/protocols/aggregator.proto

@@ -1,4 +1,4 @@
-// Copyright (C) 2020-2023 Intel Corporation
+// Copyright (C) 2020-2025 Intel Corporation
 // Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you.
 
 syntax = "proto3";
@@ -9,9 +9,9 @@ import "openfl/protocols/base.proto";
 
 
 service Aggregator {
-  rpc SendTaskResults(TaskResultsRequest) returns (TaskResultsResponse) {}
-  rpc GetTasks(GetTasksRequest) returns (GetTasksResponse) {}
-  rpc CallCheckpoint(CheckpointRequest) returns (CheckpointResponse) {}
+  rpc SendTaskResults(stream DataStream) returns (TaskResultsResponse) {}
+  rpc GetTasks(GetTasksRequest) returns (stream DataStream) {}
+  rpc CallCheckpoint(stream DataStream) returns (CheckpointResponse) {}
 }
 
 message MessageHeader {

openfl/experimental/workflow/protocols/director.proto

@@ -1,10 +1,11 @@
-// Copyright 2020-2024 Intel Corporation
+// Copyright 2020-2025 Intel Corporation
 // SPDX-License-Identifier: Apache-2.0
 
 syntax = "proto3";
 
 package openfl.experimental.workflow.director;
 
+import "openfl/protocols/base.proto";
 import "google/protobuf/timestamp.proto";
 import "google/protobuf/duration.proto";
 
@@ -18,7 +19,7 @@ service Director {
     //Runtime RPCs
     rpc SetNewExperiment(stream ExperimentInfo) returns (SetNewExperimentResponse) {}
     rpc GetEnvoys(GetEnvoysRequest) returns (GetEnvoysResponse) {}
-    rpc GetFlowState(GetFlowStateRequest) returns (GetFlowStateResponse) {}
+    rpc GetFlowState(GetFlowStateRequest) returns (stream DataStream) {}
     rpc ConnectRuntime(SendRuntimeRequest) returns (RuntimeRequestResponse) {}
     rpc GetExperimentStdout(GetExperimentStdoutRequest) returns (stream GetExperimentStdoutResponse) {}
 }

openfl/experimental/workflow/transport/grpc/aggregator_client.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2024 Intel Corporation
+# Copyright 2020-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 
@@ -12,6 +12,7 @@
 
 from openfl.experimental.workflow.protocols import aggregator_pb2, aggregator_pb2_grpc
 from openfl.experimental.workflow.transport.grpc.grpc_channel_options import channel_options
+from openfl.protocols.utils import datastream_to_proto, proto_to_datastream
 
 
 class ConstantBackoff:
@@ -280,7 +281,7 @@ def send_task_results(self, collaborator_name, round_number, next_step, clone_by
             execution_environment=clone_bytes,
         )
 
-        response = self.stub.SendTaskResults(request)
+        response = self.stub.SendTaskResults(proto_to_datastream(request))
         self.validate_response(response, collaborator_name)
 
         return response.header
@@ -291,8 +292,8 @@ def get_tasks(self, collaborator_name):
         """Get tasks from the aggregator."""
         self._set_header(collaborator_name)
         request = aggregator_pb2.GetTasksRequest(header=self.header)
-
-        response = self.stub.GetTasks(request)
+        response_stream = self.stub.GetTasks(request)
+        response = datastream_to_proto(aggregator_pb2.GetTasksResponse(), response_stream)
         self.validate_response(response, collaborator_name)
 
         return (
@@ -316,7 +317,7 @@ def call_checkpoint(self, collaborator_name, clone_bytes, function, stream_buffe
             stream_buffer=stream_buffer,
         )
 
-        response = self.stub.CallCheckpoint(request)
+        response = self.stub.CallCheckpoint(proto_to_datastream(request))
         self.validate_response(response, collaborator_name)
 
         return response.header

openfl/experimental/workflow/transport/grpc/aggregator_server.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2024 Intel Corporation
+# Copyright 2020-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 
@@ -14,6 +14,7 @@
 
 from openfl.experimental.workflow.protocols import aggregator_pb2, aggregator_pb2_grpc
 from openfl.experimental.workflow.transport.grpc.grpc_channel_options import channel_options
+from openfl.protocols.utils import datastream_to_proto, proto_to_datastream
 
 logger = logging.getLogger(__name__)
 
@@ -128,18 +129,20 @@ def check_request(self, request):
         )
 
     def SendTaskResults(self, request, context):  # NOQA:N802
-        """<FIND OUT WHAT COMMENT TO PUT HERE>.
+        """Processes a request from a collaborator to retrieve the results of a locally
+        executed task.
 
         Args:
             request: The gRPC message request
             context: The gRPC context
         """
-        self.validate_collaborator(request, context)
-        self.check_request(request)
-        collaborator_name = request.header.sender
-        round_number = (request.round_number,)
-        next_step = (request.next_step,)
-        execution_environment = request.execution_environment
+        proto = datastream_to_proto(aggregator_pb2.TaskResultsRequest(), request)
+        self.validate_collaborator(proto, context)
+        self.check_request(proto)
+        collaborator_name = proto.header.sender
+        round_number = (proto.round_number,)
+        next_step = (proto.next_step,)
+        execution_environment = proto.execution_environment
 
         _ = self.aggregator.send_task_results(
             collaborator_name, round_number[0], next_step, execution_environment
@@ -160,14 +163,15 @@ def GetTasks(self, request, context):  # NOQA:N802
 
         rn, f, ee, st, q = self.aggregator.get_tasks(request.header.sender)
 
-        return aggregator_pb2.GetTasksResponse(
+        response = aggregator_pb2.GetTasksResponse(
             header=self.get_header(collaborator_name),
             round_number=rn,
             function_name=f,
             execution_environment=ee,
             sleep_time=st,
             quit=q,
         )
+        return proto_to_datastream(response)
 
     def CallCheckpoint(self, request, context):  # NOQA:N802
         """Request aggregator to perform a checkpoint for a given function.
@@ -176,12 +180,13 @@ def CallCheckpoint(self, request, context):  # NOQA:N802
             request: The gRPC message request
             context: The gRPC context
         """
-        self.validate_collaborator(request, context)
-        self.check_request(request)
-        collaborator_name = request.header.sender
-        execution_environment = request.execution_environment
-        function = request.function
-        stream_buffer = request.stream_buffer
+        proto = datastream_to_proto(aggregator_pb2.CheckpointRequest(), request)
+        self.validate_collaborator(proto, context)
+        self.check_request(proto)
+        collaborator_name = proto.header.sender
+        execution_environment = proto.execution_environment
+        function = proto.function
+        stream_buffer = proto.stream_buffer
 
         self.aggregator.call_checkpoint(
             collaborator_name, execution_environment, function, stream_buffer

openfl/experimental/workflow/transport/grpc/director_client.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2024 Intel Corporation
+# Copyright 2020-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 """DirectorClient module."""
@@ -12,6 +12,7 @@
 
 from openfl.experimental.workflow.protocols import director_pb2, director_pb2_grpc
 from openfl.experimental.workflow.transport.grpc.exceptions import EnvoyNotFoundError
+from openfl.protocols.utils import datastream_to_proto
 
 from .grpc_channel_options import channel_options
 
@@ -313,8 +314,8 @@ def get_flow_state(self) -> Tuple:
                 - flspec_obj (object): The FLSpec object containing
                     details of the updated flow state.
         """
-        response = self.stub.GetFlowState(director_pb2.GetFlowStateRequest())
-
+        response_stream = self.stub.GetFlowState(director_pb2.GetFlowStateRequest())
+        response = datastream_to_proto(director_pb2.GetFlowStateResponse(), response_stream)
         return response.completed, response.flspec_obj
 
     def stream_experiment_stdout(self, experiment_name) -> Iterator[Dict[str, Any]]:

openfl/experimental/workflow/transport/grpc/director_server.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2024 Intel Corporation
+# Copyright 2020-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 """DirectorGRPCServer module."""
@@ -15,7 +15,7 @@
 from openfl.experimental.workflow.protocols import director_pb2, director_pb2_grpc
 from openfl.experimental.workflow.transport.grpc.exceptions import EnvoyNotFoundError
 from openfl.experimental.workflow.transport.grpc.grpc_channel_options import channel_options
-from openfl.protocols.utils import get_headers
+from openfl.protocols.utils import get_headers, proto_to_datastream
 
 logger = logging.getLogger(__name__)
 
@@ -326,7 +326,12 @@ async def GetFlowState(self, request, context) -> director_pb2.GetFlowStateRespo
             director_pb2.GetFlowStateResponse: The response to the request.
         """
         status, flspec_obj = await self.director.get_flow_state()
-        return director_pb2.GetFlowStateResponse(completed=status, flspec_obj=flspec_obj)
+        response = director_pb2.GetFlowStateResponse(
+            completed=status,
+            flspec_obj=flspec_obj,
+        )
+        for chunk in proto_to_datastream(response):
+            await context.write(chunk)
 
     async def GetExperimentStdout(
         self, request, context