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

[scngupta-dsp]

Is your feature request related to a problem? Please describe.

The issue related to securefederatedai#1565, where the User encountered a problem while importing user-defined modules in Jupyter notebooks while running a FederatedRuntime experiment

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 ModuleNotFoundError failure during execution on participants in a distributed infrastructure

Describe the solution you'd like

Enable users to import helper functions from a separate python script. For e.g. FL experiment tutorial: crowd_guard.ipynb is importing some helper functions / classes from user-defined script validation.py in a folder workspace

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

validation.py (contains helper class)

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):

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. Analyse all imports in Jupyter Notebook and identify user-defined imports
2. Copy user defined python scripts / folders containing these imports into the generated_workpace

This shall ensure that generated_workspace (shown below) includes all user-defined code and ensure that it works on the distributed infrastructure

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

N.A.

Additional context

This enhancement shall be based on following Requirements & Guidelines

Export Directives

• User-defined imports should be present in a notebook cell that is annotated by #| export directive as the first line
• Rationale:
  • #| export directives are required to export the user-defined imports to exported script and further processing

User-defined scripts should not install any packages:

• User-defined scripts should not install any package
• Rationale:
  • While the exported script is analyzed to identify dependencies and build the requirements.txt for the FL experiment, User-defined scripts are not analyzed by the infrastructure to identify dependencies

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 to generate_workspace
• Rationale:
  • Path of user-defined module may not exist on distributed nodes
    For e.g. Notebook contains from openfl.experimental.workflow.external.helper.validation import CrowdGuardClientValidation
    While this may work on the user node - it may not be possible for infrastructure to recreate the same path on distributed nodes
  • Another possible option is to modify the imports, but this logic could be complex and is recommended to be avoided
    For e.g. from openfl.experimental.workflow.external.helper.validation import CrowdGuardClientValidation could be modified to from helper.validation import CrowdGuardClientValidation
    Proposal is to avoid this complexity in notebook_tools

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

CrowdGuard.ipynb (Jupyter notebook for Workflow API experiment)

...
...
sys.path.append('./helper')

from validation import CrowdGuardClientValidation
...

Recommended_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.

utils.py (contains additional helper functions)

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

validation.py (contains helper functions)

from utils import calculate_accuracy

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

• Rationale:
  • In above example notebook_tools shall be able to identify that from helper.validation import CrowdGuardClientValidation in Notebook is a user defined import
  • However analyzing validation.py and identifying from utils import calculate_accuracy is a user-defined import could be complex to implement and is recommended to be avoided