From 65dc0f614a9e19d4c2b5b015af1cce10631b7628 Mon Sep 17 00:00:00 2001 From: Teodor Parvanov Date: Wed, 30 Apr 2025 08:54:58 +0200 Subject: [PATCH] Revert "End of round callback adjustment (#1568)" This reverts commit e848af18a8f5eb75716599c8e4d0adb175c38213. Signed-off-by: Teodor Parvanov --- openfl/component/aggregator/aggregator.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/openfl/component/aggregator/aggregator.py b/openfl/component/aggregator/aggregator.py index 6a77bb5a9b..6675ac240e 100644 --- a/openfl/component/aggregator/aggregator.py +++ b/openfl/component/aggregator/aggregator.py @@ -1148,6 +1148,10 @@ def _end_of_round_check(self): # Once all of the task results have been processed self._end_of_round_check_done[self.round_number] = True + # End of round callbacks. + # todo handle case when aggregator restarted before callback was successful + self.callbacks.on_round_end(self.round_number, logs) + # Save the latest model if not self.assigner.is_task_group_evaluation(): logger.info("Saving round %s model...", self.round_number) @@ -1155,10 +1159,6 @@ def _end_of_round_check(self): else: logger.info("Skipping model save for round %s in evaluation mode.", self.round_number) - # End of round callbacks. - # todo handle case when aggregator restarted before callback was successful - self.callbacks.on_round_end(self.round_number, logs) - self.round_number += 1 # resetting stragglers for task for a new round @@ -1171,8 +1171,6 @@ def _end_of_round_check(self): # TODO This needs to be fixed! if self._time_to_quit(): logger.info("Experiment Completed. Cleaning up...") - # End of experiment callbacks. - self.callbacks.on_experiment_end() else: logger.info("Starting round %s...", self.round_number) # https://github.com/securefederatedai/openfl/pull/1195#discussion_r1879479537