onnx

1
# Copyright (c) ONNX Project Contributors
2
#
3
# SPDX-License-Identifier: Apache-2.0
4

5
# NOTE: Put all metadata in pyproject.toml.
6
# Set the environment variable `ONNX_PREVIEW_BUILD=1` to build the dev preview release.
7
from __future__ import annotations
8

9
import contextlib
10
import datetime
11
import glob
12
import logging
13
import multiprocessing
14
import os
15
import platform
16
import shlex
17
import shutil
18
import subprocess
19
import sys
20
import sysconfig
21
import textwrap
22
from typing import ClassVar
23

24
import setuptools
25
import setuptools.command.build_ext
26
import setuptools.command.build_py
27
import setuptools.command.develop
28

29
TOP_DIR = os.path.realpath(os.path.dirname(__file__))
30
CMAKE_BUILD_DIR = os.path.join(TOP_DIR, ".setuptools-cmake-build")
31

32
WINDOWS = os.name == "nt"
33

34
CMAKE = shutil.which("cmake3") or shutil.which("cmake")
35

36
################################################################################
37
# Global variables for controlling the build variant
38
################################################################################
39

40
# Default value is set to TRUE\1 to keep the settings same as the current ones.
41
# However going forward the recommended way to is to set this to False\0
42
ONNX_ML = os.getenv("ONNX_ML") != "0"
43
ONNX_VERIFY_PROTO3 = os.getenv("ONNX_VERIFY_PROTO3") == "1"
44
ONNX_NAMESPACE = os.getenv("ONNX_NAMESPACE", "onnx")
45
ONNX_BUILD_TESTS = os.getenv("ONNX_BUILD_TESTS") == "1"
46
ONNX_DISABLE_EXCEPTIONS = os.getenv("ONNX_DISABLE_EXCEPTIONS") == "1"
47
ONNX_DISABLE_STATIC_REGISTRATION = os.getenv("ONNX_DISABLE_STATIC_REGISTRATION") == "1"
48
ONNX_PREVIEW_BUILD = os.getenv("ONNX_PREVIEW_BUILD") == "1"
49

50
USE_MSVC_STATIC_RUNTIME = os.getenv("USE_MSVC_STATIC_RUNTIME", "0") == "1"
51
DEBUG = os.getenv("DEBUG", "0") == "1"
52
COVERAGE = os.getenv("COVERAGE", "0") == "1"
53

54
# Customize the wheel plat-name; sometimes useful for MacOS builds.
55
# See https://github.com/onnx/onnx/pull/6117
56
ONNX_WHEEL_PLATFORM_NAME = os.getenv("ONNX_WHEEL_PLATFORM_NAME")
57

58
################################################################################
59
# Pre Check
60
################################################################################
61

62
assert CMAKE, "Could not find cmake in PATH"
63

64
################################################################################
65
# Version
66
################################################################################
67

68
try:
69
    _git_version = (
70
        subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=TOP_DIR)
71
        .decode("ascii")
72
        .strip()
73
    )
74
except (OSError, subprocess.CalledProcessError):
75
    _git_version = ""
76

77
with open(os.path.join(TOP_DIR, "VERSION_NUMBER"), encoding="utf-8") as version_file:
78
    _version = version_file.read().strip()
79
    if ONNX_PREVIEW_BUILD:
80
        # Create the dev build for weekly releases
81
        todays_date = datetime.date.today().strftime("%Y%m%d")
82
        _version += ".dev" + todays_date
83
    VERSION_INFO = {"version": _version, "git_version": _git_version}
84

85
################################################################################
86
# Utilities
87
################################################################################
88

89

90
@contextlib.contextmanager
91
def cd(path):
92
    if not os.path.isabs(path):
93
        raise RuntimeError(f"Can only cd to absolute path, got: {path}")
94
    orig_path = os.getcwd()
95
    os.chdir(path)
96
    try:
97
        yield
98
    finally:
99
        os.chdir(orig_path)
100

101

102
def get_ext_suffix():
103
    return sysconfig.get_config_var("EXT_SUFFIX")
104

105

106
################################################################################
107
# Customized commands
108
################################################################################
109

110

111
def create_version(directory: str):
112
    """Create version.py based on VERSION_INFO."""
113
    version_file_path = os.path.join(directory, "onnx", "version.py")
114
    os.makedirs(os.path.dirname(version_file_path), exist_ok=True)
115

116
    with open(version_file_path, "w", encoding="utf-8") as f:
117
        f.write(
118
            textwrap.dedent(
119
                f"""\
120
                # This file is generated by setup.py. DO NOT EDIT!
121

122

123
                version = "{VERSION_INFO['version']}"
124
                git_version = "{VERSION_INFO['git_version']}"
125
                """
126
            )
127
        )
128

129

130
class CmakeBuild(setuptools.Command):
131
    """Compiles everything when `python setup.py build` is run using cmake.
132

133
    Custom args can be passed to cmake by specifying the `CMAKE_ARGS`
134
    environment variable.
135

136
    The number of CPUs used by `make` can be specified by passing `-j<ncpus>`
137
    to `setup.py build`.  By default all CPUs are used.
138
    """
139

140
    user_options: ClassVar[list] = [
141
        ("jobs=", "j", "Specifies the number of jobs to use with make")
142
    ]
143

144
    def initialize_options(self):
145
        self.jobs = None
146

147
    def finalize_options(self):
148
        self.set_undefined_options("build", ("parallel", "jobs"))
149
        if self.jobs is None and os.getenv("MAX_JOBS") is not None:
150
            self.jobs = os.getenv("MAX_JOBS")
151
        self.jobs = multiprocessing.cpu_count() if self.jobs is None else int(self.jobs)
152

153
    def run(self):
154
        os.makedirs(CMAKE_BUILD_DIR, exist_ok=True)
155

156
        with cd(CMAKE_BUILD_DIR):
157
            build_type = "Release"
158
            # configure
159
            cmake_args = [
160
                CMAKE,
161
                f"-DPYTHON_INCLUDE_DIR={sysconfig.get_path('include')}",
162
                f"-DPYTHON_EXECUTABLE={sys.executable}",
163
                "-DBUILD_ONNX_PYTHON=ON",
164
                "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
165
                f"-DONNX_NAMESPACE={ONNX_NAMESPACE}",
166
                f"-DPY_EXT_SUFFIX={get_ext_suffix() or ''}",
167
            ]
168
            if COVERAGE:
169
                cmake_args.append("-DONNX_COVERAGE=ON")
170
            if COVERAGE or DEBUG:
171
                # in order to get accurate coverage information, the
172
                # build needs to turn off optimizations
173
                build_type = "Debug"
174
            cmake_args.append(f"-DCMAKE_BUILD_TYPE={build_type}")
175
            if WINDOWS:
176
                cmake_args.extend(
177
                    [
178
                        # we need to link with libpython on windows, so
179
                        # passing python version to window in order to
180
                        # find python in cmake
181
                        f"-DPY_VERSION={'{}.{}'.format(*sys.version_info[:2])}",
182
                    ]
183
                )
184
                if USE_MSVC_STATIC_RUNTIME:
185
                    cmake_args.append("-DONNX_USE_MSVC_STATIC_RUNTIME=ON")
186
                if platform.architecture()[0] == "64bit":
187
                    if "arm" in platform.machine().lower():
188
                        cmake_args.extend(["-A", "ARM64"])
189
                    else:
190
                        cmake_args.extend(["-A", "x64", "-T", "host=x64"])
191
                else:  # noqa: PLR5501
192
                    if "arm" in platform.machine().lower():
193
                        cmake_args.extend(["-A", "ARM"])
194
                    else:
195
                        cmake_args.extend(["-A", "Win32", "-T", "host=x86"])
196
            if ONNX_ML:
197
                cmake_args.append("-DONNX_ML=1")
198
            if ONNX_VERIFY_PROTO3:
199
                cmake_args.append("-DONNX_VERIFY_PROTO3=1")
200
            if ONNX_BUILD_TESTS:
201
                cmake_args.append("-DONNX_BUILD_TESTS=ON")
202
            if ONNX_DISABLE_EXCEPTIONS:
203
                cmake_args.append("-DONNX_DISABLE_EXCEPTIONS=ON")
204
            if ONNX_DISABLE_STATIC_REGISTRATION:
205
                cmake_args.append("-DONNX_DISABLE_STATIC_REGISTRATION=ON")
206
            if "CMAKE_ARGS" in os.environ:
207
                extra_cmake_args = shlex.split(os.environ["CMAKE_ARGS"])
208
                # prevent crossfire with downstream scripts
209
                del os.environ["CMAKE_ARGS"]
210
                logging.info("Extra cmake args: %s", extra_cmake_args)
211
                cmake_args.extend(extra_cmake_args)
212
            cmake_args.append(TOP_DIR)
213
            logging.info("Using cmake args: %s", cmake_args)
214
            if "-DONNX_DISABLE_EXCEPTIONS=ON" in cmake_args:
215
                raise RuntimeError(
216
                    "-DONNX_DISABLE_EXCEPTIONS=ON option is only available for c++ builds. Python binding require exceptions to be enabled."
217
                )
218
            subprocess.check_call(cmake_args)
219

220
            build_args = [CMAKE, "--build", os.curdir]
221
            if WINDOWS:
222
                build_args.extend(["--config", build_type])
223
                build_args.extend(["--", f"/maxcpucount:{self.jobs}"])
224
            else:
225
                build_args.extend(["--", "-j", str(self.jobs)])
226
            subprocess.check_call(build_args)
227

228

229
class BuildPy(setuptools.command.build_py.build_py):
230
    def run(self):
231
        if self.editable_mode:
232
            dst_dir = TOP_DIR
233
        else:
234
            dst_dir = self.build_lib
235
        create_version(dst_dir)
236
        return super().run()
237

238

239
class Develop(setuptools.command.develop.develop):
240
    def run(self):
241
        create_version(TOP_DIR)
242
        return super().run()
243

244

245
class BuildExt(setuptools.command.build_ext.build_ext):
246
    def run(self):
247
        self.run_command("cmake_build")
248
        return super().run()
249

250
    def build_extensions(self):
251
        # We override this method entirely because the actual building is done
252
        # by cmake_build. Here we just copy the built extensions to the final
253
        # destination.
254
        build_lib = self.build_lib
255
        extension_dst_dir = os.path.join(build_lib, "onnx")
256
        os.makedirs(extension_dst_dir, exist_ok=True)
257

258
        for ext in self.extensions:
259
            fullname = self.get_ext_fullname(ext.name)
260
            filename = os.path.basename(self.get_ext_filename(fullname))
261

262
            if not WINDOWS:
263
                lib_dir = CMAKE_BUILD_DIR
264
            else:
265
                # Windows compiled extensions are stored in Release/Debug subfolders
266
                debug_lib_dir = os.path.join(CMAKE_BUILD_DIR, "Debug")
267
                release_lib_dir = os.path.join(CMAKE_BUILD_DIR, "Release")
268
                if os.path.exists(debug_lib_dir):
269
                    lib_dir = debug_lib_dir
270
                elif os.path.exists(release_lib_dir):
271
                    lib_dir = release_lib_dir
272
            src = os.path.join(lib_dir, filename)
273
            dst = os.path.join(extension_dst_dir, filename)
274
            self.copy_file(src, dst)
275

276
        # Copy over the generated python files to build/source dir depending on editable mode
277
        if self.editable_mode:
278
            dst_dir = TOP_DIR
279
        else:
280
            dst_dir = build_lib
281

282
        generated_py_files = glob.glob(os.path.join(CMAKE_BUILD_DIR, "onnx", "*.py"))
283
        generated_pyi_files = glob.glob(os.path.join(CMAKE_BUILD_DIR, "onnx", "*.pyi"))
284
        assert generated_py_files, "Bug: No generated python files found"
285
        assert generated_pyi_files, "Bug: No generated python stubs found"
286
        for src in (*generated_py_files, *generated_pyi_files):
287
            dst = os.path.join(dst_dir, os.path.relpath(src, CMAKE_BUILD_DIR))
288
            os.makedirs(os.path.dirname(dst), exist_ok=True)
289
            self.copy_file(src, dst)
290

291

292
CMD_CLASS = {
293
    "cmake_build": CmakeBuild,
294
    "build_py": BuildPy,
295
    "build_ext": BuildExt,
296
    "develop": Develop,
297
}
298

299
################################################################################
300
# Extensions
301
################################################################################
302

303
EXT_MODULES = [setuptools.Extension(name="onnx.onnx_cpp2py_export", sources=[])]
304

305

306
################################################################################
307
# Final
308
################################################################################
309

310
setuptools.setup(
311
    ext_modules=EXT_MODULES,
312
    cmdclass=CMD_CLASS,
313
    version=VERSION_INFO["version"],
314
    options=(
315
        {"bdist_wheel": {"plat_name": ONNX_WHEEL_PLATFORM_NAME}}
316
        if ONNX_WHEEL_PLATFORM_NAME is not None
317
        else {}
318
    ),
319
)
320