Comments improved

Signed-off-by: noopur <[email protected]>

.github/workflows/task_runner_e2e.yml

@@ -35,15 +35,19 @@ env:
 
 jobs:
   test_run:
+    name: Task Runner E2E Test
     runs-on: ubuntu-22.04
+
     strategy:
       matrix:
         # There are open issues for some of the models, so excluding them for now:
         # 1. https://github.com/securefederatedai/openfl/issues/1126
         # 2. https://github.com/securefederatedai/openfl/issues/1127 
         # model_name: [ "torch_cnn_mnist", "keras_cnn_mnist", "torch_cnn_histology", "tf_2dunet", "tf_cnn_histology" ]
         model_name: [ "torch_cnn_mnist", "keras_cnn_mnist" ]
+        python_version: [ "3.8", "3.9", "3.10" ]
       fail-fast: false # do not immediately fail if one of the combinations fail
+
     env:
       MODEL_NAME: ${{ inputs.model_name == 'all' && matrix.model_name || inputs.model_name }} # run all models if model_name is not provided else run the provided model
 
@@ -58,7 +62,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v3
       with:
-        python-version: "3.9"
+        python-version: ${{ matrix.python_version }}
 
     - name: Install dependencies
       run: |

tests/openfl_e2e/conftest.py

@@ -211,42 +211,45 @@ def fx_federation(request, pytestconfig):
         federation_fixture: Named tuple containing the objects for model owner, aggregator, and collaborators
     """
     log.info("Fixture for federation setup")
+    collaborators = []
+    # Default name for bare metal approach, modify as required.
+    agg_domain_name = "aggregator"
+
     # Parse the command line arguments
     args = parse_arguments()
     model_name = args.model_name
     results_dir = args.results_dir or pytestconfig.getini("results_dir")
     num_collaborators = args.num_collaborators
     num_rounds = args.num_rounds
 
-    collaborators = []
-    # Default name for bare metal approach, modify as required.
-    agg_domain_name = "aggregator"
+    # Validate the model name and create the workspace name
+    if not model_name.upper() in constants.ModelName._member_names_:
+        raise ValueError(f"Invalid model name: {model_name}")
 
-    log.info(
-        f"num_collaborators: {num_collaborators}, num_rounds: {num_rounds}, model_name: {model_name}"
-    )
     workspace_name = f"workspace_{model_name}"
+
+    # Create model owner object and the workspace for the model
     model_owner = participants.ModelOwner(workspace_name, model_name)
     try:
         workspace_path = model_owner.create_workspace(results_dir=results_dir)
     except Exception as e:
         log.error(f"Failed to create the workspace: {e}")
         raise e
 
-    # Modify the plan if the number of collaborators or rounds are provided
-    if num_collaborators or num_rounds:
-        try:
-            model_owner.modify_plan(new_rounds=num_rounds, num_collaborators=num_collaborators)
-        except Exception as e:
-            log.error(f"Failed to modify the plan: {e}")
-            raise e
+    # Modify and initialize the plan
+    try:
+        model_owner.modify_plan(new_rounds=num_rounds, num_collaborators=num_collaborators)
+    except Exception as e:
+        log.error(f"Failed to modify the plan: {e}")
+        raise e
 
     try:
         model_owner.initialize_plan(agg_domain_name=agg_domain_name)
     except Exception as e:
         log.error(f"Failed to initialize the plan: {e}")
         raise e
 
+    # Modify and initialize the plan
     try:
         model_owner.certify_workspace()
     except Exception as e:
@@ -266,7 +269,7 @@ def fx_federation(request, pytestconfig):
         )
         collaborators.append(collaborator)
 
-    # Return the objects for model owner, aggregator, and collaborators
+    # Return the federation fixture
     return federation_fixture(
         model_owner=model_owner,
         aggregator=aggregator,

tests/openfl_e2e/utils/constants.py

@@ -1,6 +1,21 @@
 # Copyright 2024-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from enum import Enum
+
+# Define the model names. This is a non exhaustive list of models that can be used in the tests
+class ModelName(Enum):
+    """
+    Enum class to define the model names.
+    """
+    # IMP - The model name must be same (and in uppercase) as the model value.
+    # This is used to identify the model in the tests.
+    TORCH_CNN_MNIST = "torch_cnn_mnist"
+    KERAS_CNN_MNIST = "keras_cnn_mnist"
+    TORCH_CNN_HISTOLOGY = "torch_cnn_histology"
+    TF_2DUNET = "tf_2dunet"
+    TF_CNN_HISTOLOGY = "tf_cnn_histology"
+
 NUM_COLLABORATORS = 2
 NUM_ROUNDS = 5
 WORKSPACE_NAME = "my_federation"