Merge pull request #2102 from FedML-AI/alexleung/dev_branch_latest

[CoreEngine] make the deployment work.
FedML-AI · May 15, 2024 · 5d8bbc5 · 5d8bbc5
2 parents c3436fe + 496732a
commit 5d8bbc5
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 4 deletions.
diff --git a/python/fedml/computing/scheduler/master/base_master_job_runner_manager.py b/python/fedml/computing/scheduler/master/base_master_job_runner_manager.py
@@ -47,6 +47,17 @@ def stop_job_runner(
                 target=FedMLCloudServerManager.stop_cloud_server, args=(run_id, server_id, args.agent_config))
             stopping_process.start()
 
+    def complete_job_runner(
+            self, run_id, args=None, server_id=None, request_json=None,
+            run_as_cloud_agent=False, run_as_cloud_server=False
+    ):
+        super().complete_job_runner(run_id)
+
+        if run_as_cloud_agent or run_as_cloud_server:
+            stopping_process = Process(
+                target=FedMLCloudServerManager.stop_cloud_server, args=(run_id, server_id, args.agent_config))
+            stopping_process.start()
+
     def _start_cloud_server(
             self, args, run_id, request_json, edge_id=None,
             use_local_process_as_cloud_server=False

diff --git a/python/fedml/computing/scheduler/master/base_master_protocol_manager.py b/python/fedml/computing/scheduler/master/base_master_protocol_manager.py
@@ -292,7 +292,7 @@ def callback_exit_train(self, topic, payload):
             server_id = request_json.get("server_id", None)
 
         # Stop the job runner
-        self._get_job_runner_manager().stop_job_runner(
+        self._get_job_runner_manager().complete_job_runner(
             run_id, args=self.args, server_id=server_id, request_json=request_json,
             run_as_cloud_agent=self.run_as_cloud_agent, run_as_cloud_server=self.run_as_cloud_server)
 

diff --git a/python/fedml/computing/scheduler/model_scheduler/master_job_runner.py b/python/fedml/computing/scheduler/model_scheduler/master_job_runner.py
@@ -361,8 +361,7 @@ def process_deployment_result_message(self, topic=None, payload=None):
         # Update the global deployment result mapping
         self.slave_deployment_results_map[str(device_id)] = model_status
 
-        logging.info("callback_deployment_result_message: topic {}, payload {}, result mapping {}.".format(
-            topic, payload, self.slave_deployment_results_map))
+        logging.info("callback_deployment_result_message: topic {}, payload {}.".format(topic, payload))
 
         request_json = self.request_json
         if request_json is None:

diff --git a/python/fedml/computing/scheduler/model_scheduler/master_protocol_manager.py b/python/fedml/computing/scheduler/model_scheduler/master_protocol_manager.py
@@ -100,9 +100,12 @@ def print_connected_info(self):
         pass
 
     def callback_deployment_result_message(self, topic=None, payload=None):
-        logging.info(f"Received deployment result: {self}")
+        logging.info(f"Received deployment result")
         FedMLDeployJobRunnerManager.get_instance().save_deployment_result(topic, payload)
 
+    def callback_exit_train(self, topic, payload):
+        pass
+
     def callback_delete_deployment(self, topic, payload):
         logging.info("[Master] callback_delete_deployment")
         # Parse payload as the model message object.