Review comments

Signed-off-by: Chaurasiya, Payal <[email protected]>

.github/workflows/task_runner_eval_dws_e2e.yml

@@ -1,5 +1,5 @@
 ---
-# Task Runner E2E tests for dockerized approach
+# Task Runner E2E Evaluation tests for dockerized approach
 
 name: Task_Runner_Eval_DWS_E2E  # Please do not modify the name as it is used in the composite action
 

tests/end_to_end/models/aggregator.py

@@ -26,6 +26,7 @@ def __init__(self, agg_domain_name=None, workspace_path=None, container_id=None,
             agg_domain_name (str): Aggregator domain name
             workspace_path (str): Workspace path
             container_id (str): Container ID
+            eval_scope (bool, optional): Scope of aggregator is evaluation. Default is False.
         """
         self.name = "aggregator"
         self.agg_domain_name = agg_domain_name

tests/end_to_end/models/model_owner.py

@@ -154,15 +154,16 @@ def modify_plan(self, param_config, plan_path, eval_scope=False):
             data["network"]["settings"]["use_tls"] = param_config.use_tls
 
             if eval_scope:
-                # Remove all existing task_groups
-                data['assigner']['settings']['task_groups'] = []
+                # Remove all existing task_groups and set num_rounds to 1
+                data["assigner"]["settings"]["task_groups"] = []
                 # Add new task_groups for evaluation scope with task as aggregated_model_validation
                 new_task_group = {
                     "name": "evaluation",
                     "percentage": 1.0,
                     "tasks": ["aggregated_model_validation"]
                 }
-                data['assigner']['settings']['task_groups'].append(new_task_group)
+                data["assigner"]["settings"]["task_groups"].append(new_task_group)
+                data["aggregator"]["settings"]["rounds_to_train"] = 1
 
             with open(plan_file, "w+") as write_file:
                 yaml.dump(data, write_file)

tests/end_to_end/test_suites/tr_with_eval_tests.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2023 Intel Corporation
+# Copyright 2020-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 import pytest
@@ -10,7 +10,7 @@
     fx_federation_tr_dws,
 )
 from tests.end_to_end.utils import federation_helper as fed_helper
-from tests.end_to_end.utils.tr_workspace import create_tr_workspace, create_tr_workspace_dws
+from tests.end_to_end.utils.tr_workspace import create_tr_workspace, create_tr_dws_workspace
 from tests.end_to_end.utils.summary_helper import get_aggregated_accuracy
 
 log = logging.getLogger(__name__)
@@ -35,7 +35,7 @@ def test_eval_federation_via_native(request, fx_federation_tr):
         num_rounds=request.config.num_rounds,
     ), "Federation completion failed"
 
-    # Set the best model path
+    # Set the best model path in request. It is used during plan initialization for evaluation step
     request.config.best_model_path = os.path.join(fx_federation_tr.aggregator.workspace_path, "save", "best.pbuf")
     metric_file_path = os.path.join(fx_federation_tr.aggregator.workspace_path, "logs", "aggregator_metrics.txt")
     model_accuracy = get_aggregated_accuracy(metric_file_path)
@@ -81,14 +81,14 @@ def test_eval_federation_via_dockerized_workspace(request, fx_federation_tr_dws)
         num_rounds=request.config.num_rounds,
     ), "Federation completion failed"
 
-    # Set the best model path
+    # Set the best model path in request. It is used during plan initialization for evaluation step
     request.config.best_model_path = os.path.join(fx_federation_tr_dws.aggregator.workspace_path, "save", "best.pbuf")
     metric_file_path = os.path.join(fx_federation_tr_dws.aggregator.workspace_path, "logs", "aggregator_metrics.txt")
     model_accuracy = get_aggregated_accuracy(metric_file_path)
     log.info(f"Model accuracy post {request.config.num_rounds} rounds: {model_accuracy}")
 
     # Create new workspace with evaluation scope
-    new_fed_obj = create_tr_workspace_dws(request, eval_scope=True)
+    new_fed_obj = create_tr_dws_workspace(request, eval_scope=True)
 
     results_new = fed_helper.run_federation_for_dws(new_fed_obj, use_tls=request.config.use_tls)
     # Verify the completion of the federation run

tests/end_to_end/utils/tr_common_fixtures.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from tests.end_to_end.utils.tr_workspace import create_tr_workspace, create_tr_workspace_dws
+from tests.end_to_end.utils.tr_workspace import create_tr_workspace, create_tr_dws_workspace
 
 
 @pytest.fixture(scope="function")
@@ -37,4 +37,4 @@ def fx_federation_tr_dws(request):
     Note: As this is a function level fixture, thus no import is required at test level.
     """
     request.config.test_env = "task_runner_dockerized_ws"
-    return create_tr_workspace_dws(request)
+    return create_tr_dws_workspace(request)

tests/end_to_end/utils/tr_workspace.py

@@ -23,6 +23,15 @@
 
 
 def create_tr_workspace(request, eval_scope=False):
+    """Summary: Create a task runner workspace.
+
+    Args:
+        request (object): Pytest request object.
+        eval_scope (bool, optional): If True, sets up the evaluation scope for a single round. Defaults to False.
+
+    Returns:
+        tuple : A named tuple containing the objects for model owner, aggregator, and collaborators.
+    """
     collaborators = []
     executor = concurrent.futures.ThreadPoolExecutor()
 
@@ -48,10 +57,6 @@ def create_tr_workspace(request, eval_scope=False):
     plan_path = constants.AGG_PLAN_PATH.format(local_bind_path)
     param_config=request.config
 
-    if eval_scope:
-        log.info("Setting up evaluation scope, so update the plan for 1 round")
-        param_config.num_rounds = 1
-
     model_owner.modify_plan(param_config, plan_path=plan_path, eval_scope=eval_scope)
 
     # Initialize the plan
@@ -117,27 +122,15 @@ def create_tr_workspace(request, eval_scope=False):
     )
 
 
-def create_tr_workspace_dws(request, eval_scope=False):
+def create_tr_dws_workspace(request, eval_scope=False):
     """
-    Create a task runner workspace for distributed workload simulation.
-
-    This function sets up the necessary environment and configurations for a
-    task runner workspace, including model owner and collaborator workspaces,
-    plan modifications, and Docker container setups.
-
+    Run task runner experiment thru dockerized workspace.
     Args:
-        request: A request object containing configuration details.
-        eval_scope (bool, optional): If True, sets up the evaluation scope for
-            a single round. Defaults to False.
+        request (object): Pytest request object.
+        eval_scope (bool, optional): If True, sets up the evaluation scope for a single round. Defaults to False.
 
     Returns:
-        federation_details: An object containing details about the federation
-        setup, including model owner, aggregator, collaborators, workspace path,
-        and local bind path.
-
-    Raises:
-        Exception: If there is an error during the creation of the tar for the
-        aggregator.
+        tuple: A named tuple containing the objects for model owner, aggregator, and collaborators.
     """
 
     collaborators = []