[Task Runner API] [Flower Interoperability] Integrate changes for flower-app-pytorch workspace's custom Collaborator into core component (#1598)

* integrate flower collaborator into core collaborator

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

* formatting

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

* change task name interop -> prepare_for_interop, change references to local_grpc_server to interop_server

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

* add prepare_interop_server as callback to run on_experiment_begin

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

* fix typo

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

* change to on_round_begin()

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

* return empty dict

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

* set self.prepare_interop_server() to run on experiment begin

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

* remove self-explanatory comment, added another comment for clarity

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

---------

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

Author: kminhta

openfl-workspace/flower-app-pytorch/README.md

@@ -91,16 +91,26 @@ task_runner:
 
 3. `FlowerDataLoader` with similar high-level functionality to other dataloaders.
 
-4. `Task` - we introduce a `tasks_connector.yaml` that will allow the collaborator to connect to Flower framework via the local gRPC server. It also handles the task runner's `start_client_adapter` method, which actually starts the Flower component and local gRPC server. By setting `local_server_port` to 0, the port is dynamically allocated. This is mainly for local experiments to avoid overlapping the ports.
+4. `Task` - we introduce a `tasks_connector.yaml` that will allow the collaborator to connect to Flower framework via the interop server. It also handles the task runner's `start_client_adapter` method, which actually starts the Flower component and interop server. By setting `local_server_port` to 0, the port is dynamically allocated. This is mainly for local experiments to avoid overlapping the ports.
 
 ```yaml
 tasks:
-  settings:
-    connect_to: Flower
-  start_client_adapter:
+  prepare_for_interop:
     function: start_client_adapter
     kwargs:
-      local_server_port: 0
+      interop_server_port: 0
+  settings:
+    interop_server: src.grpc.connector.flower.interop_server
+```
+
+5.`Collaborator` has an additional setting `interop_mode` which will invoke a callback to prepare the interop server that'll eventually be started by the Task Runner
+
+```yaml
+collaborator :
+  defaults : plan/defaults/collaborator.yaml
+  template : openfl.component.Collaborator
+  settings: 
+    interop_mode : True
 ```
 
 > **Note**: `aggregator.settings.rounds_to_train` is set to 1. __Do not edit this__. The actual number of rounds for the experiment is controlled by Flower logic inside of `./app-pytorch/pyproject.toml`. The entirety of the Flower experiment will run in a single OpenFL round. Increasing this will cause OpenFL to attempt to run the experiment again. The aggregator round is there to stop the OpenFL components at the completion of the experiment.

openfl-workspace/flower-app-pytorch/plan/plan.yaml

@@ -25,7 +25,9 @@ connector :
 
 collaborator :
   defaults : plan/defaults/collaborator.yaml
-  template : src.collaborator.CollaboratorFlower
+  template : openfl.component.Collaborator
+  settings: 
+    interop_mode : True
 
 data_loader :
   defaults : plan/defaults/data_loader.yaml
@@ -51,12 +53,12 @@ assigner :
       - name  : Connector_Flower
         percentage : 1.0
         tasks : 
-          - start_client_adapter
+          - prepare_for_interop
 
 tasks :
   defaults : plan/defaults/tasks_connector.yaml
   settings :
-    connect_to : Flower
+    interop_server: src.grpc.connector.flower.interop_server
 
 compression_pipeline :
   defaults : plan/defaults/compression_pipeline.yaml

openfl-workspace/flower-app-pytorch/src/collaborator.py

@@ -21,16 +21,14 @@ class FlowerInteropServer(grpcadapter_pb2_grpc.GrpcAdapterServicer):
     request handling issues.
     """
 
-    def __init__(self, openfl_client, collaborator_name):
+    def __init__(self, send_message_to_client):
         """
         Initialize.
 
         Args:
-            openfl_client: An instance of the OpenFL Client.
-            collaborator_name: The name of the collaborator.
+            send_message_to_client (Callable): A callable function to send messages to the OpenFL client.
         """
-        self.openfl_client = openfl_client
-        self.collaborator_name = collaborator_name
+        self.send_message_to_client = send_message_to_client
         self.end_experiment_callback = None
         self.request_queue = queue.Queue()
         self.processing_thread = threading.Thread(target=self.process_queue)
@@ -86,7 +84,7 @@ def process_queue(self):
             openfl_request = flower_to_openfl_message(request)
 
             # Send request to the OpenFL server
-            openfl_response = self.openfl_client.send_message_to_server(openfl_request, self.collaborator_name)
+            openfl_response = self.send_message_to_client(openfl_request)
 
             # Check to end experiment
             if hasattr(openfl_response, 'metadata'):

openfl-workspace/flower-app-pytorch/src/grpc/connector/flower/interop_server.py

@@ -51,30 +51,40 @@ def __init__(self, **kwargs):
 
         self.shutdown_requested = False  # Flag to signal shutdown
 
-    def start_client_adapter(self, interop_server, **kwargs):
+    def start_client_adapter(self,
+                             col_name=None,
+                             round_num=None,
+                             input_tensor_dict=None,
+                             **kwargs):
         """
-        Start the local gRPC server and the Flower SuperNode.
+        Start the FlowerInteropServer and the Flower SuperNode.
 
         Args:
-            interop_server: The local gRPC server instance.
-            **kwargs: Additional parameters, including 'local_server_port'.
+            col_name (str, optional): The collaborator name. Defaults to None.
+            round_num (int, optional): The current round number. Defaults to None.
+            input_tensor_dict (dict, optional): The input tensor dictionary. Defaults to None.
+            **kwargs: Additional parameters for configuration.
+                includes:
+                    interop_server (object): The FlowerInteropServer instance.
+                    interop_server_port (int): The port for the interop server.
         """
-        local_server_port = kwargs.get('local_server_port')
 
         def message_callback():
             self.shutdown_requested = True
 
-        # Set the callback for ending the experiment
+        interop_server = kwargs.get('interop_server')
+        interop_server_port = kwargs.get('interop_server_port')
         interop_server.set_end_experiment_callback(message_callback)
-        interop_server.start_server(local_server_port)
+        interop_server.start_server(interop_server_port)
 
-        local_server_port = interop_server.get_port()
+        # interop server sets port dynamically
+        interop_server_port = interop_server.get_port()
 
         command = [
             "flower-supernode",
             "--insecure",
             "--grpc-adapter",
-            "--superlink", f"127.0.0.1:{local_server_port}",
+            "--superlink", f"127.0.0.1:{interop_server_port}",
             "--clientappio-api-address", f"127.0.0.1:{self.client_port}",
             "--node-config", f"data-path='{self.data_path}'"
         ]
@@ -93,7 +103,7 @@ def message_callback():
 
         if self.sgx_enabled:
             # Check if port is open before starting the client app
-            while not is_port_open('127.0.0.1', local_server_port):
+            while not is_port_open('127.0.0.1', interop_server_port):
                 time.sleep(0.5)
 
             time.sleep(1) # Add a small delay after confirming the port is open
@@ -116,6 +126,14 @@ def message_callback():
                 interop_server.stop_server()
             time.sleep(0.1)
 
+        # Collaborator expects these dictionaries, but they are not used in this context
+        # as Flower will handle the tensors internally.
+        global_output_tensor_dict = {}
+        local_output_tensor_dict = {}
+
+        return global_output_tensor_dict, local_output_tensor_dict
+
+
 
     def set_tensor_dict(self, tensor_dict, with_opt_vars=False):
         """
@@ -160,6 +178,9 @@ def initialize_tensorkeys_for_functions(self, with_opt_vars=False):
         """Initialize tensor keys for functions. Currently not implemented."""
         pass
 
+    def get_required_tensorkeys_for_function(self, func_name, **kwargs):
+        """Get tensor keys for functions. Return empty dict."""
+        return {}
 
 def install_flower_FAB(flwr_app_name):
     """

openfl-workspace/flower-app-pytorch/src/runner.py

@@ -1,4 +1,4 @@
-start_client_adapter:
+prepare_for_interop:
   function : start_client_adapter
   kwargs   :
-    local_server_port : 0 # local grpc server, 0 to dynamically allocate
+    interop_server_port : 0 # interop server port, 0 to dynamically allocate

openfl-workspace/workspace/plan/defaults/tasks_connector.yaml

@@ -4,6 +4,7 @@
 
 """Collaborator module."""
 
+import importlib
 import logging
 from enum import Enum
 from time import sleep
@@ -75,6 +76,7 @@ def __init__(
         write_logs=False,
         callbacks: Optional[List] = [],
         secure_aggregation=False,
+        interop_mode=False,
     ):
         """Initialize the Collaborator object.
 
@@ -143,6 +145,15 @@ def __init__(
             else:
                 callbacks = [secure_aggregation_callback]
 
+        # Interoperability mode
+        self._interop_mode_enabled = interop_mode
+        if self._interop_mode_enabled:
+            callbacks.append(
+                callbacks_module.LambdaCallback(
+                    on_experiment_begin=lambda logs=None: self.prepare_interop_server()
+                )
+            )
+
         # Callbacks
         self.callbacks = callbacks_module.CallbackList(
             callbacks,
@@ -584,3 +595,26 @@ def _apply_masks(
                 continue
             masked_metric = np.add(self._private_mask, tensor_dict[tensor_key])
             tensor_dict[tensor_key] = np.add(masked_metric, self._shared_mask)
+
+    def prepare_interop_server(self):
+        """
+        Prepare the interoperability server.
+
+        This function initializes the interoperability server and sets up
+        the callback for receiving messages from the interop server.
+        It also sets the interop server in the task configuration to be used
+        by the Task Runner.
+        """
+
+        # Initialize the interop server
+        framework = self.task_config["settings"]["interop_server"]
+        module = importlib.import_module(framework)
+
+        def receive_message_from_interop(message):
+            """Receive message from interop server."""
+            # Process the request and return a response
+            response = self.client.send_message_to_server(message, self.collaborator_name)
+            return response
+
+        interop_server = module.FlowerInteropServer(receive_message_from_interop)
+        self.task_config["prepare_for_interop"]["kwargs"]["interop_server"] = interop_server