[Workflow API] Resolve issues in setting unserializable objects as private attributes (#1605)

* Delayed aggregator initialization

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

* Added test case

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

* Added test case

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

* Documentation update

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

* Incorporated review comments

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

* Incorporated review comments.

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

* Incorporated review comments

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

* Incorporated review comments.

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

---------

Signed-off-by: Sachin Gupta <sachin.gupta.dsp@gmail.com>

Author: scngupta-dsp

docs/about/features_index/workflowinterface.rst

@@ -236,6 +236,9 @@ Some important points to remember while creating callback function and private a
     - If no Callback Function or private attributes is specified then the Participant shall not have any *private attributes*
     - In above example multiple collaborators have the same callback function or private attributes. Depending on the Federated Learning requirements, user can specify unique callback function or private attributes for each Participant
     - *Private attributes* needs to be set after instantiating the participant.
+    - **Known Limitations**: When using a `callable` to initialize *private attributes* that are **not serializable**, users should be aware of following limitations:
+        * `checkpoint` should not be enabled with `LocalRuntime`. Users should ensure that default (disabled) setting of checkpoint is used or it is explicitly disabled :code:`flow = FederatedFlow( ..., checkpoint = false)`
+        * filtering of attributes (via `include` or `exclude`) cannot be used during the  transition from aggregator step to collaborator steps. This limitation applies to **all attributes** if any non-serializable private attribute is present in aggregator. The flow logic must be updated to avoid filtering in steps that transition control from aggregator to collaborators
 
 Now let's see how the runtime for a flow is assigned, and the flow gets run:
 
@@ -558,6 +561,7 @@ In a distributed environment consisting of Director, Envoys and User Node (where
 
 **IMPORTANT**: While this information is useful for debugging, depending on your workflow it may require significant disk space. For this reason, checkpoint is disabled by default.
 
+
 Future Plans
 ==============
 Following functionalities are planned for inclusion in future releases of the Workflow Interface:

openfl/experimental/workflow/interface/fl_spec.py

@@ -170,13 +170,14 @@ def _setup_initial_state(self) -> None:
         """
         self._metaflow_interface = MetaflowInterface(self.__class__, self.runtime.backend)
         self._run_id = self._metaflow_interface.create_run()
-        # Initialize aggregator private attributes
-        self.runtime.initialize_aggregator()
-        self._foreach_methods = []
         FLSpec._reset_clones()
         FLSpec._create_clones(self, self.runtime.collaborators)
-        # Initialize collaborator private attributes
+
+        # Initialize participant private attributes
+        self.runtime.initialize_aggregator()
         self.runtime.initialize_collaborators()
+        self._foreach_methods = []
+
         if self._checkpoint:
             print(f"Created flow {self.__class__.__name__}")
 

tests/end_to_end/test_suites/wf_local_func_tests.py

@@ -8,21 +8,32 @@
 import random
 from metaflow import Step
 
-from tests.end_to_end.utils.wf_common_fixtures import fx_local_federated_workflow, fx_local_federated_workflow_prvt_attr
+from tests.end_to_end.utils.wf_common_fixtures import (
+    fx_local_federated_workflow,
+    fx_local_federated_workflow_prvt_attr,
+    fx_local_fed_wf_unserializable_pvt_attrs,
+)
+
 from tests.end_to_end.workflow.exclude_flow import TestFlowExclude
 from tests.end_to_end.workflow.include_exclude_flow import TestFlowIncludeExclude
 from tests.end_to_end.workflow.include_flow import TestFlowInclude
 from tests.end_to_end.workflow.internal_loop import TestFlowInternalLoop
 from tests.end_to_end.workflow.reference_flow import TestFlowReference
 from tests.end_to_end.workflow.subset_flow import TestFlowSubsetCollaborators
-from tests.end_to_end.workflow.private_attr_wo_callable import TestFlowPrivateAttributesWoCallable
+from tests.end_to_end.workflow.private_attr_wo_callable import (
+    TestFlowPrivateAttributesWoCallable,
+)
 from tests.end_to_end.workflow.private_attributes_flow import TestFlowPrivateAttributes
 from tests.end_to_end.workflow.private_attr_both import TestFlowPrivateAttributesBoth
+from tests.end_to_end.workflow.unserializable_private_attr import (
+    TestFlowUnserializablePrivateAttributes,
+)
 
 from tests.end_to_end.utils import wf_helper as wf_helper
 
 log = logging.getLogger(__name__)
 
+
 def test_exclude_flow(request, fx_local_federated_workflow):
     """
     Test if variable is excluded, variables not show in next step
@@ -73,7 +84,9 @@ def test_internal_loop(request, fx_local_federated_workflow):
     model = None
     optimizer = None
 
-    flflow = TestFlowInternalLoop(model, optimizer, request.config.num_rounds, checkpoint=True)
+    flflow = TestFlowInternalLoop(
+        model, optimizer, request.config.num_rounds, checkpoint=True
+    )
     flflow.runtime = fx_local_federated_workflow.runtime
     flflow.run()
 
@@ -87,25 +100,37 @@ def test_internal_loop(request, fx_local_federated_workflow):
         "end",
     ]
 
-    steps_present_in_cli, missing_steps_in_cli, extra_steps_in_cli = wf_helper.validate_flow(
-            flflow, expected_flow_steps
-        )
-
-    assert len(steps_present_in_cli) == len(expected_flow_steps), "Number of steps fetched from Datastore through CLI do not match the Expected steps provided"
-    assert len(missing_steps_in_cli) == 0, f"Following steps missing from Datastore: {missing_steps_in_cli}"
-    assert len(extra_steps_in_cli) == 0, f"Following steps are extra in Datastore: {extra_steps_in_cli}"
+    steps_present_in_cli, missing_steps_in_cli, extra_steps_in_cli = (
+        wf_helper.validate_flow(flflow, expected_flow_steps)
+    )
+
+    assert len(steps_present_in_cli) == len(
+        expected_flow_steps
+    ), "Number of steps fetched from Datastore through CLI do not match the Expected steps provided"
+    assert (
+        len(missing_steps_in_cli) == 0
+    ), f"Following steps missing from Datastore: {missing_steps_in_cli}"
+    assert (
+        len(extra_steps_in_cli) == 0
+    ), f"Following steps are extra in Datastore: {extra_steps_in_cli}"
     assert flflow.end_count == 1, "End function called more than one time"
 
-    log.info("\n  Summary of internal flow testing \n"
-             "No issues found and below are the tests that ran successfully\n"
-             "1. Number of training completed is equal to training rounds\n"
-             "2. CLI steps and Expected steps are matching\n"
-             "3. Number of tasks are aligned with number of rounds and number of collaborators\n"
-             "4. End function executed one time")
+    log.info(
+        "\n  Summary of internal flow testing \n"
+        "No issues found and below are the tests that ran successfully\n"
+        "1. Number of training completed is equal to training rounds\n"
+        "2. CLI steps and Expected steps are matching\n"
+        "3. Number of tasks are aligned with number of rounds and number of collaborators\n"
+        "4. End function executed one time"
+    )
     log.info("Successfully ended test_internal_loop")
 
 
-@pytest.mark.parametrize("fx_local_federated_workflow", [("init_collaborator_private_attr_index", "int", None )], indirect=True)
+@pytest.mark.parametrize(
+    "fx_local_federated_workflow",
+    [("init_collaborator_private_attr_index", "int", None)],
+    indirect=True,
+)
 def test_reference_flow(request, fx_local_federated_workflow):
     """
     Test reference variables matched through out the flow
@@ -118,7 +143,12 @@ def test_reference_flow(request, fx_local_federated_workflow):
         flflow.run()
     log.info("Successfully ended test_reference_flow")
 
-@pytest.mark.parametrize("fx_local_federated_workflow", [("init_collaborator_private_attr_name", "str", None )], indirect=True)
+
+@pytest.mark.parametrize(
+    "fx_local_federated_workflow",
+    [("init_collaborator_private_attr_name", "str", None)],
+    indirect=True,
+)
 def test_subset_collaborators(request, fx_local_federated_workflow):
     """
     Test the subset of collaborators in a federated workflow.
@@ -158,16 +188,16 @@ def test_subset_collaborators(request, fx_local_federated_workflow):
         )
 
         assert len(list(step)) == len(subset_collaborators), (
-                f"...Flow only ran for {len(list(step))} "
-                + f"instead of the {len(subset_collaborators)} expected "
-                + f"collaborators- Testcase Failed."
-            )
+            f"...Flow only ran for {len(list(step))} "
+            + f"instead of the {len(subset_collaborators)} expected "
+            + f"collaborators- Testcase Failed."
+        )
         log.info(
             f"Found {len(list(step))} tasks for each of the "
             + f"{len(subset_collaborators)} collaborators"
         )
-        log.info(f'subset_collaborators = {subset_collaborators}')
-        log.info(f'collaborators_ran = {collaborators_ran}')
+        log.info(f"subset_collaborators = {subset_collaborators}")
+        log.info(f"collaborators_ran = {collaborators_ran}")
         for collaborator_name in subset_collaborators:
             assert collaborator_name in collaborators_ran, (
                 f"...Flow did not execute for "
@@ -177,7 +207,8 @@ def test_subset_collaborators(request, fx_local_federated_workflow):
 
     log.info(
         f"Testing FederatedFlow - Ending test for validating "
-        + f"the subset of collaborators.")
+        + f"the subset of collaborators."
+    )
     log.info("Successfully ended test_subset_collaborators")
 
 
@@ -194,7 +225,11 @@ def test_private_attr_wo_callable(request, fx_local_federated_workflow_prvt_attr
     log.info("Successfully ended test_private_attr_wo_callable")
 
 
-@pytest.mark.parametrize("fx_local_federated_workflow", [("init_collaborate_pvt_attr_np", "int", "init_agg_pvt_attr_np" )], indirect=True)
+@pytest.mark.parametrize(
+    "fx_local_federated_workflow",
+    [("init_collaborate_pvt_attr_np", "int", "init_agg_pvt_attr_np")],
+    indirect=True,
+)
 def test_private_attributes(request, fx_local_federated_workflow):
     """
     Set private attribute through callable function
@@ -208,7 +243,11 @@ def test_private_attributes(request, fx_local_federated_workflow):
     log.info("Successfully ended test_private_attributes")
 
 
-@pytest.mark.parametrize("fx_local_federated_workflow_prvt_attr", [("init_collaborate_pvt_attr_np", "int", "init_agg_pvt_attr_np" )], indirect=True)
+@pytest.mark.parametrize(
+    "fx_local_federated_workflow_prvt_attr",
+    [("init_collaborate_pvt_attr_np", "int", "init_agg_pvt_attr_np")],
+    indirect=True,
+)
 def test_private_attr_both(request, fx_local_federated_workflow_prvt_attr):
     """
     Set private attribute through callable function and direct assignment
@@ -220,3 +259,25 @@ def test_private_attr_both(request, fx_local_federated_workflow_prvt_attr):
         log.info(f"Starting round {i}...")
         flflow.run()
     log.info("Successfully ended test_private_attr_both")
+
+
+@pytest.mark.parametrize(
+    "fx_local_fed_wf_unserializable_pvt_attrs",
+    [
+        ("callable_to_init_collab_unserializable_pvt_attrs",
+        "int",
+        "callable_to_init_agg_unserializable_pvt_attrs")
+    ],
+    indirect=True,
+)
+def test_unserializable_private_attr(
+    request, fx_local_fed_wf_unserializable_pvt_attrs
+):
+    """
+    Validate unserializable objects are accessible as private attributes
+    """
+    log.info("Starting Test for unserializable private attributes")
+    flflow = TestFlowUnserializablePrivateAttributes(rounds=request.config.num_rounds, checkpoint=False)
+    flflow.runtime = fx_local_fed_wf_unserializable_pvt_attrs.runtime
+    flflow.run()
+    log.info("Successfully ended Test for unserializable private attributes")

tests/end_to_end/utils/wf_common_fixtures.py

@@ -148,3 +148,62 @@ def fx_local_federated_workflow_prvt_attr(request):
         collaborators=collaborators_list,
         runtime=local_runtime,
     )
+
+
+@pytest.fixture(scope="function")
+def fx_local_fed_wf_unserializable_pvt_attrs(request):
+    """
+    Fixture to set up a local federated workflow for testing.
+    This fixture initializes an `Aggregator` and sets up a list of collaborators
+    based on the number specified in the test configuration. It also configures
+    a `LocalRuntime` with the aggregator, collaborators, and an optional backend
+    if specified in the test configuration.
+    Args:
+        request (FixtureRequest): The pytest request object that provides access
+                                to the test configuration.
+    Yields:
+        LocalRuntime: An instance of `LocalRuntime` configured with the aggregator,
+                    collaborators, and backend.
+    """
+    # Inline import
+    from tests.end_to_end.utils.wf_helper import (
+        callable_to_init_agg_unserializable_pvt_attrs,
+        callable_to_init_collab_unserializable_pvt_attrs
+    )
+    collab_callback_func = request.param[0] if hasattr(request, 'param') and request.param else None
+    collab_value = request.param[1] if hasattr(request, 'param') and request.param else None
+    agg_callback_func = request.param[2] if hasattr(request, 'param') and request.param else None
+
+    # Get the callback functions from the locals using string
+    collab_callback_func_name = locals()[collab_callback_func] if collab_callback_func else None
+    agg_callback_func_name = locals()[agg_callback_func] if agg_callback_func else None
+    collaborators_list = []
+
+    # Setup aggregator
+    if agg_callback_func_name:
+        aggregator = Aggregator(name="agg",
+                                private_attributes_callable=agg_callback_func_name)
+    else:
+        aggregator = Aggregator()
+
+    # Setup collaborators
+    for i in range(request.config.num_collaborators):
+        func_var = i if collab_value == "int" else f"collaborator{i}" if collab_value == "str" else None
+        collab = Collaborator(
+                name=f"collaborator{i}",
+                private_attributes_callable=collab_callback_func_name
+            )
+        collaborators_list.append(collab)
+
+    workflow_backend = request.config.workflow_backend if hasattr(request.config, 'workflow_backend') else None
+    if workflow_backend:
+        local_runtime = LocalRuntime(aggregator=aggregator, collaborators=collaborators_list, backend=workflow_backend)
+    else:
+        local_runtime = LocalRuntime(aggregator=aggregator, collaborators=collaborators_list)
+
+    # Return the federation fixture
+    return workflow_local_fixture(
+        aggregator=aggregator,
+        collaborators=collaborators_list,
+        runtime=local_runtime,
+    )

tests/end_to_end/utils/wf_helper.py

@@ -4,6 +4,8 @@
 from metaflow import Flow
 import logging
 import numpy as np
+from openfl.databases import TensorDB
+from openfl.utilities import TensorKey
 
 import tests.end_to_end.utils.exceptions as ex
 
@@ -116,6 +118,20 @@ def init_agg_pvt_attr_np():
     return {"test_loader": np.random.rand(10, 28, 28)}
 
 
+def callable_to_init_collab_unserializable_pvt_attrs():
+    """
+    Create and return a TensorDB
+    """
+    return {"col_tensor_db": TensorDB()}
+
+
+def callable_to_init_agg_unserializable_pvt_attrs():
+    """
+    Create and return a TensorDB
+    """
+    return {"agg_tensor_db": TensorDB()}
+
+
 def run_notebook(notebook_path, output_notebook_path):
     """
     Function to run the notebook.