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

Alexleung/dev branch latest

Author: fedml-alex

python/fedml/computing/scheduler/master/base_master_protocol_manager.py

@@ -33,7 +33,7 @@ def __init__(self, args, agent_config=None):
         self.agent_config = agent_config
         self.topic_start_train = None
         self.topic_stop_train = None
-        self.topic_exit_train = None
+        self.topic_complete_job = None
         self.topic_report_status = None
         self.topic_ota_msg = None
         self.topic_response_device_info = None
@@ -62,8 +62,8 @@ def generate_topics(self):
         # The topi for stopping training
         self.topic_stop_train = "mlops/flserver_agent_" + str(self.edge_id) + "/stop_train"
 
-        # The topi for exiting training
-        self.topic_exit_train = GeneralConstants.get_topic_exit_train(self.edge_id)
+        # The topic for completing job
+        self.topic_complete_job = GeneralConstants.get_topic_complete_job(self.edge_id)
 
         # The topic for reporting current device status.
         self.topic_report_status = "mlops/report_device_status"
@@ -93,7 +93,7 @@ def generate_topics(self):
         self.subscribed_topics.clear()
         self.add_subscribe_topic(self.topic_start_train)
         self.add_subscribe_topic(self.topic_stop_train)
-        self.add_subscribe_topic(self.topic_exit_train)
+        self.add_subscribe_topic(self.topic_complete_job)
         self.add_subscribe_topic(self.topic_report_status)
         self.add_subscribe_topic(self.topic_ota_msg)
         self.add_subscribe_topic(self.topic_response_device_info)
@@ -108,7 +108,7 @@ def add_protocol_handler(self):
         # Add the message listeners for all topics
         self.add_message_listener(self.topic_start_train, self.callback_start_train)
         self.add_message_listener(self.topic_stop_train, self.callback_stop_train)
-        self.add_message_listener(self.topic_exit_train, self.callback_exit_train)
+        self.add_message_listener(self.topic_complete_job, self.callback_complete_job)
         self.add_message_listener(self.topic_ota_msg, FedMLBaseMasterProtocolManager.callback_server_ota_msg)
         self.add_message_listener(self.topic_report_status, self.callback_report_current_status)
         self.add_message_listener(self.topic_response_device_info, self.callback_response_device_info)
@@ -281,7 +281,7 @@ def callback_stop_train(self, topic, payload, use_payload=None):
             run_id, args=self.args, server_id=server_id, request_json=request_json,
             run_as_cloud_agent=self.run_as_cloud_agent)
 
-    def callback_exit_train(self, topic, payload):
+    def callback_complete_job(self, topic, payload):
         # Parse the parameters.
         request_json = json.loads(payload)
         run_id = request_json.get("runId", None)
@@ -291,10 +291,10 @@ def callback_exit_train(self, topic, payload):
         if server_id is None:
             server_id = request_json.get("server_id", None)
 
-        # Stop the 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)
+        self._process_job_complete_status(run_id, server_id, request_json)
+
+    def _process_job_complete_status(self, run_id, server_id, complete_payload):
+        pass
 
     def callback_run_logs(self, topic, payload):
         run_id = str(topic).split('/')[-1]

python/fedml/computing/scheduler/master/master_protocol_manager.py

@@ -34,3 +34,10 @@ def _init_extra_items(self):
     # Override
     def print_connected_info(self):
         super().print_connected_info()
+
+    # Override
+    def _process_job_complete_status(self, run_id, server_id, complete_payload):
+        # Complete the job runner
+        self._get_job_runner_manager().complete_job_runner(
+            run_id, args=self.args, server_id=server_id, request_json=complete_payload,
+            run_as_cloud_agent=self.run_as_cloud_agent, run_as_cloud_server=self.run_as_cloud_server)

python/fedml/computing/scheduler/model_scheduler/master_protocol_manager.py

@@ -103,9 +103,6 @@ def callback_deployment_result_message(self, topic=None, payload=None):
         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.

python/fedml/computing/scheduler/scheduler_core/general_constants.py

@@ -208,6 +208,11 @@ def get_ip_address(request_json, infer_host=None):
         return ip
 
     @staticmethod
-    def get_topic_exit_train(server_id):
-        topic_exit_train = f"status_center/flserver_agent_{server_id}/exit_train"
-        return topic_exit_train
+    def get_topic_complete_job(server_id):
+        topic_complete_job = f"status_center/master_agent_{server_id}/complete_job"
+        return topic_complete_job
+
+    @staticmethod
+    def get_payload_complete_job(run_id, server_id):
+        payload_complete_job = {"runId": run_id, "serverId": server_id}
+        return payload_complete_job

python/fedml/computing/scheduler/scheduler_core/status_manager_protocols.py

@@ -67,9 +67,9 @@ def process_job_completed_status(self, master_id, status):
         # self.stop_cloud_server()
         # self.remove_listener_for_run_metrics(self.run_id)
         # self.remove_listener_for_run_logs(self.run_id)
-        payload_exit_train = {"runId": self.run_id, "serverId": master_id}
         self.message_center.receive_message(
-            GeneralConstants.get_topic_exit_train(master_id), json.dumps(payload_exit_train))
+            GeneralConstants.get_topic_complete_job(master_id),
+            json.dumps(GeneralConstants.get_payload_complete_job(self.run_id, master_id)))
 
         if self.status_center.is_deployment_status_center and status == ServerConstants.MSG_MLOPS_SERVER_STATUS_FAILED:
             self.report_deployment_status(self.run_id, GeneralConstants.MSG_MODELOPS_DEPLOYMENT_STATUS_FAILED)