Merge pull request #2099 from FedML-AI/alexleung/dev_v070_for_refactor

Alexleung/dev v070 for refactor

Author: fedml-alex

python/examples/launch/serve_mnist/fedml_model_config.yaml

@@ -1,6 +1,8 @@
 workspace: "./"
 entry_point: "mnist_serve_main.py"
 
+auto_detect_public_ip: true
+
 data_cache_dir: ""
 bootstrap: ""
 

python/fedml/__init__.py

@@ -34,7 +34,7 @@
 _global_training_type = None
 _global_comm_backend = None
 
-__version__ = "0.8.30"
+__version__ = "0.8.31"
 
 
 # This is the deployment environment used for different roles (RD/PM/BD/Public Developers). Potential VALUE: local, dev, test, release

python/fedml/computing/scheduler/comm_utils/job_monitor.py

@@ -223,7 +223,7 @@ def monitor_replicas_number():
             curr_version = fedml.get_env_version()
             num_replica_url_path = "fedmlModelServer/api/v1/endpoint/replica-info"
             mlops_prefix = fedml._get_backend_service()
-            url = f"{mlops_prefix}{num_replica_url_path}"
+            url = f"{mlops_prefix}/{num_replica_url_path}"
 
             cached_token = FedMLModelCache.get_instance().get_end_point_token_with_eid(endpoint_id)
             if cached_token is None:

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

@@ -17,18 +17,18 @@ def __init__(self):
         self.master_api_process = None
         self.mlops_metrics = MLOpsMetrics()
         self.status_reporter = None
-        self.enable_simulation_cloud_agent = True
+        self.enable_simulation_cloud_agent = False
         self.use_local_process_as_cloud_server = False
         self.protocol_mgr = None
 
     def login(
             self, user_id, api_key=None, device_id=None,
-            os_name=None, role=None
+            os_name=None, role=None, runner_cmd=None
     ):
         # Login account
         login_result = FedMLAccountManager.get_instance().login(
             user_id, api_key=api_key, device_id=device_id,
-            os_name=os_name, role=role
+            os_name=os_name, role=role, runner_cmd=runner_cmd
         )
         if login_result is not None:
             self.agent_args = login_result

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

@@ -136,6 +136,10 @@ def run_impl(
 
         logging.info("Detect all status of Edge ids: " + str(edge_ids))
 
+        self.status_reporter.report_server_id_status(
+            self.run_id, ServerConstants.MSG_MLOPS_SERVER_STATUS_STARTING, edge_id=self.edge_id,
+            server_id=self.edge_id, server_agent_id=self.edge_id)
+
         status_ok, active_edge_info_dict, inactivate_edges = self.detect_edges_status(
             edge_device_info_queue, edge_device_info_global_queue=edge_device_info_global_queue,
             callback_when_edges_ready=self.send_training_request_to_edges)

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

@@ -36,6 +36,28 @@ def start_job_runner(
             status_center_queue=status_center_queue
         )
 
+    def stop_job_runner(
+            self, run_id, args=None, server_id=None, request_json=None,
+            run_as_cloud_agent=False, run_as_cloud_server=False
+    ):
+        super().stop_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 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

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

@@ -33,6 +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_complete_job = None
         self.topic_report_status = None
         self.topic_ota_msg = None
         self.topic_response_device_info = None
@@ -44,7 +45,7 @@ def __init__(self, args, agent_config=None):
         self.run_as_cloud_server = False
         self.run_as_edge_server_and_agent = False
         self.run_as_cloud_server_and_agent = False
-        self.enable_simulation_cloud_agent = True
+        self.enable_simulation_cloud_agent = False
         self.use_local_process_as_cloud_server = False
         self.ota_upgrade = FedMLOtaUpgrade(edge_id=args.edge_id)
         self.running_request_json = dict()
@@ -61,6 +62,9 @@ def generate_topics(self):
         # The topi for stopping training
         self.topic_stop_train = "mlops/flserver_agent_" + str(self.edge_id) + "/stop_train"
 
+        # 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"
 
@@ -89,6 +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_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)
@@ -103,6 +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_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)
@@ -140,12 +146,6 @@ def callback_start_train(self, topic=None, payload=None):
         except Exception:
             pass
 
-        # Parse the message when running in the cloud server mode.
-        if self.run_as_cloud_server:
-            message_bytes = payload.encode("ascii")
-            base64_bytes = base64.b64decode(message_bytes)
-            payload = base64_bytes.decode("ascii")
-
         # Parse the parameters
         # [NOTES] Example Request JSON:
         # https://fedml-inc.larksuite.com/wiki/ScnIwUif9iupbjkYS0LuBrd6sod#WjbEdhYrvogmlGxKTOGu98C6sSb
@@ -264,6 +264,9 @@ 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)
+        server_id = request_json.get("serverId", None)
+        if server_id is None:
+            server_id = request_json.get("server_id", None)
 
         # Broadcast the job status to all edges
         self.rebuild_status_center(self.get_status_queue())
@@ -274,7 +277,24 @@ def callback_stop_train(self, topic, payload, use_payload=None):
             self.running_request_json.pop(run_id_str)
 
         # Stop the job runner
-        self._get_job_runner_manager().stop_job_runner(run_id)
+        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)
+
+    def callback_complete_job(self, topic, payload):
+        # Parse the parameters.
+        request_json = json.loads(payload)
+        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)
+        server_id = request_json.get("serverId", None)
+        if server_id is None:
+            server_id = request_json.get("server_id", None)
+
+        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]
@@ -498,6 +518,11 @@ def send_training_stop_request_to_specific_edge(self, edge_id, payload):
         logging.info("stop_train: send topic " + topic_stop_train)
         self.message_center.send_message(topic_stop_train, payload)
 
+    def send_training_stop_request_to_cloud_server(self, edge_id, payload):
+        topic_stop_train = "mlops/flserver_agent_" + str(edge_id) + "/stop_train"
+        logging.info("stop_train: send topic " + topic_stop_train)
+        self.message_center.send_message(topic_stop_train, payload)
+
     def send_status_check_msg(self, run_id, edge_id, server_id, context=None):
         topic_status_check = f"server/client/request_device_info/{edge_id}"
         payload = {"server_id": server_id, "run_id": run_id}

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

@@ -3,13 +3,16 @@
 import logging
 import os
 import traceback
+
+import fedml
 from fedml.computing.scheduler.comm_utils.sys_utils import get_python_program
 
 
 class FedMLCloudServerManager:
     FEDML_CLOUD_SERVER_PREFIX = "fedml-server-run-"
     LOCAL_RUNNER_INFO_DIR_NAME = 'runner_infos'
     STATUS_IDLE = "IDLE"
+    FEDML_SERVER_BASE_IMAGE = "/fedml-device-image:"
 
     def __init__(self, args, run_id=None, edge_id=None, request_json=None, agent_config=None, version=None):
         self.server_docker_image = None
@@ -18,11 +21,13 @@ def __init__(self, args, run_id=None, edge_id=None, request_json=None, agent_con
         self.edge_id = edge_id
         self.request_json = request_json
         self.agent_config = agent_config
+        if version is None:
+            version = fedml.get_env_version()
         self.version = version
         image_version = self.version
         if image_version == "local":
-            image_version = "dev"
-        self.server_docker_base_image = "/fedml-device-image:" + image_version
+            image_version = "test"
+        self.server_docker_base_image = FedMLCloudServerManager._get_server_base_image(image_version)
         self.cloud_server_name = None
 
     @staticmethod
@@ -121,44 +126,52 @@ def start_cloud_server(self, packages_config):
         logging.info("start run with k8s: " + run_deployment_cmd)
         os.system(run_deployment_cmd)
 
-    def stop_cloud_server(self):
-        self.cloud_server_name = FedMLCloudServerManager.FEDML_CLOUD_SERVER_PREFIX + str(self.run_id) \
-                                 + "-" + str(self.edge_id)
-        self.server_docker_image = (
-                self.agent_config["docker_config"]["registry_server"]
-                + self.agent_config["docker_config"]["registry_dir"]
-                + self.server_docker_base_image
+    @staticmethod
+    def stop_cloud_server(run_id, server_id, agent_config):
+        cloud_server_name = FedMLCloudServerManager._get_cloud_server_name(run_id, server_id)
+        server_docker_image = (
+                agent_config["docker_config"]["registry_server"]
+                + agent_config["docker_config"]["registry_dir"]
+                + FedMLCloudServerManager._get_server_base_image(fedml.get_env_version())
         )
         delete_deployment_cmd = (
                 "export FEDML_AGGREGATOR_NAME="
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";export FEDML_AGGREGATOR_SVC="
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";export FEDML_AGGREGATOR_VERSION="
-                + self.version
+                + fedml.get_env_version()
                 + ';export FEDML_AGGREGATOR_IMAGE_PATH="'
-                + self.server_docker_image
+                + server_docker_image
                 + '"'
                 + ";export FEDML_CONF_ID="
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";export FEDML_DATA_PV_ID="
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";export FEDML_DATA_PVC_ID="
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";export FEDML_REGISTRY_SECRET_SUFFIX="
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";kubectl -n fedml-devops-aggregator-"
-                + self.version
+                + fedml.get_env_version()
                 + " delete deployment "
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";kubectl -n fedml-devops-aggregator-"
-                + self.version
+                + fedml.get_env_version()
                 + " delete svc "
-                + self.cloud_server_name
+                + cloud_server_name
                 + ";kubectl -n fedml-devops-aggregator-"
-                + self.version
+                + fedml.get_env_version()
                 + " delete secret secret-"
-                + self.cloud_server_name
+                + cloud_server_name
         )
         logging.info("stop run with k8s: " + delete_deployment_cmd)
         os.system(delete_deployment_cmd)
+
+    @staticmethod
+    def _get_server_base_image(version):
+        return f"{FedMLCloudServerManager.FEDML_SERVER_BASE_IMAGE}{version}"
+
+    @staticmethod
+    def _get_cloud_server_name(run_id, server_id):
+        return f"{FedMLCloudServerManager.FEDML_CLOUD_SERVER_PREFIX}{run_id}-{server_id}"

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/master/server_login.py

@@ -39,6 +39,6 @@ def logout():
     master_agent = FedMLLaunchMasterAgent()
     if args.type == 'login':
         master_agent.login(args.api_key, api_key=args.api_key, device_id=args.device_id,
-                           os_name=args.os_name, role=args.role)
+                           os_name=args.os_name, role=args.role, runner_cmd=args.runner_cmd)
     else:
         master_agent.logout()