From 6f2307ba6ac802f4c4650cd21869e49d1845188d Mon Sep 17 00:00:00 2001 From: Peter Goldsborough Date: Tue, 27 Nov 2018 17:33:54 -0800 Subject: [PATCH] Allow building libraries with setuptools that dont have abi suffix (#14130) Summary: When using `setuptools` to build a Python extension, setuptools will automatically add an ABI suffix like `cpython-37m-x86_64-linux-gnu` to the shared library name when using Python 3. This is required for extensions meant to be imported as Python modules. When we use setuptools to build shared libraries not meant as Python modules, for example libraries that define and register TorchScript custom ops, having your library called `my_ops.cpython-37m-x86_64-linux-gnu.so` is a bit annoying compared to just `my_ops.so`, especially since you have to reference the library name when loading it with `torch.ops.load_library` in Python. This PR fixes this by adding a `with_options` class method to the `torch.utils.cpp_extension.BuildExtension` which allows configuring the `BuildExtension`. In this case, the first option we add is `no_python_abi_suffix`, which we then use in `get_ext_filename` (override from `setuptools.build_ext`) to throw away the ABI suffix. I've added a test `setup.py` in a `no_python_abi_suffix_test` folder. Fixes https://github.com/pytorch/pytorch/issues/14188 t-vi fmassa soumith Pull Request resolved: https://github.com/pytorch/pytorch/pull/14130 Differential Revision: D13216575 Pulled By: goldsborough fbshipit-source-id: 67dc345c1278a1a4ee4ca907d848bc1fb4956cfa --- .../no_python_abi_suffix_test.cpp | 1 + .../no_python_abi_suffix_test/setup.py | 10 ++++ test/cpp_extensions/setup.py | 4 +- test/custom_operator/model.py | 20 ++++---- test/run_test.py | 8 ++- test/test_cpp_extensions.py | 20 ++++++-- torch/utils/cpp_extension.py | 57 ++++++++++++++++++---- 7 files changed, 94 insertions(+), 26 deletions(-) create mode 100644 test/cpp_extensions/no_python_abi_suffix_test/no_python_abi_suffix_test.cpp create mode 100644 test/cpp_extensions/no_python_abi_suffix_test/setup.py diff --git a/test/cpp_extensions/no_python_abi_suffix_test/no_python_abi_suffix_test.cpp b/test/cpp_extensions/no_python_abi_suffix_test/no_python_abi_suffix_test.cpp new file mode 100644 index 0000000..b4ff740 --- /dev/null +++ b/test/cpp_extensions/no_python_abi_suffix_test/no_python_abi_suffix_test.cpp @@ -0,0 +1 @@ +void dummy(int) { } diff --git a/test/cpp_extensions/no_python_abi_suffix_test/setup.py b/test/cpp_extensions/no_python_abi_suffix_test/setup.py new file mode 100644 index 0000000..9b39cd5 --- /dev/null +++ b/test/cpp_extensions/no_python_abi_suffix_test/setup.py @@ -0,0 +1,10 @@ +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CppExtension + +setup( + name="no_python_abi_suffix_test", + ext_modules=[ + CppExtension("no_python_abi_suffix_test", ["no_python_abi_suffix_test.cpp"]) + ], + cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)}, +) diff --git a/test/cpp_extensions/setup.py b/test/cpp_extensions/setup.py index e7a61ec..82f990a 100644 --- a/test/cpp_extensions/setup.py +++ b/test/cpp_extensions/setup.py @@ -1,7 +1,7 @@ import sys import torch.cuda from setuptools import setup -from torch.utils.cpp_extension import CppExtension, CUDAExtension +from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension from torch.utils.cpp_extension import CUDA_HOME CXX_FLAGS = [] if sys.platform == 'win32' else ['-g', '-Werror'] @@ -27,4 +27,4 @@ setup( name='torch_test_cpp_extension', packages=['torch_test_cpp_extension'], ext_modules=ext_modules, - cmdclass={'build_ext': torch.utils.cpp_extension.BuildExtension}) + cmdclass={'build_ext': BuildExtension}) diff --git a/test/custom_operator/model.py b/test/custom_operator/model.py index 990ecbf..5131b4a 100644 --- a/test/custom_operator/model.py +++ b/test/custom_operator/model.py @@ -4,16 +4,15 @@ import sys import torch -SHARED_LIBRARY_NAMES = { - 'linux': 'libcustom_ops.so', - 'darwin': 'libcustom_ops.dylib', - 'win32': 'custom_ops.dll' -} - def get_custom_op_library_path(): - path = os.path.abspath('build/{}'.format( - SHARED_LIBRARY_NAMES[sys.platform])) + if sys.platform.startswith("win32"): + library_filename = "custom_ops.dll" + elif sys.platform.startswith("darwin"): + library_filename = "libcustom_ops.dylib" + else: + library_filename = "libcustom_ops.so" + path = os.path.abspath("build/{}".format(library_filename)) assert os.path.exists(path), path return path @@ -30,7 +29,8 @@ class Model(torch.jit.ScriptModule): def main(): parser = argparse.ArgumentParser( - description="Serialize a script module with custom ops") + description="Serialize a script module with custom ops" + ) parser.add_argument("--export-script-module-to", required=True) options = parser.parse_args() @@ -40,5 +40,5 @@ def main(): model.save(options.export_script_module_to) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/test/run_test.py b/test/run_test.py index 9991411..547aab9 100644 --- a/test/run_test.py +++ b/test/run_test.py @@ -148,10 +148,16 @@ def test_cpp_extensions(executable, test_module, test_directory, options): except RuntimeError: print(CPP_EXTENSIONS_ERROR) return 1 + cpp_extensions_test_dir = os.path.join(test_directory, 'cpp_extensions') return_code = shell([sys.executable, 'setup.py', 'install', '--root', './install'], - os.path.join(test_directory, 'cpp_extensions')) + cwd=cpp_extensions_test_dir) if return_code != 0: return return_code + if sys.platform != 'win32': + return_code = shell([sys.executable, 'setup.py', 'install', '--root', './install'], + cwd=os.path.join(cpp_extensions_test_dir, 'no_python_abi_suffix_test')) + if return_code != 0: + return return_code python_path = os.environ.get('PYTHONPATH', '') try: diff --git a/test/test_cpp_extensions.py b/test/test_cpp_extensions.py index 9dec09a..a032b4d 100755 --- a/test/test_cpp_extensions.py +++ b/test/test_cpp_extensions.py @@ -29,7 +29,7 @@ IS_WINDOWS = sys.platform == 'win32' class TestCppExtension(common.TestCase): def setUp(self): - if sys.platform != 'win32': + if not IS_WINDOWS: default_build_root = torch.utils.cpp_extension.get_default_build_root() if os.path.exists(default_build_root): shutil.rmtree(default_build_root) @@ -121,7 +121,7 @@ class TestCppExtension(common.TestCase): @unittest.skipIf(not TEST_CUDNN, "CuDNN not found") def test_jit_cudnn_extension(self): # implementation of CuDNN ReLU - if sys.platform == 'win32': + if IS_WINDOWS: extra_ldflags = ['cudnn.lib'] else: extra_ldflags = ['-lcudnn'] @@ -381,7 +381,6 @@ class TestCppExtension(common.TestCase): self.assertEqual(len(net.parameters()), 4) p = net.named_parameters() - self.assertEqual(type(p), dict) self.assertEqual(len(p), 4) self.assertIn('fc.weight', p) self.assertIn('fc.bias', p) @@ -402,6 +401,21 @@ class TestCppExtension(common.TestCase): is_python_module=False) self.assertEqual(torch.ops.test.func(torch.eye(5)), torch.eye(5)) + @unittest.skipIf(IS_WINDOWS, "Not available on Windows") + def test_no_python_abi_suffix_sets_the_correct_library_name(self): + # For this test, run_test.py will call `python setup.py install` in the + # cpp_extensions/no_python_abi_suffix_test folder, where the + # `BuildExtension` class has a `no_python_abi_suffix` option set to + # `True`. This *should* mean that on Python 3, the produced shared + # library does not have an ABI suffix like + # "cpython-37m-x86_64-linux-gnu" before the library suffix, e.g. "so". + # On Python 2 there is no ABI suffix anyway. + root = os.path.join("cpp_extensions", "no_python_abi_suffix_test", "build") + print(list(os.walk(os.path.join("cpp_extensions", "no_python_abi_suffix_test")))) + matches = [f for _, _, fs in os.walk(root) for f in fs if f.endswith("so")] + self.assertEqual(len(matches), 1, str(matches)) + self.assertEqual(matches[0], "no_python_abi_suffix_test.so", str(matches)) + if __name__ == '__main__': common.run_tests() diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py index 3a67a27..7c855a3 100644 --- a/torch/utils/cpp_extension.py +++ b/torch/utils/cpp_extension.py @@ -84,7 +84,6 @@ with compiling PyTorch from source. !! WARNING !! ''' -ACCEPTED_COMPILERS_FOR_PLATFORM = {'darwin': ['clang++', 'clang'], 'linux': ['g++', 'gcc']} CUDA_HOME = _find_cuda_home() CUDNN_HOME = os.environ.get('CUDNN_HOME') or os.environ.get('CUDNN_PATH') # PyTorch releases have the version pattern major.minor.patch, whereas when @@ -106,6 +105,10 @@ def _is_binary_build(): return not BUILT_FROM_SOURCE_VERSION_PATTERN.match(torch.version.__version__) +def _accepted_compilers_for_platform(): + return ['clang++', 'clang'] if sys.platform.startswith('darwin') else ['g++', 'gcc'] + + def get_default_build_root(): ''' Returns the path to the root folder under which extensions will built. @@ -135,8 +138,7 @@ def check_compiler_ok_for_platform(compiler): which = subprocess.check_output(['which', compiler], stderr=subprocess.STDOUT) # Use os.path.realpath to resolve any symlinks, in particular from 'c++' to e.g. 'g++'. compiler_path = os.path.realpath(which.decode().strip()) - accepted_compilers = ACCEPTED_COMPILERS_FOR_PLATFORM[sys.platform] - return any(name in compiler_path for name in accepted_compilers) + return any(name in compiler_path for name in _accepted_compilers_for_platform()) def check_compiler_abi_compatibility(compiler): @@ -160,15 +162,15 @@ def check_compiler_abi_compatibility(compiler): if not check_compiler_ok_for_platform(compiler): warnings.warn(WRONG_COMPILER_WARNING.format( user_compiler=compiler, - pytorch_compiler=ACCEPTED_COMPILERS_FOR_PLATFORM[sys.platform][0], + pytorch_compiler=_accepted_compilers_for_platform()[0], platform=sys.platform)) return False - if sys.platform == 'darwin': + if sys.platform.startswith('darwin'): # There is no particular minimum version we need for clang, so we're good here. return True try: - if sys.platform == 'linux': + if sys.platform.startswith('linux'): minimum_required_version = MINIMUM_GCC_VERSION version = subprocess.check_output([compiler, '-dumpfullversion', '-dumpversion']) version = version.decode().strip().split('.') @@ -191,7 +193,11 @@ def check_compiler_abi_compatibility(compiler): return False -class BuildExtension(build_ext): +# See below for why we inherit BuildExtension from object. +# https://stackoverflow.com/questions/1713038/super-fails-with-error-typeerror-argument-1-must-be-type-not-classobj-when + + +class BuildExtension(build_ext, object): ''' A custom :mod:`setuptools` build extension . @@ -206,6 +212,22 @@ class BuildExtension(build_ext): the C++ and CUDA compiler during mixed compilation. ''' + @classmethod + def with_options(cls, **options): + ''' + Returns an alternative constructor that extends any original keyword + arguments to the original constructor with the given options. + ''' + def init_with_options(*args, **kwargs): + kwargs = kwargs.copy() + kwargs.update(options) + return cls(*args, **kwargs) + return init_with_options + + def __init__(self, *args, **kwargs): + super(BuildExtension, self).__init__(*args, **kwargs) + self.no_python_abi_suffix = kwargs.get("no_python_abi_suffix", False) + def build_extensions(self): self._check_abi() for extension in self.extensions: @@ -261,9 +283,7 @@ class BuildExtension(build_ext): extra_postargs = None def spawn(cmd): - orig_cmd = cmd # Using regex to match src, obj and include files - src_regex = re.compile('/T(p|c)(.*)') src_list = [ m.group(2) for m in (src_regex.match(elem) for elem in cmd) @@ -322,6 +342,23 @@ class BuildExtension(build_ext): build_ext.build_extensions(self) + def get_ext_filename(self, ext_name): + # Get the original shared library name. For Python 3, this name will be + # suffixed with ".so", where will be something like + # cpython-37m-x86_64-linux-gnu. On Python 2, there is no such ABI name. + # The final extension, .so, would be .lib/.dll on Windows of course. + ext_filename = super(BuildExtension, self).get_ext_filename(ext_name) + # If `no_python_abi_suffix` is `True`, we omit the Python 3 ABI + # component. This makes building shared libraries with setuptools that + # aren't Python modules nicer. + if self.no_python_abi_suffix and sys.version_info >= (3, 0): + # The parts will be e.g. ["my_extension", "cpython-37m-x86_64-linux-gnu", "so"]. + ext_filename_parts = ext_filename.split('.') + # Omit the second to last element. + without_abi = ext_filename_parts[:-2] + ext_filename_parts[-1:] + ext_filename = '.'.join(without_abi) + return ext_filename + def _check_abi(self): # On some platforms, like Windows, compiler_cxx is not available. if hasattr(self.compiler, 'compiler_cxx'): @@ -998,7 +1035,7 @@ def _write_ninja_file(path, else: ldflags = ['-shared'] + extra_ldflags # The darwin linker needs explicit consent to ignore unresolved symbols. - if sys.platform == 'darwin': + if sys.platform.startswith('darwin'): ldflags.append('-undefined dynamic_lookup') elif IS_WINDOWS: ldflags = _nt_quote_args(ldflags) -- 2.7.4