Use azcopy to download test data (#4221)

Use azcopy from download_e2e_test_data.py, add helper function for downloading azcopy.
Update download_test_data.py to use helper function.

orttraining/tools/ci_test/download_e2e_test_data.py

@@ -5,17 +5,26 @@
 import argparse
 import hashlib
 import os
+import shutil
+import subprocess
 import sys
 import tempfile
 import urllib.request
 import zipfile
 
+SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
+REPO_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..", ".."))
+
+sys.path.append(os.path.join(REPO_DIR, "tools", "python"))
+
+import get_azcopy  # noqa: E402
+
 # update these if the E2E test data changes
 ARCHIVE_BLOB_URL = "https://onnxruntimetestdata.blob.core.windows.net/training/onnxruntime_training_data.zip?snapshot=2020-06-15T23:17:35.8314853Z"
 ARCHIVE_SHA256_DIGEST = "B01C169B6550D1A0A6F1B4E2F34AE2A8714B52DBB70AC04DA85D371F691BDFF9"
 
-def _download(url, local_path):
-  urllib.request.urlretrieve(url, local_path)
+def _download(azcopy_path, url, local_path):
+  subprocess.run([azcopy_path, "cp", "--log-level", "NONE", url, local_path], check=True)
 
 def _get_sha256_digest(file_path):
   alg = hashlib.sha256()
@@ -36,22 +45,19 @@ def _check_file_sha256_digest(path, expected_digest):
     raise RuntimeError(
         "SHA256 digest mismatch, expected: {}, actual: {}".format(expected_digest.lower(), actual_digest.lower()))
 
-def _extract_archive(archive_path, target_dir):
-  with zipfile.ZipFile(archive_path) as archive:
-    archive.extractall(target_dir)
-
 def main():
   parser = argparse.ArgumentParser(description="Downloads training end-to-end test data.")
   parser.add_argument("target_dir", help="The test data destination directory.")
   args = parser.parse_args()
 
-  with tempfile.TemporaryDirectory() as temp_dir:
+  with tempfile.TemporaryDirectory() as temp_dir, \
+       get_azcopy.get_azcopy() as azcopy_path:
     archive_path = os.path.join(temp_dir, "archive.zip")
     print("Downloading E2E test data from '{}'...".format(ARCHIVE_BLOB_URL))
-    _download(ARCHIVE_BLOB_URL, archive_path)
+    _download(azcopy_path, ARCHIVE_BLOB_URL, archive_path)
     _check_file_sha256_digest(archive_path, ARCHIVE_SHA256_DIGEST)
     print("Extracting to '{}'...".format(args.target_dir))
-    _extract_archive(archive_path, args.target_dir)
+    shutil.unpack_archive(archive_path, args.target_dir)
     print("Done.")
 
 if __name__ == "__main__":

tools/ci_build/github/azure-pipelines/linux-ort-srv-nightly-pipeline.yml

@@ -14,14 +14,6 @@ jobs:
       continueOnError: true
       condition: always()
 
-    - task: CmdLine@2
-      displayName: 'Download azcopy'
-      inputs:
-        script: |
-          curl -so azcopy.tar.gz -L 'https://aka.ms/downloadazcopy-v10-linux'
-          tar -zxvf azcopy.tar.gz --strip 1
-        workingDirectory: $(Build.BinariesDirectory)
-
     - task: PythonScript@0
       displayName: 'Download test data'
       inputs:

tools/ci_build/github/azure-pipelines/nuget/templates/mac-set-variables-and-download.yml

@@ -1,12 +1,4 @@
 steps:
-- task: CmdLine@2
-  displayName: 'Download azcopy'
-  inputs:
-    script: |
-      curl -so azcopy.tar.gz -L 'https://aka.ms/downloadazcopy-v10-mac'
-      tar -zxvf azcopy.tar.gz --strip 1
-    workingDirectory: $(Build.BinariesDirectory)
-
 - task: PythonScript@0
   displayName: 'Download test data'
   inputs:

tools/ci_build/github/azure-pipelines/nuget/templates/win-download-test-data.yml

@@ -1,4 +1,4 @@
-# Assumes AZCopy and Python download is already done
+# Assumes Python download is already done
 steps:
 - task: PythonScript@0
   displayName: 'Download test data'

tools/ci_build/github/download_test_data.py

@@ -10,6 +10,13 @@ from urllib.parse import urlparse
 from urllib.parse import urljoin
 from urllib.parse import urlsplit
 
+SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
+REPO_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..", ".."))
+sys.path.append(os.path.join(REPO_DIR, "tools", "python"))
+
+from get_azcopy import get_azcopy  # noqa: E402
+
+
 # Hardcoded map of storage account to azure region endpoint
 storage_account_to_endpoint_map = {
     'onnxruntimetestdata.blob.core.windows.net': {
@@ -47,7 +54,7 @@ def get_azure_region():
 
 def parse_arguments():
     parser = argparse.ArgumentParser(description="ONNXRuntime Data Downloader.")
-    parser.add_argument("--test_data_url", help="Test data URL.")
+    parser.add_argument("--test_data_url", required=True, help="Test data URL.")
     parser.add_argument("--azure_region", help="Azure region")
     parser.add_argument("--build_dir", required=True, help="Path to the build directory.")
     parser.add_argument("--edge_device", action="store_true", help="Edge device with limit disk space.")
@@ -80,7 +87,7 @@ def get_region_based_url(url, azure_location):
     return url
 
 
-def download_and_unzip(build_dir, url, dest_folder, use_token=True):
+def download_and_unzip(azcopy_path, build_dir, url, dest_folder, use_token=True):
     dest_folder = os.path.join(build_dir, dest_folder)
     # attach the SAS token to the url. Note DO NOT print the url with the token in any logs.
     token = os.environ.get('Test_Data_Download_Key')
@@ -90,14 +97,13 @@ def download_and_unzip(build_dir, url, dest_folder, use_token=True):
         url_with_token = url
 
     # Download data using AZCopy tool
-    # Our linux CI build machine has azcopy in /usr/bin but the version is too old
-    azcopy_exe = \
-        'azcopy.exe' if sys.platform.startswith("win") and shutil.which('azcopy') else os.path.join(build_dir, 'azcopy')
     try:
-        subprocess.run([azcopy_exe, 'cp', '--log-level', 'ERROR', '--recursive', url_with_token, build_dir], check=True)
+        subprocess.run(
+            [azcopy_path, 'cp', '--log-level', 'ERROR', '--recursive', url_with_token, build_dir],
+            check=True)
     except Exception as e:
         print(e)
-        print(azcopy_exe)
+        print(azcopy_path)
         raise Exception("Downloading data failed. Source: " + url + " Destination: " + build_dir)
 
     os.makedirs(dest_folder, exist_ok=True)
@@ -116,29 +122,6 @@ def download_and_unzip(build_dir, url, dest_folder, use_token=True):
     os.unlink(local_file_name)
 
 
-def download_additional_data(build_dir, azure_region):
-    additional_data_url = 'https://onnxruntimetestdata.blob.core.windows.net/models/'
-    # url = get_region_based_url(args.test_data_url, azure_region)
-    if not shutil.which('cmake'):
-        cmake_url = urljoin(additional_data_url, 'cmake-3.15.1-win64-x64.zip')
-        print("Starting download for cmake : " + cmake_url)
-        download_and_unzip(build_dir, cmake_url, 'cmake_temp', False)
-        dest_dir = os.path.join(build_dir, 'cmake')
-        if os.path.exists(dest_dir):
-            print('deleting %s' % dest_dir)
-            shutil.rmtree(dest_dir)
-        shutil.move(os.path.join(build_dir, 'cmake_temp', 'cmake-3.15.1-win64-x64'), dest_dir)
-
-    # Download OpenCPPCoverageSetup.exe
-    opencpp_url = urljoin(additional_data_url, 'OpenCppCoverageSetup-x64-0.9.7.0.exe')
-    print("Starting download for opencppcoverage " + opencpp_url)
-    dest_folder = os.path.join(build_dir, 'installer', 'opencppcoverage')
-    os.makedirs(dest_folder, exist_ok=True)
-    azcopy_exe = 'azcopy.exe' if shutil.which('azcopy') else os.path.join(build_dir, 'azcopy')
-    subprocess.run([azcopy_exe, 'cp', '--log-level', 'ERROR', opencpp_url, os.path.join(dest_folder, 'installer.exe')],
-                   check=True)
-
-
 args = parse_arguments()
 models_folder = 'models'
 
@@ -157,9 +140,10 @@ else:
         azure_region = get_azure_region()
     try:
         # Download test data
-        url = get_region_based_url(args.test_data_url, azure_region)
-        print("Starting test data download %s" % url)
-        download_and_unzip(args.build_dir, url, models_folder)
+        with get_azcopy(os.path.join(args.build_dir, "azcopy")) as azcopy_path:
+            url = get_region_based_url(args.test_data_url, azure_region)
+            print("Starting test data download %s" % url)
+            download_and_unzip(azcopy_path, args.build_dir, url, models_folder)
 
         all_downloads_done = True
 

tools/python/get_azcopy.py

@@ -0,0 +1,78 @@
+import contextlib
+import os
+import platform
+import re
+import shutil
+import stat
+import subprocess
+import tempfile
+import urllib.parse
+import urllib.request
+
+AZCOPY_VERSION = "10.4.3"
+
+# See here for instructions on getting stable download links:
+# https://docs.microsoft.com/en-us/azure/storage/common/storage-use-azcopy-v10#obtain-a-static-download-link
+_AZCOPY_DOWNLOAD_URLS = {
+    "Linux": "https://azcopyvnext.azureedge.net/release20200501/azcopy_linux_amd64_10.4.3.tar.gz",
+    "Darwin": "https://azcopyvnext.azureedge.net/release20200501/azcopy_darwin_amd64_10.4.3.zip",
+    "Windows": "https://azcopyvnext.azureedge.net/release20200501/azcopy_windows_amd64_10.4.3.zip",
+}
+
+
+def _check_version(azcopy_path):
+    proc = subprocess.run(
+        [azcopy_path, "--version"],
+        stdout=subprocess.PIPE, universal_newlines=True)
+    match = re.search(r"\d+(?:\.\d+)+", proc.stdout)
+
+    if not match:
+        raise RuntimeError("Failed to determine azcopy version.")
+
+    return match.group(0) == AZCOPY_VERSION
+
+
+def _find_azcopy(start_dir):
+    for root, _, file_names in os.walk(start_dir):
+        for file_name in file_names:
+            if file_name == "azcopy" or file_name == "azcopy.exe":
+                return os.path.join(root, file_name)
+    raise RuntimeError("Failed to azcopy in '{}'.".format(start_dir))
+
+
+@contextlib.contextmanager
+def get_azcopy(local_azcopy_path="azcopy"):
+    """
+    Creates a context manager that returns a path to a particular version of
+    azcopy (specified in AZCOPY_VERSION). Downloads a temporary copy if needed.
+
+    :param local_azcopy_path: Path to a local azcopy to try first.
+
+    Example usage:
+        with get_azcopy() as azcopy_path:
+            subprocess.run([azcopy_path, "--version"])
+    """
+    with contextlib.ExitStack() as context_stack:
+        azcopy_path = shutil.which(local_azcopy_path)
+
+        if azcopy_path is None or not _check_version(azcopy_path):
+            temp_dir = context_stack.enter_context(
+                tempfile.TemporaryDirectory())
+
+            download_url = _AZCOPY_DOWNLOAD_URLS[platform.system()]
+            download_basename = urllib.parse.urlsplit(
+                download_url).path.rsplit("/", 1)[-1]
+            assert len(download_basename) > 0
+            downloaded_path = os.path.join(temp_dir, download_basename)
+
+            print("Downloading azcopy from '{}'...".format(download_url))
+            urllib.request.urlretrieve(download_url, downloaded_path)
+
+            extracted_path = os.path.join(temp_dir, "azcopy")
+            shutil.unpack_archive(downloaded_path, extracted_path)
+
+            azcopy_path = _find_azcopy(extracted_path)
+
+            os.chmod(azcopy_path, stat.S_IXUSR)
+
+        yield azcopy_path