From dceb40c774c4193df24dd3426ea9bb6f8dbb7996 Mon Sep 17 00:00:00 2001
From: Jeff <jeffreyr@securitydept.com>
Date: Fri, 22 Nov 2024 13:38:05 -0500
Subject: [PATCH] Macfaiss (#818)

* macOS platforms faiss OMP Thread patch

* Add TODO

* remove TODO, as build checks for it

* Apply recomendations

* Build Error
---
 docs/embeddings/configuration/ann.md | 2 ++
 docs/faq.md                          | 3 ++-
 src/python/txtai/ann/faiss.py        | 9 +++++++++
 test/python/testann.py               | 8 ++++++++
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/docs/embeddings/configuration/ann.md b/docs/embeddings/configuration/ann.md
index 533918754..162cda345 100644
--- a/docs/embeddings/configuration/ann.md
+++ b/docs/embeddings/configuration/ann.md
@@ -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:
diff --git a/docs/faq.md b/docs/faq.md
index 33c3baf9c..80db87044 100644
--- a/docs/faq.md
+++ b/docs/faq.md
@@ -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__
diff --git a/src/python/txtai/ann/faiss.py b/src/python/txtai/ann/faiss.py
index 72be55470..75d0958e9 100644
--- a/src/python/txtai/ann/faiss.py
+++ b/src/python/txtai/ann/faiss.py
@@ -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):
     """
diff --git a/test/python/testann.py b/test/python/testann.py
index f313cb9ca..c932be541 100644
--- a/test/python/testann.py
+++ b/test/python/testann.py
@@ -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):
         """