Adding Secure Aggregation Test (#1419)

* Secure Aggregation Tests

Signed-off-by: Chaurasiya, Payal <[email protected]>

* Secure Aggregation Tests

Signed-off-by: Chaurasiya, Payal <[email protected]>

* Handle success and exception messages

Signed-off-by: Chaurasiya, Payal <[email protected]>

* Handle success and exception messages

Signed-off-by: Chaurasiya, Payal <[email protected]>

* Format issue

Signed-off-by: Chaurasiya, Payal <[email protected]>

* Put Secure Agg

Signed-off-by: Chaurasiya, Payal <[email protected]>

* Review comments

Signed-off-by: Chaurasiya, Payal <[email protected]>

---------

Signed-off-by: Chaurasiya, Payal <[email protected]>

Author: payalcha

.github/workflows/pq_pipeline.yml

@@ -65,6 +65,13 @@ jobs:
     needs: task_runner_e2e
     uses: ./.github/workflows/task_runner_fedeval_e2e.yml
 
+  task_runner_secure_agg_e2e:
+    if: |
+      (github.event_name == 'schedule' && github.repository_owner == 'securefederatedai') ||
+      (github.event_name == 'workflow_dispatch')
+    name: TaskRunner Secure Aggregation E2E
+    uses: ./.github/workflows/task_runner_secure_agg_e2e.yml
+
   task_runner_straggler_e2e:
     if: |
       (github.event_name == 'schedule' && github.repository_owner == 'securefederatedai') ||

.github/workflows/pr_pipeline.yml

@@ -59,10 +59,14 @@ jobs:
     name: TaskRunner E2E
     uses: ./.github/workflows/task_runner_basic_e2e.yml
 
-  task_runner_e2e_resiliency:
+  task_runner_resiliency_e2e:
     name: TaskRunner Resiliency E2E
     uses: ./.github/workflows/task_runner_resiliency_e2e.yml
 
+  task_runner_secure_agg_e2e:
+    name: TaskRunner Secure Aggregation E2E
+    uses: ./.github/workflows/task_runner_secure_agg_e2e.yml
+
   tr_docker_gramine_direct:
     name: TaskRunner (docker/gramine-direct)
     uses: ./.github/workflows/tr_docker_gramine_direct.yml

.github/workflows/task_runner_secure_agg_e2e.yml

@@ -0,0 +1,52 @@
+---
+# Task Runner E2E tests with Secure aggregation with bare metal approach
+
+name: Task_Runner_Secure_Agg_E2E  # Please do not modify the name as it is used in the composite action
+
+on:
+  workflow_call:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  test_secure_aggregation:
+    name: Secure Aggregation (torch/mnist, 3.10)
+    runs-on: ubuntu-22.04
+    timeout-minutes: 30
+    env:
+      MODEL_NAME: 'torch/mnist'
+      PYTHON_VERSION: '3.10'
+      NUM_ROUNDS: 5
+      NUM_COLLABORATORS: 2
+    steps:
+      - name: Checkout OpenFL repository
+        id: checkout_openfl
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 2 # needed for detecting changes
+          submodules: "true"
+          token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Install Secure Aggregation dependies
+        run: |
+          python -m pip install pycryptodome
+          
+      - name: Pre test run
+        uses: ./.github/actions/tr_pre_test_run
+        if: ${{ always() }}
+
+      - name: Run Task Runner with secure aggregation
+        id: run_tests
+        run: |
+          python -m pytest -s tests/end_to_end/test_suites/task_runner_tests.py \
+          -m task_runner_basic --model_name ${{ env.MODEL_NAME }} \
+          --num_rounds ${{ env.NUM_ROUNDS }} --num_collaborators ${{ env.NUM_COLLABORATORS }} --secure_agg
+          echo "Task runner end to end test run completed"
+
+      - name: Post test run
+        uses: ./.github/actions/tr_post_test_run
+        if: ${{ always() }}
+        with:
+          test_type: "With_Secure_Agg"

tests/end_to_end/conftest.py

@@ -29,6 +29,7 @@ def pytest_addoption(parser):
     parser.addoption("--disable_client_auth", action="store_true")
     parser.addoption("--disable_tls", action="store_true")
     parser.addoption("--log_memory_usage", action="store_true")
+    parser.addoption("--secure_agg", action="store_true")
 
 
 def pytest_configure(config):
@@ -46,6 +47,7 @@ def pytest_configure(config):
     config.require_client_auth = not args.disable_client_auth
     config.use_tls = not args.disable_tls
     config.log_memory_usage = args.log_memory_usage
+    config.secure_agg = args.secure_agg
     config.results_dir = config.getini("results_dir")
 
 

tests/end_to_end/models/model_owner.py

@@ -159,8 +159,8 @@ def modify_plan(self, param_config, plan_path):
             data["data_loader"]["settings"]["collaborator_count"] = int(self.num_collaborators)
             data["network"]["settings"]["require_client_auth"] = param_config.require_client_auth
             data["network"]["settings"]["use_tls"] = param_config.use_tls
-
-
+            if param_config.secure_agg:
+                data["aggregator"]["settings"]["secure_aggregation"] = True
             with open(plan_file, "w+") as write_file:
                 yaml.dump(data, write_file)
             log.info(f"Modified the plan with provided parameters.")

tests/end_to_end/utils/conftest_helper.py

@@ -33,6 +33,7 @@ def parse_arguments():
         parser.add_argument("--disable_client_auth", action="store_true", help="Disable client authentication. Default is False")
         parser.add_argument("--disable_tls", action="store_true", help="Disable TLS for communication. Default is False")
         parser.add_argument("--log_memory_usage", action="store_true", help="Enable Memory leak logs. Default is False")
+        parser.add_argument("--secure_agg", action="store_true", help="Enable secure aggregation. Default is False")
         args = parser.parse_known_args()[0]
         return args
 

tests/end_to_end/utils/constants.py

@@ -54,3 +54,4 @@ class ModelName(Enum):
 
 COL_CERTIFY_CMD = "fx collaborator certify --import 'agg_to_col_{}_signed_cert.zip'"
 DFLT_DOCKERIZE_IMAGE_NAME = "workspace"
+EXCEPTION = "Exception"

tests/end_to_end/utils/federation_helper.py

@@ -171,12 +171,11 @@ def _create_tarball(collaborator_name, data_file_path, local_bind_path, add_data
     return True
 
 
-def import_pki_for_collaborators(collaborators, local_bind_path):
+def import_pki_for_collaborators(collaborators):
     """
     Import and certify the CSR for the collaborators
     """
     executor = concurrent.futures.ThreadPoolExecutor()
-    local_agg_ws_path = constants.AGG_WORKSPACE_PATH.format(local_bind_path)
     try:
         results = [
             executor.submit(
@@ -359,13 +358,21 @@ def _verify_completion_for_participant(
     ):
         with open(participant.res_file, "r") as file:
             lines = [line.strip() for line in file.readlines()]
-        content = list(filter(str.rstrip, lines))[-1:]
+        # Below change is done to incorporate warnings coming in end of runs
+        content = list(filter(str.rstrip, lines))[-7:] if len(lines) >= 7 else lines
 
         # Print last line of the log file on screen to track the progress
-        log.info(f"Last line in {participant.name} log: {content}")
+        log.info(f"Last line in {participant.name} log: {lines[-1:]}")
         if constants.SUCCESS_MARKER in content:
             break
         log.info(f"Process is yet to complete for {participant.name}")
+        # If in logs Exception is encountered, throw Exception and stop the process
+        if constants.EXCEPTION in content:
+            log.error(
+                f"Process {participant.name} is throwing Exception. Check the logs for more details"
+            )
+            raise Exception(f"Process failed for {participant.name}")
+
         time.sleep(45)
 
     if constants.SUCCESS_MARKER not in content:
@@ -430,6 +437,7 @@ def federation_env_setup_and_validate(request, eval_scope=False):
         f"\tModel name: {request.config.model_name}\n"
         f"\tClient authentication: {request.config.require_client_auth}\n"
         f"\tTLS: {request.config.use_tls}\n"
+        f"\tSecure Aggregation: {request.config.secure_agg}\n"
         f"\tMemory Logs: {request.config.log_memory_usage}\n"
         f"\tResults directory: {request.config.results_dir}\n"
         f"\tWorkspace path: {workspace_path}"

tests/end_to_end/utils/tr_workspace.py

@@ -145,7 +145,7 @@ def create_tr_workspace(request, eval_scope=False):
 
     if request.config.use_tls:
         fh.setup_pki_for_collaborators(collaborators, model_owner, local_bind_path)
-        fh.import_pki_for_collaborators(collaborators, local_bind_path)
+        fh.import_pki_for_collaborators(collaborators)
 
     fh.remove_stale_processes(request.config.num_collaborators)