microsoft/onnxruntime-extensions

Public

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

CodeCommitsIssuesPull requestsActionsInsightsSecurity
b661d5f22f396e757eb1de6e1ab28f2a50f0e81b

Branches

Tags

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

Clone

HTTPS

Download ZIP

.pyproject/cmdclass.py

327lines · 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, OSError):
34 pass
35
36 return None
37
38
39def _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
53def _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
71def prepare_env(project_root):
72 if sys.platform == "win32":
73 _load_vsdevcmd(project_root)
74
75
76def 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
96class 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
119class CmdDevelop(CommandMixin, _develop):
120 user_options = getattr(_develop, 'user_options', []
121 ) + CommandMixin.user_options
122
123
124class 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' / 'win-arm64' in Windows to fix that.
133 if sys.platform == "win32":
134 if "arm" in sys.version.lower():
135 self.plat_name = "win-arm64"
136 else:
137 self.plat_name = "win-amd64"
138 if os.environ.get('OCOS_SCB_DEBUG', None) == '1':
139 self.debug = True
140 super().finalize_options()
141
142
143class CmdBuildCMakeExt(_build_ext):
144
145 # noinspection PyAttributeOutsideInit
146 def initialize_options(self):
147 super().initialize_options()
148 self.use_cuda = None
149 self.no_azure = True
150 self.no_opencv = True
151 self.cc_debug = None
152 self.pp_api = True
153 self.cuda_archs = None
154 self.ort_pkg_dir = None
155
156 def _parse_options(self, options):
157 for segment in options.split(','):
158 if not segment:
159 continue
160 key = segment
161 if '=' in segment:
162 key, value = segment.split('=')
163 else:
164 value = 1
165
166 key = key.replace('-', '_')
167 if not hasattr(self, key):
168 raise RuntimeError(
169 f"Unknown {ORTX_USER_OPTION} option value: {key}")
170 setattr(self, key, value)
171 return self
172
173 def finalize_options(self) -> None:
174 if CommandMixin.config_settings is not None:
175 self._parse_options(
176 CommandMixin.config_settings.get(ORTX_USER_OPTION, ""))
177 if self.cc_debug:
178 self.debug = True
179 super().finalize_options()
180
181 def run(self):
182 """
183 Perform build_cmake before doing the 'normal' stuff
184 """
185 for extension in self.extensions:
186 if extension.name == 'onnxruntime_extensions._extensions_pydll':
187 self.build_cmake(extension)
188
189 def build_cmake(self, extension):
190 project_dir = pathlib.Path().absolute()
191 build_temp = pathlib.Path(self.build_temp)
192 build_temp.mkdir(parents=True, exist_ok=True)
193 ext_fullpath = pathlib.Path(
194 self.get_ext_fullpath(extension.name)).absolute()
195
196 config = 'RelWithDebInfo' if self.debug else 'Release'
197 cmake_args = [
198 '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' +
199 str(ext_fullpath.parent.absolute()),
200 '-DCMAKE_POLICY_VERSION_MINIMUM=3.5',
201 '-DOCOS_ENABLE_CTEST=OFF',
202 '-DOCOS_BUILD_PYTHON=ON',
203 '-DOCOS_PYTHON_MODULE_PATH=' + str(ext_fullpath),
204 '-DCMAKE_BUILD_TYPE=' + config
205 ]
206
207 if self.ort_pkg_dir:
208 cmake_args += ['-DONNXRUNTIME_PKG_DIR=' + self.ort_pkg_dir]
209
210 if self.no_opencv:
211 # Disabling openCV can drastically reduce the build time.
212 cmake_args += [
213 '-DOCOS_ENABLE_OPENCV_CODECS=OFF',
214 '-DOCOS_ENABLE_CV2=OFF']
215
216 if self.pp_api:
217 if not self.no_opencv:
218 raise RuntimeError(
219 "Cannot enable PP C API Python Wrapper without disabling OpenCV.")
220 cmake_args += ['-DOCOS_ENABLE_C_API=ON']
221
222 if self.no_azure is not None:
223 azure_flag = "OFF" if self.no_azure == 1 else "ON"
224 cmake_args += ['-DOCOS_ENABLE_AZURE=' + azure_flag]
225 print("=> AzureOp build flag: " + azure_flag)
226
227 if self.use_cuda is not None:
228 cuda_flag = "OFF" if self.use_cuda == 0 else "ON"
229 cmake_args += ['-DOCOS_USE_CUDA=' + cuda_flag]
230 print("=> CUDA build flag: " + cuda_flag)
231 if cuda_flag == "ON":
232 cuda_ver = _load_cuda_version()
233 if cuda_ver is None:
234 raise RuntimeError("Cannot find nvcc in your env:path, use-cuda doesn't work")
235 if sys.platform == "win32":
236 cuda_path = os.environ.get("CUDA_PATH")
237 cmake_args += [f'-T cuda={cuda_path}']
238 # TODO: temporarily add a flag for MSVC 19.40
239 cmake_args += ['-DCMAKE_CUDA_FLAGS_INIT=-allow-unsupported-compiler']
240 f_ver = ext_fullpath.parent / "_version.py"
241 with f_ver.open('a') as _f:
242 _f.writelines(["\n", f"cuda = \"{cuda_ver}\"", "\n"])
243
244 if self.cuda_archs is not None:
245 cmake_args += ['-DCMAKE_CUDA_ARCHITECTURES=' + self.cuda_archs]
246 else:
247 smi = _load_nvidia_smi()
248 if not smi:
249 raise RuntimeError(
250 "Cannot detect the CUDA archs from your machine, please specify it manually.")
251 cmake_args += ['-DCMAKE_CUDA_ARCHITECTURES=' + smi]
252
253 # CMake lets you override the generator - we need to check this.
254 # Can be set with Conda-Build, for example.
255 cmake_generator = os.environ.get("CMAKE_GENERATOR", "")
256 # Adding CMake arguments set as environment variable
257 # (needed e.g. to build for ARM OSx on conda-forge)
258 if "CMAKE_ARGS" in os.environ:
259 cmake_args += [
260 item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
261
262 if sys.platform != "win32":
263 # Using Ninja-build since it a) is available as a wheel and b)
264 # multithread automatically. MSVC would require all variables be
265 # exported for Ninja to pick it up, which is a little tricky to do.
266 # Users can override the generator with CMAKE_GENERATOR in CMake
267 # 3.15+.
268 if not cmake_generator or cmake_generator == "Ninja":
269 try:
270 import ninja # noqa: F401
271
272 ninja_executable_path = os.path.join(
273 ninja.BIN_DIR, "ninja")
274 cmake_args += [
275 "-GNinja",
276 f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}",
277 ]
278 except ImportError:
279 pass
280
281 if sys.platform.startswith("darwin"):
282 cmake_args += ["-DCMAKE_OSX_DEPLOYMENT_TARGET=10.15"]
283 # Cross-compile support for macOS - respect ARCHFLAGS if set
284 archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
285 if archs:
286 cmake_args += [
287 "-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
288
289 # overwrite the Python module info if the auto-detection doesn't work.
290 # export Python3_INCLUDE_DIRS=/opt/python/cp38-cp38
291 # export Python3_LIBRARIES=/opt/python/cp38-cp38
292 for env in ['Python3_INCLUDE_DIRS', 'Python3_LIBRARIES']:
293 if env in os.environ:
294 cmake_args.append("-D%s=%s" % (env, os.environ[env]))
295
296 if self.debug:
297 cmake_args += ['-DCC_OPTIMIZE=OFF']
298
299 # the parallel build has to be limited on some Linux VM machine.
300 cpu_number = os.environ.get('CPU_NUMBER')
301 build_args = [
302 '--config', config,
303 '--parallel' + ('' if cpu_number is None else ' ' + cpu_number)
304 ]
305 cmake_exe = 'cmake'
306 # if sys.platform == "win32":
307 # # unlike Linux/macOS, cmake pip package on Windows fails to build some 3rd party dependencies.
308 # # so we have to use the cmake from a standalone installation or the one from Visual Studio.
309 # standalone_cmake = os.path.join(os.environ.get("ProgramFiles"), "\\CMake\\bin\\cmake.exe")
310 # if os.path.exists(standalone_cmake):
311 # cmake_exe = standalone_cmake
312 # elif os.environ.get(VSINSTALLDIR_NAME):
313 # cmake_exe = os.environ[VSINSTALLDIR_NAME] + \
314 # 'Common7\\IDE\\CommonExtensions\\Microsoft\\CMake\\CMake\\bin\\cmake.exe'
315 # # Add this cmake directory into PATH to make sure the child-process still find it.
316 # os.environ['PATH'] = os.path.dirname(
317 # cmake_exe) + os.pathsep + os.environ['PATH']
318
319 self.spawn([cmake_exe, '-S', str(project_dir),
320 '-B', str(build_temp)] + cmake_args)
321 if not self.dry_run:
322 self.spawn([cmake_exe, '--build', str(build_temp)] + build_args)
323
324
325ortx_cmdclass = dict(build=CmdBuild,
326 develop=CmdDevelop,
327 build_ext=CmdBuildCMakeExt)
328