Add faiss-gpu pip wheel packaging (facebookresearch#5131)

alibeklfc · facebook-github-bot · commit 4adf26018b61 · 2026-04-20T14:22:35.000-07:00
Summary:

Add pip wheel packaging for `faiss-gpu`, enabling `pip install faiss-gpu`
from PyPI. Builds GPU wheels for Linux x86_64, Python 3.10-3.13, with
CUDA 12.6 on manylinux_2_28.

**New files:**
- `pyproject-gpu.toml`: scikit-build-core config for faiss-gpu. Enables
  CUDA GPU support, uses Dynamic Dispatch, OpenBLAS, and excludes CUDA
  runtime libs (provided by nvidia-cuda-runtime-cu12 pip dependency).
- `.github/workflows/build-pip-gpu.yml`: CI workflow for GPU wheels using
  cibuildwheel on GPU runners (4-core-ubuntu-gpu-t4).
- `tests/test_wheel_smoke_gpu.py`: GPU smoke tests validating GPU
  detection, StandardGpuResources, GPU FlatL2/IP search, IVF+PQ on GPU,
  CPU&lt;-&gt;GPU transfer roundtrip, and GPU serialization.

**Modified files:**
- `python/__init__.py`: Added `_preload_gpu_libs()` to pre-load CUDA
  shared libraries from pip packages before SWIG import. No-op on
  CPU-only installs.
- `.github/workflows/build.yml`: Added `build-pip-gpu` workflow call.
- `tests/BUCK`: Added `test_wheel_smoke_gpu` target for Buck CI.

Differential Revision: D97507525
diff --git a/.github/workflows/build-pip-gpu.yml b/.github/workflows/build-pip-gpu.yml
@@ -0,0 +1,70 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+name: Build GPU pip wheel
+
+on:
+  workflow_call:
+  workflow_dispatch:
+
+concurrency:
+  group: gpu-pip-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  build-faiss-gpu:
+    name: Build & test faiss-gpu
+    runs-on: 4-core-ubuntu-gpu-t4
+    timeout-minutes: 90
+    steps:
+      - name: Checkout
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          fetch-depth: 0
+          fetch-tags: true
+
+      - name: Activate GPU pyproject.toml
+        # cibuildwheel + scikit-build-core both require literal pyproject.toml.
+        # Workspace is fresh per job (actions/checkout creates new tree).
+        run: cp pyproject-gpu.toml pyproject.toml
+
+      - name: Set up Python
+        # Pin host Python explicitly: the 4-core-ubuntu-gpu-t4 runner image
+        # ships a broken Python 3.14.4 in its tool cache (libpython3.14.so.1.0
+        # missing), and pypa/cibuildwheel picks that up by default.
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Build wheel (cibuildwheel)
+        uses: pypa/cibuildwheel@ee02a1537ce3071a004a6b08c41e72f0fdc42d9a  # v3.4.0
+        env:
+          CIBW_BUILD_VERBOSITY: 1
+
+      - name: Upload wheel
+        if: always()
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02  # v4.6.2
+        with:
+          name: wheels-gpu
+          path: wheelhouse/*.whl
+
+  publish-gpu:
+    name: Publish faiss-gpu to PyPI
+    needs: [build-faiss-gpu]
+    runs-on: ubuntu-latest
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
+    environment:
+      name: pypi
+      url: https://pypi.org/p/faiss-gpu
+    permissions:
+      id-token: write
+    steps:
+      - name: Download wheel
+        uses: actions/download-artifact@65a9edc5881444af0b9093a5e628f2fe47ea3b2e  # v4.1.7
+        with:
+          name: wheels-gpu
+          path: dist
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc  # release/v1
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -17,3 +17,5 @@ jobs:
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
   build-pip:
     uses: ./.github/workflows/build-pip.yml
+  build-pip-gpu:
+    uses: ./.github/workflows/build-pip-gpu.yml
diff --git a/faiss/python/__init__.py b/faiss/python/__init__.py
@@ -13,6 +13,59 @@
 import sys
 import inspect
 
+
+def _preload_gpu_libs():
+    """Pre-load CUDA shared libraries from nvidia-*-cu12 pip packages.
+
+    The libcudart.so / libcublas.so files live in
+    site-packages/nvidia/{cuda_runtime,cublas}/lib/ — not on the default
+    ld.so search path. Pre-load them with RTLD_GLOBAL before the SWIG
+    extension dlopen()s libfaiss.so (which has CUDA undefined symbols).
+
+    Behavior:
+    - faiss-cpu install (no nvidia.cuda_runtime package): silent no-op.
+    - faiss-gpu install with partial nvidia-* deps (cuda_runtime present
+      but cublas missing): raises RuntimeError with fix-it message.
+    """
+    import ctypes
+    import os
+
+    # Gate on cuda_runtime first: its absence signals a faiss-cpu install.
+    try:
+        import nvidia.cuda_runtime
+    except ImportError:
+        return  # faiss-cpu install, no GPU libs expected
+
+    # cublas is mandatory once cuda_runtime is present; partial install is an error.
+    try:
+        import nvidia.cublas
+    except ImportError as e:
+        raise RuntimeError(
+            "faiss-gpu installed but nvidia-cublas-cu12 is missing — "
+            "pip install nvidia-cublas-cu12==12.6.*"
+        ) from e
+
+    # Order matters: libcudart provides symbols that libcublas* depend on.
+    _LIBS = [
+        (nvidia.cuda_runtime, "libcudart.so.12"),
+        (nvidia.cublas,       "libcublas.so.12"),
+        (nvidia.cublas,       "libcublasLt.so.12"),
+    ]
+    for pkg, soname in _LIBS:
+        so_path = os.path.join(os.path.dirname(pkg.__file__), "lib", soname)
+        try:
+            ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)
+        except OSError as e:
+            raise RuntimeError(
+                f"faiss-gpu: failed to load {soname} from {so_path} — "
+                f"corrupt nvidia-*-cu12 wheel? Try: pip install --force-reinstall "
+                f"nvidia-cuda-runtime-cu12==12.6.* nvidia-cublas-cu12==12.6.*"
+            ) from e
+
+
+_preload_gpu_libs()
+del _preload_gpu_libs
+
 # We import * so that the symbol foo can be accessed as faiss.foo.
 from .loader import *
 
diff --git a/pyproject-gpu.toml b/pyproject-gpu.toml
@@ -0,0 +1,112 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+[build-system]
+requires = ["scikit-build-core>=0.10", "swig>=4.2,<5", "numpy>=2.0", "cmake>=3.24", "ninja"]
+build-backend = "scikit_build_core.build"
+
+[project]
+name = "faiss-gpu"
+dynamic = ["version"]
+description = "A library for efficient similarity search and clustering of dense vectors (GPU support)."
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE", "THIRD_PARTY_NOTICES"]
+requires-python = ">=3.10"
+authors = [
+  { name = "Matthijs Douze" },
+  { name = "Jeff Johnson" },
+  { name = "Hervé Jégou" },
+  { name = "Lucas Hosseini" },
+]
+keywords = ["similarity-search", "vector-search", "nearest-neighbors", "clustering", "gpu", "cuda"]
+classifiers = [
+  "Development Status :: 5 - Production/Stable",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: POSIX :: Linux",
+  "Programming Language :: C++",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: Implementation :: CPython",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Software Development :: Libraries :: Python Modules",
+  "Environment :: GPU :: NVIDIA CUDA :: 12",
+]
+dependencies = [
+  "numpy>=1.25",
+  "packaging",
+  "nvidia-cuda-runtime-cu12==12.6.*",
+  "nvidia-cublas-cu12==12.6.*",
+]
+
+[project.urls]
+Homepage = "https://github.com/facebookresearch/faiss"
+Repository = "https://github.com/facebookresearch/faiss"
+Issues = "https://github.com/facebookresearch/faiss/issues"
+Documentation = "https://github.com/facebookresearch/faiss/wiki"
+
+[tool.scikit-build]
+cmake.build-type = "Release"
+metadata.version.provider = "scikit_build_core.metadata.regex"
+metadata.version.input = "CMakeLists.txt"
+metadata.version.regex = '^\s+VERSION\s+(?P<value>[0-9]+\.[0-9]+\.[0-9]+)'
+wheel.exclude = ["include/**", "share/**", "lib/cmake/**", "lib/pkgconfig/**"]
+
+[tool.scikit-build.cmake.define]
+FAISS_OPT_LEVEL = "dd"
+FAISS_ENABLE_GPU = "ON"
+FAISS_ENABLE_CUVS = "OFF"
+FAISS_ENABLE_PYTHON = "ON"
+FAISS_ENABLE_MKL = "OFF"
+FAISS_ENABLE_C_API = "OFF"
+FAISS_ENABLE_EXTRAS = "OFF"
+FAISS_ENABLE_SVS = "OFF"
+FAISS_USE_LTO = "ON"
+CMAKE_CUDA_ARCHITECTURES = "70-real;75-real;80-real;86-real;89-real;90"
+CMAKE_CUDA_FLAGS_RELEASE = "-O3 -Xfatbin=-compress-all --threads 4"
+CMAKE_C_FLAGS = "-fdata-sections -ffunction-sections"
+CMAKE_CXX_FLAGS = "-fdata-sections -ffunction-sections"
+CMAKE_SHARED_LINKER_FLAGS = "-Wl,--gc-sections -Wl,--strip-all"
+BUILD_TESTING = "OFF"
+BUILD_SHARED_LIBS = "ON"
+
+# abi3 opt-in (mirrors CPU pyproject.toml).
+[[tool.scikit-build.overrides]]
+if.platform-system = "linux"
+if.abi-flags = "^$"
+wheel.py-api = "cp310"
+
+[tool.cibuildwheel]
+build = "cp310-manylinux_x86_64"
+skip = "*-win* *-macosx* *-manylinux_i686 *-manylinux_aarch64 *-musllinux_*"
+test-requires = [
+  "numpy>=2.0,<3.0",
+  "pytest",
+  "nvidia-cuda-runtime-cu12==12.6.*",
+  "nvidia-cublas-cu12==12.6.*",
+]
+test-command = "python -m pytest {project}/tests/test_wheel_smoke_gpu.py -v"
+
+[tool.cibuildwheel.linux]
+manylinux-x86_64-image = "manylinux_2_28"
+before-all = "dnf install -y epel-release && dnf install -y openblas-devel openblas-openmp swig && dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && dnf install -y cuda-toolkit-12-6"
+# auditwheel: exclude libcuda.so.1 (host driver, not pip-installable)
+# and the libcudart/libcublas* shipped via nvidia-*-cu12 wheels (preloaded at import time).
+repair-wheel-command = "auditwheel repair -w {dest_dir} {wheel} --exclude libcudart.so.12 --exclude libcublas.so.12 --exclude libcublasLt.so.12 --exclude libcuda.so.1"
+
+[[tool.cibuildwheel.overrides]]
+select = "*-manylinux*"
+inherit.repair-wheel-command = "append"
+repair-wheel-command = "pipx run abi3audit --strict --report {wheel}"
+
+[tool.cibuildwheel.linux.environment]
+CUDA_HOME = "/usr/local/cuda-12.6"
+CUDACXX = "/usr/local/cuda-12.6/bin/nvcc"
diff --git a/tests/test_wheel_smoke_gpu.py b/tests/test_wheel_smoke_gpu.py
