Skip to content

Commit 37a4813

Browse files
alibeklfcfacebook-github-bot
alibeklfc
authored and
facebook-github-bot
committed
Add faiss-gpu pip wheel packaging (facebookresearch#5131)
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<->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 8061b24 commit 37a4813

5 files changed

Lines changed: 421 additions & 0 deletions

File tree

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
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: true
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+
with:
25+
fetch-depth: 0
26+
fetch-tags: true
27+
28+
- name: Activate GPU pyproject.toml
29+
# cibuildwheel + scikit-build-core both require literal pyproject.toml.
30+
# Workspace is fresh per job (actions/checkout creates new tree).
31+
run: cp pyproject-gpu.toml pyproject.toml
32+
33+
- name: Set up Python
34+
# Pin host Python explicitly: the 4-core-ubuntu-gpu-t4 runner image
35+
# ships a broken Python 3.14.4 in its tool cache (libpython3.14.so.1.0
36+
# missing), and pypa/cibuildwheel picks that up by default.
37+
uses: actions/setup-python@v5
38+
with:
39+
python-version: '3.12'
40+
41+
- name: Build wheel (cibuildwheel)
42+
uses: pypa/cibuildwheel@ee02a1537ce3071a004a6b08c41e72f0fdc42d9a # v3.4.0
43+
env:
44+
CIBW_BUILD_VERBOSITY: 1
45+
46+
- name: Upload wheel
47+
if: always()
48+
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
49+
with:
50+
name: wheels-gpu
51+
path: wheelhouse/*.whl
52+
53+
publish-gpu:
54+
name: Publish faiss-gpu to PyPI
55+
needs: [build-faiss-gpu]
56+
runs-on: ubuntu-latest
57+
if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
58+
environment:
59+
name: pypi
60+
url: https://pypi.org/p/faiss-gpu
61+
permissions:
62+
id-token: write
63+
steps:
64+
- name: Download wheel
65+
uses: actions/download-artifact@65a9edc5881444af0b9093a5e628f2fe47ea3b2e # v4.1.7
66+
with:
67+
name: wheels-gpu
68+
path: dist
69+
- name: Publish to PyPI
70+
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/__init__.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,59 @@
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+
Behavior:
26+
- faiss-cpu install (no nvidia.cuda_runtime package): silent no-op.
27+
- faiss-gpu install with partial nvidia-* deps (cuda_runtime present
28+
but cublas missing): raises RuntimeError with fix-it message.
29+
"""
30+
import ctypes
31+
import os
32+
33+
# Gate on cuda_runtime first: its absence signals a faiss-cpu install.
34+
try:
35+
import nvidia.cuda_runtime
36+
except ImportError:
37+
return # faiss-cpu install, no GPU libs expected
38+
39+
# cublas is mandatory once cuda_runtime is present; partial install is an error.
40+
try:
41+
import nvidia.cublas
42+
except ImportError as e:
43+
raise RuntimeError(
44+
"faiss-gpu installed but nvidia-cublas-cu12 is missing — "
45+
"pip install nvidia-cublas-cu12==12.6.*"
46+
) from e
47+
48+
# Order matters: libcudart provides symbols that libcublas* depend on.
49+
_LIBS = [
50+
(nvidia.cuda_runtime, "libcudart.so.12"),
51+
(nvidia.cublas, "libcublas.so.12"),
52+
(nvidia.cublas, "libcublasLt.so.12"),
53+
]
54+
for pkg, soname in _LIBS:
55+
so_path = os.path.join(os.path.dirname(pkg.__file__), "lib", soname)
56+
try:
57+
ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)
58+
except OSError as e:
59+
raise RuntimeError(
60+
f"faiss-gpu: failed to load {soname} from {so_path} — "
61+
f"corrupt nvidia-*-cu12 wheel? Try: pip install --force-reinstall "
62+
f"nvidia-cuda-runtime-cu12==12.6.* nvidia-cublas-cu12==12.6.*"
63+
) from e
64+
65+
66+
_preload_gpu_libs()
67+
del _preload_gpu_libs
68+
1669
# We import * so that the symbol foo can be accessed as faiss.foo.
1770
from .loader import *
1871

‎pyproject-gpu.toml‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
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+
skip = "*-win* *-macosx* *-manylinux_i686 *-manylinux_aarch64 *-musllinux_*"
89+
test-requires = [
90+
"numpy>=2.0,<3.0",
91+
"pytest",
92+
"nvidia-cuda-runtime-cu12==12.6.*",
93+
"nvidia-cublas-cu12==12.6.*",
94+
]
95+
test-command = "python -m pytest {project}/tests/test_wheel_smoke_gpu.py -v"
96+
97+
[tool.cibuildwheel.linux]
98+
manylinux-x86_64-image = "manylinux_2_28"
99+
# manylinux_2_28 ships gcc-toolset-14; CUDA 12.6 only supports GCC <=13, so install
100+
# gcc-toolset-13 explicitly and force the build to use it via CC/CXX/CUDAHOSTCXX below.
101+
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-toolkit-12-6"
102+
# auditwheel: exclude libcuda.so.1 (host driver, not pip-installable)
103+
# and the libcudart/libcublas* shipped via nvidia-*-cu12 wheels (preloaded at import time).
104+
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"
105+
106+
[[tool.cibuildwheel.overrides]]
107+
select = "*-manylinux*"
108+
inherit.repair-wheel-command = "append"
109+
repair-wheel-command = "pipx run abi3audit --strict --report {wheel}"
110+
111+
[tool.cibuildwheel.linux.environment]
112+
CUDA_HOME = "/usr/local/cuda-12.6"
113+
CUDACXX = "/usr/local/cuda-12.6/bin/nvcc"
114+
CC = "/opt/rh/gcc-toolset-13/root/usr/bin/gcc"
115+
CXX = "/opt/rh/gcc-toolset-13/root/usr/bin/g++"
116+
CUDAHOSTCXX = "/opt/rh/gcc-toolset-13/root/usr/bin/g++"

0 commit comments

Comments
 (0)