#------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. #-------------------------------------------------------------------------- from setuptools import setup, find_packages, Extension from distutils import log as logger from distutils.command.build_ext import build_ext as _build_ext from glob import glob from os import path, getcwd, environ, remove, walk, makedirs, listdir from shutil import copyfile, copytree, rmtree import platform import subprocess import sys import datetime nightly_build = False featurizers_build = False package_name = 'onnxruntime' wheel_name_suffix = None def parse_arg_remove_boolean(argv, arg_name): arg_value = False if arg_name in sys.argv: arg_value = True argv.remove(arg_name) return arg_value def parse_arg_remove_string(argv, arg_name_equal): arg_value = None for arg in sys.argv[1:]: if arg.startswith(arg_name_equal): arg_value = arg[len(arg_name_equal):] sys.argv.remove(arg) break return arg_value # Any combination of the following arguments can be applied featurizers_build = parse_arg_remove_boolean(sys.argv, '--use_featurizers') if parse_arg_remove_boolean(sys.argv, '--nightly_build'): package_name = 'ort-nightly' nightly_build = True wheel_name_suffix = parse_arg_remove_string(sys.argv, '--wheel_name_suffix=') cuda_version = None rocm_version = None # The following arguments are mutually exclusive if parse_arg_remove_boolean(sys.argv, '--use_tensorrt'): package_name = 'onnxruntime-gpu-tensorrt' if not nightly_build else 'ort-trt-nightly' elif parse_arg_remove_boolean(sys.argv, '--use_cuda'): package_name = 'onnxruntime-gpu' if not nightly_build else 'ort-gpu-nightly' cuda_version = parse_arg_remove_string(sys.argv, '--cuda_version=') elif parse_arg_remove_boolean(sys.argv, '--use_rocm'): package_name = 'onnxruntime-rocm' if not nightly_build else 'ort-rocm-nightly' rocm_version = parse_arg_remove_string(sys.argv, '--rocm_version=') elif parse_arg_remove_boolean(sys.argv, '--use_openvino'): package_name = 'onnxruntime-openvino' elif parse_arg_remove_boolean(sys.argv, '--use_dnnl'): package_name = 'onnxruntime-dnnl' elif parse_arg_remove_boolean(sys.argv, '--use_nuphar'): package_name = 'onnxruntime-nuphar' elif parse_arg_remove_boolean(sys.argv, '--use_vitisai'): package_name = 'onnxruntime-vitisai' elif parse_arg_remove_boolean(sys.argv, '--use_acl'): package_name = 'onnxruntime-acl' elif parse_arg_remove_boolean(sys.argv, '--use_armnn'): package_name = 'onnxruntime-armnn' elif parse_arg_remove_boolean(sys.argv, '--use_dml'): package_name = 'onnxruntime-dml' # PEP 513 defined manylinux1_x86_64 and manylinux1_i686 # PEP 571 defined manylinux2010_x86_64 and manylinux2010_i686 # PEP 599 defines the following platform tags: # manylinux2014_x86_64 # manylinux2014_i686 # manylinux2014_aarch64 # manylinux2014_armv7l # manylinux2014_ppc64 # manylinux2014_ppc64le # manylinux2014_s390x manylinux_tags = [ 'manylinux1_x86_64', 'manylinux1_i686', 'manylinux2010_x86_64', 'manylinux2010_i686', 'manylinux2014_x86_64', 'manylinux2014_i686', 'manylinux2014_aarch64', 'manylinux2014_armv7l', 'manylinux2014_ppc64', 'manylinux2014_ppc64le', 'manylinux2014_s390x', ] is_manylinux = environ.get('AUDITWHEEL_PLAT', None) in manylinux_tags class build_ext(_build_ext): def build_extension(self, ext): dest_file = self.get_ext_fullpath(ext.name) logger.info('copying %s -> %s', ext.sources[0], dest_file) copyfile(ext.sources[0], dest_file) try: from wheel.bdist_wheel import bdist_wheel as _bdist_wheel class bdist_wheel(_bdist_wheel): def finalize_options(self): _bdist_wheel.finalize_options(self) if not is_manylinux: self.root_is_pure = False def _rewrite_ld_preload(self, to_preload): with open('onnxruntime/capi/_ld_preload.py', 'rt') as f: ld_preload = f.read().splitlines() with open('onnxruntime/capi/_ld_preload.py', 'wt') as f: for line in ld_preload: f.write(line) f.write('\n') if 'LD_PRELOAD_BEGIN_MARK' in line: break if len(to_preload) > 0: f.write('from ctypes import CDLL, RTLD_GLOBAL\n') for library in to_preload: f.write('_{} = CDLL("{}", mode=RTLD_GLOBAL)\n'.format(library.split('.')[0], library)) def run(self): if is_manylinux: source = 'onnxruntime/capi/onnxruntime_pybind11_state.so' dest = 'onnxruntime/capi/onnxruntime_pybind11_state_manylinux1.so' logger.info('copying %s -> %s', source, dest) copyfile(source, dest) result = subprocess.run(['patchelf', '--print-needed', dest], check=True, stdout=subprocess.PIPE, universal_newlines=True) dependencies = ['librccl.so', 'libamdhip64.so', 'librocblas.so', 'libMIOpen.so', 'libhsa-runtime64.so', 'libhsakmt.so'] to_preload = [] args = ['patchelf', '--debug'] for line in result.stdout.split('\n'): for dependency in dependencies: if dependency in line: to_preload.append(line) args.extend(['--remove-needed', line]) args.append(dest) if len(args) > 3: subprocess.run(args, check=True, stdout=subprocess.PIPE) dest = 'onnxruntime/capi/libonnxruntime_providers_cuda.so' if path.isfile(dest): result = subprocess.run(['patchelf', '--print-needed', dest], check=True, stdout=subprocess.PIPE, universal_newlines=True) cuda_dependencies = ['libcublas.so', 'libcublasLt.so', 'libcudnn.so', 'libcudart.so', 'libcurand.so', 'libcufft.so', 'libnvToolsExt.so'] args = ['patchelf', '--debug'] for line in result.stdout.split('\n'): for dependency in cuda_dependencies: if dependency in line: if not dependency in to_preload: to_preload.append(line) args.extend(['--remove-needed', line]) args.append(dest) if len(args) > 3: subprocess.run(args, check=True, stdout=subprocess.PIPE) self._rewrite_ld_preload(to_preload) _bdist_wheel.run(self) if is_manylinux: file = glob(path.join(self.dist_dir, '*linux*.whl'))[0] logger.info('repairing %s for manylinux1', file) try: subprocess.run(['auditwheel', 'repair', '-w', self.dist_dir, file], check=True, stdout=subprocess.PIPE) finally: logger.info('removing %s', file) remove(file) except ImportError as error: print("Error importing dependencies:") print(error) bdist_wheel = None # Additional binaries if platform.system() == 'Linux': libs = ['onnxruntime_pybind11_state.so', 'libdnnl.so.2', 'libmklml_intel.so', 'libmklml_gnu.so', 'libiomp5.so', 'mimalloc.so'] dl_libs = ['libonnxruntime_providers_shared.so', 'libonnxruntime_providers_cuda.so'] # DNNL, TensorRT & OpenVINO EPs are built as shared libs libs.extend(['libonnxruntime_providers_shared.so']) libs.extend(['libonnxruntime_providers_dnnl.so']) libs.extend(['libonnxruntime_providers_tensorrt.so']) libs.extend(['libonnxruntime_providers_openvino.so']) libs.extend(['libonnxruntime_providers_cuda.so']) # Nuphar Libs libs.extend(['libtvm.so.0.5.1']) if nightly_build: libs.extend(['libonnxruntime_pywrapper.so']) elif platform.system() == "Darwin": libs = ['onnxruntime_pybind11_state.so', 'libdnnl.2.dylib', 'mimalloc.so'] # TODO add libmklml and libiomp5 later. # DNNL & TensorRT EPs are built as shared libs libs.extend(['libonnxruntime_providers_shared.dylib']) libs.extend(['libonnxruntime_providers_dnnl.dylib']) libs.extend(['libonnxruntime_providers_tensorrt.dylib']) libs.extend(['libonnxruntime_providers_cuda.dylib']) if nightly_build: libs.extend(['libonnxruntime_pywrapper.dylib']) else: libs = ['onnxruntime_pybind11_state.pyd', 'dnnl.dll', 'mklml.dll', 'libiomp5md.dll'] # DNNL, TensorRT & OpenVINO EPs are built as shared libs libs.extend(['onnxruntime_providers_shared.dll']) libs.extend(['onnxruntime_providers_dnnl.dll']) libs.extend(['onnxruntime_providers_tensorrt.dll']) libs.extend(['onnxruntime_providers_openvino.dll']) libs.extend(['onnxruntime_providers_cuda.dll']) # DirectML Libs libs.extend(['DirectML.dll']) # Nuphar Libs libs.extend(['tvm.dll']) if nightly_build: libs.extend(['onnxruntime_pywrapper.dll']) if is_manylinux: data = ['capi/libonnxruntime_pywrapper.so'] if nightly_build else [] data += [path.join('capi', x) for x in dl_libs if path.isfile(path.join('onnxruntime', 'capi', x))] ext_modules = [ Extension( 'onnxruntime.capi.onnxruntime_pybind11_state', ['onnxruntime/capi/onnxruntime_pybind11_state_manylinux1.so'], ), ] else: data = [path.join('capi', x) for x in libs if path.isfile(path.join('onnxruntime', 'capi', x))] ext_modules = [] # Additional examples examples_names = ["mul_1.onnx", "logreg_iris.onnx", "sigmoid.onnx"] examples = [path.join('datasets', x) for x in examples_names] # Extra files such as EULA and ThirdPartyNotices extra = ["LICENSE", "ThirdPartyNotices.txt", "Privacy.md"] # Description README = path.join(getcwd(), "docs/python/README.rst") if not path.exists(README): this = path.dirname(__file__) README = path.join(this, "docs/python/README.rst") if not path.exists(README): raise FileNotFoundError("Unable to find 'README.rst'") with open(README) as f: long_description = f.read() packages = [ 'onnxruntime', 'onnxruntime.backend', 'onnxruntime.capi', 'onnxruntime.capi.training', 'onnxruntime.datasets', 'onnxruntime.tools', 'onnxruntime.tools.ort_format_model', 'onnxruntime.tools.ort_format_model.ort_flatbuffers_py', 'onnxruntime.tools.ort_format_model.ort_flatbuffers_py.experimental', 'onnxruntime.tools.ort_format_model.ort_flatbuffers_py.experimental.fbs', 'onnxruntime.quantization', 'onnxruntime.quantization.operators', 'onnxruntime.quantization.CalTableFlatBuffers', 'onnxruntime.transformers', 'onnxruntime.transformers.longformer', ] requirements_file = "requirements.txt" local_version = None enable_training = parse_arg_remove_boolean(sys.argv, '--enable_training') if enable_training: packages.extend(['onnxruntime.training', 'onnxruntime.training.amp', 'onnxruntime.training.optim', 'onnxruntime.training.ortmodule']) requirements_file = "requirements-training.txt" # with training, we want to follow this naming convention: # stable: # onnxruntime-training-1.7.0+cu111-cp36-cp36m-linux_x86_64.whl # nightly: # onnxruntime-training-1.7.0.dev20210408+cu111-cp36-cp36m-linux_x86_64.whl # this is needed immediately by pytorch/ort so that the user is able to # install an onnxruntime training package with matching torch cuda version. package_name = 'onnxruntime-training' if cuda_version: # removing '.' to make local Cuda version number in the same form as Pytorch. local_version = '+cu' + cuda_version.replace('.', '') if rocm_version: # removing '.' to make Cuda version number in the same form as Pytorch. rocm_version = rocm_version.replace('.', '') local_version = '+rocm' + rocm_version package_data = {} data_files = [] if package_name == 'onnxruntime-nuphar': packages += ["onnxruntime.nuphar"] extra += [path.join('nuphar', 'NUPHAR_CACHE_VERSION')] if featurizers_build: # Copy the featurizer data from its current directory into the onnx runtime directory so that the # content can be included as module data. # Apparently, the root_dir is different based on how the script is invoked source_root_dir = None dest_root_dir = None for potential_source_prefix, potential_dest_prefix in [ (getcwd(), getcwd()), (path.dirname(__file__), path.dirname(__file__)), (path.join(getcwd(), ".."), getcwd()), ]: potential_dir = path.join(potential_source_prefix, "external", "FeaturizersLibrary", "Data") if path.isdir(potential_dir): source_root_dir = potential_source_prefix dest_root_dir = potential_dest_prefix break if source_root_dir is None: raise Exception("Unable to find the build root dir") assert dest_root_dir is not None featurizer_source_dir = path.join(source_root_dir, "external", "FeaturizersLibrary", "Data") assert path.isdir(featurizer_source_dir), featurizer_source_dir featurizer_dest_dir = path.join(dest_root_dir, "onnxruntime", "FeaturizersLibrary", "Data") if path.isdir(featurizer_dest_dir): rmtree(featurizer_dest_dir) for item in listdir(featurizer_source_dir): this_featurizer_source_fullpath = path.join(featurizer_source_dir) assert path.isdir(this_featurizer_source_fullpath), this_featurizer_source_fullpath copytree(this_featurizer_source_fullpath, featurizer_dest_dir) packages.append("onnxruntime.FeaturizersLibrary.Data.{}".format(item)) package_data[packages[-1]] = listdir(path.join(featurizer_dest_dir, item)) package_data["onnxruntime"] = data + examples + extra version_number = '' with open('VERSION_NUMBER') as f: version_number = f.readline().strip() if nightly_build: # https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables build_suffix = environ.get('BUILD_BUILDNUMBER') if build_suffix is None: # The following line is only for local testing build_suffix = str(datetime.datetime.now().date().strftime("%Y%m%d")) else: build_suffix = build_suffix.replace('.', '') if enable_training: from packaging import version from packaging.version import Version # with training package, we need to bump up version minor number so that # nightly releases take precedence over the latest release when --pre is used during pip install. # eventually this shall be the behavior of all onnxruntime releases. # alternatively we may bump up version number right after every release. ort_version = version.parse(version_number) if isinstance(ort_version, Version): version_number = '{major}.{minor}.{macro}'.format( major=ort_version.major, minor=ort_version.minor + 1, macro=ort_version.micro) version_number = version_number + ".dev" + build_suffix if local_version: version_number = version_number + local_version if wheel_name_suffix: package_name = "{}_{}".format(package_name, wheel_name_suffix) cmd_classes = {} if bdist_wheel is not None: cmd_classes['bdist_wheel'] = bdist_wheel cmd_classes['build_ext'] = build_ext requirements_path = path.join(getcwd(), requirements_file) if not path.exists(requirements_path): this = path.dirname(__file__) requirements_path = path.join(this, requirements_file) if not path.exists(requirements_path): raise FileNotFoundError("Unable to find " + requirements_file) with open(requirements_path) as f: install_requires = f.read().splitlines() if enable_training: def save_build_and_package_info(package_name, version_number, cuda_version): sys.path.append(path.join(path.dirname(__file__), 'onnxruntime', 'python')) from onnxruntime_collect_build_info import find_cudart_versions version_path = path.join('onnxruntime', 'capi', 'build_and_package_info.py') with open(version_path, 'w') as f: f.write("package_name = '{}'\n".format(package_name)) f.write("__version__ = '{}'\n".format(version_number)) if cuda_version: f.write("cuda_version = '{}'\n".format(cuda_version)) # cudart_versions are integers cudart_versions = find_cudart_versions(build_env=True) if cudart_versions and len(cudart_versions) == 1: f.write("cudart_version = {}\n".format(cudart_versions[0])) else: print( "Error getting cudart version. ", "did not find any cudart library" if not cudart_versions or len(cudart_versions) == 0 else "found multiple cudart libraries") else: # TODO: rocm pass save_build_and_package_info(package_name, version_number, cuda_version) # Setup setup( name=package_name, version=version_number, description='ONNX Runtime is a runtime accelerator for Machine Learning models', long_description=long_description, author='Microsoft Corporation', author_email='onnxruntime@microsoft.com', cmdclass=cmd_classes, license="MIT License", packages=packages, ext_modules=ext_modules, package_data=package_data, url="https://onnxruntime.ai", download_url='https://github.com/microsoft/onnxruntime/tags', data_files=data_files, install_requires=install_requires, keywords='onnx machine learning', entry_points= { 'console_scripts': [ 'onnxruntime_test = onnxruntime.tools.onnxruntime_test:main', ] }, classifiers=[ 'Development Status :: 5 - Production/Stable', 'Intended Audience :: Developers', 'License :: OSI Approved :: MIT License', 'Operating System :: POSIX :: Linux', 'Operating System :: Microsoft :: Windows', 'Operating System :: MacOS', 'Topic :: Scientific/Engineering', 'Topic :: Scientific/Engineering :: Mathematics', 'Topic :: Scientific/Engineering :: Artificial Intelligence', 'Topic :: Software Development', 'Topic :: Software Development :: Libraries', 'Topic :: Software Development :: Libraries :: Python Modules', 'Programming Language :: Python', 'Programming Language :: Python :: 3 :: Only', 'Programming Language :: Python :: 3.6', 'Programming Language :: Python :: 3.7', 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9'], )