add method to flex base class for acquiring local grpc client

Signed-off-by: kta-intel <[email protected]>

openfl/component/aggregator/aggregator.py

@@ -705,7 +705,7 @@ def stop_flex(self):
     def get_local_grpc_client(self):
         if not self.is_flex_available():
             raise RuntimeError("Federated Learning exchange as not been enabled.")
-        return self.flex.local_grpc_client
+        return self.flex.get_local_grpc_client()
 
     def _end_of_round_with_stragglers_check(self):
         """

openfl/component/interoperability/flex.py

@@ -14,6 +14,7 @@ def __init__(self, command: list[str], **kwargs):
         Args:
             command (list[str]): The command to run the server as a subprocess.
         """
+        self.local_grpc_client = None
         self._command = command
         self._process = None
 
@@ -40,3 +41,10 @@ def stop(self):
             logger.info("[FLEX] Subprocess stopped.")
         else:
             logger.info("[FLEX] No subprocess is currently running.")
+
+
+    def get_local_grpc_client(self):
+        """
+        Get the local gRPC client.
+        """
+        return self.local_grpc_client

openfl/federated/task/runner_flower.py

@@ -3,7 +3,6 @@
 from flwr.proto import grpcadapter_pb2_grpc
 from multiprocessing import cpu_count
 from openfl.federated.task.runner import TaskRunner
-from openfl.transport import AggregatorGRPCClient
 from openfl.transport.grpc.flex.flower.local_grpc_server import LocalGRPCServer
 import subprocess
 

openfl/transport/grpc/aggregator_client.py

@@ -503,6 +503,7 @@ def send_message_to_server(self, flower_message, collaborator_name):
             The response from the OpenFL server, converted back to a Flower message.
         """
         self._set_header(collaborator_name)
+        #TODO: use a general to/from openfl_message function with "add_header" option, do the message conversion before calling send_message_to_server
         openfl_message = flower_to_openfl_message(flower_message, 
                                                   header=self.header)
         openfl_response = self.stub.PelicanDrop(openfl_message)

openfl/transport/grpc/flex/flower/local_grpc_server.py

@@ -52,6 +52,7 @@ def process_queue(self):
         while True:
             request, response_queue = self.request_queue.get()
             # Send request to the OpenFL server
+            # TODO: do message conversions here
             flower_response = self.openfl_client.send_message_to_server(request, self.collaborator_name)
             # Send response to Flower client
             response_queue.put(flower_response)