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