GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/onnxruntime-extensions

Public

mirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
sayanshaw/genai-tutorial

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/.pyproject

.pyproject/cmdclass.py

275lines · modecode

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