revert to initialize, add super call to handle_event

NVIDIA · Mar 15, 2024 · 1cb0079 · 1cb0079
1 parent 82cd200
commit 1cb0079
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 4 deletions.
diff --git a/nvflare/apis/impl/controller.py b/nvflare/apis/impl/controller.py
@@ -35,7 +35,7 @@ def __init__(self, task_check_period=0.2):
         self._task_check_period = task_check_period
         self.communicator = None
 
-    def initialize_controller(self, fl_ctx: FLContext):
+    def initialize(self, fl_ctx: FLContext):
         engine = fl_ctx.get_engine()
         if not engine:
             self.system_panic(f"Engine not found. {self.__class__.__name__} exiting.", fl_ctx)

diff --git a/nvflare/app_common/workflows/cyclic_ctl.py b/nvflare/app_common/workflows/cyclic_ctl.py
@@ -266,6 +266,7 @@ def restore(self, state_data: dict, fl_ctx: FLContext):
             pass
 
     def handle_event(self, event_type, fl_ctx):
+        super().handle_event(event_type, fl_ctx)
         if event_type == EventType.JOB_DEAD:
             client_name = fl_ctx.get_prop(FLContextKey.DEAD_JOB_CLIENT_NAME)
             new_client_list = []

diff --git a/nvflare/app_common/workflows/model_controller.py b/nvflare/app_common/workflows/model_controller.py
@@ -144,7 +144,7 @@ def start_controller(self, fl_ctx: FLContext) -> None:
         self.event(AppEventType.INITIAL_MODEL_LOADED)
 
         self.engine = self.fl_ctx.get_engine()
-        self.initialize()
+        FLComponentWrapper.initialize(self)
 
     def _build_shareable(self, data: FLModel = None) -> Shareable:
         if not data:  # if no data is given, send self.model

diff --git a/nvflare/private/fed/server/server_runner.py b/nvflare/private/fed/server/server_runner.py
@@ -127,7 +127,7 @@ def _execute_run(self):
                     fl_ctx.set_prop(FLContextKey.WORKFLOW, wf.id, sticky=True)
 
                     wf.controller.communicator.initialize_run(fl_ctx)
-                    wf.controller.initialize_controller(fl_ctx)
+                    wf.controller.initialize(fl_ctx)
 
                     self.log_info(fl_ctx, "Workflow {} ({}) started".format(wf.id, type(wf.controller)))
                     self.log_debug(fl_ctx, "firing event EventType.START_WORKFLOW")

diff --git a/tests/unit_test/apis/impl/controller_test.py b/tests/unit_test/apis/impl/controller_test.py
@@ -126,7 +126,7 @@ def _setup_system(num_clients=1):
     fl_ctx = mock_server_engine.new_context()
     communicator = WFCommServer()
     controller.set_communicator(communicator)
-    controller.initialize_controller(fl_ctx)
+    controller.initialize(fl_ctx)
     controller.communicator.initialize_run(fl_ctx=fl_ctx)
     return controller, mock_server_engine, fl_ctx, clients_list