#------------------------------------------------------------------------- # 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 from shutil import copyfile import platform import subprocess import sys import datetime nightly_build = False package_name = 'onnxruntime' if '--use_tensorrt' in sys.argv: package_name = 'onnxruntime-gpu-tensorrt' sys.argv.remove('--use_tensorrt') if '--nightly_build' in sys.argv: package_name = 'ort-trt-nightly' nightly_build = True sys.argv.remove('--nightly_build') elif '--use_cuda' in sys.argv: package_name = 'onnxruntime-gpu' sys.argv.remove('--use_cuda') if '--nightly_build' in sys.argv: package_name = 'ort-gpu-nightly' nightly_build = True sys.argv.remove('--nightly_build') elif '--use_ngraph' in sys.argv: package_name = 'onnxruntime-ngraph' sys.argv.remove('--use_ngraph') elif '--use_openvino' in sys.argv: package_name = 'onnxruntime-openvino' elif '--use_nuphar' in sys.argv: package_name = 'onnxruntime-nuphar' sys.argv.remove('--use_nuphar') if '--nightly_build' in sys.argv: package_name = 'ort-nightly' nightly_build = True sys.argv.remove('--nightly_build') is_manylinux2010 = False if environ.get('AUDITWHEEL_PLAT', None) == 'manylinux2010_x86_64': is_manylinux2010 = True 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_manylinux2010: 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_manylinux2010: source = 'onnxruntime/capi/onnxruntime_pybind11_state.so' dest = 'onnxruntime/capi/onnxruntime_pybind11_state_manylinux2010.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) cuda_dependencies = ['libcublas.so', 'libcudnn.so', 'libcudart.so'] to_preload = [] args = ['patchelf', '--debug'] for line in result.stdout.split('\n'): for dependency in cuda_dependencies: if dependency in line: to_preload.append(line) args.extend(['--remove-needed', line]) args.append(dest) if len(to_preload) > 0: subprocess.run(args, check=True, stdout=subprocess.PIPE) self._rewrite_ld_preload(to_preload) _bdist_wheel.run(self) if is_manylinux2010: file = glob(path.join(self.dist_dir, '*linux*.whl'))[0] logger.info('repairing %s for manylinux2010', file) try: subprocess.run(['auditwheel', 'repair', '--plat', 'manylinux2010_x86_64', '-w', self.dist_dir, file], check=True, stdout=subprocess.PIPE) finally: logger.info('removing %s', file) remove(file) except ImportError: bdist_wheel = None # Additional binaries if platform.system() == 'Linux': libs = ['onnxruntime_pybind11_state.so', 'libmkldnn.so.0', 'libmklml_intel.so', 'libiomp5.so'] # nGraph Libs libs.extend(['libngraph.so', 'libcodegen.so', 'libcpu_backend.so', 'libmkldnn.so', 'libtbb_debug.so', 'libtbb_debug.so.2', 'libtbb.so', 'libtbb.so.2']) # Nuphar Libs libs.extend(['libtvm.so']) elif platform.system() == "Darwin": libs = ['onnxruntime_pybind11_state.so', 'libmkldnn.0.dylib'] # TODO add libmklml and libiomp5 later. else: libs = ['onnxruntime_pybind11_state.pyd', 'mkldnn.dll', 'mklml.dll', 'libiomp5md.dll'] libs.extend(['ngraph.dll', 'cpu_backend.dll', 'tbb.dll']) # Nuphar Libs libs.extend(['tvm.dll']) if is_manylinux2010: data = [] ext_modules = [ Extension( 'onnxruntime.capi.onnxruntime_pybind11_state', ['onnxruntime/capi/onnxruntime_pybind11_state_manylinux2010.so'], ), ] else: data = [path.join('capi', x) for x in libs if path.isfile(path.join('onnxruntime', 'capi', x))] ext_modules = [] python_modules_list = list() if '--use_openvino' in sys.argv: #Adding python modules required for openvino ep python_modules_list.extend(['openvino_mo', 'openvino_emitter']) sys.argv.remove('--use_openvino') # 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"] # 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() version_number = '' with open('VERSION_NUMBER') as f: version_number = f.readline().strip() if nightly_build: date_suffix = str(datetime.datetime.now().date().strftime("%m%d")) version_number = version_number + ".dev" + date_suffix # Setup setup( name=package_name, version=version_number, description='ONNX Runtime Python bindings', long_description=long_description, author='Microsoft Corporation', author_email='onnx@microsoft.com', cmdclass={'bdist_wheel': bdist_wheel, 'build_ext': build_ext}, license="MIT License", packages=['onnxruntime', 'onnxruntime.backend', 'onnxruntime.capi', 'onnxruntime.datasets', 'onnxruntime.tools', ], ext_modules=ext_modules, package_data={ 'onnxruntime': data + examples + extra, }, py_modules=python_modules_list, extras_require={ 'backend': ['onnx>=1.2.3'], 'numpy': ['numpy>=1.15.0'] }, entry_points= { 'console_scripts': [ 'onnxruntime_test = onnxruntime.tools.onnxruntime_test:main', ] }, classifiers=[ 'Development Status :: 4 - Beta', 'Environment :: Console', 'Intended Audience :: Developers', 'Operating System :: POSIX :: Linux', 'Programming Language :: Python', 'Programming Language :: Python :: 3 :: Only', 'Programming Language :: Python :: 3.5', 'Programming Language :: Python :: 3.6', 'Programming Language :: Python :: 3.7'], )