File size: 6,034 Bytes
f0c91ec | 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | from __future__ import annotations
import argparse
from pathlib import Path
def main() -> None:
parser = argparse.ArgumentParser(
description="Prepare kernel-builder output for a platform wheel build."
)
parser.add_argument("project", type=Path)
parser.add_argument("--version", required=True)
parser.add_argument(
"--torch-requirement",
default="torch>=2.11",
help=(
"torch dependency for the wheel metadata; non-stable-ABI variants "
'must pin the torch minor they were built against (e.g. "torch>=2.9,<2.10")'
),
)
args = parser.parse_args()
pyproject = args.project / "pyproject.toml"
text = pyproject.read_text(encoding="utf-8")
old_version = 'version = "0.1.0"'
if text.count(old_version) != 1:
raise RuntimeError("generated pyproject must contain one stub version")
text = text.replace(old_version, f'version = "{args.version}"', 1)
requires_python = 'requires-python = ">=3.9"'
if text.count(requires_python) != 1:
raise RuntimeError("generated pyproject must contain one Python requirement")
text = text.replace(
requires_python,
f'{requires_python}\ndependencies = ["{args.torch_requirement}"]',
1,
)
pyproject.write_text(text, encoding="utf-8")
cmake = args.project / "CMakeLists.txt"
cmake_text = cmake.read_text(encoding="utf-8")
for required in (False, True):
marker = " REQUIRED" if required else ""
development = (
f"find_package(Python3{marker} COMPONENTS Development "
"Development.SABIModule Interpreter)"
)
if cmake_text.count(development) != 1:
raise RuntimeError(
"generated CMake must contain one Python development lookup"
)
cmake_text = cmake_text.replace(
development,
f"find_package(Python3{marker} COMPONENTS Development.SABIModule Interpreter)",
1,
)
cmake.write_text(cmake_text, encoding="utf-8")
setup = args.project / "setup.py"
setup_text = setup.read_text(encoding="utf-8")
shutil_import = "from shutil import which, move\n"
if setup_text.count(shutil_import) != 1:
raise RuntimeError("generated setup must contain one shutil import")
setup_text = setup_text.replace(
shutil_import,
"from shutil import copy2, move, which\n",
1,
)
ninja_path = 'ninja_executable_path = Path(ninja.BIN_DIR) / "ninja"'
if setup_text.count(ninja_path) != 1:
raise RuntimeError("generated setup must contain one Ninja executable path")
setup_text = setup_text.replace(
ninja_path,
'ninja_executable_path = Path(ninja.BIN_DIR) / '
'("ninja.exe" if os.name == "nt" else "ninja")',
1,
)
for cache_tool in ("sccache", "ccache"):
availability = f'return which("{cache_tool}") is not None'
if setup_text.count(availability) != 1:
raise RuntimeError(
f"generated setup must contain one {cache_tool} availability check"
)
setup_text = setup_text.replace(
availability,
f'return os.name != "nt" and which("{cache_tool}") is not None',
1,
)
cmake_args_hook = (
' if "CMAKE_ARGS" in os.environ:\n'
' cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") '
"if item]\n"
)
if setup_text.count(cmake_args_hook) != 1:
raise RuntimeError("generated setup must contain one CMAKE_ARGS hook")
setup_text = setup_text.replace(
cmake_args_hook,
cmake_args_hook
+ ' cmake_make_program = os.environ.get("ORBITQUANT_CMAKE_MAKE_PROGRAM")\n'
+ " if cmake_make_program:\n"
+ ' cmake_args.append(f"-DCMAKE_MAKE_PROGRAM:FILEPATH={cmake_make_program}")\n',
1,
)
build_temp = " build_temp = Path(self.build_temp) / ext.name"
if setup_text.count(build_temp) != 1:
raise RuntimeError("generated setup must contain one extension build temp")
setup_text = setup_text.replace(
build_temp,
' build_temp_root = os.environ.get("ORBITQUANT_BUILD_TEMP", '
"self.build_temp)\n"
" build_temp = (Path(build_temp_root) / ext.name).resolve()",
1,
)
windows_multi_config = (
' if sys.platform == "win32":\n'
" # Move the dylib one folder up for discovery."
)
if setup_text.count(windows_multi_config) != 1:
raise RuntimeError("generated setup must contain one Windows output move")
setup_text = setup_text.replace(
windows_multi_config,
' if sys.platform == "win32" and (extdir / cfg).is_dir():\n'
" # Move the dylib one folder up for discovery.",
1,
)
build_call = (
" subprocess.run(\n"
' ["cmake", "--build", str(build_temp), *build_args], '
"cwd=build_temp, check=True\n"
" )\n"
)
if setup_text.count(build_call) != 1:
raise RuntimeError("generated setup must contain one wheel CMake build call")
setup_text = setup_text.replace(
build_call,
build_call
+ "\n"
+ ' package_name = ext.name.split(".", 1)[0]\n'
+ " generated_ops = (\n"
+ ' Path(ext.sourcedir) / "torch-ext" / package_name / "_ops.py"\n'
+ " )\n"
+ ' copy2(generated_ops, extdir / "_ops.py")\n',
1,
)
zip_safe = " zip_safe=False,\n"
if setup_text.count(zip_safe) != 1:
raise RuntimeError("generated setup must contain one zip-safe option")
setup_text = setup_text.replace(
zip_safe,
' options={"bdist_wheel": {"py_limited_api": "cp39"}},\n' + zip_safe,
1,
)
setup.write_text(setup_text, encoding="utf-8")
if __name__ == "__main__":
main()
|