ARROW-112 Make PyArrow Binding more Robust (#95)

blink1073 · web-flow · commit 6e2bc1d18b24 · 2022-08-31T14:57:35.000-05:00
diff --git a/.github/workflows/release-python.yml b/.github/workflows/release-python.yml
@@ -8,6 +8,10 @@ concurrency:
   group: wheels-${{ github.ref }}
   cancel-in-progress: true
 
+defaults:
+  run:
+    working-directory: ./bindings/python
+
 jobs:
   build_wheels:
     name: Build wheel for ${{ matrix.python }}-${{ matrix.buildplat[1] }}
@@ -37,17 +41,13 @@ jobs:
 
       - name: Build MacOS Py38 Wheel
         if: ${{ matrix.python == 'cp38' && matrix.buildplat[0] == 'macos-10.15' }}
-        working-directory: ./bindings/python
-        shell: bash
         env:
           CIBW_BUILD: cp38-macosx_x86_64
           MACOSX_DEPLOYMENT_TARGET: "10.13"
         run: python -m cibuildwheel --output-dir wheelhouse
 
       - name: Build wheels
         if: ${{ matrix.python != 'cp38' || matrix.buildplat[0] != 'macos-10.15' }}
-        working-directory: ./bindings/python
-        shell: bash
         env:
           CIBW_BUILD: ${{ matrix.python }}-${{ matrix.buildplat[1] }}
           MACOSX_DEPLOYMENT_TARGET: "10.13"
@@ -92,7 +92,7 @@ jobs:
       - uses: actions/upload-artifact@v2
         with:
           name: "sdist"
-          path: ./bindings/python/dist/*.tar.gz
+          path: ./bindings/python/wheelhouse/*.tar.gz
 
   collect-dist:
     runs-on: ubuntu-latest
@@ -102,6 +102,7 @@ jobs:
       - name: Download all workflow run artifacts
         uses: actions/download-artifact@v2
       - name: Flatten directory
+        working-directory: .
         run: |
           find . -mindepth 2 -type f -exec mv {} . \;
           find . -type d -empty -delete
diff --git a/bindings/python/MANIFEST.in b/bindings/python/MANIFEST.in
@@ -7,6 +7,7 @@ exclude THIRD-PARTY-NOTICES
 exclude release.sh
 exclude addtags.py
 exclude benchmark.py
+exclude repair_wheel.py
 exclude .flake8
 exclude RELEASE.rst
 exclude asv.conf.json
diff --git a/bindings/python/cibw_before_build.sh b/bindings/python/cibw_before_build.sh
@@ -19,9 +19,8 @@ then
   fi
 
   # Install pyarrow with the appropriate platform.
-  pip install --platform $platform --target $HOME/wheels --no-deps --only-binary=:all: pyarrow
+  pip install --platform $platform --upgrade --target $HOME/wheels --no-deps --only-binary=:all: pyarrow
 fi
 
-
 # Build libbson with the appropriate arch.
 ./build-libbson.sh
diff --git a/bindings/python/pyproject.toml b/bindings/python/pyproject.toml
@@ -12,6 +12,7 @@ skip = "pp* *-manylinux_i686 *_ppc64le *_s390x *-musllinux*"
 before-build = "bash ./cibw_before_build.sh"
 build-frontend = "build"
 test-command = "python -c \"from pymongoarrow.lib import process_bson_stream\""
+repair-wheel-command = "python repair_wheel.py {dest_dir} {wheel} {delocate_archs}"
 
 [tool.cibuildwheel.environment]
 LIBBSON_INSTALL_DIR = "./libbson"
diff --git a/bindings/python/release.sh b/bindings/python/release.sh
@@ -10,6 +10,9 @@ then
   exit 1
 fi
 
+# Clean up
+rm -rf dist wheelhouse build pymongoarrow/*.so pymongoarrow/*.dll pymongoarrow/*.dylib
+
 # Platform-dependent actions:
 PYTHON=${PYTHON_BINARY:-"python"}
 if [ "Linux" = "$(uname -s)" ]
@@ -38,3 +41,7 @@ then
   $PYTHON -m pip install auditwheel
   $PYTHON addtags.py dist/*.whl "$PLAT" ./wheelhouse
 fi
+
+# Repair the wheel, copying shared libraries as needed
+MACHINE=$(python -c "import platform;print(platform.machine())")
+python repair_wheel.py "./wheelhouse" dist/*.whl $MACHINE
diff --git a/bindings/python/repair_wheel.py b/bindings/python/repair_wheel.py
@@ -0,0 +1,94 @@
+import atexit
+import glob
+import os
+import sys
+import tempfile
+from subprocess import run
+
+HERE = os.path.abspath(os.path.dirname(__file__))
+wheel_dir, wheel_file, delocate_args = sys.argv[1:]
+wheel_dir = wheel_dir.replace(os.sep, "/")
+wheel_file = wheel_file.replace(os.sep, "/")
+
+# Ensure pyarrow.
+if "universal2" in wheel_file:
+    # pip selects the most specific platform by default,
+    # so we have to tell it to install the universal2 version.
+    # See https://github.com/pypa/packaging/issues/381
+    macos_ver = os.environ.get("MACOSX_DEPLOYMENT_TARGET", "10.3")
+    macos_ver = macos_ver.replace(".", "_")
+    wheel_temp_dir = tempfile.TemporaryDirectory()
+    atexit.register(wheel_temp_dir.cleanup)
+    run(
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "--platform",
+            f"macosx_{macos_ver}_universal2",
+            "--upgrade",
+            "--target",
+            wheel_temp_dir.name,
+            "--only-binary=:all:",
+            "pyarrow",
+        ]
+    )
+    # Allow the installed pyarrow library to be imported.
+    sys.path.insert(0, wheel_temp_dir.name)
+else:
+    run([sys.executable, "-m", "pip", "install", "pyarrow"])
+import pyarrow as pa  # noqa
+
+libbson = os.environ.get("LIBBSON_INSTALL_DIR", os.path.join(HERE, "libbson"))
+libbson = os.path.abspath(libbson)
+if os.name == "nt":
+    libbson_lib = glob.glob(os.path.join(libbson, "bin"))
+else:
+    libbson_lib = glob.glob(os.path.join(libbson, "lib*"))
+extra_paths = pa.get_library_dirs() + libbson_lib
+extra_path = os.path.pathsep.join([a.replace(os.sep, "/") for a in extra_paths])
+
+
+def append_os_variable(name, extra_path):
+    if os.environ.get(name):
+        os.environ[name] = os.environ[name] + os.pathsep + extra_path
+    else:
+        os.environ[name] = extra_path
+    print(f"{name}: {os.environ[name]}")
+
+
+if os.name == "nt":
+    append_os_variable("PATH", extra_path)
+    run([sys.executable, "-m", "pip", "install", "delvewheel"])
+    # Do not mangle ucrtbased.dll to avoid:
+    # "RuntimeError: Unable to rename the dependencies of vcruntime140d.dll
+    # because this DLL has trailing data. If this DLL was created with MinGW,
+    # run the strip utility. Otherwise, include ucrtbased.dll in the
+    # --no-mangle flag. In addition, if you believe that delvewheel should
+    # avoid name-mangling a specific DLL by default, open an issue at
+    # https://github.com/adang1345/delvewheel/issues and include this error
+    # message."
+    run(["delvewheel", "repair", "--no-mangle", "ucrtbased.dll", "-w", wheel_dir, wheel_file])
+
+elif sys.platform == "darwin":
+    # FIXME: We should not have to do this.
+    site_pkgs = sys.base_prefix
+    dylib = glob.glob(f"{sys.base_prefix}/lib/python*/lib-dynload")[0]
+    extra_path = f"{dylib}:{extra_path}"
+    append_os_variable("DYLD_LIBRARY_PATH", extra_path)
+    run([sys.executable, "-m", "pip", "install", "delocate"])
+    run(
+        [
+            "delocate-wheel",
+            "--require-archs",
+            delocate_args,
+            "-w",
+            wheel_dir,
+            wheel_file,
+        ]
+    )
+else:
+    append_os_variable("LD_LIBRARY_PATH", extra_path)
+    run([sys.executable, "-m", "pip", "install", "auditwheel"])
+    run(["auditwheel", "repair", "-w", wheel_dir, wheel_file])
diff --git a/bindings/python/setup.py b/bindings/python/setup.py
@@ -11,15 +11,6 @@
 BUILD_DIR = os.path.join(HERE, "pymongoarrow")
 IS_WIN = platform == "win32"
 
-# Find and copy the binary arrow files, unless
-# MONGO_NO_COPY_ARROW_LIB is set (for instance in a conda build).
-# Wheels are meant to be self-contained, per PEP 513.
-# https://www.python.org/dev/peps/pep-0513/#id40
-# Conda has the opposite philosphy, where libraries are meant to be
-# shared.  For instance, there is an arrow-cpp library available on conda-forge
-# that provides the libarrow files.
-COPY_LIBARROW = not os.environ.get("MONGO_NO_COPY_LIBARROW", False)
-
 # Find and copy the binary libbson file, unless
 # MONGO_NO_COPY_LIBBSON is set (for instance in a conda build).
 COPY_LIBBSON = not os.environ.get("MONGO_NO_COPY_LIBBSON", False)
@@ -47,6 +38,7 @@ def append_libbson_flags(module):
     install_dir = os.environ.get("LIBBSON_INSTALL_DIR")
     if install_dir:
         install_dir = os.path.abspath(install_dir)
+
         # Handle the copy-able library file if applicable.
         if COPY_LIBBSON:
             if platform == "darwin":
@@ -62,6 +54,14 @@ def append_libbson_flags(module):
                 if os.path.exists(lib_file):
                     shutil.copy(lib_file, BUILD_DIR)
 
+            # Ensure our Cython extension can dynamically link to libraries
+            # - https://blog.krzyzanowskim.com/2018/12/05/rpath-what/
+            # - https://nehckl0.medium.com/creating-relocatable-linux-executables-by-setting-rpath-with-origin-45de573a2e98
+            if platform == "darwin":
+                module.extra_link_args += ["-rpath", "@loader_path"]
+            elif platform == "linux":
+                module.extra_link_args += ["-Wl,-rpath,$ORIGIN"]
+
         # Find the linkable library file, and explicity add it to the linker if on Windows.
         lib_dirs = glob.glob(os.path.join(install_dir, "lib*"))
         if len(lib_dirs) != 1:
@@ -125,64 +125,33 @@ def append_libbson_flags(module):
     module.libraries.extend(libnames)
 
 
-def append_arrow_flags(module):
+def append_arrow_flags(ext):
     import numpy as np
     import pyarrow as pa
 
-    if IS_WIN:
-        module.include_dirs.append(np.get_include())
-        module.include_dirs.append(pa.get_include())
-    else:
-        module.extra_compile_args.append("-isystem" + pa.get_include())
-        module.extra_compile_args.append("-isystem" + np.get_include())
-        # Arrow's manylinux{2010, 2014} binaries are built with gcc < 4.8 which predates CXX11 ABI
-        # - https://uwekorn.com/2019/09/15/how-we-build-apache-arrows-manylinux-wheels.html
-        # - https://arrow.apache.org/docs/python/extending.html#example
-        if "std=" not in os.environ.get("CXXFLAGS", ""):
-            module.extra_compile_args.append("-std=c++11")
-            module.extra_compile_args.append("-D_GLIBCXX_USE_CXX11_ABI=0")
-
-    # Handle the arrow library files manually.
-    # Alternative to using pyarrow.create_library_symlinks().
-    # You can use MONGO_LIBARROW_DIR to explicitly set the location of the
-    # arrow libraries (for instance in a conda build).
-    # We first check for an unmodified path to the library,
-    # then look for a library file with a version modifier, e.g. libarrow.600.dylib.
-    arrow_lib = os.environ.get("MONGO_LIBARROW_DIR", pa.get_library_dirs()[0])
-    if platform == "darwin":
-        exts = [".dylib", ".*.dylib"]
-    elif platform == "linux":
-        exts = [".so", ".so.*"]
-    else:
-        # Windows is handled differently (see below)
-        pass
-
-    # Find the appropriate library file and optionally copy it locally.
-    # Explicitly handle "parquet" library as a workaround for
-    # https://issues.apache.org/jira/browse/ARROW-17327
-    for name in pa.get_libraries() + ["parquet"]:
-        if IS_WIN:
-            if COPY_LIBARROW:
-                lib_file = os.path.join(arrow_lib, f"{name}.dll")
-                if not os.path.exists(lib_file):
-                    if name == "parquet":
-                        continue
-                    raise ValueError("Could not find compiled arrow library")
-                shutil.copy(lib_file, BUILD_DIR)
-            lib_file = os.path.join(arrow_lib, f"{name}.lib")
-            module.extra_link_args.append(lib_file)
-            continue
-
-        for ext in exts:
-            files = glob.glob(os.path.join(arrow_lib, f"lib{name}{ext}"))
-            if not files:
-                continue
-            path = files[0]
-            if COPY_LIBARROW:
-                shutil.copy(path, BUILD_DIR)
-                path = os.path.join(BUILD_DIR, os.path.basename(path))
-            module.extra_link_args.append(path)
-            break
+    # From https://arrow.apache.org/docs/python/integration/extending.html#example
+    # The Numpy C headers are currently required
+    ext.include_dirs.append(np.get_include())
+    ext.include_dirs.append(pa.get_include())
+    ext.libraries.extend(pa.get_libraries())
+    ext.library_dirs.extend(pa.get_library_dirs())
+
+    if os.name != "nt":
+        # On Linux and MacOS, we must run pyarrow.create_library_symlinks()
+        # as a user with write access to the directory where pyarrow is
+        # installed.
+        # See https://arrow.apache.org/docs/python/integration/extending.html#building-extensions-against-pypi-wheels.
+        pa.create_library_symlinks()
+
+    if os.name == "posix":
+        ext.extra_compile_args.append("-std=c++11")
+
+    # Arrow's manylinux{2010, 2014} binaries are built with gcc < 4.8 which predates CXX11 ABI
+    # - https://uwekorn.com/2019/09/15/how-we-build-apache-arrows-manylinux-wheels.html
+    # - https://arrow.apache.org/docs/python/extending.html#example
+    if "std=" not in os.environ.get("CXXFLAGS", ""):
+        ext.extra_compile_args.append("-std=c++11")
+        ext.extra_compile_args.append("-D_GLIBCXX_USE_CXX11_ABI=0")
 
 
 def get_extension_modules():
@@ -197,13 +166,6 @@ def get_extension_modules():
     for module in modules:
         append_libbson_flags(module)
         append_arrow_flags(module)
-        # Ensure our Cython extension can dynamically link to libraries
-        # - https://blog.krzyzanowskim.com/2018/12/05/rpath-what/
-        # - https://nehckl0.medium.com/creating-relocatable-linux-executables-by-setting-rpath-with-origin-45de573a2e98
-        if platform == "darwin":
-            module.extra_link_args += ["-rpath", "@loader_path"]
-        elif platform == "linux":
-            module.extra_link_args += ["-Wl,-rpath,$ORIGIN"]
 
     return modules
 
