Proposal for branching in Vertical Federated learning

[nisha987]

Summary

Branching in Metaflow can be explained with parallel steps. it supports any number of parallel branches. This proposal is to enable branching support for Vertical Federated Learning (VFL) in OpenFL, providing a flexible workflow management system that allows for complex, multi-path federated learning scenarios.

Motivation

1. Allow flexible implementation of Vertical Federated Learning scenarios
2. Support different processing paths based on collaborator characteristics or data requirements

Proposal

Vertical FL

Vertical Federated Learning (VFL) is a federated machine learning method that allows multiple clients, each possessing complementary and non-overlapping features about the same entity, to collaboratively train machine learning models without disclosing their raw data. This approach is particularly valuable in situations where organizations have distinct yet highly informative data subsets. By combining these subsets, VFL can significantly improve the predictive performance of machine learning models.

Core Features

We propose branching support for the following use-case:

• Support branching within collaborator steps

2. Branching in collaborator steps
   This use-case addresses the branching within collaborator steps. collaborator can branch into multiple steps and join at the end. Example mentioned below is for understanding how branching works and not a usage example.

   @aggregator
   def start(self):
       self.next(self.collaborator_branch1, self.collaborator_branch2)
   
   @collaborator
   def collaborator_branch1(self):
       self.next(self.join)
   
   @collaborator
   def collaborator_branch2(self):
       self.next(self.join)
   
   @aggregator
   def join(self, inputs):
       self.next(self.end)
   

Implementation Requirements

• All branches must eventually converge at join points

Scheduling the collaborator

• 1:1 mapping of number of branches with number of list of collaborator.

Parallel Execution:

• Each of the branch is executed parallely.
• For localruntime with "ray" we execute each collaborator with different/assigned CPU.
• For federated runtime, each collaborator uses different envoy.

Implementation:

Components which will require changes:

FLSpec:

The FLSpec class needs to be updated to support branch selection and routing in the next() method. Currently, the next() method supports transitioning to the next task with foreach parameter for collaborator selection. It needs to be extended to handle branching.

1. Enhanced next() method: Modify to handle a new keyword parameter branches which will specify different execution paths based on branch names.
2. Branch Mapping : Add a mechanism to map branch names to specific collaborators, maintaining the 1:1 relationship between branches and collaborators.
3. Branch State Management : Implement proper state tracking to ensure that different branches can join later at a common aggregator step.

def next(self, f, **kwargs) -> None:
    """Specifies the next task in the flow to execute.

    Args:
        f (Callable): The next task that will be executed in the flow.
        **kwargs: Additional keyword arguments, which can include:
            - foreach: Specifies to run the task on each item in the given collection
            - branches: Specifies different execution branches for different collaborators
            - include/exclude: Filter attributes to include or exclude in transitions
    """
    # Get the name and reference to the calling function
    parent = inspect.stack()[1][3]
    parent_func = getattr(self, parent)

    # Take back-up of current state of self for transitions
    agg_to_collab_ss = None
    if aggregator_to_collaborator(f, parent_func):
        agg_to_collab_ss = self._capture_instance_snapshot(kwargs=kwargs)

    # Handle branching if specified
    if "branches" in kwargs:
        # Each branch maps to exactly one collaborator
        branches = kwargs.get("branches", {})
        self._branch_mapping = branches
        # Store branch information for runtime execution
        # Rest of branching logic...
    
    # Existing foreach and other logic
    # ...

Runtime Changes

Both LocalRuntime and FederatedRuntime need updates to handle branch execution and joining. The primary focus is on ensuring that branching works in both simulation and distributed environments.

Aggregator Changes

The aggregator component is central to the branching feature and requires significant changes.

1. Branch Task Assignment: Modify how tasks are assigned to collaborators based on branch specifications.
2. Branch Task Execution: Enable execution of different branches with different collaborators.
3. Result Collection: Enhance result collection logic to handle results from different branches.
4. Join Operation: Implement a robust join operation that can handle results from multiple branches.

def do_task(self, f_name: str) -> Any:
    """Execute aggregator steps until transition.

    Args:
        f_name (str): Aggregator step

    Returns:
        string / None: Next collaborator function or None end of the flow.
    """
    # Set aggregator private attributes to flow object
    self.__set_private_attrs_to_clone(self.flow)

    not_at_transition_point = True
    # Run a loop to execute flow steps until not_at_transition_point is False
    while not_at_transition_point:
        f = getattr(self.flow, f_name)
        # Get the list of parameters of function f
        args = inspect.signature(f)._parameters

        # Handle branch-specific logic
        if hasattr(self.flow, "_branch_mapping"):
            # Process branch mappings
            # Each collaborator is assigned exactly one branch
            # ...

        # Existing logic continues
        # ...

The sequence diagram below illustrates the enhanced internal flow of the aggregator component when handling branching in Vertical Federated Learning. The key steps include:

1. Initialization: Client sets up the workflow and runtime
2. Start Task: The aggregator executes the start task and initializes parameters
3. Branch Dispatch: The aggregator uses next() with the branches parameter to dispatch different tasks to different collaborators
4. Parallel Branch Execution: Each branch executes on its assigned collaborator
5. Results Collection: Results from all branches are collected
6. Join Operation: The aggregator combines results from different branches
7. Next Round or Completion: Either start the next round or complete the workflow

sequenceDiagram
    title OpenFL Vertical Federated Learning with Branching - Aggregator Flow
    
    participant Client as Client Application
    participant FLSpec as FLSpec
    participant Runtime as Runtime
    participant Agg as Aggregator
    participant Branch1 as "Branch1 Collaborator"
    participant Branch2 as "Branch2 Collaborator"
    
    Note over Client,Branch2: Initialization Phase
    Client->>FLSpec: Create FLSpec workflow
    Client->>Runtime: Create runtime(aggregator, collaborators)
    Client->>FLSpec: Set runtime
    Client->>FLSpec: Run()
    
    Note over FLSpec,Branch2: Start Task Execution
    FLSpec->>Runtime: execute_task(self, start_function)
    Runtime->>Agg: execute_agg_task()
    
    Note over Agg: Execute start() function
    Agg->>Agg: Initialize model, data, parameters
    
    Note over Agg: Branch Dispatch
    Agg->>FLSpec: next() with branches={"branch1": ["collab1"], "branch2": ["collab2"]}
    FLSpec->>FLSpec: Store branch mapping
    
    Note over FLSpec,Branch2: Branch Execution Phase
    FLSpec->>Runtime: execute_task() with branch information
    
    Note over Runtime: Execute Branch 1
    Runtime->>Runtime: Map branch1 to collab1
    Runtime->>Branch1: execute_branch_task(branch1_task)
    Branch1->>Branch1: Process data in branch1
    Branch1->>Branch1: Generate partial results
    
    Note over Runtime: Execute Branch 2 (in parallel)
    Runtime->>Runtime: Map branch2 to collab2
    Runtime->>Branch2: execute_branch_task(branch2_task)
    Branch2->>Branch2: Process data in branch2
    Branch2->>Branch2: Generate partial results
    
    Note over Branch1,Agg: Results Collection
    Branch1-->>Runtime: Return branch1 results
    Branch2-->>Runtime: Return branch2 results
    
    Note over Runtime,Agg: Join Operation
    Runtime->>Agg: execute_agg_task(join) with results from branches
    
    Note over Agg: Join Logic
    Agg->>Agg: Merge results from different branches
    Agg->>Agg: Integrate different feature components
    Agg->>Agg: Update model with combined results
    
    Note over Agg: Next Round Decision
    alt More rounds needed
        Agg->>FLSpec: next() to branches or next task
    else Final round completed
        Agg->>FLSpec: next(self.end)
        FLSpec->>Runtime: execute_task(end_function)
        Runtime->>Agg: execute_agg_task(end)
        Agg->>Agg: Finalize model and metrics
    end
    
    Note over FLSpec,Agg: Completion
    Runtime->>FLSpec: Return final results
    FLSpec->>Client: Return completed workflow

Loading

Sequence diagram for federated runtime is added below. Key differences from the LocalRuntime approach include:

1. Director Orchestration: The Director orchestrates the distributed components, including the Aggregator and Envoys, each running as separate processes.
2. Envoy Communication: Each branch task is dispatched to the appropriate Envoy which communicates with its Collaborator process.
3. Distributed Task Execution: Branch tasks execute on separate physical or virtual machines via the Envoy-Collaborator connection.
4. Asynchronous Result Collection: Results from branches are collected asynchronously by the Director as they become available.
5. State Management: The distributed nature requires careful state management between components to ensure proper synchronization.

The 1:1 mapping between branches and collaborators is maintained, with each branch executing on its dedicated collaborator via its associated envoy.

sequenceDiagram
    title OpenFL Vertical Federated Learning with Branching - Federated Runtime Flow
    
    participant Client as Client Application
    participant FLSpec as FLSpec
    participant Director as Director
    participant Agg as Aggregator
    participant Envoy1 as "Envoy (Collab1)"
    participant Collab1 as "Collaborator 1"
    participant Envoy2 as "Envoy (Collab2)"
    participant Collab2 as "Collaborator 2"
    
    Note over Client,Collab2: Initialization Phase
    Client->>FLSpec: Create FLSpec workflow
    Client->>FLSpec: Create FederatedRuntime(director, envoys)
    Client->>FLSpec: Set runtime
    Client->>FLSpec: Run()
    
    Note over FLSpec,Director: Runtime Initialization
    FLSpec->>Director: start_experiment_execution_loop()
    Director->>Agg: Initialize aggregator coroutine
    Director->>Envoy1: Connect to envoy
    Director->>Envoy2: Connect to envoy
    Envoy1->>Collab1: Initialize collaborator
    Envoy2->>Collab2: Initialize collaborator
    
    Note over Director,Agg: Start Task Execution
    Director->>Agg: Execute start function
    
    Note over Agg: Execute start() function
    Agg->>Agg: Initialize model, data, parameters
    
    Note over Agg: Branch Dispatch
    Agg->>FLSpec: next() with branches={"branch1": ["collab1"], "branch2": ["collab2"]}
    FLSpec->>FLSpec: Store branch mapping
    
    Note over Director,Collab2: Branch Execution Phase (Parallel)
    
    par Branch 1 and Branch 2 Execute in Parallel
        Note over Director,Envoy1: Branch 1 Execution
        Director->>Envoy1: Dispatch branch1_task to envoy
        Envoy1->>Collab1: Execute branch1_task
        Collab1->>Collab1: Process data in branch1
        Collab1->>Collab1: Generate partial results
        Collab1-->>Envoy1: Return branch1 results
    and
        Note over Director,Envoy2: Branch 2 Execution
        Director->>Envoy2: Dispatch branch2_task to envoy
        Envoy2->>Collab2: Execute branch2_task
        Collab2->>Collab2: Process data in branch2
        Collab2->>Collab2: Generate partial results
        Collab2-->>Envoy2: Return branch2 results
    end
    
    Note over Envoy1,Director: Results Collection (May happen asynchronously)
    par Asynchronous Results Collection
        Envoy1-->>Director: Forward branch1 results
    and
        Envoy2-->>Director: Forward branch2 results
    end
    
    Note over Director,Agg: Join Operation
    Director->>Agg: Provide results from all branches
    
    Note over Agg: Join Logic
    Agg->>Agg: Merge results from different branches
    Agg->>Agg: Integrate different feature components
    Agg->>Agg: Update model with combined results
    
    Note over Agg,Director: Next Round Decision
    alt More rounds needed
        Agg->>FLSpec: next() to continue training
        FLSpec->>Director: Continue with next round
    else Final round completed
        Agg->>FLSpec: next(self.end)
        FLSpec->>Director: Execute end function
        Director->>Agg: Execute end function
        Agg->>Agg: Finalize model and metrics
    end
    
    Note over Director,Client: Completion
    Director->>FLSpec: Return final experiment state
    FLSpec->>Client: Return completed workflow results

Loading

Implementation Approach

1. FLSpec and Utility Function Updates

1.1. Add branching support to FLSpec.next() method
1.2. Create helper utilities for branch-to-collaborator mapping
1.3. Add branch-specific decorators (if needed)

2. Runtime Support for Branching

2.1. Update LocalRuntime to handle branching execution
2.2. Test with simulation-based workflows
2.3. Update FederatedRuntime to support distributed branching

3. Aggregator Integration

3.1. Modify aggregator to handle branch task assignment and execution
3.2. Implement efficient join operation for branches
3.3. Update task tracking and result collection

Sample Usage Example

@aggregator
def start(self):
    self.next(self.collaborator_branch1, self.collaborator_branch2, 
             branches={"branch1": ["collaborator1"], "branch2": ["collaborator2"]})

@collaborator
def collaborator_branch1(self):
    # Branch 1 specific processing
    self.next(self.join)

@collaborator
def collaborator_branch2(self):
    # Branch 2 specific processing  
    self.next(self.join)

@aggregator
def join(self, inputs):
    # Process results from both branches
    self.next(self.end)

Limitations and Scope

• All branches must adhere to Metaflow validation rules
• Branch synchronization points must be clearly defined

Current Design and Limitations

• Current implementation for VFL executes all steps sequentially with each aggregator initiating workflow for collaborator after aggregating results from previous collaborator. Proposed approach simplifies the code while supporting multiple branches/execution flow for aggregator and collaborator.

Next steps and other information.

• To be implemented for Local Runtime initially and then Federated Runtime.

[theakshaypant]

For all the usage code blocks, the join steps are incorrectly labelled; shouldn't they be joining at the @aggregator?

[nisha987]

yes. that's right. updated.

[teoparvanov]

Thanks for adding more implementation details, @nisha987, I think it's a very good starting point!
From an FLSpec interface perspective it looks good.

On the other hand, I think the sequence diagram you've shared mostly corresponds to the LocalRuntime, which acts as an orchestrator among the aggregator and the collaborators (running within the same process). On the other hand, FederatedRuntime does not have a symmetric role (f.e. it does not implement the execute_task method). You can start tracking the difference by inspecting _run_local() and _run_federated in FLSpec's run() method.

In fact, for FederatedRuntime, the Director is the entity that orchestrates start-up of the aggregator it manages (as a coroutine), and the envoy+collaborator processes (see the director's start_experiment_execution_loop).

As a next step, I would therefore suggest creating a separate sequence diagram for FederatedRuntime.

[nisha987]

Added a sequence diagram for federated runtime.

[teoparvanov]

PS: in terms of implementation sequence, it could make sense to deliver branching support for LocalRuntime independently (before) support in FederatedRuntime, to improve local simulation parallelism.

[nisha987]

agreed.

[teoparvanov]

Another small question - what does "All branches must adhere to Metaflow validation rules" mean? Could you reference those rules, presumably from Metaflow's documentation?

[nisha987]

It just means that we support branching as supported in metaflow. (join all branches at the end, create any number of branches.)
metaflow link: https://docs.metaflow.org/metaflow/basics#branch