Update decorators

python/fedml/api/__init__.py

@@ -13,6 +13,7 @@
             print(f"run status {run_log_result.run_status}, total log nums {log_result.total_log_lines}, "
                   f"total log pages {log_result.total_log_pages}, log list {log_result.log_line_list}")
 """
+from io import BytesIO
 from typing import List, Tuple
 
 from fedml.api.constants import RunStatus
@@ -27,6 +28,7 @@
 from fedml.computing.scheduler.model_scheduler.device_server_constants import ServerConstants
 from fedml.computing.scheduler.model_scheduler.device_client_constants import ClientConstants
 
+from .fedml_response import ResponseCode
 
 def fedml_login(api_key: str = None):
     """
@@ -182,10 +184,11 @@ def cluster_killall(api_key=None) -> bool:
 
 
 def upload(data_path, api_key=None, tag_list=[], service="R2", name=None, description=None, metadata=None, show_progress=False,
-           out_progress_to_err=True, progress_desc=None) -> FedMLResponse:
+           out_progress_to_err=True, progress_desc=None, byte_data: BytesIO = None, encrypted_api_key_flag=False) -> FedMLResponse:
     return storage.upload(data_path=data_path, api_key=api_key, name=name, description=description, tag_list =tag_list,
                           service=service, progress_desc=progress_desc, show_progress=show_progress,
-                          out_progress_to_err=out_progress_to_err, metadata=metadata)
+                          out_progress_to_err=out_progress_to_err, metadata=metadata, byte_data=byte_data,
+                          encrypted_api_key_flag=encrypted_api_key_flag)
 
 def get_storage_user_defined_metadata(data_name, api_key=None) -> FedMLResponse:
     return storage.get_user_metadata(data_name=data_name, api_key=api_key)

python/fedml/api/modules/utils.py

@@ -9,33 +9,33 @@
 FEDML_MLOPS_BUILD_PRE_IGNORE_LIST = 'dist-packages,client-package.zip,server-package.zip,__pycache__,*.pyc,*.git'
 
 
-def fedml_login(api_key):
-    api_key_is_valid, api_key = _check_api_key(api_key=api_key)
+def fedml_login(api_key, encrypted_api_key_flag=False):
+    api_key_is_valid, api_key = _check_api_key(api_key=api_key, encrypted_api_key_flag=encrypted_api_key_flag)
     if api_key_is_valid:
         return 0, api_key
 
     return -1, api_key
 
 
-def _check_api_key(api_key=None):
+def _check_api_key(api_key=None, encrypted_api_key_flag=False):
     if api_key is None or api_key == "":
         saved_api_key = get_api_key()
         if saved_api_key is None or saved_api_key == "":
             api_key = click.prompt("FedML® Launch API Key is not set yet, please input your API key")
         else:
             api_key = saved_api_key
 
-    is_valid_heartbeat = FedMLResourceManager.get_instance().check_heartbeat(api_key)
+    is_valid_heartbeat = FedMLResourceManager.get_instance().check_heartbeat(api_key, encrypted_api_key_flag)
     if not is_valid_heartbeat:
         return False, api_key
     else:
         save_api_key(api_key)
         return True, api_key
 
 
-def authenticate(api_key):
+def authenticate(api_key, encrypted_api_key_flag=False):
 
-    error_code, api_key = fedml_login(api_key)
+    error_code, api_key = fedml_login(api_key, encrypted_api_key_flag)
 
     # Exit if not able to authenticate successfully
     if error_code:

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

@@ -16,7 +16,7 @@
 from ....core.mlops.mlops_utils import MLOpsUtils
 from ..scheduler_core.log_manager import LogsManager
 from ..scheduler_core.metrics_manager import MetricsManager
-from fedml.utils.debugging import debug
+from fedml.utils.decorators import debug
 from ..scheduler_core.status_center import JobStatus
 from ..scheduler_core.compute_cache_manager import ComputeCacheManager
 from multiprocessing import Process, Queue

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

@@ -251,7 +251,7 @@ def get_user_url():
 
     @staticmethod
     def get_dataset_url():
-        create_dataset_url = "{}/fedmlOpsServer/api/v1/cli/dataset".format(
+        create_dataset_url = "{}/system/api/v1/cli/storage".format(
             ServerConstants.get_mlops_url())
         return create_dataset_url
 
@@ -271,13 +271,13 @@ def get_complete_multipart_upload_url():
 
     @staticmethod
     def list_dataset_url():
-        list_dataset_url = "{}/fedmlOpsServer/api/v1/cli/dataset/list".format(
+        list_dataset_url = "{}/system/api/v1/cli/storage/list".format(
             ServerConstants.get_mlops_url())
         return list_dataset_url
 
     @staticmethod
     def get_dataset_metadata_url():
-        get_dataset_metadata_url = "{}/fedmlOpsServer/api/v1/cli/dataset/meta".format(
+        get_dataset_metadata_url = "{}/system/api/v1/cli/storage/meta".format(
             ServerConstants.get_mlops_url())
         return get_dataset_metadata_url
 

python/fedml/computing/scheduler/scheduler_entry/resource_manager.py

@@ -15,9 +15,10 @@ def __init__(self):
     def get_instance():
         return FedMLResourceManager()
 
-    def check_heartbeat(self, api_key):
+    def check_heartbeat(self, api_key, encrypted_api_key_flag=False):
         heartbeat_url = ServerConstants.get_heartbeat_url()
-        heartbeat_api_headers = {'Content-Type': 'application/json', 'Connection': 'close'}
+        heartbeat_api_headers = {'Content-Type': 'application/json', 'Connection': 'close',
+                                 'Encrypted': str(encrypted_api_key_flag)}
         heartbeat_json = {
             "apiKey": api_key
         }

python/fedml/utils/debugging.py → python/fedml/utils/decorators.py

@@ -1,9 +1,10 @@
 from typing import Callable
 import os
+import time
 import functools
 
-def debug(_func: Callable=None, *, output_file="output.txt"):
 
+def debug(_func: Callable = None, *, output_file="output.txt"):
     def decorator(func: Callable):
 
         @functools.wraps(func)
@@ -36,3 +37,18 @@ def wrapper(*args, **kwargs):
         return decorator
 
     return decorator(_func)
+
+
+def timeit(func: Callable):
+    """Print the runtime of the decorated function"""
+    functools.wraps(func)
+
+    def wrapper(*args, **kwargs):
+        start = time.perf_counter()
+        value = func(*args, **kwargs)
+        end = time.perf_counter()
+        run_time = end - start
+        print(f"Finished {func.__name__!r} in {run_time:.4f} seconds")
+        return value
+
+    return wrapper