Task Runner API: Secure Aggregation (securefederatedai#1346)

* feat(secagg): add utility functions Signed-off-by: Pant, Akshay <[email protected]> * feat(secagg): add callbacks for setup Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): used callbacks to use aggregator client and tensor db Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): read flag from plan and enable respective callbacks at the component Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): modify function to save tensors sent by collaborators to handle secure aggregation setup phase Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): add depedency in requirements list Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): task runner exampel workspace Signed-off-by: Pant, Akshay <[email protected]> * fix(secure aggregation): change utility functions after testing Signed-off-by: Pant, Akshay <[email protected]> * fix(secure aggregation): change to only use callback for collaborator Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): add setup steps for the server Signed-off-by: Pant, Akshay <[email protected]> * fix(secure aggregation): read enabling flag from plan.yaml Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): add server side unmasking Signed-off-by: Pant, Akshay <[email protected]> * feat(secure aggregation): add client side masking Signed-off-by: Pant, Akshay <[email protected]> * fix(secure aggregation): add AggregationFunction * restructure: formatting changes Signed-off-by: Pant, Akshay <[email protected]> * restructure: formatting changes Signed-off-by: Pant, Akshay <[email protected]> * doc(secure aggregation): change workspace header Signed-off-by: Pant, Akshay <[email protected]> * fix(secure aggregation): remove redundant flag init and use conditional for agg function Signed-off-by: Pant, Akshay <[email protected]> * restructure: formatting changes Signed-off-by: Pant, Akshay <[email protected]> * fix(secure aggregation): add secure aggregation to collaborator callback list Signed-off-by: Pant, Akshay <[email protected]> * docs(secure aggregation): add documentation for taskrunner and workflow Signed-off-by: Pant, Akshay <[email protected]> * fix(sec agg): remove mnist example Signed-off-by: Pant, Akshay <[email protected]> * restructure(secagg): renamed setup fucntion to a more appropriate name Signed-off-by: Pant, Akshay <[email protected]> * style: pre-commit changes Signed-off-by: Pant, Akshay <[email protected]> * doc(secagg): fix typo Signed-off-by: Pant, Akshay <[email protected]> * restructure(secagg): rename methods and update docstrings Signed-off-by: Pant, Akshay <[email protected]> * restructure(secagg): rename callback Signed-off-by: Pant, Akshay <[email protected]> * fix(secagg): change agg_function selection in the collaborator Signed-off-by: Pant, Akshay <[email protected]> * style(secagg): move conditional outside method to where it is called Signed-off-by: Pant, Akshay <[email protected]> * fix(secagg): read only aggregator setting to check if secagg is enabled Signed-off-by: Pant, Akshay <[email protected]> * fix(secagg): read masks from db only once per collaborator Signed-off-by: Pant, Akshay <[email protected]> * feat(secagg): store col pvt seeds in tensor db insted of masks during bootstrap Signed-off-by: Pant, Akshay <[email protected]> * feat(secagg): subtract weighted avg of masks when taking a weighted average of tensors sent by collaborators Signed-off-by: Pant, Akshay <[email protected]> * restructure(secagg): reduce complexity of _generate_masks by splitting it into another method Signed-off-by: Pant, Akshay <[email protected]> * doc(secagg): fix the flags to be set in plan for enabling secagg Signed-off-by: Pant, Akshay <[email protected]> * restructure(aggregation function): rename SecureAggregation to SecureWeightedAverage Signed-off-by: Pant, Akshay <[email protected]> * docs(secagg): change to mention secure aggregation as a core security feature Signed-off-by: Pant, Akshay <[email protected]> * docs(secagg): emphasize the change required in plan Signed-off-by: Pant, Akshay <[email protected]> * docs(secagg): update docstring for aggregation function Signed-off-by: Pant, Akshay <[email protected]> * docs(secagg): update code-block language Signed-off-by: Pant, Akshay <[email protected]> --------- Signed-off-by: Pant, Akshay <[email protected]>
payalcha · Feb 26, 2025 · e7a8adf · e7a8adf
1 parent bc227e5
commit e7a8adf
Show file tree

Hide file tree

Showing 12 changed files with 1,042 additions and 8 deletions.
diff --git a/docs/about/features.rst b/docs/about/features.rst
@@ -97,3 +97,16 @@ Quantitatively audit data privacy in statistical and machine learning algorithms
 
     features_index/privacy_meter
 
+.. _secure_aggregation:
+
+---------------------
+Secure Aggregation
+---------------------
+
+In Federated Learning (FL), Secure Aggregation (SecAgg) is a technique that allows the participants to collaborate on the central model without revealing their individual contributions (local model updates). For more info see :doc:`features_index/secure_aggregation`
+
+.. toctree::
+    :hidden:
+
+    features_index/secure_aggregation
+
diff --git a/docs/about/features_index/secure_aggregation.rst b/docs/about/features_index/secure_aggregation.rst
@@ -0,0 +1,120 @@
+.. # Copyright (C) 2020-2024 Intel Corporation
+.. # SPDX-License-Identifier: Apache-2.0
+
+Secure Aggregation
+=======================================
+
+In Federated Learning (FL), Secure Aggregation (SecAgg) restricts the aggregator to only learn the summation or average of the updates from collaborators.
+
+OpenFL integrates `SecAgg <https://eprint.iacr.org/2017/281.pdf>`_ into TaskRunner API as well as the Workflow API.
+
+TaskRunner API
+-------------------------------------
+
+OpenFL treats SecAgg as a core security feature and can be enabled for any experiment by simply modifying the plan.
+
+The following plan shows secure aggregation being enabled on `keras/mnist <https://github.com/securefederatedai/openfl/tree/develop/openfl-workspace/keras/mnist>`_ workspace by simply modifying the plan.
+
+.. code-block:: yaml
+    :emphasize-lines: 10
+
+    aggregator:
+        settings:
+            best_state_path: save/best.pbuf
+            db_store_rounds: 2
+            init_state_path: save/init.pbuf
+            last_state_path: save/last.pbuf
+            persist_checkpoint: true
+            persistent_db_path: local_state/tensor.db
+            rounds_to_train: 1
+            secure_aggregation: true
+        template: openfl.component.Aggregator
+    assigner:
+        settings:
+            task_groups:
+            - name: learning
+            percentage: 1.0
+            tasks:
+            - aggregated_model_validation
+            - train
+            - locally_tuned_model_validation
+            - name: evaluation
+            percentage: 0
+            tasks:
+            - aggregated_model_validation
+        template: openfl.component.RandomGroupedAssigner
+    collaborator:
+        settings:
+            db_store_rounds: 1
+            delta_updates: false
+            opt_treatment: RESET
+        template: openfl.component.Collaborator
+    compression_pipeline:
+        settings: {}
+        template: openfl.pipelines.NoCompressionPipeline
+    data_loader:
+        settings:
+            batch_size: 256
+            collaborator_count: 2
+            data_group_name: mnist
+        template: src.dataloader.KerasMNISTInMemory
+    network:
+        settings:
+            agg_addr: localhost
+            agg_port: 53788
+            cert_folder: cert
+            client_reconnect_interval: 5
+            hash_salt: auto
+            require_client_auth: true
+            use_tls: true
+        template: openfl.federation.Network
+    task_runner:
+        settings: {}
+        template: src.taskrunner.KerasCNN
+    tasks:
+        aggregated_model_validation:
+            function: validate_task
+            kwargs:
+            apply: global
+            batch_size: 32
+            metrics:
+            - accuracy
+        locally_tuned_model_validation:
+            function: validate_task
+            kwargs:
+            apply: local
+            batch_size: 32
+            metrics:
+            - accuracy
+        settings: {}
+        train:
+            function: train_task
+            kwargs:
+            batch_size: 32
+            epochs: 1
+            metrics:
+            - loss
+        
+As can be seen in the above plan, by only enabling ``aggregator.settings.secure_aggregation`` in the workspace plan, one can enable SecAgg.
+
+After the flags have been set in plan.yml and the setup for the experiment is completed, one can verify that SecAgg was enabled by looking at the aggregator logs
+
+.. code-block:: bash
+
+    [21:55:01] INFO     SecAgg: recreated secrets successfully                                                                                          setup.py:281
+               INFO     SecAgg: setup completed, saved required tensors to db
+
+Similarly, in the collaborator logs
+
+.. code-block:: bash
+
+               INFO     Secure aggregation is enabled, starting setup...                                                                    secure_aggregation.py:48
+    [21:55:01] INFO     SecAgg: setup completed, saved required tensors to db.
+
+
+Workflow API
+-------------------------------------
+
+OpenFL provides `utility functions <https://github.com/securefederatedai/openfl/tree/develop/openfl/utilities/secagg>`_ that can be utilised to perform SecAgg in Workflow API.
+
+An example notebook can be found `here <https://github.com/securefederatedai/openfl/tree/develop/openfl-tutorials/experimental/workflow/SecAgg>`_ that showcases how the secure aggregation flow can be achieved in Workflow API using both, LocalRuntime and FederatedRuntime.
diff --git a/openfl/callbacks/__init__.py b/openfl/callbacks/__init__.py
@@ -5,3 +5,4 @@
 from openfl.callbacks.lambda_callback import LambdaCallback
 from openfl.callbacks.memory_profiler import MemoryProfiler
 from openfl.callbacks.metric_writer import MetricWriter
+from openfl.callbacks.secure_aggregation import SecAggBootstrapping
diff --git a/openfl/callbacks/callback_list.py b/openfl/callbacks/callback_list.py
@@ -29,7 +29,7 @@ def __init__(
         **params,
     ):
         super().__init__()
-        self.callbacks = _flatten(callbacks) if callbacks else []
+        self.callbacks = list(_flatten(callbacks)) if callbacks else []
 
         self._add_default_callbacks(add_memory_profiler, add_metric_writer)
 
@@ -77,7 +77,7 @@ def on_round_end(self, round_num: int, logs=None):
 
     def on_experiment_begin(self, logs=None):
         for callback in self.callbacks:
-            callback.on_experiment_begin(logs)
+            callback.on_experiment_begin(logs=logs)
 
     def on_experiment_end(self, logs=None):
         for callback in self.callbacks: