mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-16 21:00:14 +00:00
85 lines
3 KiB
Python
Executable file
85 lines
3 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
|
|
import argparse
|
|
|
|
import os
|
|
import sys
|
|
import shutil
|
|
import subprocess
|
|
import hashlib
|
|
from os.path import expanduser
|
|
|
|
|
|
|
|
test_data_url = 'https://onnxruntimetestdata.blob.core.windows.net/models/cmake-3.13.2-win64-x64.zip'
|
|
test_data_checksum = '4cbaf72047d20bc84742327a5eafffd1'
|
|
|
|
def check_md5(filename, expected_md5):
|
|
if not os.path.exists(filename):
|
|
return False
|
|
hash_md5 = hashlib.md5()
|
|
BLOCKSIZE = 1024*64
|
|
with open(filename, "rb") as f:
|
|
buf = f.read(BLOCKSIZE)
|
|
while len(buf) > 0:
|
|
hash_md5.update(buf)
|
|
buf = f.read(BLOCKSIZE)
|
|
hex = hash_md5.hexdigest()
|
|
if hex != expected_md5:
|
|
print('md5 mismatch, expect %s, got %s' % (expected_md5, hex))
|
|
os.remove(filename)
|
|
return False
|
|
return True
|
|
|
|
def is_windows():
|
|
return sys.platform.startswith("win")
|
|
|
|
#the last part of src_url should be unique, across all the builds
|
|
def download_test_data(models_dir, src_url, expected_md5):
|
|
cache_dir = os.path.join(expanduser("~"), '.cache','onnxruntime')
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
local_zip_file = os.path.join(cache_dir, os.path.basename(src_url))
|
|
if not check_md5(local_zip_file, expected_md5):
|
|
print("Downloading test data")
|
|
if is_windows():
|
|
subprocess.run(['powershell', '-Command', 'Invoke-WebRequest %s -OutFile %s' % (src_url, local_zip_file)], check=True)
|
|
elif shutil.which('aria2c'):
|
|
subprocess.run(['aria2c','-x', '5', '-j',' 5', '-q', src_url, '-d', cache_dir], check=True)
|
|
elif shutil.which('curl'):
|
|
subprocess.run(['curl', '-s', src_url, '-o', local_zip_file], check=True)
|
|
else:
|
|
import urllib.request
|
|
urllib.request.urlretrieve(src_url, local_zip_file)
|
|
if not check_md5(local_zip_file, expected_md5):
|
|
print('Download failed')
|
|
exit(-1)
|
|
if os.path.exists(models_dir):
|
|
print('deleting %s' % models_dir)
|
|
shutil.rmtree(models_dir)
|
|
if is_windows():
|
|
subprocess.run(['powershell', '-Command', 'Expand-Archive -LiteralPath "%s" -DestinationPath "%s" -Force' % (local_zip_file, models_dir)], check=True)
|
|
else:
|
|
subprocess.run(['unzip','-qd', models_dir, local_zip_file], check=True)
|
|
return True
|
|
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--build_dir", required=True, help="Path to the build directory.")
|
|
return parser.parse_args()
|
|
|
|
|
|
args = parse_arguments()
|
|
os.makedirs(args.build_dir,exist_ok=True)
|
|
|
|
download_test_data(os.path.join(args.build_dir,'cmake_temp'), test_data_url, test_data_checksum)
|
|
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.13.2-win64-x64'),dest_dir)
|
|
if not os.path.exists(os.path.join(dest_dir,'bin','cmake.exe')):
|
|
print('download failed')
|
|
exit(-1)
|