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

Alexleung/dev branch latest

Author: fedml-alex

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

@@ -264,23 +264,26 @@ def callback_stop_train(self, topic, payload, use_payload=None):
         run_id = request_json.get("runId", None)
         run_id = request_json.get("id", None) if run_id is None else run_id
         run_id_str = str(run_id)
+        edge_ids = request_json.get("edgeids", None)
         server_id = request_json.get("serverId", None)
         if server_id is None:
             server_id = request_json.get("server_id", None)
-        edge_ids = request_json.get("edgeids", None)
-
-        # Stop the job runner
-        self._get_job_runner_manager().stop_job_runner(
-            run_id, args=self.args, server_id=server_id, request_json=request_json,
-            run_as_cloud_agent=self.run_as_cloud_agent)
+        server_agent_id = server_id
 
         # Cleanup the cached object
         if self.running_request_json.get(run_id_str, None) is not None:
             self.running_request_json.pop(run_id_str)
 
+        # If it is the cloud agent, then forward the stopping request to the corresponding cloud server.
+        if self.run_as_cloud_agent:
+            server_agent_id = self.edge_id
+            topic_stop_train_to_cloud_server = f"mlops/flserver_agent_{server_id}/stop_train"
+            self.message_center.send_message(topic_stop_train_to_cloud_server, payload)
+            return
+
         # Reset all edge status and server status
         for iter_edge_id in edge_ids:
-            self.generate_status_report(run_id, iter_edge_id, server_agent_id=server_id).\
+            self.generate_status_report(run_id, iter_edge_id, server_agent_id=server_agent_id).\
                 report_client_id_status(iter_edge_id, GeneralConstants.MSG_MLOPS_SERVER_STATUS_KILLED,
                                         run_id=run_id, server_id=server_id)
 

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

@@ -23,7 +23,7 @@ def __init__(self, run_id=None, edge_id=None, server_id=None,
         self.edge_id = edge_id
         self.server_id = server_id
         self.edge_id_list = edge_id_list
-        self.client_agent_active_list = dict()
+        self.edge_status_dict = None
         self.running_scheduler_contract = running_scheduler_contract if running_scheduler_contract is not None else dict()
         self.message_reporter = MLOpsMetrics()
         self.message_reporter.set_messenger(message_center)
@@ -163,6 +163,8 @@ def status_center_process_master_status(self, topic, payload):
         status = request_json["status"]
         edge_id = request_json["edge_id"]
         server_id = request_json.get("server_id", None)
+        if server_id is None or str(server_id) == "0":
+            server_id = self.server_id
         run_id_str = str(run_id)
 
         # Process the job status
@@ -185,8 +187,7 @@ def process_job_status_consensus(self, run_id, master_id, status):
         status = self.get_entire_job_status()
 
         # Set the device status based on the job status
-        edge_id_status_dict = self.client_agent_active_list.get(f"{run_id}", {})
-        for edge_id_item, edge_status_item in edge_id_status_dict.items():
+        for edge_id_item, edge_status_item in self.edge_status_dict.items():
             if edge_id_item == "server":
                 continue
 
@@ -233,31 +234,29 @@ def status_center_process_slave_status(self, topic, payload):
         init_edge_id_list = payload_json.get("init_all_edge_id_list", None)
         init_server_id = payload_json.get("init_server_id", None)
 
-        active_item_dict = self.client_agent_active_list.get(f"{run_id}", None)
-        if active_item_dict is None:
-            self.client_agent_active_list[f"{run_id}"] = dict()
+        if self.edge_status_dict is None:
+            self.edge_status_dict = dict()
 
         if init_edge_id_list is not None:
-            self.client_agent_active_list[f"{run_id}"][f"server"] = init_server_id
+            self.edge_status_dict[f"server"] = init_server_id
             for edge_id_item in init_edge_id_list:
-                self.client_agent_active_list[f"{run_id}"][f"{edge_id_item}"] = \
+                self.edge_status_dict[f"{edge_id_item}"] = \
                     ClientConstants.MSG_MLOPS_CLIENT_STATUS_IDLE
 
         if run_id is not None and edge_id is not None:
-            self.client_agent_active_list[f"{run_id}"][f"{edge_id}"] = status
+            self.edge_status_dict[f"{edge_id}"] = status
 
             self.process_device_status(run_id, edge_id, status)
 
     def process_device_status(self, run_id, edge_id, status):
         number_of_failed_edges = 0
         number_of_finished_edges = 0
         number_of_killed_edges = 0
-        edge_id_status_dict = self.client_agent_active_list.get(f"{run_id}", {})
-        server_id = edge_id_status_dict.get("server", 0)
+        server_id = self.edge_status_dict.get("server", 0)
         enable_fault_tolerance, fault_tolerance_rate = self.parse_fault_tolerance_params(run_id)
         running_edges_list = list()
         edge_nums = 0
-        for edge_id_item, status_item in edge_id_status_dict.items():
+        for edge_id_item, status_item in self.edge_status_dict.items():
             if edge_id_item == "server":
                 continue