onnxruntime/tools/ci_build/github/download_test_data.py
2019-04-12 22:23:16 -07:00

43 lines
No EOL
1.4 KiB
Python
Executable file

#!/usr/bin/python3
import urllib.request
import json
import subprocess
import os
import argparse
from urllib.parse import urlparse
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.")
return parser.parse_args()
def get_server_hostname():
#should be northcentralus or centralus
azure_location=get_azure_region()
print(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
args = parse_arguments()
hostname=get_server_hostname()
url=args.test_data_url.replace('onnxruntimetestdata', hostname)
print('data url=%s' % url)
subprocess.run(['./azcopy','cp', '--log-level','ERROR', url,'.'],check=True)
os.makedirs('models',exist_ok=True)
local_file_name = os.path.basename(urlparse(url).path)
subprocess.run(['unzip', '-qd','models','20190327.zip'])