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
rebased 26-Mar.2

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

Author: ishaileshpant

openfl/federated/plan/plan.py

@@ -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,
@@ -623,6 +634,16 @@ def get_server(
         Returns:
             AggregatorGRPCServer: gRPC server of the aggregator instance.
         """
+        server_args = self.get_server_args(root_certificate, private_key, certificate, kwargs)
+
+        server_args["aggregator"] = self.get_aggregator()
+
+        if self.server_ is None:
+            self.server_ = AggregatorGRPCServer(**server_args)
+
+        return self.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 +659,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/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.proto

@@ -11,8 +11,8 @@ import "openfl/protocols/base.proto";
 service Aggregator {
   rpc GetTasks(GetTasksRequest) returns (GetTasksResponse) {}
   rpc GetAggregatedTensor(GetAggregatedTensorRequest) returns (GetAggregatedTensorResponse) {}
-  rpc SendLocalTaskResults(stream DataStream) returns (SendLocalTaskResultsResponse) {}
   rpc InteropRelay(InteropMessage) returns (InteropMessage) {}
+  rpc SendLocalTaskResults(stream DataStream) returns (SendLocalTaskResultsResponse) {}
 }
 
 message MessageHeader {

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

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,5 @@
+# Copyright 2020-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+from openfl.transport.rest.aggregator_client import AggregatorRESTClient