Make LocalRuntime working with support for the exclusion of prohibited data types.

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

Author: yuliasherman

openfl/experimental/workflow/interface/fl_spec.py

@@ -22,6 +22,7 @@
     filter_attributes,
     generate_artifacts,
     should_transfer,
+    validate_data_types,
 )
 
 
@@ -343,8 +344,14 @@ 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 prohibited attributes from the next task
+        if kwargs:
+            filter_attributes(self, f, **kwargs)
+            if self._runtime._prohibited_data_types:
+                # try:
+                validate_data_types(self._runtime._prohibited_data_types, **kwargs)
+                # except Exception as exception:
+                #     print(exception)
 
         if str(self._runtime) == "FederatedRuntime":
             if f.collaborator_step and not f.aggregator_step:

openfl/experimental/workflow/runtime/federated_runtime.py

@@ -42,6 +42,7 @@ def __init__(
         collaborators: Optional[List[str]] = None,
         director: Optional[Dict[str, Any]] = None,
         notebook_path: Optional[str] = None,
+        prohibited_data_types: Optional[List[str]] = None,
         tls: bool = False,
     ) -> None:
         """Initializes the FederatedRuntime object.
@@ -51,9 +52,11 @@ def __init__(
                 Defaults to None.
             director (Optional[Dict[str, Any]]): Director information. Defaults to None
             notebook_path (Optional[str]): Jupyter notebook path
+            prohibited_data_types (List[str]): A list of data types that are not allowed to be sent
+                through the network.
             tls (bool): Whether to use TLS for the connection.
         """
-        super().__init__()
+        super().__init__(prohibited_data_types=prohibited_data_types)
         self.__collaborators = collaborators
 
         self.tls = tls

openfl/experimental/workflow/runtime/local_runtime.py

@@ -302,6 +302,7 @@ def __init__(
         aggregator: Dict = None,
         collaborators: Dict = None,
         backend: str = "single_process",
+        prohibited_data_types: Optional[List[str]] = None,
         **kwargs,
     ) -> None:
         """Initializes the LocalRuntime object to run the flow on a single
@@ -314,6 +315,8 @@ def __init__(
             collaborators (List[Type[Collaborator]], optional): A list of
                 collaborators; each with their own private attributes.
             backend (str, optional): The backend that will execute the tasks.
+            prohibited_data_types (List[str]): A list of data types that are not allowed to be sent
+                through the network.
             Defaults to "single_process".
                 Available options are:
                 - 'single_process': (default) Executes every task within the
@@ -324,7 +327,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.
 
@@ -346,7 +349,7 @@ def some_collaborator_task(self):
             access. If the system has one GPU, collaborator tasks will run
             sequentially.
         """
-        super().__init__()
+        super().__init__(prohibited_data_types=prohibited_data_types)
         if backend not in ["ray", "single_process"]:
             raise ValueError(
                 f"Invalid 'backend' value '{backend}', accepted values are "
@@ -737,7 +740,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=super().prohibited_data_types
+            )
 
             # write the clone to the object store
             # ensure clone is getting latest _metaflow_interface

openfl/experimental/workflow/runtime/runtime.py

@@ -4,19 +4,37 @@
 
 """openfl.experimental.workflow.runtime module Runtime class."""
 
-from typing import Callable, List
+from typing import Callable, List, Optional
 
 from openfl.experimental.workflow.interface.fl_spec import FLSpec
 from openfl.experimental.workflow.interface.participants import Aggregator, Collaborator
 
 
 class Runtime:
-    def __init__(self):
+    def __init__(self, prohibited_data_types: Optional[List[str]] = None):
         """Initializes the Runtime object.
 
         This serves as a base interface for runtimes that can run FLSpec flows.
+
+        Args:
+            prohibited_data_types (Optional[List[str]]): A list of data types that are not allowed to be sent
+                through the network. Defaults to an empty list if not provided.
         """
-        pass
+        self.prohibited_data_types = prohibited_data_types or []
+
+    @property
+    def prohibited_data_types(self) -> List[str]:
+        """Return the prohibited data types for the runtime."""
+        return self._prohibited_data_types
+
+    @prohibited_data_types.setter
+    def prohibited_data_types(self, value: List[str]):
+        """Set the prohibited data types for the runtime."""
+        if not isinstance(value, list):
+            raise TypeError("prohibited_data_types must be a list of strings.")
+        if not all(isinstance(item, str) for item in value):
+            raise ValueError("All items in prohibited_data_types must be strings.")
+        self._prohibited_data_types = value
 
     @property
     def aggregator(self):

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