[FEATURE] [Build-TimeReduction V1] Merge Optimization: Streaming vectors from Java Layer to JNI layer to avoid OOM/Circuit Breaker in native engines #1506
Assignees
Labels
enhancement
indexing-improvements
This label should be attached to all the github issues which will help improving the indexing time.
v2.14.0
Comments
vamshin
changed the title
This was referenced Mar 28, 2024
Merged
|
The issue is updated with all the details on the solution and benchmarking |
navneet1v
added
the
indexing-improvements
This was referenced Apr 9, 2024
Merged
|
Closing the PR as changes are merged and will be released with 2.14 version of Opensearch |
github-project-automation
bot
moved this from 2.14.0
to ✅ Done
in Vector Search RoadMap
navneet1v
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Problem Statement
While testing on with below details, I was seeing CB exceptions.
Logs:
If CB is enabled
If CB is not enabled
Root Cause
The reason why this OOM exception/CB is tripping because while creating the Faiss/nmslib index at a segment level we first load all the vectors(floats) in JVM heap. As vectors are 4byte floats this lead to an array of size ~28.4GB ((4 * 1536 * 5000000)/2^30) and then OOM. Ref: https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java#L42-L59
Solution
The solution I am proposing here is while reading the vectors from doc values we will stream/transfer the vectors to a memory address in Native memory(RAM) and then pass that address to JNI layer while creating indices for native libraries( Faiss and Nmslib), rather than accumulating the vectors in a list in heap and then pass this to JNI layer. I have done a POC implementation for this same here. This will help resolve the issue described at the start reason being we are just keeping a finite amount of vectors in heap hence no OOM or CB will happen.
Critical Design Choices
How many vectors we should be streaming at once from Java to JNI Layer?
This is an interesting choice to take, as we don’t want to stream a lot of vectors to JNI layer at once because it can lead to GC getting triggered when Heap Memory is under stress, we also don’t want to steam too less which can lead to this context switch and more fragmentation in Native memory.
Approach 1
So what I am proposing here is may be instead of Number of vectors we should focus on amount of data we should be streaming because this is what actually being sent. Considering a typical JVM size as 32GB for production workloads streaming 100Mb which is 0.003% of the whole heap.
Table 1: Providing details around segment size and vectors
Considering the above table we can see that number of trips to JNI will be increased as the dimension increase if we keep a constant data that can be sent to JNI.
The concern here is not the number of trips we are making to JNI, problem is every time we go to JNI we will be adding the floats in c++ stl vectors. If there is not enough memory to expand the vector in place then c++ will copy this whole vector to new memory location and then all data to it(ref). This will add latency in the overall system. Check below benchmarks which shows that if you send all data at once and if you send in batch how much extra latency gets added.
So what I am proposing here is may be instead of Number of vectors we should focus on amount of data we should be streaming because this is what actually being sent. Considering a typical JVM size as 32GB for production workloads streaming 100Mb which is 0.003% of the whole heap.
Table 2: Benchmarking results when a fixed number of vectors are transferred without initial capacity being set.
Approach 2 (Recommended)
To ensure that we are not adding any extra latency that is coming due to sizing and re-sizing of the stl::vector we can set an initial capacity for the vector. If we do so then there is no re-sizing happening and we can avoid the extra latency. Below benchmarks run the same experiment except now the stl::vector expansion is not happening.Table 3: Benchmarking results when a fixed number of vectors are transferred with initial capacity being set.
As we can see by removing the expansion out of picture and applying few more optimization which is only possible by setting the initial capacity(directly converting java float array to vectors rather than an intermediate storage), we are able to reduce the time taken to move vectors from Java to JNI layer > 50%.
How do we find out accurately the number of vectors which we want to stream to JNI layer without reading all the vectors in the segment?
On examining closely the DocValues interface provides a function called as cost() which returns the max documents present. But this includes the deleted documents too.
The segment creation/graph creation happens in 2 scenarios:
For #1, the number of vectors won’t will be that high which can cause the OOM issues, hence we can stream all the vectors to JNI layer directly and we don’t need to depend on cost() to determine the size of vectors upfront.
For #2, as the merges happen for large segments and there will be deleted docs then cost() function cannot be used, as this will lead to creating large chunk of memory which we will not use for graph creation. To accurately find the number of docs in BinaryDocValues we can use the LiveDoc bits. Please refer this POC code ref. Sample sudo code below:
So, Approach 2 and setting the initial capacity for the stl::vector we can remove the decision of selecting the accurate size of vectors to be transferred and become more memory oriented transfers. Below are the benchmarks that shows with different memory size like 100MB, 200MB .. 500MB what will be the impact on latency. We can start with a default value of 1% of heap size by setting this in cluster setting. This will provide user flexibility to change it in future.
The benefit of using a memory based limit on the transfer of vectors over fixed number based limit is even when number of shards increase on the node, the system will be stable. Example if we decide to use lets say 100K as a limit then having 50 shards on the node can lead to heap CBs as total heap required will ~29GB((586*50)/1024).
Table 4: Shows the details on how heap usage changes based dimension for a fixed number of vectors.
Benchmarks with different size of data transfers
Table 5: Benchmarking results when fixed memory of vectors are transferred with initial capacity being set.
If we compare table 5 with table 3, we can see that the difference in the latencies are very minimal < 100ms. This provides a strong evidence that we should use a fixed amount of memory for data transfer which is more flexible and robust as compared to transferring a fixed amount of vectors.
Apart from vectors should we also stream the docIds?
I think we should not stream docids. Here are some stats, In Lucene we can have at 2^(31)-1 docIds and an Integer takes 4 bytes so total required memory to hold all docids is 0.5GB((2^(31)-1)/2^30). We do this we will doing over engineering in our solution.In native memory should we store vectors as 1-D array or create a 2-D array to store the vectors?
If we look at create index interfaces of Faiss and Nmslib we can see that both of them takes a vectors in 1D array. Hence it make sense to use a 1-D array otherwise we will need to construct the 1-D array from 2D array which will consume computations for a larger datasets.Test Plan
Below are the list of tests/Benchmarks that will performed apart from Unit test and Integration tests.Correctness Testing
The below tests will ensure that we are able to merge the large dataset even when heap size is less that what Opensearch can accommodate. The below configurations errors out on Opensearch version 2.13. The reason for choosing these configurations because we have seen errors in these configurations recently.A/B Testing with Opensearch version 2.13
We are going to use nightly benchmarks to validate if there are any regression happening in the system with this change. The main parameters we will be looking for is the indexing time, force merge time and refresh time. Along with that recall should remain intact.Currently nightly benchmarks doesn’t run training related workloads for those workloads we are going to run them separately. Below are the details.
The experiment results will be compared with results here: #1473
Tasks
The text was updated successfully, but these errors were encountered: