[Workflow API] Enable user's to import helper functions / utilities from a separate python script in FederatedRuntime Notebooks

[refai06]

Is your feature request related to a problem? Please describe.
The issue related to <1565>, where the User encountered a problem while importing user-defined modules in Jupyter notebooks while running an experiment based on Workflow API & FederatedRuntime

In the current implementation of Workflow API, jupyter notebook is expected to define the Federated Learning experiment in it's entirety. If the user attempts to import a user defined module from a different python script it will fail due to following reasons

• When the notebook is exported, the script inside the generated_workspace does not contain the user defined code
• This will lead to ImportError/ModuleNotFoundError failure during execution on participants in a distributed infrastructure

Describe the solution you'd like
Use case: Enable user's to import the helper functions from a separate python script. For e.g.

Filename: validation.py (contains helper functions)

class CrowdGuardClientValidation:

    def __distance_global_model_final_metric(distance_type: str, prediction_matrix,
                                             prediction_global_model, sample_indices_by_label,
                                             own_index):

    def __predict_for_single_model(model, local_data, device):

    def __do_predictions(models, global_model, local_data, device):

    def __prune_poisoned_models(num_layers, total_number_of_clients, own_client_index,
                                distances_by_metric, verbose=False):

    def validate_models(global_model, models, own_client_index, local_data, device):

Filename: CrowdGuard.ipynb (Jupyter notebook for Workflow API experiment)

#| export

from validation import CrowdGuardClientValidation

class FederatedFlow_CrowdGuard(FLSpec):

    @aggregator
    def start(self):

    @collaborator
    def train(self):

    @collaborator
    def local_validation(self):
        ...
        detected_suspicious_models = CrowdGuardClientValidation.validate_models(self.global_model,
                                                                                all_models,
                                                                                own_client_index,
                                                                                self.train_loader,
                                                                                self.device)
	...
	

    @aggregator
    def end(self): 

To support this use case, existing export process in notebook_tools needs to be enhanced to

1. Identify user-defined imports

• Use the Python ast module to parse the exported script (i.e experiment.py)
• Extract all import statements:

  import module
  from module import ...
  

2. Classify imports and filter only user-defined modules present in the workspace

3. Copy user-defined modules into the generated_workpace

• Locate its corresponding .py file or folder within the workspace, the same directory where the Jupyter notebook is located.
• Copy the module into the generated_workspace as follows
  • Directly copy, if it's a single file (e.g. validation.py)
  • If it's a folder (e.g. helper/validation.py), the entire folder is copied

This shall ensure that the FL experiment and workspace includes all user-defined module code

To enable the user defined imports from a different python file (workspace/validation.py), the generated_workspace shall also contain the additional python scrips. For e.g. existing workspace

workspace
├── crowd_guard.ipynb 
└── validation.py

generated_workspace after the workspace is exported

workspace
├── generated_workspace
│   ├── src
│	│   ├── __init__.py
│	│   ├── experiment.py
│	│   └── validation.py 
│	├── .workspace
│	├── plan
│	│   └── plan.yaml
│	└── requirements.txt
├── crowd_guard.ipynb 
└── validation.py

Describe alternatives you've considered

• Current implementation assumes that the Jupyter notebook contains all the user-defined logic required for the experiment
  Additional context: This enhancement is based on following Requirements & Guidelines

Export Directives

• User defined utility / helper functions should be present in a notebook cell that is annotated by #| export directive as the first line
• This is required to enable infrastructure to export the logic to exported script and analyze the exported script for any user-defined imports

No installations in User defined imports:

• User defined utility should not install any packages as this will not be analyzed by the infrastructure to build the requirements.txt for the FL experiment

Location of User defined python scripts:

• User-defined modules must be present inside the same directory as the Jupyter Notebook to enable infrastructure to copy and package in generate_workspace
• Rationale: Path of user-defined module may not exist on distributed nodes and it may not be possible to copy the user-defined script to the given path

  • For e.g. If user creates scripts in openfl directory and uses them in the jupyter notebook (for e.g. from openfl.experimental.workflow.external.helper.validation import CrowdGuardClientValidation)

  • During AST module analysis of the exported script, the import statements could not be resolved.

  • It requires system path modification and rewriting of imports in the exported script to align with the generated_workspace
    For e.g.
    from openfl.experimental.workflow.external.helper.validation import CrowdGuardClientValidation would need to be modified to
    from helper.validation import CrowdGuardClientValidation

  • Modification of imports is not recommended as it would introduce complexity

Restrictions on User-defined imports
User defined code should not modify the sys.path to enable python to find the scripts to import
For e.g.

workspace
├── crowd_guard.ipynb 
└── helper
          └── validaiton.py

Filename:CrowdGuard.ipynb (Jupyter notebook for Workflow API experiment)

...
...
sys.path.append('./helper')
from validation import CrowdGuardClientValidation
...

Correct_Usage

from helper.validation import CrowdGuardClientValidation

User-defined imports should be self-contained:

• User defined code should not import other user-defined code from different python scripts. For e.g.

Filename: utils.py (contains additional helper functions)

def calculate_accuracy(predictions, labels):
      correct = (predications == labels).sum()

      return correct / len(labels)

def compute_precision(predictions, labels):
      ...

      return 0.85

Filename: validation.py (contains helper functions)

from utils import calculate_accuracy, compute_precision

class CrowdGuardClientValidation:
      def validate_models(global_model, models, own_client_index, local_data, device):
      ...
      ...
      accuracy = caculate_accuracy(predictions, labels)

      precision = compute_precision(predictions, labels)

• Rationale:
  • Analysing whether ALL imports in user-defined code are user-defined could be quite complex
  • In above example notebook_tools shall be able to identify that from helper.validation import CrowdGuardClientValidation is a user defined import. But it would be quite complex to analyze and keep track if CrowdGuardClientValidation has other user defined imports, such as from utils import calculate_accuracy, compute_precision and whether they are in same python script or different

User defined module imports:

• Used defined imports for the helper module (e.g.helper/validation) with or without an __init__.py are supported, because entire user-defined helper folders are copied to the generated_workspace, which would resolve the imports correctly in the exported script (e.g. experiment.py)

  • For e.g., Import statements like from helper.validation import CrowdGuardClientValidation and from helper import CrowdGuardClientValidation

• If an __init__.py file is present inside the same directory (for e.g. workspace) as the Jupyter Notebook along with user-defined module validation.py. For e.g.

workspace
├── __init__.py
├── crowd_guard.ipynb 
└── validation.py

• Using import statement in Jupyter notebook (for e.g. from validation import CrowdGuardClientValidation) remains valid in the exported script (experiment.py), as validation.py is copied to the generated_workspace.