microsoft/onnxruntime-extensions
Publicmirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable
.pyproject/cmdclass.py
325lines · modecode
| 1 | # -*- coding: utf-8 -*- |
| 2 | # Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | # Licensed under the MIT License. See License.txt in the project root for |
| 4 | # license information. |
| 5 | ########################################################################### |
| 6 | |
| 7 | import re |
| 8 | import os |
| 9 | import sys |
| 10 | import pathlib |
| 11 | import subprocess |
| 12 | |
| 13 | from textwrap import dedent |
| 14 | from setuptools.command.build import build as _build |
| 15 | from setuptools.command.build_ext import build_ext as _build_ext |
| 16 | from setuptools.command.develop import develop as _develop |
| 17 | |
| 18 | VSINSTALLDIR_NAME = 'VSINSTALLDIR' |
| 19 | ORTX_USER_OPTION = 'ortx-user-option' |
| 20 | |
| 21 | |
| 22 | def _load_cuda_version(): |
| 23 | nvcc_path = 'nvcc' |
| 24 | cuda_path = os.environ.get('CUDA_PATH') |
| 25 | if cuda_path is not None: |
| 26 | nvcc_path = os.path.join(cuda_path, 'bin', 'nvcc') |
| 27 | try: |
| 28 | output = subprocess.check_output([nvcc_path, "--version"], stderr=subprocess.STDOUT).decode("utf-8") |
| 29 | pattern = r"\bV(\d+\.\d+\.\d+)\b" |
| 30 | match = re.search(pattern, output) |
| 31 | if match: |
| 32 | return match.group(1) |
| 33 | except (subprocess.CalledProcessError, OSError): |
| 34 | pass |
| 35 | |
| 36 | return None |
| 37 | |
| 38 | |
| 39 | def _load_nvidia_smi(): |
| 40 | try: |
| 41 | outputs = subprocess.check_output( |
| 42 | ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader,nounits"], |
| 43 | stderr=subprocess.STDOUT).decode("utf-8").splitlines() |
| 44 | output = outputs[0] if outputs else "" |
| 45 | arch = output.strip().replace('.', '') |
| 46 | return arch if arch.isdigit() else None |
| 47 | except (subprocess.CalledProcessError, OSError): |
| 48 | pass |
| 49 | |
| 50 | return None |
| 51 | |
| 52 | |
| 53 | def _load_vsdevcmd(project_root): |
| 54 | if os.environ.get(VSINSTALLDIR_NAME) is None: |
| 55 | stdout, _ = subprocess.Popen([ |
| 56 | 'powershell', ' -noprofile', '-executionpolicy', |
| 57 | 'bypass', '-f', project_root + '/tools/get_vsdevcmd.ps1', '-outputEnv', '1'], |
| 58 | stdout=subprocess.PIPE, shell=False, universal_newlines=True).communicate() |
| 59 | for line in stdout.splitlines(): |
| 60 | kv_pair = line.split('=') |
| 61 | if len(kv_pair) == 2: |
| 62 | os.environ[kv_pair[0]] = kv_pair[1] |
| 63 | else: |
| 64 | import shutil |
| 65 | if shutil.which('cmake') is None: |
| 66 | raise SystemExit( |
| 67 | "Cannot find cmake in the executable path, " |
| 68 | "please run this script under Developer Command Prompt for VS.") |
| 69 | |
| 70 | |
| 71 | def prepare_env(project_root): |
| 72 | if sys.platform == "win32": |
| 73 | _load_vsdevcmd(project_root) |
| 74 | |
| 75 | |
| 76 | def read_git_refs(project_root): |
| 77 | release_branch = False |
| 78 | stdout, _ = subprocess.Popen( |
| 79 | ['git'] + ['log', '-1', '--format=%H'], |
| 80 | cwd=project_root, |
| 81 | stdout=subprocess.PIPE, universal_newlines=True).communicate() |
| 82 | HEAD = dedent(stdout.splitlines()[0]).strip('\n\r') |
| 83 | stdout, _ = subprocess.Popen( |
| 84 | ['git'] + ['show-ref', '--head'], |
| 85 | cwd=project_root, |
| 86 | stdout=subprocess.PIPE, universal_newlines=True).communicate() |
| 87 | for _ln in stdout.splitlines(): |
| 88 | _ln = dedent(_ln).strip('\n\r') |
| 89 | if _ln.startswith(HEAD): |
| 90 | _, _2 = _ln.split(' ') |
| 91 | if _2.startswith('refs/remotes/origin/rel-'): |
| 92 | release_branch = True |
| 93 | return release_branch, HEAD |
| 94 | |
| 95 | |
| 96 | class CommandMixin: |
| 97 | user_options = [ |
| 98 | (ORTX_USER_OPTION + '=', None, "extensions options for kernel building") |
| 99 | ] |
| 100 | config_settings = None |
| 101 | |
| 102 | # noinspection PyAttributeOutsideInit |
| 103 | def initialize_options(self) -> None: |
| 104 | super().initialize_options() |
| 105 | self.ortx_user_option = None |
| 106 | |
| 107 | def finalize_options(self) -> None: |
| 108 | if self.ortx_user_option is not None: |
| 109 | if CommandMixin.config_settings is None: |
| 110 | CommandMixin.config_settings = { |
| 111 | ORTX_USER_OPTION: self.ortx_user_option} |
| 112 | else: |
| 113 | raise RuntimeError( |
| 114 | f"Cannot pass {ORTX_USER_OPTION} several times, like as the command args and in backend API.") |
| 115 | |
| 116 | super().finalize_options() |
| 117 | |
| 118 | |
| 119 | class CmdDevelop(CommandMixin, _develop): |
| 120 | user_options = getattr(_develop, 'user_options', [] |
| 121 | ) + CommandMixin.user_options |
| 122 | |
| 123 | |
| 124 | class CmdBuild(CommandMixin, _build): |
| 125 | user_options = getattr(_build, 'user_options', []) + \ |
| 126 | CommandMixin.user_options |
| 127 | |
| 128 | # noinspection PyAttributeOutsideInit |
| 129 | def finalize_options(self) -> None: |
| 130 | # There is a bug in setuptools that prevents the build get the right platform name from arguments. |
| 131 | # So, it cannot generate the correct wheel with the right arch in Official release pipeline. |
| 132 | # Force plat_name to be 'win-amd64' in Windows to fix that, |
| 133 | # since extensions cmake is only available on x64 for Windows now, it is not a problem to hardcode it. |
| 134 | if sys.platform == "win32" and "arm" not in sys.version.lower(): |
| 135 | self.plat_name = "win-amd64" |
| 136 | if os.environ.get('OCOS_SCB_DEBUG', None) == '1': |
| 137 | self.debug = True |
| 138 | super().finalize_options() |
| 139 | |
| 140 | |
| 141 | class CmdBuildCMakeExt(_build_ext): |
| 142 | |
| 143 | # noinspection PyAttributeOutsideInit |
| 144 | def initialize_options(self): |
| 145 | super().initialize_options() |
| 146 | self.use_cuda = None |
| 147 | self.no_azure = True |
| 148 | self.no_opencv = True |
| 149 | self.cc_debug = None |
| 150 | self.pp_api = True |
| 151 | self.cuda_archs = None |
| 152 | self.ort_pkg_dir = None |
| 153 | |
| 154 | def _parse_options(self, options): |
| 155 | for segment in options.split(','): |
| 156 | if not segment: |
| 157 | continue |
| 158 | key = segment |
| 159 | if '=' in segment: |
| 160 | key, value = segment.split('=') |
| 161 | else: |
| 162 | value = 1 |
| 163 | |
| 164 | key = key.replace('-', '_') |
| 165 | if not hasattr(self, key): |
| 166 | raise RuntimeError( |
| 167 | f"Unknown {ORTX_USER_OPTION} option value: {key}") |
| 168 | setattr(self, key, value) |
| 169 | return self |
| 170 | |
| 171 | def finalize_options(self) -> None: |
| 172 | if CommandMixin.config_settings is not None: |
| 173 | self._parse_options( |
| 174 | CommandMixin.config_settings.get(ORTX_USER_OPTION, "")) |
| 175 | if self.cc_debug: |
| 176 | self.debug = True |
| 177 | super().finalize_options() |
| 178 | |
| 179 | def run(self): |
| 180 | """ |
| 181 | Perform build_cmake before doing the 'normal' stuff |
| 182 | """ |
| 183 | for extension in self.extensions: |
| 184 | if extension.name == 'onnxruntime_extensions._extensions_pydll': |
| 185 | self.build_cmake(extension) |
| 186 | |
| 187 | def build_cmake(self, extension): |
| 188 | project_dir = pathlib.Path().absolute() |
| 189 | build_temp = pathlib.Path(self.build_temp) |
| 190 | build_temp.mkdir(parents=True, exist_ok=True) |
| 191 | ext_fullpath = pathlib.Path( |
| 192 | self.get_ext_fullpath(extension.name)).absolute() |
| 193 | |
| 194 | config = 'RelWithDebInfo' if self.debug else 'Release' |
| 195 | cmake_args = [ |
| 196 | '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + |
| 197 | str(ext_fullpath.parent.absolute()), |
| 198 | '-DCMAKE_POLICY_VERSION_MINIMUM=3.5', |
| 199 | '-DOCOS_ENABLE_CTEST=OFF', |
| 200 | '-DOCOS_BUILD_PYTHON=ON', |
| 201 | '-DOCOS_PYTHON_MODULE_PATH=' + str(ext_fullpath), |
| 202 | '-DCMAKE_BUILD_TYPE=' + config |
| 203 | ] |
| 204 | |
| 205 | if self.ort_pkg_dir: |
| 206 | cmake_args += ['-DONNXRUNTIME_PKG_DIR=' + self.ort_pkg_dir] |
| 207 | |
| 208 | if self.no_opencv: |
| 209 | # Disabling openCV can drastically reduce the build time. |
| 210 | cmake_args += [ |
| 211 | '-DOCOS_ENABLE_OPENCV_CODECS=OFF', |
| 212 | '-DOCOS_ENABLE_CV2=OFF'] |
| 213 | |
| 214 | if self.pp_api: |
| 215 | if not self.no_opencv: |
| 216 | raise RuntimeError( |
| 217 | "Cannot enable PP C API Python Wrapper without disabling OpenCV.") |
| 218 | cmake_args += ['-DOCOS_ENABLE_C_API=ON'] |
| 219 | |
| 220 | if self.no_azure is not None: |
| 221 | azure_flag = "OFF" if self.no_azure == 1 else "ON" |
| 222 | cmake_args += ['-DOCOS_ENABLE_AZURE=' + azure_flag] |
| 223 | print("=> AzureOp build flag: " + azure_flag) |
| 224 | |
| 225 | if self.use_cuda is not None: |
| 226 | cuda_flag = "OFF" if self.use_cuda == 0 else "ON" |
| 227 | cmake_args += ['-DOCOS_USE_CUDA=' + cuda_flag] |
| 228 | print("=> CUDA build flag: " + cuda_flag) |
| 229 | if cuda_flag == "ON": |
| 230 | cuda_ver = _load_cuda_version() |
| 231 | if cuda_ver is None: |
| 232 | raise RuntimeError("Cannot find nvcc in your env:path, use-cuda doesn't work") |
| 233 | if sys.platform == "win32": |
| 234 | cuda_path = os.environ.get("CUDA_PATH") |
| 235 | cmake_args += [f'-T cuda={cuda_path}'] |
| 236 | # TODO: temporarily add a flag for MSVC 19.40 |
| 237 | cmake_args += ['-DCMAKE_CUDA_FLAGS_INIT=-allow-unsupported-compiler'] |
| 238 | f_ver = ext_fullpath.parent / "_version.py" |
| 239 | with f_ver.open('a') as _f: |
| 240 | _f.writelines(["\n", f"cuda = \"{cuda_ver}\"", "\n"]) |
| 241 | |
| 242 | if self.cuda_archs is not None: |
| 243 | cmake_args += ['-DCMAKE_CUDA_ARCHITECTURES=' + self.cuda_archs] |
| 244 | else: |
| 245 | smi = _load_nvidia_smi() |
| 246 | if not smi: |
| 247 | raise RuntimeError( |
| 248 | "Cannot detect the CUDA archs from your machine, please specify it manually.") |
| 249 | cmake_args += ['-DCMAKE_CUDA_ARCHITECTURES=' + smi] |
| 250 | |
| 251 | # CMake lets you override the generator - we need to check this. |
| 252 | # Can be set with Conda-Build, for example. |
| 253 | cmake_generator = os.environ.get("CMAKE_GENERATOR", "") |
| 254 | # Adding CMake arguments set as environment variable |
| 255 | # (needed e.g. to build for ARM OSx on conda-forge) |
| 256 | if "CMAKE_ARGS" in os.environ: |
| 257 | cmake_args += [ |
| 258 | item for item in os.environ["CMAKE_ARGS"].split(" ") if item] |
| 259 | |
| 260 | if sys.platform != "win32": |
| 261 | # Using Ninja-build since it a) is available as a wheel and b) |
| 262 | # multithread automatically. MSVC would require all variables be |
| 263 | # exported for Ninja to pick it up, which is a little tricky to do. |
| 264 | # Users can override the generator with CMAKE_GENERATOR in CMake |
| 265 | # 3.15+. |
| 266 | if not cmake_generator or cmake_generator == "Ninja": |
| 267 | try: |
| 268 | import ninja # noqa: F401 |
| 269 | |
| 270 | ninja_executable_path = os.path.join( |
| 271 | ninja.BIN_DIR, "ninja") |
| 272 | cmake_args += [ |
| 273 | "-GNinja", |
| 274 | f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}", |
| 275 | ] |
| 276 | except ImportError: |
| 277 | pass |
| 278 | |
| 279 | if sys.platform.startswith("darwin"): |
| 280 | cmake_args += ["-DCMAKE_OSX_DEPLOYMENT_TARGET=10.15"] |
| 281 | # Cross-compile support for macOS - respect ARCHFLAGS if set |
| 282 | archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", "")) |
| 283 | if archs: |
| 284 | cmake_args += [ |
| 285 | "-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))] |
| 286 | |
| 287 | # overwrite the Python module info if the auto-detection doesn't work. |
| 288 | # export Python3_INCLUDE_DIRS=/opt/python/cp38-cp38 |
| 289 | # export Python3_LIBRARIES=/opt/python/cp38-cp38 |
| 290 | for env in ['Python3_INCLUDE_DIRS', 'Python3_LIBRARIES']: |
| 291 | if env in os.environ: |
| 292 | cmake_args.append("-D%s=%s" % (env, os.environ[env])) |
| 293 | |
| 294 | if self.debug: |
| 295 | cmake_args += ['-DCC_OPTIMIZE=OFF'] |
| 296 | |
| 297 | # the parallel build has to be limited on some Linux VM machine. |
| 298 | cpu_number = os.environ.get('CPU_NUMBER') |
| 299 | build_args = [ |
| 300 | '--config', config, |
| 301 | '--parallel' + ('' if cpu_number is None else ' ' + cpu_number) |
| 302 | ] |
| 303 | cmake_exe = 'cmake' |
| 304 | # if sys.platform == "win32": |
| 305 | # # unlike Linux/macOS, cmake pip package on Windows fails to build some 3rd party dependencies. |
| 306 | # # so we have to use the cmake from a standalone installation or the one from Visual Studio. |
| 307 | # standalone_cmake = os.path.join(os.environ.get("ProgramFiles"), "\\CMake\\bin\\cmake.exe") |
| 308 | # if os.path.exists(standalone_cmake): |
| 309 | # cmake_exe = standalone_cmake |
| 310 | # elif os.environ.get(VSINSTALLDIR_NAME): |
| 311 | # cmake_exe = os.environ[VSINSTALLDIR_NAME] + \ |
| 312 | # 'Common7\\IDE\\CommonExtensions\\Microsoft\\CMake\\CMake\\bin\\cmake.exe' |
| 313 | # # Add this cmake directory into PATH to make sure the child-process still find it. |
| 314 | # os.environ['PATH'] = os.path.dirname( |
| 315 | # cmake_exe) + os.pathsep + os.environ['PATH'] |
| 316 | |
| 317 | self.spawn([cmake_exe, '-S', str(project_dir), |
| 318 | '-B', str(build_temp)] + cmake_args) |
| 319 | if not self.dry_run: |
| 320 | self.spawn([cmake_exe, '--build', str(build_temp)] + build_args) |
| 321 | |
| 322 | |
| 323 | ortx_cmdclass = dict(build=CmdBuild, |
| 324 | develop=CmdDevelop, |
| 325 | build_ext=CmdBuildCMakeExt) |
| 326 | |