Resolve Visibility of Aggregator Private Attributes in Experimental Workflow

openfl/experimental/component/aggregator/aggregator.py

@@ -153,6 +153,9 @@ def run_flow(self) -> None:
         """Start the execution and run flow until transition."""
         # Start function will be the first step if any flow
         f_name = "start"
+        # Creating a clones from the flow object
+        FLSpec._reset_clones()
+        FLSpec._create_clones(self.flow, self.flow.runtime.collaborators)
 
         self.logger.info(f"Starting round {self.current_round}...")
         while True:
@@ -363,22 +366,21 @@ def do_task(self, f_name: str) -> Any:
 
             # Transition check
             if aggregator_to_collaborator(f, parent_func):
-                # Extract clones, instance snapshot and kwargs when reached
-                # foreach loop first time
-                if len(self.flow.execute_task_args) > 4:
-                    temp = self.flow.execute_task_args[3:]
-                    self.clones_dict, self.instance_snapshot, self.kwargs = temp
-
+                # Delete aggregator private attribute from flow object
+                self.__delete_agg_attrs_from_clone(self.flow)
+                # Unpack execute_task_args - clones_dict, instance snapshot and kwargs
+                self.clones_dict, self.instance_snapshot, self.kwargs = self.flow.execute_task_args[
+                    3:
+                ]
+                if "foreach" in self.kwargs:
+                    self.flow.filter_exclude_include(f, **self.kwargs)
                     self.selected_collaborators = getattr(self.flow, self.kwargs["foreach"])
                 else:
-                    self.kwargs = self.flow.execute_task_args[3]
+                    self.kwargs = self.flow.execute_task_args[-1]
 
                 # Transition encountered, break the loop
                 not_at_transition_point = False
 
-        # Delete private attributes from flow object
-        self.__delete_agg_attrs_from_clone(self.flow)
-
         return f_name if f_name != "end" else None
 
     def send_task_results(

openfl/experimental/interface/fl_spec.py

@@ -228,22 +228,6 @@ def restore_instance_snapshot(self, ctx: FLSpec, instance_snapshot: List[FLSpec]
                 if not hasattr(ctx, name):
                     setattr(ctx, name, attr)
 
-    def get_clones(self, kwargs):
-        """Create, and prepare clones."""
-        FLSpec._reset_clones()
-        FLSpec._create_clones(self, self.runtime.collaborators)
-        selected_collaborators = self.__getattribute__(kwargs["foreach"])
-
-        for col in selected_collaborators:
-            clone = FLSpec._clones[col]
-            clone.input = col
-            artifacts_iter, _ = generate_artifacts(ctx=clone)
-            attributes = artifacts_iter()
-            for name, attr in attributes:
-                setattr(clone, name, deepcopy(attr))
-            clone._foreach_methods = self._foreach_methods
-            clone._metaflow_interface = self._metaflow_interface
-
     def next(self, f, **kwargs):
         """Specifies the next task in the flow to execute.
 
@@ -264,30 +248,21 @@ def next(self, f, **kwargs):
         if aggregator_to_collaborator(f, parent_func):
             agg_to_collab_ss = self._capture_instance_snapshot(kwargs=kwargs)
 
-            if str(self._runtime) == "FederatedRuntime":
-                if len(FLSpec._clones) == 0:
-                    self.get_clones(kwargs)
-
         # Remove included / excluded attributes from next task
         filter_attributes(self, f, **kwargs)
 
         if str(self._runtime) == "FederatedRuntime":
             if f.collaborator_step and not f.aggregator_step:
                 self._foreach_methods.append(f.__name__)
 
-            if "foreach" in kwargs:
-                self.filter_exclude_include(f, **kwargs)
-                # if "foreach" in kwargs:
-                self.execute_task_args = (
-                    self,
-                    f,
-                    parent_func,
-                    FLSpec._clones,
-                    agg_to_collab_ss,
-                    kwargs,
-                )
-            else:
-                self.execute_task_args = (self, f, parent_func, kwargs)
+            self.execute_task_args = (
+                self,
+                f,
+                parent_func,
+                FLSpec._clones,
+                agg_to_collab_ss,
+                kwargs,
+            )
 
         elif str(self._runtime) == "LocalRuntime":
             # update parameters required to execute execute_task function

tests/github/experimental/workspace/testcase_private_attributes/src/aggregator_private_attrs.py

@@ -4,4 +4,4 @@
 
 
 def aggregator_private_attrs():
-    return {"test_loader": np.random.rand(10, 28, 28)}  # Random data
+    return {"test_loader_agg": np.random.rand(10, 28, 28)}  # Random data

tests/github/experimental/workspace/testcase_private_attributes/src/testflow_privateattributes.py

@@ -37,7 +37,7 @@ def start(self):
         )
         self.collaborators = self.runtime.collaborators
 
-        validate_collab_private_attr(self, "test_loader", "start")
+        validate_agg_private_attr(self, "start", aggr = ["test_loader_agg"], collabs =["train_loader", "test_loader"])
 
         self.exclude_agg_to_agg = 10
         self.include_agg_to_agg = 100
@@ -49,7 +49,7 @@ def aggregator_step(self):
         Testing whether Agg private attributes are accessible in next agg step.
         Collab private attributes should not be accessible here
         """
-        validate_collab_private_attr(self, "test_loader", "aggregator_step")
+        validate_agg_private_attr(self, "aggregator_step", aggr = ["test_loader_agg"], collabs =["train_loader", "test_loader"])
 
         self.include_agg_to_collab = 42
         self.exclude_agg_to_collab = 40
@@ -65,8 +65,8 @@ def collaborator_step_a(self):
         Testing whether Collab private attributes are accessible in collab step
         Aggregator private attributes should not be accessible here
         """
-        validate_agg_private_attrs(
-            self, "train_loader", "test_loader", "collaborator_step_a"
+        validate_collab_private_attrs(
+            self, "collaborator_step_a", aggr = ["test_loader_agg"], collabs =["train_loader", "test_loader"]
         )
 
         self.exclude_collab_to_collab = 2
@@ -80,8 +80,8 @@ def collaborator_step_b(self):
         Aggregator private attributes should not be accessible here
         """
 
-        validate_agg_private_attrs(
-            self, "train_loader", "test_loader", "collaborator_step_b"
+        validate_collab_private_attrs(
+            self, "collaborator_step_b", aggr = ["test_loader_agg"], collabs =["train_loader", "test_loader"]
         )
         self.exclude_collab_to_agg = 10
         self.include_collab_to_agg = 12
@@ -93,7 +93,7 @@ def join(self, inputs):
         Testing whether attributes are excluded from collab to agg
         """
         # Aggregator should only be able to access its own attributes
-        if hasattr(self, "test_loader") is False:
+        if hasattr(self, "test_loader_agg") is False:
             TestFlowPrivateAttributes.ERROR_LIST.append(
                 "aggregator_join_aggregator_attributes_missing"
             )
@@ -143,19 +143,45 @@ def end(self):
         TestFlowPrivateAttributes.ERROR_LIST = []
 
 
-def validate_collab_private_attr(self, private_attr, step_name):
+def validate_agg_private_attr(self, step_name, **private_attrs_kwargs):
+
+    """
+        Validate that aggregator can only access their own attributes 
+
+        Args:
+            step_name: Name of the step being validated 
+            private_attr_kwargs: Keyword arguments specifying the names of private attributes for the aggregator and collaborators.
+    """
+    agg_attrs = private_attrs_kwargs.get('aggr',[])
+    collab_attrs = private_attrs_kwargs.get('collabs', [])
     # Aggregator should only be able to access its own attributes
-    if hasattr(self, private_attr) is False:
+
+    # check for missing aggregator attributes 
+    inaccessible_agg_attrs = [attr for attr in agg_attrs if not hasattr(self, attr)]
+    if inaccessible_agg_attrs:
         TestFlowPrivateAttributes.ERROR_LIST.append(
-            step_name + "_aggregator_attributes_missing"
+            step_name + " aggregator_attributes_missing"
         )
         print(
             f"{bcolors.FAIL} ...Failed in {step_name} - aggregator private attributes not "
             + f"accessible {bcolors.ENDC}"
         )
 
+    # check for collaborator private attributes that should not be accessible 
+    breached_collab_attrs = [attr for attr in collab_attrs if hasattr(self,attr) is True]
+    if breached_collab_attrs:
+        TestFlowPrivateAttributes.ERROR_LIST.append(
+            step_name + "_collaborator_attributes_found"
+        )
+        print(
+            f"{bcolors.FAIL} ... Attribute test failed in {step_name} - collaborator"
+            + f"private attributes accessible:{','.join(breached_collab_attrs)} {bcolors.ENDC}"
+        )
+
+
+
     for idx, collab in enumerate(self.collaborators):
-        # Collaborator private attributes should not be accessible
+        # Collaborator attributes should not be accessible in aggregator step
         if (
             type(self.collaborators[idx]) is not str
             or hasattr(self.runtime, "_collaborators") is True
@@ -171,17 +197,44 @@ def validate_collab_private_attr(self, private_attr, step_name):
             )
 
 
-def validate_agg_private_attrs(self, private_attr_1, private_attr_2, step_name):
+def validate_collab_private_attrs(self, step_name, **private_attrs_kwargs):
+
+    """
+        Validate that collaborators can only access their own attributes
+
+        Args:
+            step_name: Name of the step being validated 
+            private_attr_kwargs: Keyword arguments specifying the names of private attributes for the aggregator and collaborators.
+    """
+    agg_attrs = private_attrs_kwargs.get('aggr',[])
+    collab_attrs = private_attrs_kwargs.get('collabs', [])
+
     # Collaborator should only be able to access its own attributes
-    if not hasattr(self, private_attr_1) or not hasattr(self, private_attr_2):
+
+    # check for missing collaborators attributes 
+    inaccessible_collab_attrs = [attr for attr in collab_attrs if not hasattr(self,attr)]
+
+    if inaccessible_collab_attrs:
         TestFlowPrivateAttributes.ERROR_LIST.append(
-            step_name + "collab_attributes_not_found"
+            step_name + " collab_attributes_not_found"
         )
         print(
             f"{bcolors.FAIL} ... Attribute test failed in {step_name} - Collab "
             + f"private attributes not accessible {bcolors.ENDC}"
         )
+    # check for aggregator private attributes that should not be accessible 
+    breached_agg_attr = [attr for attr in agg_attrs if hasattr(self, attr) is True]
+    if breached_agg_attr:
+        TestFlowPrivateAttributes.ERROR_LIST.append(
+            step_name + "_aggregator_attributes_found"
+        )
+
+        print(
+            f"{bcolors.FAIL} ... Attribute test failed in {step_name} - Aggregator" 
+            + f" private attributes accessible: {','.join(breached_agg_attr)} {bcolors.ENDC}"
+        )
 
+    # Aggregator attributes should not be accessible in collaborator step
     if hasattr(self.runtime, "_aggregator") and isinstance(self.runtime._aggregator, Aggregator):
         # Error - we are able to access aggregator attributes
         TestFlowPrivateAttributes.ERROR_LIST.append(

tests/github/experimental/workspace/testcase_private_attributes_initialization_with_both_options/src/aggregator_private_attrs.py

@@ -4,7 +4,7 @@
 
 
 def aggregator_private_attrs():
-    return {"test_loader_via_callable": np.random.rand(10, 28, 28)}  # Random data
+    return {"test_loader_agg_via_callable": np.random.rand(10, 28, 28)}  # Random data
 
 
-aggregator_private_attributes = {"test_loader": np.random.rand(10, 28, 28)}  # Random data
+aggregator_private_attributes = {"test_loader_agg": np.random.rand(10, 28, 28)}  # Random data