Spaces:
Sleeping
Sleeping
| import argparse | |
| import os | |
| import shutil | |
| import pathlib | |
| import platform | |
| from typing import NamedTuple | |
| import setuptools | |
| from wheel.bdist_wheel import bdist_wheel | |
| # Parse --build-option arguments meant for the bdist_wheel command. We have to parse these | |
| # ourselves because when bdist_wheel runs it's too late to select a subset of libraries for package_data. | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("command") | |
| parser.add_argument( | |
| "--platform", "-P", type=str, default="", help="Wheel platform: windows|linux|macos-x86_64|aarch64|universal" | |
| ) | |
| args = parser.parse_known_args()[0] | |
| # return a canonical machine architecture string | |
| # - "x86_64" for x86-64, aka. AMD64, aka. x64 | |
| # - "aarch64" for AArch64, aka. ARM64 | |
| def machine_architecture() -> str: | |
| machine = platform.machine() | |
| if machine == "x86_64" or machine == "AMD64": | |
| return "x86_64" | |
| if machine == "aarch64" or machine == "arm64": | |
| return "aarch64" | |
| raise RuntimeError(f"Unrecognized machine architecture {machine}") | |
| class Platform(NamedTuple): | |
| os: str | |
| arch: str | |
| fancy_name: str | |
| extension: str | |
| tag: str | |
| def name(self) -> str: | |
| return self.os + "-" + self.arch | |
| platforms = [ | |
| Platform("windows", "x86_64", "Windows x86-64", ".dll", "win_amd64"), | |
| Platform("linux", "x86_64", "Linux x86-64", ".so", "manylinux2014_x86_64"), | |
| Platform("linux", "aarch64", "Linux AArch64", ".so", "manylinux2014_aarch64"), | |
| Platform("macos", "universal", "macOS universal", ".dylib", "macosx_10_13_universal2"), | |
| ] | |
| class Library(NamedTuple): | |
| file: str | |
| directory: str | |
| platform: Platform | |
| # Enumerate warp/bin libraries | |
| def detect_warp_libraries(): | |
| detected_libraries = set() | |
| warp_bin = pathlib.Path("warp/bin") | |
| for file in warp_bin.rglob("*.*"): | |
| for p in platforms: | |
| if os.path.splitext(file.name)[1] == p.extension: | |
| # If this is a local build, assume we want a wheel for this machine's architecture | |
| if file.parent.name == "bin" and (p.arch == machine_architecture() or p.arch == "universal"): | |
| detected_libraries.add(Library(file.name, "bin/", p)) | |
| else: | |
| # Excpect libraries to be in a subdirectory named after the wheel platform | |
| platform_name = p.name() | |
| if file.parent.name == platform_name: | |
| detected_libraries.add(Library(file.name, "bin/" + platform_name + "/", p)) | |
| if len(detected_libraries) == 0: | |
| raise Exception("No libraries found in warp/bin. Please run build_lib.py first.") | |
| return detected_libraries | |
| detected_libraries = detect_warp_libraries() | |
| detected_platforms = set([lib.platform for lib in detected_libraries]) | |
| wheel_platform = None # The one platform for which we're building a wheel | |
| if args.command == "bdist_wheel": | |
| if args.platform != "": | |
| for p in platforms: | |
| if args.platform == p.name(): | |
| wheel_platform = p | |
| print(f"Platform argument specified for building {p.fancy_name} wheel") | |
| break | |
| if wheel_platform is None: | |
| print(f"Platform argument '{args.platform}' not recognized") | |
| elif wheel_platform not in detected_platforms: | |
| print(f"No libraries found for {wheel_platform.fancy_name}") | |
| print("Falling back to auto-detection") | |
| wheel_platform = None | |
| if wheel_platform is None: | |
| if len(detected_platforms) > 1: | |
| print("Libraries for multiple platforms were detected. Picking the first one.") | |
| print( | |
| "Run `python -m build --wheel -C--build-option=-P[windows|linux|macos]-[x86_64|aarch64|universal]` to select a specific one." | |
| ) | |
| wheel_platform = next(iter(detected_platforms)) | |
| print("Creating Warp wheel for " + wheel_platform.fancy_name) | |
| # Binary wheel distribution builds assume that the platform you're building on will be the platform | |
| # of the package. This class overrides the platform tag. | |
| # https://packaging.python.org/en/latest/specifications/platform-compatibility-tags | |
| class WarpBDistWheel(bdist_wheel): | |
| # Even though we parse the platform argument ourselves, we need to declare it here as well so | |
| # setuptools.Command can validate the command line options. | |
| user_options = bdist_wheel.user_options + [ | |
| ("platform=", "P", "Wheel platform: windows|linux|macos-x86_64|aarch64|universal"), | |
| ] | |
| def initialize_options(self): | |
| super().initialize_options() | |
| self.platform = "" | |
| def get_tag(self): | |
| if wheel_platform is not None: | |
| # The wheel's complete tag format is {python tag}-{abi tag}-{platform tag}. | |
| return "py3", "none", wheel_platform.tag | |
| else: | |
| # The target platform was not overridden. Fall back to base class behavior. | |
| return bdist_wheel.get_tag(self) | |
| def run(self): | |
| super().run() | |
| # Clean up so we can re-invoke `py -m build --wheel -C--build-option=--platform=...` | |
| # See https://github.com/pypa/setuptools/issues/1871 for details. | |
| shutil.rmtree("./build", ignore_errors=True) | |
| shutil.rmtree("./warp_lang.egg-info", ignore_errors=True) | |
| # Distributions are identified as non-pure (i.e. containing non-Python code, or binaries) if the | |
| # setuptools.setup() `ext_modules` parameter is not empty, but this assumes building extension | |
| # modules from source through the Python build. This class provides an override for prebuilt binaries: | |
| class BinaryDistribution(setuptools.Distribution): | |
| def has_ext_modules(self): | |
| return True | |
| def get_warp_libraries(platform): | |
| libraries = [] | |
| for library in detected_libraries: | |
| if library.platform == platform: | |
| src = "warp/" + library.directory + library.file | |
| dst = "warp/bin/" + library.file | |
| if src != dst: | |
| shutil.copyfile(src, dst) | |
| libraries.append("bin/" + library.file) | |
| return libraries | |
| if wheel_platform is not None: | |
| warp_binary_libraries = get_warp_libraries(wheel_platform) | |
| else: | |
| warp_binary_libraries = [] # Not needed during egg_info command | |
| setuptools.setup( | |
| package_data={ | |
| "": [ | |
| "native/*.cpp", | |
| "native/*.cu", | |
| "native/*.h", | |
| "native/clang/*.cpp", | |
| "native/nanovdb/*.h", | |
| "tests/assets/*", | |
| ] | |
| + warp_binary_libraries, | |
| }, | |
| distclass=BinaryDistribution, | |
| cmdclass={ | |
| "bdist_wheel": WarpBDistWheel, | |
| }, | |
| ) | |