From 63bf587623f75cd4299f569bf24290a1628a934b Mon Sep 17 00:00:00 2001 From: edgchen1 <18449977+edgchen1@users.noreply.github.com> Date: Tue, 16 Jun 2020 10:14:34 -0700 Subject: [PATCH] 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. --- .../tools/ci_test/download_e2e_test_data.py | 24 +++--- .../linux-ort-srv-nightly-pipeline.yml | 8 -- .../mac-set-variables-and-download.yml | 8 -- .../templates/win-download-test-data.yml | 2 +- tools/ci_build/github/download_test_data.py | 50 ++++-------- tools/python/get_azcopy.py | 78 +++++++++++++++++++ 6 files changed, 111 insertions(+), 59 deletions(-) create mode 100644 tools/python/get_azcopy.py diff --git a/orttraining/tools/ci_test/download_e2e_test_data.py b/orttraining/tools/ci_test/download_e2e_test_data.py index 39eb2b7057..0978145ebf 100755 --- a/orttraining/tools/ci_test/download_e2e_test_data.py +++ b/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__": diff --git a/tools/ci_build/github/azure-pipelines/linux-ort-srv-nightly-pipeline.yml b/tools/ci_build/github/azure-pipelines/linux-ort-srv-nightly-pipeline.yml index 7fbd8e56c4..b6de177541 100644 --- a/tools/ci_build/github/azure-pipelines/linux-ort-srv-nightly-pipeline.yml +++ b/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: diff --git a/tools/ci_build/github/azure-pipelines/nuget/templates/mac-set-variables-and-download.yml b/tools/ci_build/github/azure-pipelines/nuget/templates/mac-set-variables-and-download.yml index 20127f44d7..d689be215e 100644 --- a/tools/ci_build/github/azure-pipelines/nuget/templates/mac-set-variables-and-download.yml +++ b/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: diff --git a/tools/ci_build/github/azure-pipelines/nuget/templates/win-download-test-data.yml b/tools/ci_build/github/azure-pipelines/nuget/templates/win-download-test-data.yml index dbaf52ef28..4ba4595458 100644 --- a/tools/ci_build/github/azure-pipelines/nuget/templates/win-download-test-data.yml +++ b/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' diff --git a/tools/ci_build/github/download_test_data.py b/tools/ci_build/github/download_test_data.py index 28e8340a14..0be773b51e 100755 --- a/tools/ci_build/github/download_test_data.py +++ b/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 diff --git a/tools/python/get_azcopy.py b/tools/python/get_azcopy.py new file mode 100644 index 0000000000..520d9b17cf --- /dev/null +++ b/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