Skip to content

Commit

Permalink
Macfaiss (#818)
Browse files Browse the repository at this point in the history
* macOS platforms faiss OMP Thread patch

* Add TODO

* remove TODO, as build checks for it

* Apply recomendations

* Build Error
  • Loading branch information
yukiman76 authored Nov 22, 2024
1 parent 4f37080 commit dceb40c
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 1 deletion.
2 changes: 2 additions & 0 deletions docs/embeddings/configuration/ann.md
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ See the following Faiss documentation links for more information.
- [Binary Indexes](https://github.com/facebookresearch/faiss/wiki/Binary-indexes)
- [Search Tuning](https://github.com/facebookresearch/faiss/wiki/Faster-search)

Note: For macOS users, an existing bug in an upstream package restricts the number of processing threads to 1. This limitation is managed internally to prevent system crashes.

### hnsw
```yaml
hnsw:
Expand Down
3 changes: 2 additions & 1 deletion docs/faq.md
Original file line number Diff line number Diff line change
Expand Up @@ -103,10 +103,11 @@ __Solution__

Set the following environment parameters.

- Disable OpenMP threading via the environment variable `export OMP_NUM_THREADS=1`
- OpenMP threading is handled internally on macOS platforms, but it can be disabled by setting the environment variable OMP_NUM_THREADS to 1, e.g., export OMP_NUM_THREADS=1. For more details, refer to the related discussion on GitHub: kyamagu/faiss-wheels#100.
- Disable PyTorch MPS device via `export PYTORCH_MPS_DISABLE=1`
- Disable llama.cpp metal via `export LLAMA_NO_METAL=1`


----------

__Issue__
Expand Down
9 changes: 9 additions & 0 deletions src/python/txtai/ann/faiss.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,23 @@
"""

import math
import platform

import numpy as np

from faiss import omp_set_num_threads
from faiss import index_factory, IO_FLAG_MMAP, METRIC_INNER_PRODUCT, read_index, write_index
from faiss import index_binary_factory, read_index_binary, write_index_binary, IndexBinaryIDMap

from .base import ANN

if platform.system() == "Darwin":
# Workaround for an open bug on macOS causing segmentation faults in FAISS.
# Setting the number of threads in OpenMP to 1 avoids the issue for now.
# Ref: https://github.com/kyamagu/faiss-wheels/issues/100
# Remove this workaround once the upstream bug is patched.
omp_set_num_threads(1)


class Faiss(ANN):
"""
Expand Down
8 changes: 8 additions & 0 deletions test/python/testann.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,14 @@ def testNumPyLegacy(self):
# Validate count
self.assertEqual(ann.count(), 100)

@patch("platform.system")
def testFaissMacOS(self, mock_platform):
"""
Test backend with OS Darwin
"""
mock_platform.return_value = "Darwin"
self.runTests("faiss")

@patch("sqlalchemy.orm.Query.limit")
def testPGVector(self, query):
"""
Expand Down

0 comments on commit dceb40c

Please sign in to comment.