[Workflow API] Optimize flow state being transferred between participants in LocalRuntime & FederatedRuntime (#1589)

* add execute_task_args to reserved_keywords

Signed-off-by: Ishant Thakare <[email protected]>

* optimized flow state

Signed-off-by: Ishant Thakare <[email protected]>

---------

Signed-off-by: Ishant Thakare <[email protected]>
Co-authored-by: Payal Chaurasiya <[email protected]>
Co-authored-by: teoparvanov <[email protected]>

Author: 3 people

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

@@ -8,6 +8,7 @@
 import inspect
 import queue
 import time
+from copy import deepcopy
 from logging import getLogger
 from threading import Event
 from typing import Any, Callable, Dict, List, Tuple
@@ -16,7 +17,11 @@
 
 from openfl.experimental.workflow.interface import FLSpec
 from openfl.experimental.workflow.runtime import FederatedRuntime
-from openfl.experimental.workflow.utilities import aggregator_to_collaborator, checkpoint
+from openfl.experimental.workflow.utilities import (
+    aggregator_to_collaborator,
+    checkpoint,
+    generate_artifacts,
+)
 from openfl.experimental.workflow.utilities.metaflow_utils import MetaflowInterface
 
 logger = getLogger(__name__)
@@ -43,7 +48,8 @@ class Aggregator:
         collaborator_task_results (Event): Event to inform aggregator that
             collaborators have sent the results.
         __collaborator_tasks_queue (Dict[Queue]): queue for each collaborator.
-        flow (Any): Flow class.
+        flow (FLSpec): Flow class.
+        final_flow_state (FLSpec): Final flow state.
         name (str): aggregator in string format.
         checkpoint (bool): Whether to save checkpoint or not (default=False).
         private_attrs_callable (Callable): Function for Aggregator private
@@ -124,6 +130,7 @@ def __init__(
         self.__collaborator_tasks_queue = {collab: queue.Queue() for collab in self.authorized_cols}
 
         self.flow = flow
+        self.final_flow_state = deepcopy(flow)
         self.checkpoint = checkpoint
         self.flow._foreach_methods = []
         logger.info("MetaflowInterface creation.")
@@ -181,6 +188,12 @@ def _log_big_warning(self) -> None:
             f" WARNED!!!"
         )
 
+    def _update_final_flow(self) -> None:
+        """Update the final flow state with current flow artifacts."""
+        artifacts_iter, _ = generate_artifacts(ctx=self.flow)
+        for name, attr in artifacts_iter():
+            setattr(self.final_flow_state, name, deepcopy(attr))
+
     @staticmethod
     def _get_sleep_time() -> int:
         """Sleep 10 seconds.
@@ -250,7 +263,8 @@ async def run_flow(self) -> FLSpec:
                 self.flow.restore_instance_snapshot(self.flow, list(self.instance_snapshot))
                 delattr(self, "instance_snapshot")
 
-        return self.flow
+        self._update_final_flow()
+        return self.final_flow_state
 
     def call_checkpoint(
         self, name: str, ctx: Any, f: Callable, stream_buffer: bytes = None

openfl/experimental/workflow/utilities/runtime_utils.py

@@ -45,14 +45,16 @@ def parse_attrs(ctx, exclude=[], reserved_words=["next", "runtime", "input"]):
     return cls_attrs, valid_artifacts
 
 
-def generate_artifacts(ctx, reserved_words=["next", "runtime", "input", "checkpoint"]):
+def generate_artifacts(
+    ctx, reserved_words=["next", "runtime", "input", "checkpoint", "execute_task_args"]
+):
     """Generates artifacts from the given context, excluding specified reserved
     words.
 
     Args:
         ctx (any): The context to generate artifacts from.
         reserved_words (list, optional): A list of reserved words to exclude.
-            Defaults to ["next", "runtime", "input", "checkpoint"].
+            Defaults to ["next", "runtime", "input", "checkpoint", "execute_task_args"].
 
     Returns:
         tuple: A tuple containing a generator of artifacts and a list of
@@ -152,14 +154,14 @@ 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, chkpnt_reserved_words=["next", "runtime", "execute_task_args"]):
     """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
-            exclude from checkpointing. Defaults to ["next", "runtime"].
+            exclude from checkpointing. Defaults to ["next", "runtime", "execute_task_args"].
 
     Returns:
         step_stdout (io.StringIO): parent_func stdout