Support for the exclusion of prohibited data types.

Signed-off-by: yuliasherman <[email protected]>

Author: yuliasherman

openfl/experimental/workflow/component/aggregator/aggregator.py

@@ -198,6 +198,9 @@ async def run_flow(self) -> FLSpec:
         Returns:
             flow (FLSpec): Updated instance.
         """
+        # As an aggregator is created before a runtime, set prohibited data types for the runtime
+        # before running the flow.
+        self.flow.runtime.prohibited_data_types = FederatedRuntime.prohibited_data_types
         # Start function will be the first step if any flow
         f_name = "start"
         # Creating a clones from the flow object
@@ -394,7 +397,7 @@ def do_task(self, f_name: str) -> Any:
                     # Create list of selected collaborator clones
                     selected_clones = ([],)
                     for name, clone in self.clones_dict.items():
-                        # Check if collaboraotr is in the list of selected
+                        # Check if collaborator is in the list of selected
                         # collaborators
                         if name in self.selected_collaborators:
                             selected_clones[0].append(clone)

openfl/experimental/workflow/component/collaborator/collaborator.py

@@ -23,7 +23,7 @@ class Collaborator:
         Aggregator Server.
 
         private_attrs_callable (Callable): Function for Collaborator
-        private attriubtes.
+        private attributes.
         private_attrs_kwargs (Dict): Arguments to call private_attrs_callable.
 
     Note:

openfl/experimental/workflow/interface/fl_spec.py

@@ -22,6 +22,7 @@
     filter_attributes,
     generate_artifacts,
     should_transfer,
+    validate_data_types,
 )
 
 
@@ -164,6 +165,7 @@ def _run_local(self) -> None:
         for name, attr in final_attributes:
             setattr(self, name, attr)
 
+    # Runs only for LocalRuntime
     def _setup_initial_state(self) -> None:
         """
         Sets up the flow's initial state, initializing private attributes for
@@ -176,6 +178,7 @@ def _setup_initial_state(self) -> None:
         self._foreach_methods = []
         FLSpec._reset_clones()
         FLSpec._create_clones(self, self.runtime.collaborators)
+
         # Initialize collaborator private attributes
         self.runtime.initialize_collaborators()
         if self._checkpoint:
@@ -343,8 +346,11 @@ def next(self, f, **kwargs) -> None:
         if aggregator_to_collaborator(f, parent_func):
             agg_to_collab_ss = self._capture_instance_snapshot(kwargs=kwargs)
 
-        # Remove included / excluded attributes from next task
-        filter_attributes(self, f, **kwargs)
+        # Remove unwanted attributes from next task
+        if kwargs:
+            filter_attributes(self, f, **kwargs)
+            if self._runtime.prohibited_data_types:
+                validate_data_types(self._runtime.prohibited_data_types, **kwargs)
 
         if str(self._runtime) == "FederatedRuntime":
             if f.collaborator_step and not f.aggregator_step:

openfl/experimental/workflow/runtime/local_runtime.py

@@ -324,7 +324,7 @@ def __init__(
                   The RayGroups run concurrently while participants in the
                   group run serially.
                   The default is 1 RayGroup and can be changed by using the
-                  num_actors=1 kwarg. By using more RayGroups more concurency
+                  num_actors=1 kwarg. By using more RayGroups more concurrency
                   is allowed with the trade off being that each RayGroup has
                   extra memory overhead in the form of extra CUDA CONTEXTS.
 
@@ -737,7 +737,9 @@ def execute_collab_task(self, flspec_obj, f, parent_func, instance_snapshot, **k
             # Set new LocalRuntime for clone as it is required
             # new runtime object will not contain private attributes of
             # aggregator or other collaborators
-            clone.runtime = LocalRuntime(backend="single_process")
+            clone.runtime = LocalRuntime(
+                backend="single_process", prohibited_data_types=self.prohibited_data_types
+            )
 
             # write the clone to the object store
             # ensure clone is getting latest _metaflow_interface

openfl/experimental/workflow/runtime/runtime.py

@@ -10,7 +10,19 @@
 from openfl.experimental.workflow.interface.participants import Aggregator, Collaborator
 
 
-class Runtime:
+class AttributeValidationMeta(type):
+    def __setattr__(cls, name, value):
+        if name == "prohibited_data_types":
+            if not isinstance(value, list):
+                raise TypeError(f"'{name}' must be a list, got {type(value).__name__}")
+            if not all(isinstance(item, str) for item in value):
+                raise TypeError(f"All elements of '{name}' must be strings")
+        super().__setattr__(name, value)
+
+
+class Runtime(metaclass=AttributeValidationMeta):
+    prohibited_data_types = []
+
     def __init__(self):
         """Initializes the Runtime object.
 

openfl/experimental/workflow/utilities/__init__.py

@@ -17,6 +17,7 @@
     filter_attributes,
     generate_artifacts,
     parse_attrs,
+    validate_data_types,
 )
 from openfl.experimental.workflow.utilities.stream_redirect import (
     RedirectStdStream,

openfl/experimental/workflow/utilities/runtime_utils.py

@@ -7,6 +7,7 @@
 import inspect
 import itertools
 from types import MethodType
+from typing import List
 
 import numpy as np
 
@@ -96,6 +97,31 @@ def filter_attributes(ctx, f, **kwargs):
         _process_exclusion(ctx, cls_attrs, kwargs["exclude"], f)
 
 
+def validate_data_types(
+    prohibited_data_types: List[str], reserved_words=["collaborators"], **kwargs
+):
+    """Validates that the types of attributes in kwargs are not among the prohibited data types.
+    Raises a TypeError if any prohibited data type is found.
+
+    Args:
+        prohibited_data_types (List[str]): A list of prohibited data type names
+            (e.g., ['int', 'float']).
+        kwargs (dict): Arbitrary keyword arguments representing attribute names and their values.
+
+    Raises:
+        TypeError: If any prohibited data types are found in kwargs.
+        ValueError: If prohibited_data_types is empty.
+    """
+    if not prohibited_data_types:
+        raise ValueError("prohibited_data_types must not be empty.")
+    for attr_name, attr_value in kwargs.items():
+        if type(attr_value).__name__ in prohibited_data_types and attr_value not in reserved_words:
+            raise TypeError(
+                f"The attribute '{attr_name}' = '{attr_value}' has a prohibited value type: "
+                f"{type(attr_value).__name__}"
+            )
+
+
 def _validate_include_exclude(kwargs, cls_attrs):
     """Validates that 'include' and 'exclude' are not both present, and that
     attributes in 'include' or 'exclude' exist in the context.
@@ -152,13 +178,13 @@ def _process_exclusion(ctx, cls_attrs, exclude_list, f):
             delattr(ctx, attr)
 
 
-def checkpoint(ctx, parent_func, chkpnt_reserved_words=["next", "runtime"]):
+def checkpoint(ctx, parent_func, checkpoint_reserved_words=["next", "runtime"]):
     """Optionally saves the current state for the task just executed.
 
     Args:
         ctx (any): The context to checkpoint.
         parent_func (function): The function that was just executed.
-        chkpnt_reserved_words (list, optional): A list of reserved words to
+        checkpoint_reserved_words (list, optional): A list of reserved words to
             exclude from checkpointing. Defaults to ["next", "runtime"].
 
     Returns:
@@ -173,7 +199,7 @@ def checkpoint(ctx, parent_func, chkpnt_reserved_words=["next", "runtime"]):
     if ctx._checkpoint:
         # all objects will be serialized using Metaflow interface
         print(f"Saving data artifacts for {parent_func.__name__}")
-        artifacts_iter, _ = generate_artifacts(ctx=ctx, reserved_words=chkpnt_reserved_words)
+        artifacts_iter, _ = generate_artifacts(ctx=ctx, reserved_words=checkpoint_reserved_words)
         task_id = ctx._metaflow_interface.create_task(parent_func.__name__)
         ctx._metaflow_interface.save_artifacts(
             artifacts_iter(),
@@ -188,15 +214,15 @@ def checkpoint(ctx, parent_func, chkpnt_reserved_words=["next", "runtime"]):
 
 def old_check_resource_allocation(num_gpus, each_participant_gpu_usage):
     remaining_gpu_memory = {}
-    # TODO for each GPU the funtion tries see if all participant usages fit
+    # TODO for each GPU the function tries see if all participant usages fit
     # into a GPU, it it doesn't it removes that participant from the
     # participant list, and adds it to the remaining_gpu_memory dict. So any
     # sum of GPU requirements above 1 triggers this.
-    # But at this point the funtion will raise an error because
+    # But at this point the function will raise an error because
     # remaining_gpu_memory is never cleared.
     # The participant list should remove the participant if it fits in the gpu
-    # and save the partipant if it doesn't and continue to the next GPU to see
-    # if it fits in that one, only if we run out of GPUs should this funtion
+    # and save the participant if it doesn't and continue to the next GPU to see
+    # if it fits in that one, only if we run out of GPUs should this function
     # raise an error.
     for gpu in np.ones(num_gpus, dtype=int):
         for i, (participant_name, participant_gpu_usage) in enumerate(
@@ -230,7 +256,7 @@ def check_resource_allocation(num_gpus, each_participant_gpu_usage):
             if gpu == 0:
                 break
             if gpu < participant_gpu_usage:
-                # participant doesn't fitm break to next GPU
+                # participant doesn't fit, break to next GPU
                 break
             else:
                 # if participant fits remove from need_assigned