Add faiss support in jni by jmazanec15 · Pull Request #28 · opensearch-project/k-NN

jmazanec15 · 2021-05-24T22:41:42Z

Description

This change refactors the jni so that it can support both nmslib and faiss in a single library. The interface revolves around a new class, JNIService, which includes several methods to make calls to the jni for creating indices, training indices, and searching indices. The interface is intended to be future proof (i.e. it will not change if we add a new library).

At the moment, we do not have the functionality to test the native code directly (i.e. for memory leaks). This can be done in the future using Google tests. However, mocking JNI specific components is very tricky. We may consider separating jni components from the rest of the native code so that we can use Google tests to test the rest of the native code.

Included changes in this PR:

Add faiss v1.6.5 as submodule
Refactor build logic to support faiss
Refactor interface to jni in java layer into JNIService
Add jni utilities to abstract out shared functionality
Refactor jni consumers so build passes
Fix existing tests so build passes
Add JNIServiceTests to test JNI
Add test data

This PR DOES NOT add faiss support to the OpenSearch interface. This will be done in a future PR.

Issues Resolved

#29

Check List

New functionality includes testing.
- All tests pass
New functionality has been documented.
- New functionality has javadoc added
Commits are signed as per the DCO using --signoff

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Signed-off-by: John Mazanec <jmazane@amazon.com>

dblock

I did a cursory look paying more attention to the C++ code. I think you want to cleanup the allocation story to make sure nothing is leaking on exceptions, then we can re-review in detail.

Lack of C++ tests is a major problem. Not sure what people use today, Google tests, CPPUNIT or something, but feels like at least a beginning of something is needed, even if it's to test utility functions.

dblock · 2021-05-25T14:45:27Z

+
+    // Train index if needed -- check if there is a case where index needs part trained but is trained
+    if(!indexWriter->is_trained) {
+        //TODO: What needs to be freed???


Everything allocated should be wrapped in an auto_ptr or similar smart pointer construct to do things like env->DeleteLocalRef. I'd ensure this before this is merged.

Makes sense. I will ensure all native allocations are wrapped in smart pointers. For calls to free references in jenv, I will work on some kind of wrapper to ensure everything is freed.

So, DeleteLocalRef is more of an optimization. Local references will automatically be freed when jni call returns.

dblock · 2021-05-25T14:46:35Z

+
+    // Write the index to disk
+    std::string indexPathCpp(ConvertJavaStringToCppString(env, indexPathJ));
+    faiss::write_index(&idMap, indexPathCpp.c_str());


Does this need CatchCppExceptionAndThrowJava(env); around the whole function?

With the future goal of separating java specific functionality from native only functionality, I left this wrapper to the jni functions here: https://github.com/jmazanec15/k-NN-1/blob/faiss-initial/jni/src/org_opensearch_knn_index_JNIService.cpp

dblock · 2021-05-25T14:47:52Z

+    for (int i = 0; i < indexBytesCount; i++) {
+        vectorIoReader.data.push_back((uint8_t) indexBytesJ[i]);
+    }
+    env->ReleaseByteArrayElements(templateIndexJ, indexBytesJ, 0);


Leaks if any of the above fails? indexBytesJ needs to be inside a smart pointer that calls ReleaseByteArrayElements when it's deleted.

For this, I am not sure smart pointer is necessary. ReleaseByteArrayElements requires the Java object to be passed as the first argument. So, we would need to construct a wrapper that includes a reference to the Java object. Because the only code between allocation and release is appending to a vector, I think it is okay to keep this as a raw pointer.

dblock · 2021-05-25T14:49:40Z

+
+    // if result is not enough, padded with -1s
+    int resultSize = kJ;
+    for(int i = 0; i < kJ; i++) {


std::find maybe a 1-liner?

Signed-off-by: John Mazanec <jmazane@amazon.com>

jmazanec15 · 2021-05-25T17:43:55Z

I did a cursory look paying more attention to the C++ code. I think you want to cleanup the allocation story to make sure nothing is leaking on exceptions, then we can re-review in detail.

Sounds good. I will work on that today.

Lack of C++ tests is a major problem. Not sure what people use today, Google tests, CPPUNIT or something, but feels like at least a beginning of something is needed, even if it's to test utility functions.

Makes sense. I was debating whether or not to try to include C++ tests on this PR. I was thinking that because this is a development branch, this can be pushed back to a future PR so as not to add too many changes in one PR. What are your thoughts @dblock?

Signed-off-by: John Mazanec <jmazane@amazon.com>

VijayanB · 2021-05-26T00:59:01Z

+/**
+ * Service to interact with plugin's jni layer. Class dependencies should be minimal
+ *
+ * In order to compile C++ header file, run:


i remember that we have to comment KNNConstants to generate. is that still the case?

No need to comment out. It will work as is.

VijayanB · 2021-05-26T01:03:45Z

+        // loadIndex from different library
+        String engineName;
+        Map<String, Object> parameters;
+        if (indexPathUrl.endsWith(KNNEngine.NMSLIB.getExtension())


shall we create a method "KNNEngine getEngineNameFromPath()", so that it will be easy to extend later? what do you think?

I think that makes sense to add. Will add.

VijayanB · 2021-05-26T01:04:54Z

+            throw new IllegalArgumentException("[KNN] Invalid engine type for path: " + indexPathUrl);
+        }
+
+        final long indexPointer = JNIService.loadIndex(indexPathUrl, parameters, engineName);


may be indexReference? i feel pointer is more of cpp jargon.

I use pointer here because it is generated in the C++ layer and used to access meemory exclusively in C++ layer. So, the value it contains is in fact a valid pointer in C++ and cannot be used to access memory in Java.

VijayanB · 2021-05-26T01:16:36Z


-        private KNNIndex getKnnIndex() {
-            return knnIndex;
+        private boolean isClosed() {


Do we need to have private getter method?

This call is used exclusively in the cache. So private is fine. I guess I prefer to use getters as opposed to direct access to the members so that variables dont accidentally get changed, but I understand that this does not prevent a section of code in the cache from accessing directly a value and changing it, its more for convention. That being said, what are your thoughts?

I understand. I prefer method than variable mostly outside the class. I guess it comes in preference category, you can keep it as it is.

VijayanB · 2021-05-26T01:17:50Z

-            indices = knnIndexCache.getIndices(getHNSWPaths(searcher.getIndexReader()), getIndexName());
-        } finally {
-            searcher.close();
+        try (Engine.Searcher searcher = indexShard.acquireSearcher("knn-warmup")) {


FAR: shall we create constant?

What does FAR stand for?

For Another review

vamshin · 2021-05-26T04:12:31Z

+    if (idsJ == nullptr) {
+        throw std::runtime_error("IDs cannot be null");
+    }
+
+    if (vectorsJ == nullptr) {
+        throw std::runtime_error("Vectors cannot be null");
+    }
+
+    if (indexPathJ == nullptr) {
+        throw std::runtime_error("Index path cannot be null");
+    }
+
+    if (parametersJ == nullptr) {
+        throw std::runtime_error("Parameters cannot be null");
+    }
+


Will these exceptions crash the jvm as these are cpp Exceptions. My understanding Java does not understand these exceptions and crash? If this crashes JVM, then we may need to type case to Java Exception. Same comment for all other places where we throw cpp Exceptions.

We wrap all of the calls to wrappers in https://github.com/jmazanec15/k-NN-1/blob/faiss-initial/jni/src/org_opensearch_knn_index_JNIService.cpp#L20. So, this will ensure that exception is handled correctly. Additionally, some of the tests test exceptions.

vamshin · 2021-05-26T04:26:49Z

+                    }
+            );
+        } catch (Exception ex) {
+            throw new RuntimeException("Unable to query the index: " + ex);


Minor: Add indexName to the Exception message to help in debugging.

Will add here

vamshin · 2021-05-26T04:32:33Z

+                    KNNConstants.HNSW_ALGO_EF_SEARCH, KNNSettings.getEfSearchParam(indexName)
+            );
+        } else {
+            throw new IllegalArgumentException("[KNN] Invalid engine type for path: " + indexPathUrl);


indexPathUrl should be hidden to the users. May be just have indexName in the message. We can
indexPathUrl to a debug log for debugging.

Sure, makes sense.

vamshin · 2021-05-26T04:37:28Z

+        readLock.lock();
+        try {
+            if (knnIndexCacheEntry.isClosed()) {
+                throw new IOException("Index is already closed");


Minor: Add indexName, engineName to the message. If possible add spaceType as well. You may want to check other Exception places as well and see if we can add information to the error messages to help debug.

Signed-off-by: John Mazanec <jmazane@amazon.com>

vamshin

LGTM! Thanks

Signed-off-by: John Mazanec <jmazane@amazon.com>

jmazanec15 · 2021-05-28T17:54:26Z

I decided to refactor jni utils into an interface so that neither wrappers make calls directly to the JNIEnv. This will make it easier to add unit tests because we will just be able to mock the interface. This comes at a minimal performance cost of 1 extra function call, but, when compared to the other operations, this cost is negligible.

Signed-off-by: Jack Mazanec <jmazane1@nd.edu>

Signed-off-by: Jack Mazanec <jmazane1@nd.edu> Signed-off-by: Martin Gaievski <gaievski@amazon.com>

Signed-off-by: Jack Mazanec <jmazane1@nd.edu>

jmazanec15 added 3 commits May 24, 2021 15:42

Initial commit to add faiss to jni

8845b9d

Signed-off-by: John Mazanec <jmazane@amazon.com>

add faiss submodule

e6690a5

Signed-off-by: John Mazanec <jmazane@amazon.com>

Minor changes to get library working

d8c299a

Signed-off-by: John Mazanec <jmazane@amazon.com>

jmazanec15 force-pushed the faiss-initial branch from 34c2c00 to d8c299a Compare May 24, 2021 22:44

jmazanec15 requested review from VijayanB, dblock and vamshin and removed request for dblock May 24, 2021 22:44

dblock reviewed May 25, 2021

View reviewed changes

jmazanec15 added 4 commits May 25, 2021 09:53

Minor style improvements

ddc331c

Signed-off-by: John Mazanec <jmazane@amazon.com>

Fix parameter null checks

498b47d

Signed-off-by: John Mazanec <jmazane@amazon.com>

Add util to get object from map or throw

276cc32

Signed-off-by: John Mazanec <jmazane@amazon.com>

Simplify faiss result parsing logic

a5d1a19

Signed-off-by: John Mazanec <jmazane@amazon.com>

jmazanec15 added 3 commits May 25, 2021 13:36

Improve memory safety

92d70c3

Signed-off-by: John Mazanec <jmazane@amazon.com>

Remove parameters parameter from query index

9d1e958

Signed-off-by: John Mazanec <jmazane@amazon.com>

Add algorithm to fix build failure

4013600

Signed-off-by: John Mazanec <jmazane@amazon.com>

VijayanB reviewed May 26, 2021

View reviewed changes

vamshin reviewed May 26, 2021

View reviewed changes

jmazanec15 added 2 commits May 26, 2021 08:55

Minor cleanup to address comments

e5a1cb5

Signed-off-by: John Mazanec <jmazane@amazon.com>

fix flaky test

2041807

Signed-off-by: John Mazanec <jmazane@amazon.com>

vamshin approved these changes May 26, 2021

View reviewed changes

jmazanec15 added 2 commits May 26, 2021 12:26

Free vector after call completes

9ffc046

Signed-off-by: John Mazanec <jmazane@amazon.com>

Refactor jni utils into interface

199d5f2

Signed-off-by: John Mazanec <jmazane@amazon.com>

VijayanB approved these changes Jun 1, 2021

View reviewed changes

jmazanec15 mentioned this pull request Jun 1, 2021

Add index for storing template models #34

Merged

5 tasks

jmazanec15 merged commit adad71b into opensearch-project:faiss-develop Jun 1, 2021

dblock mentioned this pull request Jun 2, 2021

Support faiss in jni #29

Closed

jmazanec15 added a commit to jmazanec15/k-NN-1 that referenced this pull request Oct 22, 2021

Add faiss support in jni (opensearch-project#28)

104a6ac

Signed-off-by: Jack Mazanec <jmazane1@nd.edu>

jmazanec15 added a commit that referenced this pull request Oct 22, 2021

Add faiss support in jni (#28)

6f83aea

Signed-off-by: Jack Mazanec <jmazane1@nd.edu>

martin-gaievski pushed a commit to martin-gaievski/k-NN that referenced this pull request Mar 7, 2022

Add faiss support in jni (opensearch-project#28)

3470b0f

Signed-off-by: Jack Mazanec <jmazane1@nd.edu>

martin-gaievski pushed a commit to martin-gaievski/k-NN that referenced this pull request Mar 7, 2022

Add faiss support in jni (opensearch-project#28)

164b167

Signed-off-by: Jack Mazanec <jmazane1@nd.edu> Signed-off-by: Martin Gaievski <gaievski@amazon.com>

martin-gaievski pushed a commit to martin-gaievski/k-NN that referenced this pull request Mar 30, 2022

Add faiss support in jni (opensearch-project#28)

ef9338e

Signed-off-by: Jack Mazanec <jmazane1@nd.edu>

jingqimao77-spec pushed a commit to jingqimao77-spec/k-NN that referenced this pull request Mar 15, 2026

Add faiss support in jni (opensearch-project#28)

09ee5a3

Signed-off-by: Jack Mazanec <jmazane1@nd.edu>

Conversation

jmazanec15 commented May 24, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Issues Resolved

Check List

Uh oh!

dblock left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jmazanec15 commented May 25, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VijayanB May 26, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vamshin left a comment

Choose a reason for hiding this comment

Uh oh!

jmazanec15 commented May 28, 2021

Uh oh!

Reviewers

Assignees

Labels

jmazanec15 commented May 24, 2021 •

edited

Loading

VijayanB May 26, 2021 •

edited

Loading