microsoft/onnxruntime-extensions

Public

mirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
tfmtok

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

.pyproject/cmdclass.py

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