diff --git a/_custom_build/backend.py b/_custom_build/backend.py new file mode 100644 index 000000000..170ef8757 --- /dev/null +++ b/_custom_build/backend.py @@ -0,0 +1,59 @@ +# Inspired by https://github.com/python-pillow/Pillow/pull/7171. + +import sys + +from setuptools.build_meta import * # noqa: F401, F403 + +backend_class = build_wheel.__self__.__class__ + + +class _CustomBuildMetaBackend(backend_class): + def run_setup(self, setup_script="setup.py"): + if self.config_settings: + + def config_has(value): + settings = self.config_settings.get("--build-option") + if settings: + if not isinstance(settings, list): + settings = [settings] + return value in settings + + flags = [] + for ext in ( + "cpp_ext", + "cuda_ext", + "distributed_adam", + "distributed_lamb", + "permutation_search", + "bnp", + "xentropy", + "focal_loss", + "index_mul_2d", + "deprecated_fused_adam", + "deprecated_fused_lamb", + "fast_layer_norm", + "fmha", + "fast_multihead_attn", + "transducer", + "cudnn_gbn", + "peer_memory", + "nccl_p2p", + "fast_bottleneck", + "fused_conv_bias_relu", + ): + ext = "--" + ext + if config_has(ext): + flags.append(ext) + + if flags: + sys.argv = sys.argv[:1] + ["build_ext"] + flags + sys.argv[1:] + return super().run_setup(setup_script) + + def build_wheel( + self, wheel_directory, config_settings=None, metadata_directory=None + ): + self.config_settings = config_settings + return super().build_wheel(wheel_directory, config_settings, metadata_directory) + + +build_wheel = _CustomBuildMetaBackend().build_wheel diff --git a/pyproject.toml b/pyproject.toml index 44836bb2b..7858c0021 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,7 @@ [build-system] requires = [ - "setuptools", + "setuptools >= 40.8.0", "wheel", ] -build-backend = "setuptools.build_meta" +build-backend = "backend" # custom setuptools.build_meta +backend-path = ["_custom_build"]