E2E Resiliency: More restarts & edge case handling (#1629)

* Handle scenario where restart happens on the last round

Signed-off-by: noopur <noopur@intel.com>

* Break if last round reached

Signed-off-by: noopur <noopur@intel.com>

* Corrected log msg

Signed-off-by: noopur <noopur@intel.com>

* Correction

Signed-off-by: noopur <noopur@intel.com>

* Using -9

Signed-off-by: noopur <noopur@intel.com>

* Stop if the rounds are over

Signed-off-by: noopur <noopur@intel.com>

* Better logic to break the flow on reaching final round

Signed-off-by: noopur <noopur@intel.com>

---------

Signed-off-by: noopur <noopur@intel.com>

Author: noopurintel

tests/end_to_end/test_suites/tr_resiliency_tests.py

@@ -277,43 +277,42 @@ def _perform_restart_validate_rounds(fed_obj, db_file, total_rounds):
         db_file (str): Path to the database file
         total_rounds (int): Total number of rounds
     """
-
-    init_round = fed_helper.get_current_round(db_file)
-    log.info(f"Round number is {init_round} before restarts")
-
-    # Restart aggregator
-    assert int_helper.restart_participants([fed_obj.aggregator])
-    log.info("Aggregator restarted successfully")
-
-    assert (
-        round_post_agg_restart := fed_helper.validate_round_increment(
-            init_round,
-            db_file,
-            total_rounds,
+    def _restart_and_validate(participants, current_round, db_file, total_rounds, description, attempt):
+        """
+        Restarts the participants, validates the round increment, and checks if the total rounds are reached.
+        """
+        assert int_helper.restart_participants(participants), f"Failed to restart {description} on attempt {attempt}"
+        log.info(f"{description.capitalize()} restarted successfully on attempt {attempt}.")
+
+        new_round = fed_helper.validate_round_increment(current_round, db_file, total_rounds)
+        assert new_round, f"Expected current round to be ahead of {current_round} after {description} restart on attempt {attempt}"
+        return new_round
+
+    current_round = fed_helper.get_current_round(db_file)
+    log.info(f"Round number is {current_round} before restarts")
+
+    for i in range(1, 4):
+        # Restart aggregator and validate
+        current_round = _restart_and_validate(
+            [fed_obj.aggregator], current_round, db_file, total_rounds, "aggregator", attempt=i
         )
-    ), f"Expected current round to be ahead of {init_round} after aggregator restart"
-
-    # Restart collaborators
-    assert int_helper.restart_participants(fed_obj.collaborators)
-    log.info("Collaborators restarted successfully")
+        # Stop further processing in case of final round.
+        if current_round + 1 == total_rounds:
+            break
 
-    assert (
-        round_post_collab_restart := fed_helper.validate_round_increment(
-            round_post_agg_restart,
-            db_file,
-            total_rounds,
+        # Restart collaborators and validate
+        current_round = _restart_and_validate(
+            fed_obj.collaborators, current_round, db_file, total_rounds, "collaborators", attempt=i
         )
-    ), f"Expected current round to be ahead of {round_post_agg_restart} after collaborators restart"
+        if current_round + 1 == total_rounds:
+            break
 
-    # Restart all participants
-    assert int_helper.restart_participants(fed_obj.collaborators + [fed_obj.aggregator])
-    log.info("All participants restarted successfully")
-
-    assert fed_helper.validate_round_increment(
-        round_post_collab_restart,
-        db_file,
-        total_rounds,
-    ), f"Expected current round to be ahead of {round_post_collab_restart} after all participants restart"
+        # Restart all participants and validate
+        current_round = _restart_and_validate(
+            fed_obj.collaborators + [fed_obj.aggregator], current_round, db_file, total_rounds, "all participants", attempt=i
+        )
+        if current_round + 1 == total_rounds:
+            break
 
     log.info("Current round number is increasing after every restart as expected.")
 

tests/end_to_end/utils/federation_helper.py

@@ -368,7 +368,7 @@ def _verify_completion_for_participant(
         # If None, it means the process is still running
         # This is applicable for native process only
         if participant.start_process:
-            if participant.start_process.poll():
+            if participant.start_process.poll() or not len(intr_helper.get_pids_for_active_command(participant.name)):
                 log.info(f"No processes found for participant {participant.name}")
                 break
             else:
@@ -1071,7 +1071,11 @@ def validate_round_increment(inp_round, database_file, total_rounds, timeout=300
         if current_round > inp_round + 1:
             log.info(f"Round number has increased from {inp_round} to {current_round}")
             return current_round
-        log.info(f"Round number has not increased. Retrying in {sleep_interval} seconds...")
+        # Check if already at the final round (round no. index starts with 0)
+        if current_round + 1 == total_rounds:
+            log.info(f"Already at the final round")
+            return current_round
+        log.info(f"Round number has not increased from {inp_round}. Retrying in {sleep_interval} seconds...")
         time.sleep(sleep_interval)
     log.warning(f"Round number has not increased from {inp_round} after {timeout} seconds")
     return False

tests/end_to_end/utils/interruption_helper.py

@@ -76,7 +76,7 @@ def stop_start_native_participant(participant, action):
         raise ex.ParticipantStopException(f"Invalid action {action}")
 
     # Irrespective of the action, kill the processes to ensure clean state
-    log.info(f"Kill the processes for {participant.name} if running to avoid conflicts")
+    log.debug(f"Killing the processes (if running) for {participant.name} to avoid conflicts")
     participant.kill_process()
 
     if action == "stop":