Skip to content

Commit

Permalink
PyPI Release Setup for C++ Extension Module (#81)
Browse files Browse the repository at this point in the history
* wip

* debug recall

* clean up

* cleaning up

* some refactoring

* some cmake changes

* revert back to CMake v3.14

* work in progress building library for pypi release

* wip

* add script for building wheels

* workflow for pushing linux and macos wheels to pypi

* remove unnecessary change

* remove extra pyproject.toml

* fix pyton library unit tests

* fix small typo

* increment version number

* fix workflow file

* fail fast set to false

* fix wheel key logic

* make sure xcode is installed in the macos wheel builder

* fix build step for mac

* remove -march=native flag from the macos setup

* fix project setup

* upload wheels individually

* fix workflow file

* Update docker-run.sh

* Update run-benchmark.py

---------

Co-authored-by: blaise-muhirwa <[email protected]>
Co-authored-by: Vihan Lakshman <[email protected]>
  • Loading branch information
3 people authored Jan 16, 2025
1 parent 231b423 commit 207dc6d
Show file tree
Hide file tree
Showing 49 changed files with 4,106 additions and 276 deletions.
46 changes: 1 addition & 45 deletions .github/workflows/cicd.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -61,50 +61,6 @@ jobs:
fi
test-python-library:
strategy:
# Ensure that failure for one test does not cancel other actions
fail-fast: false
matrix:
os: [ubuntu-latest, macos-latest]
python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
runs-on: ${{ matrix.os }}
steps:
- name: Checkout Repository
uses: actions/checkout@v3
with:
submodules: 'recursive'

- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}

- name: Install libomp (macOS only)
if: runner.os == 'macOS'
run: |
brew install libomp
- name: Set up Poetry
uses: snok/install-poetry@v1
with:
version: ${{ env.POETRY_VERSION }}

- name: Install dependencies
run: |
cd flatnav_python
poetry install
- name: Build flatnav
run: |
cd flatnav_python
export NO_SIMD_VECTORIZATION=1
./install_flatnav.sh
- name: Run Unit Tests
run: |
make run-python-unit-tests
run-cpp-unit-tests:
needs: test-clang-format
strategy:
Expand All @@ -130,4 +86,4 @@ jobs:
- name: Build and Run C++ Unit Tests
run: |
set -ex
make run-cpp-unit-tests
make run-cpp-unit-tests
137 changes: 137 additions & 0 deletions .github/workflows/publish-python-library.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
name: Build and Publish Wheels

on:
push:
tags:
# Trigger on version tags (e.g., v1.0.0)
- 'v*'

env:
PYTHON_VERSION: 3.11

jobs:
build_wheels:
name: Build wheels on ${{ matrix.os }} for Python ${{ matrix.python }}
runs-on: ${{ matrix.os }}
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, macos-latest]
python: ['3.8', '3.9', '3.10', '3.11', '3.12']

steps:
- uses: actions/checkout@v4
with:
submodules: 'recursive'

# Set up Python 3.11 specifically for running cibuildwheel
- name: Set up global python
uses: actions/setup-python@v4
with:
python-version: ${{ env.PYTHON_VERSION }}

- name: Install macOS dependencies
if: runner.os == 'macOS'
run: |
xcode-select --install || true
brew install cmake libomp
- name: Install Dependencies
run: |
python -m pip install --upgrade pip
pip install "cibuildwheel==2.22.0"
- name: Set up Docker (Linux only)
if: runner.os == 'Linux'
uses: docker/setup-buildx-action@v3
with:
install: true

- name: Build wheel
env:
MACOSX_DEPLOYMENT_TARGET: '10.14'
run: |
set -ex
./cibuild.sh --current-version ${{ matrix.python }}
- name: Upload wheels as artifacts
uses: actions/upload-artifact@v3
if: always()
with:
name: wheels-${{ runner.os }}-py${{ matrix.python }}
path: wheelhouse/*.whl

publish:
needs: build_wheels
if: always()
runs-on: ubuntu-latest
defaults:
run:
working-directory: python-bindings
steps:
- uses: actions/checkout@v4
with:
submodules: 'recursive'

- name: Install pipx
run: |
python -m pip install --upgrade pip
python -m pip install pipx
- name: Clear dist folder and build SDist
run: |
rm -rf dist/
pipx run build --sdist
- name: Download wheels
uses: actions/download-artifact@v3
with:
path: python-bindings/dist/

- name: Prepare dist directory and upload wheels individually
env:
TWINE_USERNAME: __token__
TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}
run: |
# Move wheels to dist directory
find dist -name "*.whl" -exec mv {} dist/ \;
rm -rf dist/wheels-*
# Verify we have wheels to publish
wheel_count=$(find dist -name "*.whl" | wc -l)
if [ "$wheel_count" -eq 0 ]; then
echo "No wheels found to publish"
exit 1
fi
echo "Found $wheel_count wheels to publish:"
ls -la dist/*.whl
# Upload each wheel individually
failed_uploads=0
for wheel in dist/*.whl; do
echo "Uploading $wheel..."
if ! pipx run twine upload "$wheel"; then
echo "Warning: Failed to upload $wheel"
failed_uploads=$((failed_uploads + 1))
fi
done
# Upload sdist last
echo "Uploading sdist..."
if ! pipx run twine upload dist/*.tar.gz; then
echo "Warning: Failed to upload sdist"
failed_uploads=$((failed_uploads + 1))
fi
# Report final status
echo "Upload complete. $failed_uploads files failed to upload."
if [ "$failed_uploads" -eq "$((wheel_count + 1))" ]; then
echo "All uploads failed"
exit 1
fi
if [ "$failed_uploads" -gt 0 ]; then
echo "Some uploads failed but others succeeded"
exit 0
fi
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -34,3 +34,6 @@ docs/doxygen_output
data/
baselines/
external/openmp

# Python wheels
wheelhouse
30 changes: 15 additions & 15 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -74,21 +74,21 @@ endif()
# TODO: Using globbing or some other command that does not require writing down
# every header file
set(HEADERS
${PROJECT_SOURCE_DIR}/flatnav/distances/InnerProductDistance.h
${PROJECT_SOURCE_DIR}/flatnav/distances/SquaredL2Distance.h
${PROJECT_SOURCE_DIR}/flatnav/distances/L2DistanceDispatcher.h
${PROJECT_SOURCE_DIR}/flatnav/distances/IPDistanceDispatcher.h
${PROJECT_SOURCE_DIR}/flatnav/util/SquaredL2SimdExtensions.h
${PROJECT_SOURCE_DIR}/flatnav/util/InnerProductSimdExtensions.h
${PROJECT_SOURCE_DIR}/flatnav/util/VisitedSetPool.h
${PROJECT_SOURCE_DIR}/flatnav/util/GorderPriorityQueue.h
${PROJECT_SOURCE_DIR}/flatnav/util/Reordering.h
${PROJECT_SOURCE_DIR}/flatnav/util/Multithreading.h
${PROJECT_SOURCE_DIR}/flatnav/util/Macros.h
${PROJECT_SOURCE_DIR}/flatnav/util/Datatype.h
${PROJECT_SOURCE_DIR}/flatnav/util/SimdUtils.h
${PROJECT_SOURCE_DIR}/flatnav/distances/DistanceInterface.h
${PROJECT_SOURCE_DIR}/flatnav/index/Index.h
${PROJECT_SOURCE_DIR}/include/flatnav/distances/InnerProductDistance.h
${PROJECT_SOURCE_DIR}/include/flatnav/distances/SquaredL2Distance.h
${PROJECT_SOURCE_DIR}/include/flatnav/distances/L2DistanceDispatcher.h
${PROJECT_SOURCE_DIR}/include/flatnav/distances/IPDistanceDispatcher.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/SquaredL2SimdExtensions.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/InnerProductSimdExtensions.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/VisitedSetPool.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/GorderPriorityQueue.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/Reordering.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/Multithreading.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/Macros.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/Datatype.h
${PROJECT_SOURCE_DIR}/include/flatnav/util/SimdUtils.h
${PROJECT_SOURCE_DIR}/include/flatnav/distances/DistanceInterface.h
${PROJECT_SOURCE_DIR}/include/flatnav/index/Index.h
${PROJECT_SOURCE_DIR}/developmental-features/quantization/ProductQuantization.h
${PROJECT_SOURCE_DIR}/developmental-features/quantization/CentroidsGenerator.h
${PROJECT_SOURCE_DIR}/developmental-features/quantization/Utils.h)
Expand Down
15 changes: 3 additions & 12 deletions Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ WORKDIR ${FLATNAV_PATH}

# Copy source code
COPY flatnav/ ./flatnav/
COPY flatnav_python/ ./flatnav_python/
COPY python-bindings/ ./python-bindings/
COPY experiments/ ./experiments/

# Copy external dependencies (for now only cereal)
Expand All @@ -90,10 +90,6 @@ COPY external/ ./external/
COPY bin/supervisord.conf /etc/supervisor/conf.d/supervisord.conf

# Install needed dependencies including flatnav.
# This installs numpy as well, which is a large dependency.
WORKDIR ${FLATNAV_PATH}/flatnav_python
RUN ./install_flatnav.sh

# Install hnwlib (from a forked repo that has extensions we need)
WORKDIR ${FLATNAV_PATH}
RUN git clone https://github.com/BlaiseMuhirwa/hnswlib-original.git \
Expand All @@ -102,13 +98,8 @@ RUN git clone https://github.com/BlaiseMuhirwa/hnswlib-original.git \
&& poetry run python setup.py bdist_wheel

# Get the wheel as an environment variable
# NOTE: This is not robust and will break if there are multiple wheels in the dist folder
ENV FLATNAV_WHEEL=${FLATNAV_PATH}/flatnav_python/dist/*.whl
ENV HNSWLIB_WHEEL=${FLATNAV_PATH}/hnswlib-original/python_bindings/dist/*.whl

# Add flatnav and hnswlib to the experiment runner
# Add hnswlib to the experiment runner
WORKDIR ${FLATNAV_PATH}/experiments
RUN poetry add ${FLATNAV_WHEEL} \
&& poetry add ${HNSWLIB_WHEEL} \
&& poetry install --no-root

RUN poetry add ${HNSWLIB_WHEEL} && poetry install --no-root
7 changes: 6 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
CPP_FILES := $(wildcard flatnav/**/*.h flatnav/**/*.cpp flatnav_python/*.cpp tools/*.cpp developmental-features/**/*.h)
CIBUILDWHEEL_VERSION := 2.22.0


format-cpp:
clang-format -i $(CPP_FILES)
Expand All @@ -13,8 +15,11 @@ run-cpp-unit-tests: build-cpp
./build/test_distances
./build/test_serialization

install-cibuildwheel:
pip install "cibuildwheel==${CIBUILDWHEEL_VERSION}"

run-python-unit-tests:
cd flatnav_python && poetry run pytest -vs unit_tests
cd python-bindings && pytest -vs unit_tests

setup-clang-cmake-libomp:
./bin/install_clang_and_libomp.sh
Expand Down
27 changes: 25 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,30 @@ We currently support SIMD extensions for certain platforms as detailed below.

### Getting Started in Python

Once you've built the python bindings and you have a dataset you want to index as a numpy array, you can construct the index as shown below. This will allocate memory and create a directed graph with vectors as nodes.
Currently, we support Python wheels for versions 3.8 through 3.12 on x86_64 architectures (Intel, AMD and MacOS). Support for
ARM wheels is a future improvement.

The python library can be installed from PyPI by using
```shell
$ pip install flatnav
```

Similarly, `flatnav` can be installed from source via [cibuildwheel](https://cibuildwheel.pypa.io/en/stable/), which
builds cross-platform wheels. Follow the following steps

```shell
$ git clone https://github.com/BlaiseMuhirwa/flatnav.git --recurse-submodules
$ cd flatnav
$ make install-cibuildwheel

# This will build flatnav for the current version in your environment. If you want to build wheels
# for all supported python versions (3.8 to 3.12), remove the --current-version flag.
$ ./cibuild.sh --current-version

$ pip install wheelhouse/flatnav*.whl --force-reinstall
```

Once you have the python library installed and you have a dataset you want to index as a numpy array, you can construct the index as shown below. This will allocate memory and create a directed graph with vectors as nodes.

```python
import numpy as np
Expand All @@ -81,7 +104,7 @@ num_build_threads = 16
index = flatnav.index.create(
distance_type=distance_type,
index_data_type=DataType.float32,
dim=dim,
dim=dataset_dimension,
dataset_size=dataset_size,
max_edges_per_node=max_edges_per_node,
verbose=True,
Expand Down
2 changes: 1 addition & 1 deletion bin/docker-run.sh
Original file line number Diff line number Diff line change
Expand Up @@ -103,4 +103,4 @@ docker run \
--volume ${DATA_DIR}:/root/data \
--volume ${METRICS_DIR}:/root/metrics \
--rm flatnav:$TAG_NAME \
/usr/bin/supervisord -c /etc/supervisor/conf.d/supervisord.conf
/usr/bin/supervisord -c /etc/supervisor/conf.d/supervisord.conf
Loading

0 comments on commit 207dc6d

Please sign in to comment.