Merge pull request #2059 from FedML-AI/dev/v0.7.0

fedml-alex · web-flow · commit 60d1deb4d904 · 2024-04-29T21:41:07.000+08:00
Dev/v0.7.0
diff --git a/python/fedml/api/__init__.py b/python/fedml/api/__init__.py
@@ -194,7 +194,7 @@ def get_storage_metadata(data_name, api_key=None) -> FedMLResponse:
     return storage.get_metadata(api_key=api_key, data_name=data_name)
 
 
-def list_storage_obects(api_key=None) -> FedMLResponse:
+def list_storage_objects(api_key=None) -> FedMLResponse:
     return storage.list_objects(api_key=api_key)
 
 
diff --git a/python/fedml/api/modules/storage.py b/python/fedml/api/modules/storage.py
@@ -10,20 +10,19 @@
 from fedml.api.fedml_response import FedMLResponse, ResponseCode
 
 
-# Todo (alaydshah): Add file size
 class StorageMetadata(object):
     def __init__(self, data: dict):
         self.dataName = data.get("datasetName", None)
         self.description = data.get("description", None)
         self.tags = data.get("description", None)
         self.createdAt = data.get("createTime", None)
         self.updatedAt = data.get("updateTime", None)
+        self.size = _get_size(data.get("fileSize",None))
 
 
-# Todo (alaydshah): Add file size while creating objects. Store service name in metadata
+# Todo (alaydshah): Store service name in metadata
 # Todo (alaydshah): If data already exists, don't upload again. Instead suggest to use update command
 
-
 def upload(data_path, api_key, name, description, service, show_progress, out_progress_to_err, progress_desc,
            metadata) -> FedMLResponse:
     api_key = authenticate(api_key)
@@ -32,6 +31,9 @@ def upload(data_path, api_key, name, description, service, show_progress, out_pr
 
     if user_id is None:
         return FedMLResponse(code=ResponseCode.FAILURE, message=message)
+    
+    if(not _check_data_path(data_path)):
+        return FedMLResponse(code=ResponseCode.FAILURE,message="Invalid data path")
 
     archive_path, message = _archive_data(data_path)
     if not archive_path:
@@ -41,6 +43,7 @@ def upload(data_path, api_key, name, description, service, show_progress, out_pr
     name = os.path.splitext(os.path.basename(archive_path))[0] if name is None else name
     file_name = name + ".zip"
     dest_path = os.path.join(user_id, file_name)
+    file_size = os.path.getsize(archive_path)
 
     file_uploaded_url = store.upload_file_with_progress(src_local_path=archive_path, dest_s3_path=dest_path,
                                                         show_progress=show_progress,
@@ -53,7 +56,7 @@ def upload(data_path, api_key, name, description, service, show_progress, out_pr
     json_data = {
         "datasetName": name,
         "description": description,
-        "fileSize": "",
+        "fileSize": file_size,
         "fileUrl": file_uploaded_url,
         "tagNameList": [],
     }
@@ -228,6 +231,11 @@ def _get_storage_service(service):
     else:
         raise NotImplementedError(f"Service {service} not implemented")
 
+def _check_data_path(data_path):
+    if not os.path.isdir(data_path) or not os.path.isfile(data_path):
+        return False
+    return True
+
 
 def _archive_data(data_path: str) -> (str, str):
     src_local_path = os.path.abspath(data_path)
@@ -350,3 +358,20 @@ def _get_data_from_response(message: str, response: requests.Response) -> (Respo
             return ResponseCode.FAILURE, message, None
 
     return ResponseCode.SUCCESS, "Successfully parsed data from response", data
+
+def _get_size(size_in_bytes:str)->str:
+    size_str = ""
+    if(size_in_bytes):
+        size = int(size_in_bytes)
+        size_in_gb = size / (1024 * 1024 * 1024)
+        size_in_mb = size / (1024 * 1024)
+        size_in_kb = size / 1024
+        if(size_in_gb >= 1):
+            size_str = f"{size_in_gb:.2f} GB"
+        elif(size_in_mb >= 1):
+            size_str = f"{size_in_mb:.2f} MB"
+        elif(size_in_kb >= 1):
+            size_str = f"{size_in_kb:.2f} KB"
+        else:
+            size_str = f"{size} B"
+    return size_str 
diff --git a/python/fedml/cli/modules/storage.py b/python/fedml/cli/modules/storage.py
@@ -90,16 +90,16 @@ def upload(data_path: str, name: str, user_metadata: str, description: str, vers
 )
 def list_data(version, api_key):
     fedml.set_env_version(version)
-    response = fedml.api.list_storage_obects(api_key=api_key)
+    response = fedml.api.list_storage_objects(api_key=api_key)
     if response.code == ResponseCode.SUCCESS:
         click.echo(f"Successfully fetched list of stored objects:")
         if not response.data:
             click.echo(f"No stored objects found for account linked with apikey: {api_key}")
             return
-        object_list_table = PrettyTable(["Data Name", "Description", "Created At", "Updated At"])
+        object_list_table = PrettyTable(["Data Name", "Data Size", "Description", "Created At", "Updated At"])
         for stored_object in response.data:
             object_list_table.add_row(
-                [stored_object.dataName, stored_object.description, stored_object.createdAt, stored_object.updatedAt])
+                [stored_object.dataName, stored_object.size, stored_object.description, stored_object.createdAt, stored_object.updatedAt])
         click.echo(object_list_table)
     else:
         click.echo(f"Failed to list stored objects for account linked with apikey {api_key}. "
@@ -157,8 +157,8 @@ def get_metadata(data_name, version, api_key):
             return
         click.echo(f"Successfully fetched metadata for object {data_name}:")
         # Todo (alaydshah): Add file size and tags
-        metadata_table = PrettyTable(["Data Name", "Description", "Created At", "Updated At"])
-        metadata_table.add_row([metadata.dataName, metadata.description, metadata.createdAt, metadata.updatedAt])
+        metadata_table = PrettyTable(["Data Name","Data Size","Description", "Created At", "Updated At"])
+        metadata_table.add_row([metadata.dataName,metadata.size,metadata.description, metadata.createdAt, metadata.updatedAt])
         click.echo(metadata_table)
         click.echo("")
     else:
