1# SPDX-License-Identifier: Apache-2.0 2 3from __future__ import absolute_import 4from __future__ import division 5from __future__ import print_function 6from __future__ import unicode_literals 7 8from distutils.spawn import find_executable 9from distutils import sysconfig, log 10import setuptools 11import setuptools.command.build_py 12import setuptools.command.develop 13import setuptools.command.build_ext 14 15from collections import namedtuple 16from contextlib import contextmanager 17from datetime import date 18import glob 19import os 20import shlex 21import subprocess 22import sys 23import platform 24from textwrap import dedent 25import multiprocessing 26 27 28TOP_DIR = os.path.realpath(os.path.dirname(__file__)) 29SRC_DIR = os.path.join(TOP_DIR, 'onnx') 30TP_DIR = os.path.join(TOP_DIR, 'third_party') 31CMAKE_BUILD_DIR = os.path.join(TOP_DIR, '.setuptools-cmake-build') 32PACKAGE_NAME = 'onnx' 33 34WINDOWS = (os.name == 'nt') 35 36CMAKE = find_executable('cmake3') or find_executable('cmake') 37MAKE = find_executable('make') 38 39install_requires = [] 40setup_requires = [] 41tests_require = [] 42extras_require = {} 43 44################################################################################ 45# Global variables for controlling the build variant 46################################################################################ 47 48# Default value is set to TRUE\1 to keep the settings same as the current ones. 49# However going forward the recomemded way to is to set this to False\0 50USE_MSVC_STATIC_RUNTIME = bool(os.getenv('USE_MSVC_STATIC_RUNTIME', '1') == '1') 51ONNX_ML = not bool(os.getenv('ONNX_ML') == '0') 52ONNX_VERIFY_PROTO3 = bool(os.getenv('ONNX_VERIFY_PROTO3') == '1') 53ONNX_NAMESPACE = os.getenv('ONNX_NAMESPACE', 'onnx') 54ONNX_BUILD_TESTS = bool(os.getenv('ONNX_BUILD_TESTS') == '1') 55ONNX_DISABLE_EXCEPTIONS = bool(os.getenv('ONNX_DISABLE_EXCEPTIONS') == '1') 56 57DEBUG = bool(os.getenv('DEBUG')) 58COVERAGE = bool(os.getenv('COVERAGE')) 59 60################################################################################ 61# Version 62################################################################################ 63 64#try: 65# git_version = subprocess.check_output(['git', 'rev-parse', 'HEAD'], 66# cwd=TOP_DIR).decode('ascii').strip() 67#except (OSError, subprocess.CalledProcessError): 68# git_version = None 69git_version = None 70 71with open(os.path.join(TOP_DIR, 'VERSION_NUMBER')) as version_file: 72 VERSION_NUMBER = version_file.read().strip() 73 if '--weekly_build' in sys.argv: 74 today_number = date.today().strftime("%Y%m%d") 75 VERSION_NUMBER += '.dev' + today_number 76 PACKAGE_NAME = 'onnx-weekly' 77 sys.argv.remove('--weekly_build') 78 VersionInfo = namedtuple('VersionInfo', ['version', 'git_version'])( 79 version=VERSION_NUMBER, 80 git_version=git_version 81 ) 82 83################################################################################ 84# Pre Check 85################################################################################ 86 87assert CMAKE, 'Could not find "cmake" executable!' 88 89################################################################################ 90# Utilities 91################################################################################ 92 93 94@contextmanager 95def cd(path): 96 if not os.path.isabs(path): 97 raise RuntimeError('Can only cd to absolute path, got: {}'.format(path)) 98 orig_path = os.getcwd() 99 os.chdir(path) 100 try: 101 yield 102 finally: 103 os.chdir(orig_path) 104 105 106################################################################################ 107# Customized commands 108################################################################################ 109 110 111class ONNXCommand(setuptools.Command): 112 user_options = [] 113 114 def initialize_options(self): 115 pass 116 117 def finalize_options(self): 118 pass 119 120 121class create_version(ONNXCommand): 122 def run(self): 123 with open(os.path.join(SRC_DIR, 'version.py'), 'w') as f: 124 f.write(dedent('''\ 125 # This file is generated by setup.py. DO NOT EDIT! 126 127 from __future__ import absolute_import 128 from __future__ import division 129 from __future__ import print_function 130 from __future__ import unicode_literals 131 132 version = '{version}' 133 git_version = '{git_version}' 134 '''.format(**dict(VersionInfo._asdict())))) 135 136 137class cmake_build(setuptools.Command): 138 """ 139 Compiles everything when `python setupmnm.py build` is run using cmake. 140 141 Custom args can be passed to cmake by specifying the `CMAKE_ARGS` 142 environment variable. 143 144 The number of CPUs used by `make` can be specified by passing `-j<ncpus>` 145 to `setup.py build`. By default all CPUs are used. 146 """ 147 user_options = [ 148 (str('jobs='), str('j'), str('Specifies the number of jobs to use with make')) 149 ] 150 151 built = False 152 153 def initialize_options(self): 154 self.jobs = None 155 156 def finalize_options(self): 157 if sys.version_info[0] >= 3: 158 self.set_undefined_options('build', ('parallel', 'jobs')) 159 if self.jobs is None and os.getenv("MAX_JOBS") is not None: 160 self.jobs = os.getenv("MAX_JOBS") 161 self.jobs = multiprocessing.cpu_count() if self.jobs is None else int(self.jobs) 162 163 def run(self): 164 if cmake_build.built: 165 return 166 cmake_build.built = True 167 if not os.path.exists(CMAKE_BUILD_DIR): 168 os.makedirs(CMAKE_BUILD_DIR) 169 170 with cd(CMAKE_BUILD_DIR): 171 build_type = 'Release' 172 # configure 173 cmake_args = [ 174 CMAKE, 175 '-DPYTHON_INCLUDE_DIR={}'.format(sysconfig.get_python_inc()), 176 '-DPYTHON_EXECUTABLE={}'.format(sys.executable), 177 '-DBUILD_ONNX_PYTHON=ON', 178 '-DCMAKE_EXPORT_COMPILE_COMMANDS=ON', 179 '-DONNX_NAMESPACE={}'.format(ONNX_NAMESPACE), 180 '-DPY_EXT_SUFFIX={}'.format(sysconfig.get_config_var('EXT_SUFFIX') or ''), 181 ] 182 if COVERAGE: 183 cmake_args.append('-DONNX_COVERAGE=ON') 184 if COVERAGE or DEBUG: 185 # in order to get accurate coverage information, the 186 # build needs to turn off optimizations 187 build_type = 'Debug' 188 cmake_args.append('-DCMAKE_BUILD_TYPE=%s' % build_type) 189 if WINDOWS: 190 cmake_args.extend([ 191 # we need to link with libpython on windows, so 192 # passing python version to window in order to 193 # find python in cmake 194 '-DPY_VERSION={}'.format('{0}.{1}'.format(*sys.version_info[:2])), 195 ]) 196 if USE_MSVC_STATIC_RUNTIME: 197 cmake_args.append('-DONNX_USE_MSVC_STATIC_RUNTIME=ON') 198 if platform.architecture()[0] == '64bit': 199 cmake_args.extend(['-A', 'x64', '-T', 'host=x64']) 200 else: 201 cmake_args.extend(['-A', 'Win32', '-T', 'host=x86']) 202 if ONNX_ML: 203 cmake_args.append('-DONNX_ML=1') 204 if ONNX_VERIFY_PROTO3: 205 cmake_args.append('-DONNX_VERIFY_PROTO3=1') 206 if ONNX_BUILD_TESTS: 207 cmake_args.append('-DONNX_BUILD_TESTS=ON') 208 if ONNX_DISABLE_EXCEPTIONS: 209 cmake_args.append('-DONNX_DISABLE_EXCEPTIONS=ON') 210 if 'CMAKE_ARGS' in os.environ: 211 extra_cmake_args = shlex.split(os.environ['CMAKE_ARGS']) 212 # prevent crossfire with downstream scripts 213 del os.environ['CMAKE_ARGS'] 214 log.info('Extra cmake args: {}'.format(extra_cmake_args)) 215 cmake_args.extend(extra_cmake_args) 216 cmake_args.append(TOP_DIR) 217 log.info('Using cmake args: {}'.format(cmake_args)) 218 if '-DONNX_DISABLE_EXCEPTIONS=ON' in cmake_args: 219 raise RuntimeError("-DONNX_DISABLE_EXCEPTIONS=ON option is only available for c++ builds. Python binding require exceptions to be enabled.") 220 subprocess.check_call(cmake_args) 221 222 build_args = [CMAKE, '--build', os.curdir] 223 if WINDOWS: 224 build_args.extend(['--config', build_type]) 225 build_args.extend(['--', '/maxcpucount:{}'.format(self.jobs)]) 226 else: 227 build_args.extend(['--', '-j', str(self.jobs)]) 228 subprocess.check_call(build_args) 229 230 231class build_py(setuptools.command.build_py.build_py): 232 def run(self): 233 self.run_command('create_version') 234 self.run_command('cmake_build') 235 236 generated_python_files = \ 237 glob.glob(os.path.join(CMAKE_BUILD_DIR, 'onnx', '*.py')) + \ 238 glob.glob(os.path.join(CMAKE_BUILD_DIR, 'onnx', '*.pyi')) 239 240 for src in generated_python_files: 241 dst = os.path.join( 242 TOP_DIR, os.path.relpath(src, CMAKE_BUILD_DIR)) 243 self.copy_file(src, dst) 244 245 return setuptools.command.build_py.build_py.run(self) 246 247 248class develop(setuptools.command.develop.develop): 249 def run(self): 250 self.run_command('build_py') 251 setuptools.command.develop.develop.run(self) 252 253 254class build_ext(setuptools.command.build_ext.build_ext): 255 def run(self): 256 self.run_command('cmake_build') 257 setuptools.command.build_ext.build_ext.run(self) 258 259 def build_extensions(self): 260 for ext in self.extensions: 261 fullname = self.get_ext_fullname(ext.name) 262 filename = os.path.basename(self.get_ext_filename(fullname)) 263 264 lib_path = CMAKE_BUILD_DIR 265 if os.name == 'nt': 266 debug_lib_dir = os.path.join(lib_path, "Debug") 267 release_lib_dir = os.path.join(lib_path, "Release") 268 if os.path.exists(debug_lib_dir): 269 lib_path = debug_lib_dir 270 elif os.path.exists(release_lib_dir): 271 lib_path = release_lib_dir 272 src = os.path.join(lib_path, filename) 273 dst = os.path.join(os.path.realpath(self.build_lib), "onnx", filename) 274 self.copy_file(src, dst) 275 276 277class mypy_type_check(ONNXCommand): 278 description = 'Run MyPy type checker' 279 280 def run(self): 281 """Run command.""" 282 onnx_script = os.path.realpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "tools/mypy-onnx.py")) 283 returncode = subprocess.call([sys.executable, onnx_script]) 284 sys.exit(returncode) 285 286 287cmdclass = { 288 'create_version': create_version, 289 'cmake_build': cmake_build, 290 'build_py': build_py, 291 'develop': develop, 292 'build_ext': build_ext, 293 'typecheck': mypy_type_check, 294} 295 296################################################################################ 297# Extensions 298################################################################################ 299 300ext_modules = [ 301 setuptools.Extension( 302 name=str('onnx.onnx_cpp2py_export'), 303 sources=[]) 304] 305 306################################################################################ 307# Packages 308################################################################################ 309 310# no need to do fancy stuff so far 311packages = setuptools.find_packages() 312 313requirements_file = "requirements.txt" 314requirements_path = os.path.join(os.getcwd(), requirements_file) 315if not os.path.exists(requirements_path): 316 this = os.path.dirname(__file__) 317 requirements_path = os.path.join(this, requirements_file) 318if not os.path.exists(requirements_path): 319 raise FileNotFoundError("Unable to find " + requirements_file) 320with open(requirements_path) as f: 321 install_requires = f.read().splitlines() 322 323################################################################################ 324# Test 325################################################################################ 326 327setup_requires.append('pytest-runner') 328tests_require.append('pytest') 329tests_require.append('nbval') 330tests_require.append('tabulate') 331 332if sys.version_info[0] == 3: 333 # Mypy doesn't work with Python 2 334 extras_require['mypy'] = ['mypy==0.600'] 335 336################################################################################ 337# Final 338################################################################################ 339 340setuptools.setup( 341 name=PACKAGE_NAME, 342 version=VersionInfo.version, 343 description="Open Neural Network Exchange", 344 long_description=open("README.md").read(), 345 long_description_content_type="text/markdown", 346 ext_modules=ext_modules, 347 cmdclass=cmdclass, 348 packages=packages, 349 license='Apache License v2.0', 350 include_package_data=True, 351 install_requires=install_requires, 352 setup_requires=setup_requires, 353 tests_require=tests_require, 354 extras_require=extras_require, 355 author='ONNX', 356 author_email='onnx-technical-discuss@lists.lfai.foundation', 357 url='https://github.com/onnx/onnx', 358 entry_points={ 359 'console_scripts': [ 360 'check-model = onnx.bin.checker:check_model', 361 'check-node = onnx.bin.checker:check_node', 362 'backend-test-tools = onnx.backend.test.cmd_tools:main', 363 ] 364 }, 365) 366