# Copyright (c) ONNX Project Contributors
#
# SPDX-License-Identifier: Apache-2.0
# NOTE: Put all metadata in pyproject.toml.
# Set the environment variable `ONNX_PREVIEW_BUILD=1` to build the dev preview release.
from __future__ import annotations
import contextlib
import datetime
import glob
import logging
import multiprocessing
import os
import platform
import shlex
import shutil
import subprocess
import sys
import sysconfig
import textwrap
from typing import ClassVar
import setuptools
import setuptools.command.build_ext
import setuptools.command.build_py
import setuptools.command.develop
TOP_DIR = os.path.realpath(os.path.dirname(__file__))
CMAKE_BUILD_DIR = os.path.join(TOP_DIR, ".setuptools-cmake-build")
WINDOWS = os.name == "nt"
CMAKE = shutil.which("cmake3") or shutil.which("cmake")
################################################################################
# Global variables for controlling the build variant
################################################################################
# Default value is set to TRUE\1 to keep the settings same as the current ones.
# However going forward the recommended way to is to set this to False\0
ONNX_ML = os.getenv("ONNX_ML") != "0"
ONNX_VERIFY_PROTO3 = os.getenv("ONNX_VERIFY_PROTO3") == "1"
ONNX_NAMESPACE = os.getenv("ONNX_NAMESPACE", "onnx")
ONNX_BUILD_TESTS = os.getenv("ONNX_BUILD_TESTS") == "1"
ONNX_DISABLE_EXCEPTIONS = os.getenv("ONNX_DISABLE_EXCEPTIONS") == "1"
ONNX_DISABLE_STATIC_REGISTRATION = os.getenv("ONNX_DISABLE_STATIC_REGISTRATION") == "1"
ONNX_PREVIEW_BUILD = os.getenv("ONNX_PREVIEW_BUILD") == "1"
USE_MSVC_STATIC_RUNTIME = os.getenv("USE_MSVC_STATIC_RUNTIME", "0") == "1"
DEBUG = os.getenv("DEBUG", "0") == "1"
COVERAGE = os.getenv("COVERAGE", "0") == "1"
# Customize the wheel plat-name; sometimes useful for MacOS builds.
# See https://github.com/onnx/onnx/pull/6117
ONNX_WHEEL_PLATFORM_NAME = os.getenv("ONNX_WHEEL_PLATFORM_NAME")
################################################################################
# Pre Check
################################################################################
assert CMAKE, "Could not find cmake in PATH"
################################################################################
# Version
################################################################################
try:
_git_version = (
subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=TOP_DIR)
.decode("ascii")
.strip()
)
except (OSError, subprocess.CalledProcessError):
_git_version = ""
with open(os.path.join(TOP_DIR, "VERSION_NUMBER"), encoding="utf-8") as version_file:
_version = version_file.read().strip()
if ONNX_PREVIEW_BUILD:
# Create the dev build for weekly releases / dev build
todays_date = datetime.date.today().strftime("%Y%m%d")
_version += ".dev" + todays_date
VERSION_INFO = {"version": _version, "git_version": _git_version}
################################################################################
# Utilities
################################################################################
@contextlib.contextmanager
def cd(path):
if not os.path.isabs(path):
raise RuntimeError(f"Can only cd to absolute path, got: {path}")
orig_path = os.getcwd()
os.chdir(path)
try:
yield
finally:
os.chdir(orig_path)
def get_ext_suffix():
return sysconfig.get_config_var("EXT_SUFFIX")
def get_python_execute():
if WINDOWS:
return sys.executable
# Try to search more accurate path, because sys.executable may return a wrong one,
# as discussed in https://github.com/python/cpython/issues/84399
python_dir = os.path.abspath(
os.path.join(sysconfig.get_path("include"), "..", "..")
)
if os.path.isdir(python_dir):
python_bin = os.path.join(python_dir, "bin", "python3")
if os.path.isfile(python_bin):
return python_bin
python_bin = os.path.join(python_dir, "bin", "python")
if os.path.isfile(python_bin):
return python_bin
return sys.executable
################################################################################
# Customized commands
################################################################################
def create_version(directory: str):
"""Create version.py based on VERSION_INFO."""
version_file_path = os.path.join(directory, "onnx", "version.py")
os.makedirs(os.path.dirname(version_file_path), exist_ok=True)
with open(version_file_path, "w", encoding="utf-8") as f:
f.write(
textwrap.dedent(
f"""\
# This file is generated by setup.py. DO NOT EDIT!
version = "{VERSION_INFO['version']}"
git_version = "{VERSION_INFO['git_version']}"
"""
)
)
class CmakeBuild(setuptools.Command):
"""Compiles everything when `python setup.py build` is run using cmake.
Custom args can be passed to cmake by specifying the `CMAKE_ARGS`
environment variable.
The number of CPUs used by `make` can be specified by passing `-j