onnxruntime/tools/ci_build/github/download_test_data.py
2019-08-05 19:53:15 -07:00

81 lines
3.4 KiB
Python
Executable file

#!/usr/bin/python3
import urllib.request
import json
import subprocess
import os
import sys
import shutil
import argparse
from urllib.parse import urlparse
def is_windows():
return sys.platform.startswith("win")
def get_azure_region():
req = urllib.request.Request('http://169.254.169.254/metadata/instance?api-version=2018-10-01')
req.add_header('Metadata', 'true')
body = urllib.request.urlopen(req).read()
body = json.loads(body.decode('utf-8'))
return body['compute']['location']
def parse_arguments():
parser = argparse.ArgumentParser(description="ONNXRuntime Data Downloader.")
parser.add_argument("--test_data_url", 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.")
return parser.parse_args()
def get_server_hostname(azure_location):
if azure_location is None:
#should be northcentralus or centralus
azure_location = get_azure_region()
print("This VM is in azure location: %s" % azure_location)
if azure_location == 'centralus':
hostname = 'onnxruntimetestdata'
elif azure_location == 'northcentralus':
hostname = 'onnxruntimetestdata2'
else:
print('warning: no local data cache for azure region %s' % azure_location)
hostname = 'onnxruntimetestdata2'
return hostname
def download_and_unzip(build_dir, url, dest_folder):
print("Downloading %s" % url)
dest_folder = os.path.join(build_dir, dest_folder)
subprocess.run([os.path.join(build_dir,'azcopy'),'cp', '--log-level','ERROR', url, build_dir],check=True)
os.makedirs(dest_folder,exist_ok=True)
local_file_name = os.path.join(build_dir, os.path.basename(urlparse(url).path))
if is_windows():
print("unzip %s" % local_file_name)
if shutil.which('7z'): # 7-Zip
subprocess.run(['7z','x', local_file_name, '-y', '-o' + dest_folder], check=True)
elif shutil.which('7za'): # 7-Zip standalone
subprocess.run(['7za', 'x', local_file_name, '-y', '-o' + dest_folder], check=True)
else:
log.error("No unzip tool for use")
sys.exit(1)
else:
subprocess.run(['unzip','-qd', dest_folder ,local_file_name], check=True)
os.unlink(local_file_name)
args = parse_arguments()
hostname = get_server_hostname(args.azure_region)
url = args.test_data_url.replace('onnxruntimetestdata', hostname)
print('data url=%s' % url)
download_and_unzip(args.build_dir, url, 'models')
if is_windows():
url = 'https://onnxruntimetestdata.blob.core.windows.net/models/cmake-3.15.1-win64-x64.zip'
url = url.replace('onnxruntimetestdata', hostname)
download_and_unzip(args.build_dir, url, 'cmake_temp')
dest_dir = os.path.join(args.build_dir,'cmake')
if os.path.exists(dest_dir):
print('deleting %s' % dest_dir)
shutil.rmtree(dest_dir)
shutil.move(os.path.join(args.build_dir,'cmake_temp','cmake-3.15.1-win64-x64'),dest_dir)
url = 'https://onnxruntimetestdata.blob.core.windows.net/models/OpenCppCoverageSetup-x64-0.9.7.0.exe'
url = url.replace('onnxruntimetestdata', hostname)
dest_folder = os.path.join(args.build_dir, 'installer','opencppcoverage')
os.makedirs(dest_folder,exist_ok=True)
subprocess.run([os.path.join(args.build_dir,'azcopy'),'cp', '--log-level','ERROR', url, os.path.join(dest_folder,'installer.exe')],check=True)