Merge branch 'develop' into karansh1/batched_fetch

Author: MasterSkepticista

openfl-workspace/flower-app-pytorch/README.md

@@ -280,31 +280,18 @@ flwr run ./src/app-pytorch
 ```
 It will run another experiment. Once you are done, you can manually shut down OpenFL's `collaborator` and Flower's `SuperNode` with `CTRL+C`. This will trigger a task-completion by the task runner that'll subsequently begin the graceful shutdown process of the OpenFL and Flower components.
 
-### Running in SGX Enclave
-Gramine does not support all Linux system calls. Flower FAB is built and installed at runtime. During this, `utime()` is called, which is an [unsupported call](https://gramine.readthedocs.io/en/latest/devel/features.html#list-of-system-calls), resulting in error or unexpected behavior. To navigate this, when running in an SGX enclave, we opt to build and install the FAB during initialization and package it alongside the OpenFL workspace. To make this work, we introduce some patches to Flower's build command, which helps circumvent the unsupported system call as well as minimize read/write access.
-
-To run these patches, simply add `patch: True` to the `Connector` and `Task Runner` settings (if not already set). For the `Task Runner` also include the name of the Flower app for building and installation.
+### Running in Intel<sup>®</sup> SGX Enclave
+Intel SGX is a set of CPU extensions for creating isolated memory regions, called enclaves, that allow secure computation of sensitive data. Applications that run within enclaves are encrypted in memory and remain isolated from the rest of the system. Executing code within an enclave requires an [Intel SGX–supported Intel CPU](https://www.intel.com/content/www/us/en/support/articles/000028173/processors.html). To run this workspace in an SGX enclave, first, set `sgx_enabled: True` in the `plan.yaml` for the `task_runner`:
 
 ```yaml
-connector : 
-  defaults : plan/defaults/connector.yaml
-  template : openfl.component.ConnectorFlower
-  settings :
-    automatic_shutdown : True
-    superlink_params :
-      insecure : True
-      serverappio-api-address : 127.0.0.1:9091 
-      fleet-api-address :  127.0.0.1:9092 
-      exec-api-address : 127.0.0.1:9093
-    flwr_run_params :
-      flwr_app_name : "app-pytorch"
-      federation_name : "local-poc"
-      patch : True
-
 task_runner :
   defaults : plan/defaults/task_runner.yaml
-  template : openfl.federated.task.runner_flower.FlowerTaskRunner
+  template : src.runner.FlowerTaskRunner
   settings :
-    patch : True
-    flwr_app_name : "app-pytorch"
-```
+    flwr_app_name: app-pytorch
+    sgx_enabled: False
+```
+
+Then, follow the [instructions](https://openfl.readthedocs.io/en/latest/about/features_index/taskrunner.html#docker-container-approach) to build/pull a base image, dockerize your workspace, and containerize each component.
+
+Enabling this flag does two things. First, it runs the Flower `ClientApp` as an isolated process (`--isolation process`). This reduces the number of processes that need to be spawned in the enclave, reducing overhead time ([see](https://gramine.readthedocs.io/en/stable/performance.html#multi-process-workloads)). Second, Gramine does not support all Linux system calls. Flower FAB is built and installed at runtime. During this, `utime()` is called, which is an [unsupported call](https://gramine.readthedocs.io/en/latest/devel/features.html#list-of-system-calls), resulting in error or unexpected behavior. To navigate this, when running in an SGX enclave, we opt to build and install the FAB during initialization and package it alongside the OpenFL workspace when the `sgx_enabled` flag is set to `True`

openfl-workspace/flower-app-pytorch/plan/plan.yaml

@@ -22,7 +22,6 @@ connector :
     flwr_run_params :
       flwr_app_name : "app-pytorch"
       federation_name : "local-poc"
-      sgx_enabled: False
 
 collaborator :
   defaults : plan/defaults/collaborator.yaml

openfl-workspace/flower-app-pytorch/src/app-pytorch/pyproject.toml

@@ -8,7 +8,7 @@ version = "1.0.0"
 description = ""
 license = "Apache-2.0"
 dependencies = [
-    "flwr>=1.15.0,<1.17.0",
+    "flwr-nightly",
     "flwr-datasets[vision]>=0.5.0",
     "torch==2.5.1",
     "torchvision==0.20.1",

openfl-workspace/flower-app-pytorch/src/connector_flower.py

@@ -154,12 +154,7 @@ def _build_flwr_run_command(self) -> list[str]:
         federation_name = self.flwr_run_params.get("federation_name")
         flwr_app_name = self.flwr_run_params.get("flwr_app_name")
 
-        os.environ["TMPDIR"] = os.environ["FLWR_HOME"]
-
-        if self.flwr_run_params.get("sgx_enabled"):
-            command = ["python", "src/patch/flwr_run_patch.py", "run", f"./src/{flwr_app_name}"]
-        else:
-            command = ["flwr", "run", f"./src/{flwr_app_name}"]
+        command = ["flwr", "run", f"./src/{flwr_app_name}"]
 
         if federation_name:
             command.append(federation_name)

openfl-workspace/flower-app-pytorch/src/patch/__init__.py

@@ -5,7 +5,6 @@
 import os
 import numpy as np
 from pathlib import Path
-import sys
 import socket
 from src.util import is_safe_path
 
@@ -34,13 +33,13 @@ def __init__(self, **kwargs):
         """
         super().__init__(**kwargs)
 
-        self.sgx_enabled = kwargs.get('sgx_enabled')
         if self.data_loader is None:
             flwr_app_name = kwargs.get('flwr_app_name')
-            if self.sgx_enabled:
-                install_flower_FAB(flwr_app_name)
+            install_flower_FAB(flwr_app_name)
             return
 
+        self.sgx_enabled = kwargs.get('sgx_enabled')
+
         self.model = None
         self.logger = getLogger(__name__)
 
@@ -170,16 +169,18 @@ def install_flower_FAB(flwr_app_name):
         flwr_app_name (str): The name of the Flower application to patch.
     """
     flwr_dir = os.environ["FLWR_HOME"]
-    os.environ["TMPDIR"] = flwr_dir
+
+    # Change the current working directory to the Flower directory
+    os.chdir(flwr_dir)
 
     # Run the build command
-    subprocess.check_call([
-        sys.executable,
-        "src/patch/flwr_run_patch.py",
+    build_command = [
+        "flwr",
         "build",
         "--app",
-        f"./src/{flwr_app_name}"
-    ])
+        os.path.join("..", "..", "src", flwr_app_name)
+    ]
+    subprocess.check_call(build_command)
 
     # List .fab files after running the build command
     fab_files = list(Path(flwr_dir).glob("*.fab"))
@@ -189,8 +190,7 @@ def install_flower_FAB(flwr_app_name):
 
     # Run the install command using the newest .fab file
     subprocess.check_call([
-        sys.executable,
-        "src/patch/flwr_run_patch.py",
+        "flwr",
         "install",
         str(newest_fab_file)
     ])

openfl-workspace/flower-app-pytorch/src/patch/flwr_run_patch.py

@@ -149,15 +149,17 @@ def __init__(
         self.quit_job_sent_to = []
 
         self.tensor_db = TensorDB()
-        if persist_checkpoint:
+        if persist_checkpoint and not self.assigner.is_task_group_evaluation():
             persistent_db_path = persistent_db_path or "tensor.db"
             logger.info(
                 "Persistent checkpoint is enabled, setting persistent db at path %s",
                 persistent_db_path,
             )
             self.persistent_db = PersistentTensorDB(persistent_db_path)
         else:
-            logger.info("Persistent checkpoint is disabled")
+            logger.info(
+                "Either persistent checkpoint is disabled or the experiment is in evaluation mode"
+            )
             self.persistent_db = None
         # FIXME: I think next line generates an error on the second round
         # if it is set to 1 for the aggregator.
@@ -224,8 +226,10 @@ def __init__(
 
             self.secagg = SecAggSetup(self.uuid, self.authorized_cols, self.tensor_db)
 
-        if self.persistent_db and self._recover():
-            logger.info("Recovered state of aggregator")
+        # Only recover from persistent DB if not in evaluation mode
+        if self.persistent_db and not self.assigner.is_task_group_evaluation():
+            if self._recover():
+                logger.info("Recovered state of aggregator")
 
         # TODO: Aggregator has no concrete notion of round_begin.
         # https://github.com/securefederatedai/openfl/pull/1195#discussion_r1879479537