Add patch to fix segfault in nmslib during ingestion (#1541)

Signed-off-by: John Mazanec <[email protected]>
opensearch-project · Mar 14, 2024 · 2b0f5a3 · 2b0f5a3
1 parent bfcf7dc
commit 2b0f5a3
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 0 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -47,6 +47,9 @@ jobs:
           rm ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
           git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0002-Custom-patch-to-support-AVX2-Linux-CI.patch
           rm ../../patches/faiss/0002-Custom-patch-to-support-AVX2-Linux-CI.patch
+          cd ../nmslib
+          git apply --ignore-space-change --ignore-whitespace --3way ../../patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch
+          rm ../../patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch
         working-directory: ${{ github.workspace }}
 
       - name: Setup Java ${{ matrix.java }}

diff --git a/.github/workflows/test_security.yml b/.github/workflows/test_security.yml
@@ -47,6 +47,9 @@ jobs:
           rm ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
           git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0002-Custom-patch-to-support-AVX2-Linux-CI.patch
           rm ../../patches/faiss/0002-Custom-patch-to-support-AVX2-Linux-CI.patch
+          cd ../nmslib
+          git apply --ignore-space-change --ignore-whitespace --3way ../../patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch
+          rm ../../patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch
         working-directory: ${{ github.workspace }}
 
       - name: Setup Java ${{ matrix.java }}

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Bug Fixes
 * Disable sdc table for HNSWPQ read-only indices [#1518](https://github.com/opensearch-project/k-NN/pull/1518)
 * Switch SpaceType.INNERPRODUCT's vector similarity function to MAXIMUM_INNER_PRODUCT [#1532](https://github.com/opensearch-project/k-NN/pull/1532)
+* Add patch to fix arm segfault in nmslib during ingestion [#1541](https://github.com/opensearch-project/k-NN/pull/1541)
 ### Infrastructure
 * Manually install zlib for win CI [#1513](https://github.com/opensearch-project/k-NN/pull/1513)
 ### Documentation

diff --git a/jni/CMakeLists.txt b/jni/CMakeLists.txt
@@ -87,6 +87,19 @@ if (${CONFIG_NMSLIB} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST}
         execute_process(COMMAND git submodule update --init -- external/nmslib WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
     endif ()
 
+    # Check if patch exist, this is to skip git apply during CI build. See CI.yml with ubuntu.
+    find_path(PATCH_FILE NAMES 0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch PATHS ${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib NO_DEFAULT_PATH)
+
+    # If it exists, apply patches
+    if (EXISTS ${PATCH_FILE})
+        message(STATUS "Applying custom patches.")
+        execute_process(COMMAND git apply --ignore-space-change --ignore-whitespace --3way ${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib ERROR_VARIABLE ERROR_MSG RESULT_VARIABLE RESULT_CODE)
+
+        if(RESULT_CODE)
+            message(FATAL_ERROR "Failed to apply patch:\n${ERROR_MSG}")
+        endif()
+    endif()
+
     add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search)
 
     add_library(${TARGET_LIB_NMSLIB} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_jni_NmslibService.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/nmslib_wrapper.cpp)

diff --git a/jni/patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch b/jni/patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch
@@ -0,0 +1,31 @@
+From aa1ca485c0ab8b79dae1fb5c1567149c5f61b533 Mon Sep 17 00:00:00 2001
+From: John Mazanec <[email protected]>
+Date: Thu, 14 Mar 2024 12:22:06 -0700
+Subject: [PATCH] Initialize maxlevel during add from enterpoint->level
+
+Signed-off-by: John Mazanec <[email protected]>
+---
+ similarity_search/src/method/hnsw.cc | 6 +++++-
+ 1 file changed, 5 insertions(+), 1 deletion(-)
+
+diff --git a/similarity_search/src/method/hnsw.cc b/similarity_search/src/method/hnsw.cc
+index 35b372c..e9a725e 100644
+--- a/similarity_search/src/method/hnsw.cc
++++ b/similarity_search/src/method/hnsw.cc
+@@ -542,8 +542,12 @@ namespace similarity {
+
+         NewElement->init(curlevel, maxM_, maxM0_);
+
+-        int maxlevelcopy = maxlevel_;
++        // Get the enterpoint at this moment and then use it to set the 
++        // max level that is used. Copying maxlevel from this->maxlevel_
++        // can lead to race conditions during concurrent insertion. See:
++        // https://github.com/nmslib/nmslib/issues/544
+         HnswNode *ep = enterpoint_;
++        int maxlevelcopy = ep->level; 
+         if (curlevel < maxlevelcopy) {
+             const Object *currObj = ep->getData();
+
+-- 
+2.39.3 (Apple Git-146)
+