Add large file download warning and add AzCopy example to upload warning (Azure#28125)

diondrapeck · web-flow · commit a7c8aa549ce5 · 2023-01-09T16:32:14.000-08:00
* fix bug by replacing 'working_dir' with 'path'

* Add file size warning for gen2 uploads

* Add file size warning for gen2 uploads

* Add upload/download warning with AzCopy example for files &gt; 100MB

* Update file size warning message

* Revert unintended change

* Address pylint and cspell errors

* Address cspell errors

* Make sure proper cloud endpoint is used in warning for URL
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_blob_storage_helper.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_blob_storage_helper.py
@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-# pylint: disable=client-accepts-api-version-keyword,too-many-instance-attributes,client-method-missing-type-annotations,missing-client-constructor-parameter-kwargs
+# pylint: disable=client-accepts-api-version-keyword,too-many-instance-attributes,client-method-missing-type-annotations,missing-client-constructor-parameter-kwargs,logging-format-interpolation
 
 import logging
 import os
@@ -31,6 +31,7 @@
     upload_directory,
     upload_file,
 )
+from azure.ai.ml._azure_environments import _get_cloud_details
 from azure.ai.ml.constants._common import STORAGE_AUTH_MISMATCH_ERROR
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, MlException, ValidationException
 from azure.core.exceptions import ResourceNotFoundError
@@ -44,6 +45,7 @@
 
 class BlobStorageClient:
     def __init__(self, credential: str, account_url: str, container_name: Optional[str] = None):
+        self.account_name = account_url.split(".")[0].split("//")[1]
         self.service_client = BlobServiceClient(account_url=account_url, credential=credential)
         self.upload_to_root_container = None
         if container_name:
@@ -90,8 +92,11 @@ def upload(
             # warn if large file (> 100 MB)
             file_size, _ = get_directory_size(source)
             file_size_in_mb = file_size / 10**6
+            cloud = _get_cloud_details()
+            cloud_endpoint = cloud['storage_endpoint']  # make sure proper cloud endpoint is used
+            full_storage_url = f"https://{self.account_name}.blob.{cloud_endpoint}/{self.container}/{dest}"
             if file_size_in_mb > 100:
-                module_logger.warning(FILE_SIZE_WARNING)
+                module_logger.warning(FILE_SIZE_WARNING.format(source=source, destination=full_storage_url))
 
             # start upload
             if os.path.isdir(source):
@@ -212,6 +217,7 @@ def download(
         """
         try:
             my_list = list(self.container_client.list_blobs(name_starts_with=starts_with, include="metadata"))
+            download_size_in_mb = 0
             for item in my_list:
                 blob_name = item.name[len(starts_with) :].lstrip("/") or Path(starts_with).name
                 target_path = Path(destination, blob_name).resolve()
@@ -221,6 +227,16 @@ def download(
                     continue
 
                 blob_content = self.container_client.download_blob(item)
+
+                # check if total size of download has exceeded 100 MB
+                # make sure proper cloud endpoint is used
+                cloud = _get_cloud_details()
+                cloud_endpoint = cloud['storage_endpoint']
+                full_storage_url = f"https://{self.account_name}.blob.{cloud_endpoint}/{self.container}/{starts_with}"
+                download_size_in_mb += (blob_content.size / 10**6)
+                if download_size_in_mb > 100:
+                    module_logger.warning(FILE_SIZE_WARNING.format(source=full_storage_url, destination=destination))
+
                 blob_content = blob_content.content_as_bytes(max_concurrency)
                 target_path.parent.mkdir(parents=True, exist_ok=True)
                 with target_path.open("wb") as file:
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_constants.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_constants.py
@@ -28,9 +28,9 @@
 CHANGED_ASSET_PATH_MSG_NO_PERSONAL_DATA = "The code asset is already linked to an asset."
 EMPTY_DIRECTORY_ERROR = "Directory {0} is empty. path or local_path must be a non-empty directory."
 FILE_SIZE_WARNING = (
-    "Your file exceeds 100 MB. If you experience low upload speeds or latency, "
-    "we recommend using the AzCopy tool for this file transfer. "
-    "See https://docs.microsoft.com/azure/storage/common/storage-use-azcopy-v10 for more information."
+    "Your file exceeds 100 MB. If you experience low speeds, latency, or broken connections, we recommend using "
+    "the AzCopyv10 tool for this file transfer.\n\nExample: azcopy copy '{source}' '{destination}' "  # cspell:disable-line
+    "\n\nSee https://docs.microsoft.com/azure/storage/common/storage-use-azcopy-v10 for more information."
 )
 INVALID_MLTABLE_METADATA_SCHEMA_MSG = "Invalid MLTable metadata schema"
 INVALID_MLTABLE_METADATA_SCHEMA_ERROR = (
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_gen2_storage_helper.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_gen2_storage_helper.py
@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-# pylint: disable=client-accepts-api-version-keyword,too-many-instance-attributes,client-method-missing-type-annotations,missing-client-constructor-parameter-kwargs
+# pylint: disable=client-accepts-api-version-keyword,too-many-instance-attributes,client-method-missing-type-annotations,missing-client-constructor-parameter-kwargs,logging-format-interpolation
 
 import logging
 import os
@@ -14,15 +14,17 @@
 
 from colorama import Fore
 
-from azure.ai.ml._artifacts._constants import UPLOAD_CONFIRMATION
+from azure.ai.ml._artifacts._constants import UPLOAD_CONFIRMATION, FILE_SIZE_WARNING
 from azure.ai.ml._utils._asset_utils import (
     AssetNotChangedError,
     IgnoreFile,
     _build_metadata_dict,
     generate_asset_id,
+    get_directory_size,
     upload_directory,
     upload_file,
 )
+from azure.ai.ml._azure_environments import _get_cloud_details
 from azure.ai.ml.constants._common import STORAGE_AUTH_MISMATCH_ERROR
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, MlException, ValidationException
 from azure.core.exceptions import ResourceExistsError
@@ -34,6 +36,7 @@
 class Gen2StorageClient:
     def __init__(self, credential: str, file_system: str, account_url: str):
         service_client = DataLakeServiceClient(account_url=account_url, credential=credential)
+        self.account_name = account_url.split(".")[0].split("//")[1]
         self.file_system = file_system
         self.file_system_client = service_client.get_file_system_client(file_system=file_system)
         try:
@@ -77,6 +80,16 @@ def upload(
             # configure progress bar description
             msg = Fore.GREEN + f"Uploading {formatted_path}"
 
+            # warn if large file (> 100 MB)
+            file_size, _ = get_directory_size(source)
+            file_size_in_mb = file_size / 10**6
+
+            cloud = _get_cloud_details()
+            cloud_endpoint = cloud['storage_endpoint']  # make sure proper cloud endpoint is used
+            full_storage_url = f"https://{self.account_name}.dfs.{cloud_endpoint}/{self.file_system}/{dest}"
+            if file_size_in_mb > 100:
+                module_logger.warning(FILE_SIZE_WARNING.format(source=source, destination=full_storage_url))
+
             # start upload
             self.directory_client = self.file_system_client.get_directory_client(asset_id)
             self.check_blob_exists()
@@ -159,6 +172,7 @@ def download(self, starts_with: str, destination: str = Path.home()) -> None:
         prefix `starts_with` to the destination folder."""
         try:
             mylist = self.file_system_client.get_paths(path=starts_with)
+            download_size_in_mb = 0
             for item in mylist:
                 file_name = item.name[len(starts_with) :].lstrip("/") or Path(starts_with).name
 
@@ -168,6 +182,15 @@ def download(self, starts_with: str, destination: str = Path.home()) -> None:
 
                 target_path = Path(destination, file_name)
                 file_client = self.file_system_client.get_file_client(item.name)
+
+                # check if total size of download has exceeded 100 MB
+                cloud = _get_cloud_details()
+                cloud_endpoint = cloud['storage_endpoint']  # make sure proper cloud endpoint is used
+                full_storage_url = f"https://{self.account_name}.dfs.{cloud_endpoint}/{self.file_system}/{starts_with}"
+                download_size_in_mb += (file_client.get_file_properties().size / 10**6)
+                if download_size_in_mb > 100:
+                    module_logger.warning(FILE_SIZE_WARNING.format(source=full_storage_url, destination=destination))
+
                 file_content = file_client.download_file().readall()
                 try:
                     os.makedirs(str(target_path.parent), exist_ok=True)
