forked from NVIDIA/cuda-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbuild_hooks.py
More file actions
144 lines (111 loc) · 5.12 KB
/
build_hooks.py
File metadata and controls
144 lines (111 loc) · 5.12 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
# SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0
# This module implements basic PEP 517 backend support, see e.g.
# - https://peps.python.org/pep-0517/
# - https://setuptools.pypa.io/en/latest/build_meta.html#dynamic-build-dependencies-and-other-build-meta-tweaks
# Specifically, there are 5 APIs required to create a proper build backend, see below.
import functools
import glob
import os
import re
import subprocess
from Cython.Build import cythonize
from setuptools import Extension
from setuptools import build_meta as _build_meta
prepare_metadata_for_build_editable = _build_meta.prepare_metadata_for_build_editable
prepare_metadata_for_build_wheel = _build_meta.prepare_metadata_for_build_wheel
build_sdist = _build_meta.build_sdist
get_requires_for_build_sdist = _build_meta.get_requires_for_build_sdist
COMPILE_FOR_COVERAGE = bool(int(os.environ.get("CUDA_PYTHON_COVERAGE", "0")))
@functools.cache
def _get_proper_cuda_bindings_major_version() -> str:
# for local development (with/without build isolation)
try:
import cuda.bindings
return cuda.bindings.__version__.split(".")[0]
except ImportError:
pass
# for custom overwrite, e.g. in CI
cuda_major = os.environ.get("CUDA_CORE_BUILD_MAJOR")
if cuda_major is not None:
return cuda_major
# also for local development
try:
out = subprocess.run("nvidia-smi", env=os.environ, capture_output=True, check=True) # noqa: S603, S607
m = re.search(r"CUDA Version:\s*([\d\.]+)", out.stdout.decode())
if m:
return m.group(1).split(".")[0]
except (FileNotFoundError, subprocess.CalledProcessError):
# the build machine has no driver installed
pass
# default fallback
return "13"
# used later by setup()
_extensions = None
def _build_cuda_core():
# Customizing the build hooks is needed because we must defer cythonization until cuda-bindings,
# now a required build-time dependency that's dynamically installed via the other hook below,
# is installed. Otherwise, cimport any cuda.bindings modules would fail!
#
# This function populates "_extensions".
global _extensions
# It seems setuptools' wildcard support has problems for namespace packages,
# so we explicitly spell out all Extension instances.
root_module = "cuda.core"
root_path = f"{os.path.sep}".join(root_module.split(".")) + os.path.sep
ext_files = glob.glob(f"{root_path}/**/*.pyx", recursive=True)
def strip_prefix_suffix(filename):
return filename[len(root_path) : -4]
module_names = (strip_prefix_suffix(f) for f in ext_files)
@functools.cache
def get_cuda_paths():
CUDA_PATH = os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME", None))
if not CUDA_PATH:
raise RuntimeError("Environment variable CUDA_PATH or CUDA_HOME is not set")
CUDA_PATH = CUDA_PATH.split(os.pathsep)
print("CUDA paths:", CUDA_PATH)
return CUDA_PATH
all_include_dirs = list(os.path.join(root, "include") for root in get_cuda_paths())
extra_compile_args = []
if COMPILE_FOR_COVERAGE:
# CYTHON_TRACE_NOGIL indicates to trace nogil functions. It is not
# related to free-threading builds.
extra_compile_args += ["-DCYTHON_TRACE_NOGIL=1", "-DCYTHON_USE_SYS_MONITORING=0"]
ext_modules = tuple(
Extension(
f"cuda.core.{mod.replace(os.path.sep, '.')}",
sources=[f"cuda/core/{mod}.pyx"],
include_dirs=all_include_dirs,
language="c++",
extra_compile_args=extra_compile_args,
)
for mod in module_names
)
nthreads = int(os.environ.get("CUDA_PYTHON_PARALLEL_LEVEL", os.cpu_count() // 2))
compile_time_env = {"CUDA_CORE_BUILD_MAJOR": int(_get_proper_cuda_bindings_major_version())}
compiler_directives = {"embedsignature": True, "warn.deprecated.IF": False, "freethreading_compatible": True}
if COMPILE_FOR_COVERAGE:
compiler_directives["linetrace"] = True
_extensions = cythonize(
ext_modules,
verbose=True,
language_level=3,
nthreads=nthreads,
compiler_directives=compiler_directives,
compile_time_env=compile_time_env,
)
return
def build_editable(wheel_directory, config_settings=None, metadata_directory=None):
_build_cuda_core()
return _build_meta.build_editable(wheel_directory, config_settings, metadata_directory)
def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
_build_cuda_core()
return _build_meta.build_wheel(wheel_directory, config_settings, metadata_directory)
def _get_cuda_bindings_require():
cuda_major = _get_proper_cuda_bindings_major_version()
return [f"cuda-bindings=={cuda_major}.*"]
def get_requires_for_build_editable(config_settings=None):
return _build_meta.get_requires_for_build_editable(config_settings) + _get_cuda_bindings_require()
def get_requires_for_build_wheel(config_settings=None):
return _build_meta.get_requires_for_build_wheel(config_settings) + _get_cuda_bindings_require()