Merge pull request #2167 from FedML-AI/raphael/hotfix-inference-port

[Deploy] Finalize Gateway/Proxy port during device binding.

Author: Raphael-Jin

python/fedml/api/__init__.py

@@ -24,6 +24,8 @@
 from fedml.computing.scheduler.scheduler_entry.cluster_manager import FedMLClusterModelList
 from fedml.computing.scheduler.scheduler_entry.run_manager import FedMLRunStartedModel, FedMLGpuDevices, \
     FedMLRunModelList, FeatureEntryPoint
+from fedml.computing.scheduler.model_scheduler.device_server_constants import ServerConstants
+from fedml.computing.scheduler.model_scheduler.device_client_constants import ClientConstants
 
 
 def fedml_login(api_key: str = None):
@@ -209,16 +211,18 @@ def fedml_build(platform, type, source_folder, entry_point, config_folder, dest_
     return build.build(platform, type, source_folder, entry_point, config_folder, dest_folder, ignore)
 
 
-def login(api_key, computing, server, supplier):
-    device_bind(api_key, computing, server, supplier)
+def login(api_key, computing, server, supplier,
+          master_inference_gateway_port: int = ServerConstants.MODEL_INFERENCE_DEFAULT_PORT,
+          worker_inference_proxy_port: int = ClientConstants.LOCAL_CLIENT_API_PORT):
+    device_bind(api_key, computing, server, supplier, master_inference_gateway_port, worker_inference_proxy_port)
 
 
 def logout(computing, server):
     device_unbind(computing, server)
 
 
-def device_bind(api_key, computing, server, supplier):
-    device.bind(api_key, computing, server, supplier)
+def device_bind(api_key, computing, server, supplier, master_inference_gateway_port, worker_inference_proxy_port):
+    device.bind(api_key, computing, server, supplier, master_inference_gateway_port, worker_inference_proxy_port)
 
 
 def device_unbind(computing, server):

python/fedml/api/modules/device.py

@@ -10,14 +10,18 @@
 from fedml.computing.scheduler.comm_utils.constants import SchedulerConstants
 from fedml.computing.scheduler.comm_utils.run_process_utils import RunProcessUtils
 from fedml.computing.scheduler.master.server_constants import ServerConstants
+from fedml.computing.scheduler.model_scheduler.device_server_constants import ServerConstants as DeviceServerConstants
 from fedml.computing.scheduler.master.server_login import logout as server_logout
 from fedml.computing.scheduler.slave.client_constants import ClientConstants
+from fedml.computing.scheduler.model_scheduler.device_client_constants import ClientConstants as DeviceClientConstants
 from fedml.computing.scheduler.slave.client_login import logout as client_logout
 from fedml.computing.scheduler.scheduler_entry.resource_manager import FedMLResourceManager
 
 
 def bind(
-        api_key, computing, server, supplier
+        api_key, computing, server, supplier,
+        master_inference_gateway_port=DeviceServerConstants.MODEL_INFERENCE_DEFAULT_PORT,
+        worker_inference_proxy_port=DeviceClientConstants.LOCAL_CLIENT_API_PORT
 ):
     userid = api_key
     runner_cmd = "{}"
@@ -43,13 +47,13 @@ def bind(
     _bind(
         userid, computing, server,
         api_key, role, runner_cmd, device_id, os_name,
-        docker)
+        docker, master_inference_gateway_port, worker_inference_proxy_port)
 
 
 def _bind(
         userid, computing, server,
         api_key, role, runner_cmd, device_id, os_name,
-        docker):
+        docker, master_inference_gateway_port, worker_inference_proxy_port):
     fedml.load_env()
     if os.getenv(ModuleConstants.ENV_FEDML_INFER_HOST) is None:
         fedml.set_env_kv(ModuleConstants.ENV_FEDML_INFER_HOST, SchedulerConstants.REDIS_INFER_HOST)
@@ -60,6 +64,9 @@ def _bind(
     if os.getenv(ModuleConstants.ENV_FEDML_INFER_REDIS_PASSWORD) is None:
         fedml.set_env_kv(ModuleConstants.ENV_FEDML_INFER_REDIS_PASSWORD, SchedulerConstants.REDIS_PASSWORD)
 
+    fedml.set_env_kv(DeviceServerConstants.ENV_MASTER_INFERENCE_PORT_KEY, str(master_inference_gateway_port))
+    fedml.set_env_kv(DeviceClientConstants.ENV_CLIENT_PROXY_PORT_KEY, str(worker_inference_proxy_port))
+
     url = fedml._get_backend_service()
     platform_name = platform.system()
     docker_config_text = None

python/fedml/cli/modules/login.py

@@ -4,6 +4,8 @@
 
 import fedml.api
 from fedml.api.modules.utils import authenticate
+from fedml.computing.scheduler.model_scheduler.device_server_constants import ServerConstants
+from fedml.computing.scheduler.model_scheduler.device_client_constants import ClientConstants
 
 
 @click.command("login", help="Login the FedML® Nexus AI Platform")
@@ -51,9 +53,25 @@
     default=80,
     help="The port for local on-premise Nexus AI Platform.",
 )
+@click.option(
+    "--master_inference_gateway_port",
+    "-mgp",
+    type=int,
+    default=ServerConstants.MODEL_INFERENCE_DEFAULT_PORT,
+    help="The port for master inference gateway.",
+)
+@click.option(
+    "--worker_inference_proxy_port",
+    "-wpp",
+    type=int,
+    default=ClientConstants.LOCAL_CLIENT_API_PORT,
+    help="The port for worker inference proxy.",
+)
 def fedml_login(
         api_key, version, compute_node, server, provider, deploy_worker_num,
-        local_on_premise_platform, local_on_premise_platform_port):
+        local_on_premise_platform, local_on_premise_platform_port,
+        master_inference_gateway_port, worker_inference_proxy_port
+):
     fedml.set_env_version(version)
     fedml.set_local_on_premise_platform_host(local_on_premise_platform)
     fedml.set_local_on_premise_platform_port(local_on_premise_platform_port)
@@ -66,4 +84,4 @@ def fedml_login(
         print(f"Maybe you are using account id to login, we will try to login with account {api_key}.")
         pass
     os.environ["FEDML_MODEL_WORKER_NUM"] = str(deploy_worker_num)
-    fedml.api.login(api_key, compute_node, server, provider)
+    fedml.api.login(api_key, compute_node, server, provider, master_inference_gateway_port, worker_inference_proxy_port)

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

@@ -74,6 +74,7 @@ class ClientConstants(object):
     K8S_DEPLOYMENT_SLAVE_MOUNT_HOME_DIR = "/home/fedml/fedml-client"
 
     LOCAL_CLIENT_API_PORT = 22030
+    ENV_CLIENT_PROXY_PORT_KEY = "FEDML_WORKER_INFERENCE_PROXY_PORT"
 
     INFERENCE_HTTP_PORT = 8000
     INFERENCE_GRPC_PORT = 8001
@@ -457,6 +458,14 @@ def get_public_ip():
             logging.info("Failed to get public ip: {}".format(e))
         return ip
 
+    @staticmethod
+    def get_inference_worker_proxy_port() -> int:
+        # Use dotenv to load the environment variables
+        fedml.load_env()
+        worker_proxy_port = int(os.getenv(ClientConstants.ENV_CLIENT_PROXY_PORT_KEY,
+                                      default=ClientConstants.LOCAL_CLIENT_API_PORT))
+        return worker_proxy_port
+
     @staticmethod
     def check_process_is_running(process_id):
         for proc in psutil.process_iter():

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

@@ -103,6 +103,7 @@ class ServerConstants(object):
 
     AUTO_DETECT_PUBLIC_IP = "auto_detect_public_ip"
     MODEL_INFERENCE_DEFAULT_PORT = 2203
+    ENV_MASTER_INFERENCE_PORT_KEY = "FEDML_MASTER_INFERENCE_GATEWAY_PORT"
     MODEL_CACHE_KEY_EXPIRE_TIME = 1 * 10
 
     INFERENCE_REQUEST_TIMEOUT_KEY = "request_timeout_sec"
@@ -348,6 +349,15 @@ def get_runner_infos():
             logging.error(f"Failed to parse runner info: {e}")
         return runner_info
 
+    @staticmethod
+    def get_inference_master_gateway_port():
+        # Use dotenv to load the environment variables
+        fedml.load_env()
+        master_inference_port = int(os.getenv(ServerConstants.ENV_MASTER_INFERENCE_PORT_KEY,
+                                            default=ServerConstants.MODEL_INFERENCE_DEFAULT_PORT))
+        return master_inference_port
+
+
     @staticmethod
     def save_runner_infos(unique_device_id, edge_id, run_id=None):
         local_pkg_data_dir = ServerConstants.get_data_dir()

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

@@ -115,7 +115,7 @@ def run_impl(
             message_center=self.message_center)
 
         # start unified inference gateway process if not started
-        FedMLDeployMasterJobRunner.start_device_inference_gateway(inference_port=inference_port)
+        FedMLDeployMasterJobRunner.start_device_inference_gateway()
 
         # start inference monitor process
         FedMLDeployMasterJobRunner.stop_device_inference_monitor(
@@ -144,7 +144,7 @@ def run_impl(
                 # No device is added, updated or removed
                 logging.info("No device is added, updated or removed. No action needed for reconciliation.")
                 ip = GeneralConstants.get_ip_address(self.request_json)
-                master_port = os.getenv("FEDML_MASTER_PORT", None)
+                master_port = ServerConstants.get_inference_master_gateway_port()
                 if master_port is not None:
                     inference_port = int(master_port)
                 model_inference_port = inference_port
@@ -299,9 +299,7 @@ def process_deployment_result_message(self, topic=None, payload=None):
                 else:
                     # This is the last worker that failed, so we should continue to "ABORTED" status
                     model_config_parameters = self.request_json["parameters"]
-                    inference_port = model_config_parameters.get("server_internal_port",
-                                                                 ServerConstants.MODEL_INFERENCE_DEFAULT_PORT)
-                    inference_port_external = model_config_parameters.get("server_external_port", inference_port)
+                    inference_port_external = ServerConstants.get_inference_master_gateway_port()
                     ip = GeneralConstants.get_ip_address(self.request_json)
                     if ip.startswith("http://") or ip.startswith("https://"):
                         model_inference_url = "{}/inference/{}".format(ip, end_point_id)
@@ -369,12 +367,7 @@ def process_deployment_result_message(self, topic=None, payload=None):
             """
             When all the devices have finished the add / delete / update operation
             """
-            # Generate one unified inference api
-            # Note that here we use the gateway port instead of the inference port that is used by the slave device
-            model_config_parameters = request_json["parameters"]
-            inference_port = model_config_parameters.get("server_internal_port",
-                                                         ServerConstants.MODEL_INFERENCE_DEFAULT_PORT)
-            inference_port_external = model_config_parameters.get("server_external_port", inference_port)
+            inference_port_external = ServerConstants.get_inference_master_gateway_port()
             ip = GeneralConstants.get_ip_address(request_json)
 
             if ip.startswith("http://") or ip.startswith("https://"):
@@ -457,12 +450,10 @@ def cleanup_runner_process(self, run_id):
         ServerConstants.cleanup_run_process(run_id, not_kill_subprocess=True)
 
     @staticmethod
-    def start_device_inference_gateway(inference_port=ServerConstants.MODEL_INFERENCE_DEFAULT_PORT):
+    def start_device_inference_gateway():
         # start unified inference server
         python_program = get_python_program()
-        master_port = os.getenv("FEDML_MASTER_PORT", None)
-        if master_port is not None:
-            inference_port = int(master_port)
+        inference_port = ServerConstants.get_inference_master_gateway_port()
         if not ServerConstants.is_running_on_k8s():
             logging.info(f"start the model inference gateway...")
             inference_gw_cmd = "fedml.computing.scheduler.model_scheduler.device_model_inference:api"
@@ -543,7 +534,7 @@ def recover_inference_and_monitor():
                 if not is_activated:
                     continue
 
-                FedMLDeployMasterJobRunner.start_device_inference_gateway(inference_port=inference_port)
+                FedMLDeployMasterJobRunner.start_device_inference_gateway()
 
                 FedMLDeployMasterJobRunner.stop_device_inference_monitor(
                     run_id, end_point_name, model_id, model_name, model_version)
@@ -757,9 +748,7 @@ def parse_model_run_params(running_json):
         model_version = model_config["model_version"]
         model_config_parameters = running_json.get("parameters", {})
 
-        inference_port = model_config_parameters.get("server_internal_port",  # Internal port is for the gateway
-                                                     ServerConstants.MODEL_INFERENCE_DEFAULT_PORT)
-        inference_port_external = model_config_parameters.get("server_external_port", inference_port)
+        inference_port = ServerConstants.get_inference_master_gateway_port()
 
         return run_id, end_point_name, token, user_id, user_name, device_ids, device_objs, model_config, model_name, \
             model_id, model_storage_url, scale_min, scale_max, inference_engine, model_is_from_open, \

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

@@ -79,6 +79,9 @@ def _init_extra_items(self):
 
         client_api_cmd = "fedml.computing.scheduler.model_scheduler.device_client_api:api"
         client_api_pids = RunProcessUtils.get_pid_from_cmd_line(client_api_cmd)
+
+        worker_proxy_port = ClientConstants.get_inference_worker_proxy_port()
+
         if client_api_pids is None or len(client_api_pids) <= 0:
             # Start local API services
             cur_dir = os.path.dirname(__file__)
@@ -88,7 +91,7 @@ def _init_extra_items(self):
                 "{} -m uvicorn {} --host 0.0.0.0 --port {} --reload --reload-delay 3 --reload-dir {} "
                 "--log-level critical".format(
                     python_program, client_api_cmd,
-                    ClientConstants.LOCAL_CLIENT_API_PORT, fedml_base_dir
+                    worker_proxy_port, fedml_base_dir
                 ),
                 should_capture_stdout=False,
                 should_capture_stderr=False