Task Runner API: Secure Aggregation #1346

theakshaypant · 2025-02-06T13:49:33Z

Introduction

This PR introduces secure aggregation to Task Runner API.

Would specially apprecaite comments on how tensors are exchanged between the participants.

Changes

Secure aggregation on client side is implemented as callback.
This callback is run on_experiment_begin.
Collaborator side steps are implemented using CollaboratorSecAgg callback.
Aggregator side steps are implemented using openfl.utilities.secagg.Setup whcih performs "aggregation" of the tensors sent by the collaborators..
send_local_task_results is modfied to save secure aggregation setup stage tensors to aggregator tensor db.
Masks are added to "metric"s when a task is completed before sending the result to the aggregator.
SecureAggregation is used for "metric"s when enabled instead of WeightedAverage.
Added a secure_aggregation: bool flag to the plan under
- aggregator.settings
- collaborator.settings

NOTE

Does not cover dropouts, will be handled as part of SecAgg+ integration in future PRs.
In this implementation, there is a 60 second wait period when collaborator tries to fetch keys from the aggregator, fails if there is no response during this time.
- This time can be increased or further retries can be added.

Testing

Aggregator logs (with extra debug prints)
Collaborator 1 logs
Collaborator 2 logs

Signed-off-by: Pant, Akshay <[email protected]>

… tensor db Signed-off-by: Pant, Akshay <[email protected]>

…allbacks at the component Signed-off-by: Pant, Akshay <[email protected]>

…laborators to handle secure aggregation setup phase Signed-off-by: Pant, Akshay <[email protected]>

…o task-runner-api/secure-aggregation

Signed-off-by: Pant, Akshay <[email protected]>

payalcha · 2025-02-10T05:17:46Z

setup.py

@@ -94,6 +94,7 @@ def run(self):
        'tensorboardX',
        'protobuf>=4.22,<6.0.0',
        'grpcio>=1.56.2,<1.66.0',
+        'pycryptodome'


If we keep this as setup, even when secure aggregation not needed by customer, they need to install it.

payalcha · 2025-02-10T05:21:39Z

openfl/utilities/secagg/crypto.py

@@ -0,0 +1,154 @@
+# Copyright 2020-2025 Intel Corporation


There is cryptography package in openfl. Isn't secagg packages more suitable inside cryptography.

openfl/utilities/secagg/crypto.py

Signed-off-by: Pant, Akshay <[email protected]>

…o task-runner-api/secure-aggregation

Signed-off-by: Pant, Akshay <[email protected]>

…o task-runner-api/secure-aggregation

Signed-off-by: Pant, Akshay <[email protected]>

openfl/utilities/secagg/crypto.py

+        shape (Tuple): Shape of the numpy array to be generated.
+
+    Returns:
+        np.ndarray: array with pseudo-randomly generated numbers.


Signed-off-by: Pant, Akshay <[email protected]>

openfl-workspace/keras/mnist_secagg/plan/plan.yaml

tanwarsh · 2025-02-11T05:06:15Z

openfl-workspace/keras/mnist_secagg/src/dataloader.py

+        """
+        super().__init__(batch_size, **kwargs)
+
+        # TODO: We should be downloading the dataset shard into a directory


remove this.

you can remove this comment - refer #1333 (comment)

theakshaypant · 2025-02-11T05:08:49Z

openfl/component/collaborator/collaborator.py

@@ -168,6 +175,18 @@ def set_available_devices(self, cuda: Tuple[str] = ()):
    def run(self):
        """Run the collaborator."""
        # Experiment begin
+
+        # FIXME: Not working when added to callbacks on line 157.


@MasterSkepticista Need your input on this.

openfl-workspace/keras/mnist_secagg/plan/plan.yaml

rahulga1 · 2025-02-11T04:59:52Z

openfl-workspace/keras/mnist_secagg/src/mnist_utils.py

@@ -0,0 +1,118 @@
+# Copyright (C) 2020-2021 Intel Corporation


is it not possible to refer the original file, if there is no sec_agg specific changes in here?
If not possible, we can change name to utils.py, mnist is already in the path.

Have followed the same format as all other workspaces where the funnction definition is in all of them.
Changed file name in 1579c05.

tanwarsh · 2025-02-11T05:14:17Z

openfl/component/collaborator/collaborator.py

@@ -168,6 +175,18 @@ def set_available_devices(self, cuda: Tuple[str] = ()):
    def run(self):
        """Run the collaborator."""
        # Experiment begin
+
+        # FIXME: Not working when added to callbacks on line 157.


please fix this or remove the comment if no longer required.

Need input on this. Refer #1346 (comment).

openfl/component/aggregator/aggregator.py

rahulga1 · 2025-02-11T06:25:54Z

openfl/component/aggregator/aggregator.py

+                tensor_name = named_tensor.name
+                # Check if all collaborators have sent their data for the
+                # current key.
+                all_collaborators_sent = self.secagg.wait_for_all_collaborators(tensor_name)


in case of staggler, will it hang indefinitely?

There is no wait in this method. Only checks if all collaborators have shared "tensor_name", will rename it appropriately.

openfl/utilities/secagg/setup.py

Signed-off-by: Pant, Akshay <[email protected]>

…al for agg function Signed-off-by: Pant, Akshay <[email protected]>

Signed-off-by: Pant, Akshay <[email protected]>

payalcha · 2025-02-13T04:49:28Z

@theakshaypant are we are storing keys and cipher text right in persistent storage?
Collaborator restart should not impact, secure aggregation.

theakshaypant added 9 commits February 5, 2025 17:46

feat(secagg): add utility functions

2bc32c2

Signed-off-by: Pant, Akshay <[email protected]>

feat(secagg): add callbacks for setup

7bd840b

Signed-off-by: Pant, Akshay <[email protected]>

feat(secure aggregation): used callbacks to use aggregator client and…

1c7ec89

… tensor db Signed-off-by: Pant, Akshay <[email protected]>

feat(secure aggregation): read flag from plan and enable respective c…

a2cef93

…allbacks at the component Signed-off-by: Pant, Akshay <[email protected]>

feat(secure aggregation): modify function to save tensors sent by col…

d2f4144

…laborators to handle secure aggregation setup phase Signed-off-by: Pant, Akshay <[email protected]>

Merge branch 'develop' of https://github.com/theakshaypant/openfl int…

42c3469

…o task-runner-api/secure-aggregation

Merge branch 'develop' of https://github.com/theakshaypant/openfl int…

b003918

…o task-runner-api/secure-aggregation

feat(secure aggregation): add depedency in requirements list

8318752

Signed-off-by: Pant, Akshay <[email protected]>

feat(secure aggregation): task runner exampel workspace

f3d0f27

Signed-off-by: Pant, Akshay <[email protected]>

payalcha reviewed Feb 10, 2025

View reviewed changes

theakshaypant added 8 commits February 10, 2025 23:22

fix(secure aggregation): change utility functions after testing

31feb1b

Signed-off-by: Pant, Akshay <[email protected]>

fix(secure aggregation): change to only use callback for collaborator

413f798

Signed-off-by: Pant, Akshay <[email protected]>

feat(secure aggregation): add setup steps for the server

3146175

Signed-off-by: Pant, Akshay <[email protected]>

Merge branch 'develop' of https://github.com/theakshaypant/openfl int…

91d6ca3

…o task-runner-api/secure-aggregation

fix(secure aggregation): read enabling flag from plan.yaml

55f919c

Signed-off-by: Pant, Akshay <[email protected]>

feat(secure aggregation): add server side unmasking

54f218e

Signed-off-by: Pant, Akshay <[email protected]>

feat(secure aggregation): add client side masking

13160c7

Signed-off-by: Pant, Akshay <[email protected]>

fix(secure aggregation): add AggregationFunction

edbfd8d

theakshaypant mentioned this pull request Feb 10, 2025

Workflow API: Secure Aggregation example #1329

Open

theakshaypant added 2 commits February 11, 2025 10:12

Merge branch 'develop' of https://github.com/theakshaypant/openfl int…

3069ce2

…o task-runner-api/secure-aggregation

restructure: formatting changes

592ccb0

Signed-off-by: Pant, Akshay <[email protected]>

theakshaypant marked this pull request as ready for review February 11, 2025 04:46

theakshaypant requested review from MasterSkepticista, kta-intel, psfoley, teoparvanov, aayushgaintel, gbikkiintel and pasokan-intel as code owners February 11, 2025 04:46

theakshaypant requested review from rahulga1, rajithkrishnegowda, ishaileshpant, tanwarsh and srikanthenugul as code owners February 11, 2025 04:46

github-advanced-security bot found potential problems Feb 11, 2025

View reviewed changes

restructure: formatting changes

9b0ffef

Signed-off-by: Pant, Akshay <[email protected]>

tanwarsh reviewed Feb 11, 2025

View reviewed changes

openfl-workspace/keras/mnist_secagg/plan/plan.yaml Outdated Show resolved Hide resolved

tanwarsh reviewed Feb 11, 2025

View reviewed changes

theakshaypant commented Feb 11, 2025

View reviewed changes

rahulga1 reviewed Feb 11, 2025

View reviewed changes

tanwarsh reviewed Feb 11, 2025

View reviewed changes

rahulga1 reviewed Feb 11, 2025

View reviewed changes

theakshaypant added 3 commits February 11, 2025 12:44

doc(secure aggregation): change workspace header

1579c05

Signed-off-by: Pant, Akshay <[email protected]>

fix(secure aggregation): remove redundant flag init and use condition…

a71ba79

…al for agg function Signed-off-by: Pant, Akshay <[email protected]>

restructure: formatting changes

ee23921

Signed-off-by: Pant, Akshay <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Task Runner API: Secure Aggregation #1346

Task Runner API: Secure Aggregation #1346

theakshaypant commented Feb 6, 2025 •

edited

Loading

payalcha Feb 10, 2025

payalcha Feb 10, 2025

tanwarsh Feb 11, 2025 •

edited

Loading

tanwarsh Feb 13, 2025

theakshaypant Feb 11, 2025

rahulga1 Feb 11, 2025

theakshaypant Feb 11, 2025

tanwarsh Feb 11, 2025

theakshaypant Feb 11, 2025

rahulga1 Feb 11, 2025

theakshaypant Feb 11, 2025

payalcha commented Feb 13, 2025 •

edited

Loading

		@@ -0,0 +1,118 @@
		# Copyright (C) 2020-2021 Intel Corporation

Task Runner API: Secure Aggregation #1346

Are you sure you want to change the base?

Task Runner API: Secure Aggregation #1346

Conversation

theakshaypant commented Feb 6, 2025 • edited Loading

Introduction

Changes

NOTE

Testing

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tanwarsh Feb 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

payalcha commented Feb 13, 2025 • edited Loading

theakshaypant commented Feb 6, 2025 •

edited

Loading

tanwarsh Feb 11, 2025 •

edited

Loading

payalcha commented Feb 13, 2025 •

edited

Loading