Enable rest protocol between Aggregator and Collaborator

- add new AggregatorClientInterface to allow switching b/w grpc and rest
- endhance existing AggregatorGRPCClient to start using AggregatorClientInterface
- added new transport package for rest with AggregatorRESTClient implementing AggregatorClientInterface
- added streaming api support with custom content-type
- added various connection flag for streaming request
- send additional header key "Sender" for better request logging at server side
- aligned Rest and gRPC client for most of the init params
- added AggregatorRESTServer and necesary changes in aggregator cli and federated/plan get_server method
- added transport_protocol settings in defaults/network.yaml
rebased 26-Mar.2

Signed-off-by: Shailesh Pant <shailesh.pant@intel.com>

Author: ishaileshpant

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

@@ -7,4 +7,5 @@ settings:
     client_reconnect_interval  : 5
     require_client_auth        : True
     cert_folder                : cert
-    enable_atomic_connections  : False
+    enable_atomic_connections  : False
+    transport_protocol         : rest

openfl/federated/plan/plan.py

@@ -15,7 +15,7 @@
 
 from openfl.interface.aggregation_functions import AggregationFunction, WeightedAverage
 from openfl.interface.cli_helper import WORKSPACE
-from openfl.transport import AggregatorGRPCClient, AggregatorGRPCServer
+from openfl.transport import AggregatorGRPCClient, AggregatorGRPCServer, AggregatorRESTServer
 from openfl.utilities.utils import getfqdn_env
 
 SETTINGS = "settings"
@@ -542,8 +542,6 @@ def get_collaborator(
         else:
             defaults[SETTINGS]["client"] = self.get_client(
                 collaborator_name,
-                self.aggregator_uuid,
-                self.federation_uuid,
                 root_certificate,
                 private_key,
                 certificate,
@@ -557,8 +555,6 @@ def get_collaborator(
     def get_client(
         self,
         collaborator_name,
-        aggregator_uuid,
-        federation_uuid,
         root_certificate=None,
         private_key=None,
         certificate=None,
@@ -579,6 +575,25 @@ def get_client(
         Returns:
             AggregatorGRPCClient: gRPC client for the specified collaborator.
         """
+        client_args = self.get_client_args(
+            collaborator_name,
+            root_certificate,
+            private_key,
+            certificate,
+        )
+
+        if self.client_ is None:
+            self.client_ = AggregatorGRPCClient(**client_args)
+
+        return self.client_
+
+    def get_client_args(
+        self,
+        collaborator_name,
+        root_certificate=None,
+        private_key=None,
+        certificate=None,
+    ):
         common_name = collaborator_name
         if not root_certificate or not private_key or not certificate:
             root_certificate = "cert/cert_chain.crt"
@@ -593,14 +608,10 @@ def get_client(
         client_args["certificate"] = certificate
         client_args["private_key"] = private_key
 
-        client_args["aggregator_uuid"] = aggregator_uuid
-        client_args["federation_uuid"] = federation_uuid
+        client_args["aggregator_uuid"] = self.aggregator_uuid
+        client_args["federation_uuid"] = self.federation_uuid
         client_args["collaborator_name"] = collaborator_name
-
-        if self.client_ is None:
-            self.client_ = AggregatorGRPCClient(**client_args)
-
-        return self.client_
+        return client_args
 
     def get_server(
         self,
@@ -609,7 +620,7 @@ def get_server(
         certificate=None,
         **kwargs,
     ):
-        """Get gRPC server of the aggregator instance.
+        """Get gRPC or REST server of the aggregator instance.
 
         Args:
             root_certificate (str, optional): Root certificate for the server.
@@ -621,8 +632,29 @@ def get_server(
             **kwargs: Additional keyword arguments.
 
         Returns:
-            AggregatorGRPCServer: gRPC server of the aggregator instance.
+            Aggregator Server: returns either gRPC or REST server of the aggregator instance.
         """
+        server_args = self.get_server_args(root_certificate, private_key, certificate, kwargs)
+
+        server_args["aggregator"] = self.get_aggregator()
+        network_cfg = self.config["network"][SETTINGS]
+        protocol = network_cfg.get("transport_protocol", "grpc").lower()
+
+        if self.server_ is None:
+            self.server_ = self._get_server(protocol, **server_args)
+
+        return self.server_
+
+    def _get_server(self, protocol, **kwargs):
+        if protocol == "rest":
+            server = AggregatorRESTServer(**kwargs)
+        elif protocol == "grpc":
+            server = AggregatorGRPCServer(**kwargs)
+        else:
+            raise ValueError(f"Unsupported transport_protocol '{protocol}'")
+        return server
+
+    def get_server_args(self, root_certificate, private_key, certificate, kwargs):
         common_name = self.config["network"][SETTINGS]["agg_addr"].lower()
 
         if not root_certificate or not private_key or not certificate:
@@ -638,13 +670,7 @@ def get_server(
         server_args["root_certificate"] = root_certificate
         server_args["certificate"] = certificate
         server_args["private_key"] = private_key
-
-        server_args["aggregator"] = self.get_aggregator()
-
-        if self.server_ is None:
-            self.server_ = AggregatorGRPCServer(**server_args)
-
-        return self.server_
+        return server_args
 
     def save_model_to_state_file(self, tensor_dict, round_number, output_path):
         """Save model weights to a protobuf state file.

openfl/interface/aggregator.py

@@ -92,8 +92,9 @@ def start_(plan, authorized_cols, task_group):
         logger.info(f"Setting aggregator to assign: {task_group} task_group")
 
     logger.info("🧿 Starting the Aggregator Service.")
-
-    parsed_plan.get_server().serve()
+    # Instantiate either gRPC or REST server based on transport configuration
+    server = parsed_plan.get_server()
+    server.serve()
 
 
 @aggregator.command(name="generate-cert-request")

openfl/interface/collaborator.py

@@ -14,6 +14,7 @@
 from shutil import copy, copytree, ignore_patterns, make_archive, unpack_archive
 from tempfile import mkdtemp
 
+from click import Choice as ClickChoice
 from click import Path as ClickPath
 from click import confirm, echo, group, option, pass_context, prompt, style
 from yaml import FullLoader, dump, load
@@ -61,7 +62,14 @@ def collaborator(context):
     required=True,
     help="The certified common name of the collaborator.",
 )
-def start_(plan, collaborator_name, data_config):
+@option(
+    "-c",
+    "--client_protocol",
+    help="The client protocol to use for communication with the aggregator.",
+    default="grpc",
+    type=ClickChoice(["grpc", "rest"]),
+)
+def start_(plan, collaborator_name, data_config, client_protocol):
     """Starts a collaborator service."""
 
     if plan and is_directory_traversal(plan):
@@ -80,8 +88,15 @@ def start_(plan, collaborator_name, data_config):
 
     echo(f"Data = {plan.cols_data_paths}")
     logger.info("🧿 Starting a Collaborator Service.")
-
-    plan.get_collaborator(collaborator_name).run()
+    collaborator = plan.get_collaborator(collaborator_name)
+    # Determine if additional client protocol needs to be created
+    if client_protocol == "rest":
+        from openfl.transport.rest.aggregator_client import AggregatorRESTClient
+
+        client_args = plan.get_client_args(collaborator_name=collaborator_name)
+        aggregator_client = AggregatorRESTClient(**client_args)
+        collaborator.client = aggregator_client
+    collaborator.run()
 
 
 @collaborator.command(name="create")

openfl/protocols/aggregator_client_interface.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+from typing import Any, List, Tuple
+
+
+class AggregatorClientInterface(ABC):
+    @abstractmethod
+    def get_tasks(self) -> Tuple[List[Any], int, int, bool]:
+        """
+        Retrieves tasks for the given collaborator client.
+        Returns a tuple: (tasks, round_number, sleep_time, time_to_quit)
+        """
+        pass
+
+    @abstractmethod
+    def get_aggregated_tensor(
+        self,
+        tensor_name: str,
+        round_number: int,
+        report: bool,
+        tags: List[str],
+        require_lossless: bool,
+    ) -> Any:
+        """
+        Retrieves the aggregated tensor.
+        """
+        pass
+
+    @abstractmethod
+    def send_local_task_results(
+        self,
+        round_number: int,
+        task_name: str,
+        data_size: int,
+        named_tensors: List[Any],
+    ) -> Any:
+        """
+        Sends local task results.
+        Parameters:
+          collaborator_name: Name of the collaborator.
+          round_number: The current round.
+          task_name: Name of the task.
+          data_size: Size of the data.
+          named_tensors: A list of tensors (or named tensor objects).
+        Returns a SendLocalTaskResultsResponse.
+        """
+        pass

openfl/transport/__init__.py

@@ -3,3 +3,4 @@
 
 
 from openfl.transport.grpc import AggregatorGRPCClient, AggregatorGRPCServer
+from openfl.transport.rest import AggregatorRESTClient, AggregatorRESTServer

openfl/transport/grpc/aggregator_client.py

@@ -11,6 +11,7 @@
 import grpc
 
 from openfl.protocols import aggregator_pb2, aggregator_pb2_grpc, utils
+from openfl.protocols.aggregator_client_interface import AggregatorClientInterface
 from openfl.transport.grpc.common import create_header, create_insecure_channel, create_tls_channel
 
 logger = logging.getLogger(__name__)
@@ -165,9 +166,12 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
-class AggregatorGRPCClient:
+class AggregatorGRPCClient(AggregatorClientInterface):
     """Collaborator-side gRPC client that talks to the aggregator.
 
+    This class implements a gRPC client for communicating with an aggregator
+    over a secure (TLS) connection.
+
     Attributes:
         agg_addr (str): Aggregator address.
         agg_port (int): Aggregator port.

openfl/transport/rest/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2020-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from openfl.transport.rest.aggregator_client import AggregatorRESTClient
+from openfl.transport.rest.aggregator_server import AggregatorRESTServer