Skip to content

Commit b2f023b

Browse files
alibeklfcmeta-codesync[bot]
alibeklfc
authored and
meta-codesync[bot]
committed
Add faiss-gpu pip wheel packaging (facebookresearch#5131)
Summary: Pull Request resolved: facebookresearch#5131 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<->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
1 parent b008961 commit b2f023b

6 files changed

Lines changed: 458 additions & 0 deletions

File tree

‎.github/workflows/build-pip-gpu.yml‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
name: Build GPU pip wheel
7+
8+
on:
9+
workflow_call:
10+
workflow_dispatch:
11+
12+
concurrency:
13+
group: gpu-pip-${{ github.ref }}
14+
cancel-in-progress: ${{ !startsWith(github.ref, 'refs/tags/v') }}
15+
16+
jobs:
17+
build-faiss-gpu:
18+
name: Build & test faiss-gpu
19+
runs-on: 4-core-ubuntu-gpu-t4
20+
timeout-minutes: 90
21+
steps:
22+
- name: Checkout
23+
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
24+
25+
- name: Activate GPU pyproject.toml
26+
# cibuildwheel + scikit-build-core both require literal pyproject.toml.
27+
# Workspace is fresh per job (actions/checkout creates new tree).
28+
run: cp pyproject-gpu.toml pyproject.toml
29+
30+
- name: Set up Python
31+
# Pin host Python explicitly: the 4-core-ubuntu-gpu-t4 runner image
32+
# ships a broken Python 3.14.4 in its tool cache (libpython3.14.so.1.0
33+
# missing), and pypa/cibuildwheel picks that up by default.
34+
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
35+
with:
36+
python-version: '3.12'
37+
38+
- name: Build wheel (cibuildwheel)
39+
uses: pypa/cibuildwheel@ee02a1537ce3071a004a6b08c41e72f0fdc42d9a # v3.4.0
40+
env:
41+
CIBW_BUILD_VERBOSITY: 1
42+
43+
- name: Upload wheel
44+
if: always()
45+
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
46+
with:
47+
name: wheels-gpu
48+
path: wheelhouse/*.whl
49+
50+
publish-gpu:
51+
name: Publish faiss-gpu to PyPI
52+
needs: [build-faiss-gpu]
53+
runs-on: ubuntu-latest
54+
if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
55+
environment:
56+
name: pypi
57+
url: https://pypi.org/p/faiss-gpu
58+
permissions:
59+
id-token: write
60+
steps:
61+
- name: Download wheel
62+
uses: actions/download-artifact@65a9edc5881444af0b9093a5e628f2fe47ea3b2e # v4.1.7
63+
with:
64+
name: wheels-gpu
65+
path: dist
66+
- name: Publish to PyPI
67+
uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc # release/v1

‎.github/workflows/build.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,5 @@ jobs:
1717
if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
1818
build-pip:
1919
uses: ./.github/workflows/build-pip.yml
20+
build-pip-gpu:
21+
uses: ./.github/workflows/build-pip-gpu.yml

‎faiss/python/CMakeLists.txt‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,18 @@ install(FILES
442442
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/__init__.pyi")
443443
install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/__init__.pyi" DESTINATION faiss)
444444
endif()
445+
446+
# GPU build marker: only present in faiss-gpu wheels. Used by __init__.py to
447+
# decide whether to preload CUDA shared libs from nvidia-*-cu12 pip packages.
448+
# Gating on this marker (instead of importability of `nvidia.cuda_runtime`)
449+
# prevents false positives in faiss-cpu environments where PyTorch or cupy
450+
# have transitively pulled in the nvidia-* packages.
451+
if(FAISS_ENABLE_GPU)
452+
file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/_gpu_build.py"
453+
"# Auto-generated when FAISS_ENABLE_GPU=ON.\n"
454+
"IS_GPU_BUILD = True\n")
455+
install(FILES "${CMAKE_CURRENT_BINARY_DIR}/_gpu_build.py" DESTINATION faiss)
456+
endif()
445457
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/py.typed")
446458
install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/py.typed" DESTINATION faiss)
447459
endif()

‎faiss/python/__init__.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,71 @@
1313
import sys
1414
import inspect
1515

16+
17+
def _preload_gpu_libs():
18+
"""Pre-load CUDA shared libraries from nvidia-*-cu12 pip packages.
19+
20+
The libcudart.so / libcublas.so files live in
21+
site-packages/nvidia/{cuda_runtime,cublas}/lib/ — not on the default
22+
ld.so search path. Pre-load them with RTLD_GLOBAL before the SWIG
23+
extension dlopen()s libfaiss.so (which has CUDA undefined symbols).
24+
25+
Gating: keyed on the `faiss._gpu_build` module, which CMake writes only
26+
when FAISS_ENABLE_GPU=ON. We can't gate on `nvidia.cuda_runtime`
27+
importability — PyTorch and cupy CUDA wheels pull that in transitively,
28+
so a faiss-cpu user with PyTorch+CUDA would otherwise hit this path and
29+
see a misleading error if any nvidia-* package were missing.
30+
31+
Behavior:
32+
- faiss-cpu install (no `_gpu_build` marker): silent no-op.
33+
- faiss-gpu install with missing nvidia-cuda-runtime / nvidia-cublas:
34+
raises RuntimeError with fix-it message.
35+
"""
36+
try:
37+
from . import _gpu_build # noqa: F401
38+
except ImportError:
39+
return # faiss-cpu install: marker absent, nothing to preload
40+
41+
import ctypes
42+
import os
43+
44+
try:
45+
import nvidia.cuda_runtime
46+
except ImportError as e:
47+
raise RuntimeError(
48+
"faiss-gpu installed but nvidia-cuda-runtime-cu12 is missing — "
49+
"pip install nvidia-cuda-runtime-cu12==12.6.*"
50+
) from e
51+
52+
try:
53+
import nvidia.cublas
54+
except ImportError as e:
55+
raise RuntimeError(
56+
"faiss-gpu installed but nvidia-cublas-cu12 is missing — "
57+
"pip install nvidia-cublas-cu12==12.6.*"
58+
) from e
59+
60+
# Order matters: libcudart provides symbols that libcublas* depend on.
61+
_LIBS = [
62+
(nvidia.cuda_runtime, "libcudart.so.12"),
63+
(nvidia.cublas, "libcublas.so.12"),
64+
(nvidia.cublas, "libcublasLt.so.12"),
65+
]
66+
for pkg, soname in _LIBS:
67+
so_path = os.path.join(os.path.dirname(pkg.__file__), "lib", soname)
68+
try:
69+
ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)
70+
except OSError as e:
71+
raise RuntimeError(
72+
f"faiss-gpu: failed to load {soname} from {so_path} — "
73+
f"corrupt nvidia-*-cu12 wheel? Try: pip install --force-reinstall "
74+
f"nvidia-cuda-runtime-cu12==12.6.* nvidia-cublas-cu12==12.6.*"
75+
) from e
76+
77+
78+
_preload_gpu_libs()
79+
del _preload_gpu_libs
80+
1681
# We import * so that the symbol foo can be accessed as faiss.foo.
1782
from .loader import *
1883

‎pyproject-gpu.toml‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
[build-system]
7+
requires = ["scikit-build-core>=0.10", "swig>=4.2,<5", "numpy>=2.0", "cmake>=3.24", "ninja"]
8+
build-backend = "scikit_build_core.build"
9+
10+
[project]
11+
name = "faiss-gpu"
12+
dynamic = ["version"]
13+
description = "A library for efficient similarity search and clustering of dense vectors (GPU support)."
14+
readme = "README.md"
15+
license = "MIT"
16+
license-files = ["LICENSE", "THIRD_PARTY_NOTICES"]
17+
requires-python = ">=3.10"
18+
authors = [
19+
{ name = "Matthijs Douze" },
20+
{ name = "Jeff Johnson" },
21+
{ name = "Hervé Jégou" },
22+
{ name = "Lucas Hosseini" },
23+
]
24+
keywords = ["similarity-search", "vector-search", "nearest-neighbors", "clustering", "gpu", "cuda"]
25+
classifiers = [
26+
"Development Status :: 5 - Production/Stable",
27+
"Intended Audience :: Developers",
28+
"Intended Audience :: Science/Research",
29+
"Operating System :: POSIX :: Linux",
30+
"Programming Language :: C++",
31+
"Programming Language :: Python :: 3",
32+
"Programming Language :: Python :: 3.10",
33+
"Programming Language :: Python :: 3.11",
34+
"Programming Language :: Python :: 3.12",
35+
"Programming Language :: Python :: 3.13",
36+
"Programming Language :: Python :: 3 :: Only",
37+
"Programming Language :: Python :: Implementation :: CPython",
38+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
39+
"Topic :: Software Development :: Libraries :: Python Modules",
40+
"Environment :: GPU :: NVIDIA CUDA :: 12",
41+
]
42+
dependencies = [
43+
"numpy>=1.25",
44+
"packaging",
45+
"nvidia-cuda-runtime-cu12==12.6.*",
46+
"nvidia-cublas-cu12==12.6.*",
47+
]
48+
49+
[project.urls]
50+
Homepage = "https://github.com/facebookresearch/faiss"
51+
Repository = "https://github.com/facebookresearch/faiss"
52+
Issues = "https://github.com/facebookresearch/faiss/issues"
53+
Documentation = "https://github.com/facebookresearch/faiss/wiki"
54+
55+
[tool.scikit-build]
56+
cmake.build-type = "Release"
57+
metadata.version.provider = "scikit_build_core.metadata.regex"
58+
metadata.version.input = "CMakeLists.txt"
59+
metadata.version.regex = '^\s+VERSION\s+(?P<value>[0-9]+\.[0-9]+\.[0-9]+)'
60+
wheel.exclude = ["include/**", "share/**", "lib/cmake/**", "lib/pkgconfig/**"]
61+
62+
[tool.scikit-build.cmake.define]
63+
FAISS_OPT_LEVEL = "dd"
64+
FAISS_ENABLE_GPU = "ON"
65+
FAISS_ENABLE_CUVS = "OFF"
66+
FAISS_ENABLE_PYTHON = "ON"
67+
FAISS_ENABLE_MKL = "OFF"
68+
FAISS_ENABLE_C_API = "OFF"
69+
FAISS_ENABLE_EXTRAS = "OFF"
70+
FAISS_ENABLE_SVS = "OFF"
71+
FAISS_USE_LTO = "ON"
72+
CMAKE_CUDA_ARCHITECTURES = "70-real;75-real;80-real;86-real;89-real;90"
73+
CMAKE_CUDA_FLAGS_RELEASE = "-O3 -Xfatbin=-compress-all --threads 4"
74+
CMAKE_C_FLAGS = "-fdata-sections -ffunction-sections"
75+
CMAKE_CXX_FLAGS = "-fdata-sections -ffunction-sections"
76+
CMAKE_SHARED_LINKER_FLAGS = "-Wl,--gc-sections -Wl,--strip-all"
77+
BUILD_TESTING = "OFF"
78+
BUILD_SHARED_LIBS = "ON"
79+
80+
# abi3 opt-in (mirrors CPU pyproject.toml).
81+
[[tool.scikit-build.overrides]]
82+
if.platform-system = "linux"
83+
if.abi-flags = "^$"
84+
wheel.py-api = "cp310"
85+
86+
[tool.cibuildwheel]
87+
build = "cp310-manylinux_x86_64"
88+
test-requires = [
89+
"numpy>=2.0,<3.0",
90+
"pytest",
91+
"nvidia-cuda-runtime-cu12==12.6.*",
92+
"nvidia-cublas-cu12==12.6.*",
93+
]
94+
test-command = "python -m pytest {project}/tests/test_wheel_smoke_gpu.py -v"
95+
96+
[tool.cibuildwheel.linux]
97+
manylinux-x86_64-image = "manylinux_2_28"
98+
# manylinux_2_28 ships gcc-toolset-14; CUDA 12.6 only supports GCC <=13, so install
99+
# gcc-toolset-13 explicitly and force the build to use it via CC/CXX/CUDAHOSTCXX below.
100+
before-all = "dnf install -y gcc-toolset-13 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-nvcc-12-6 cuda-cudart-devel-12-6 cuda-cccl-12-6 libcublas-devel-12-6"
101+
# auditwheel: exclude libcuda.so.1 (host driver, not pip-installable)
102+
# and the libcudart/libcublas* shipped via nvidia-*-cu12 wheels (preloaded at import time).
103+
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"
104+
105+
[[tool.cibuildwheel.overrides]]
106+
select = "*-manylinux*"
107+
inherit.repair-wheel-command = "append"
108+
repair-wheel-command = "pipx run abi3audit --strict --report {wheel}"
109+
110+
[tool.cibuildwheel.linux.environment]
111+
CUDA_HOME = "/usr/local/cuda-12.6"
112+
CUDACXX = "/usr/local/cuda-12.6/bin/nvcc"
113+
CC = "/opt/rh/gcc-toolset-13/root/usr/bin/gcc"
114+
CXX = "/opt/rh/gcc-toolset-13/root/usr/bin/g++"
115+
CUDAHOSTCXX = "/opt/rh/gcc-toolset-13/root/usr/bin/g++"

0 commit comments

Comments
 (0)