Integrate index state sharing into mem management #1545

jmazanec15 · 2024-03-15T16:20:37Z

Description

Last PR in series to fix #1507

Adds the ability to share index state amongst indices during index load operations into the plugins memory management system. Introduces a manager of the shared state that will properly manage the lifecycle of the shared state.

Additionally, there was a bug in clear cache that had to be fixed to get this change working as well. Previously, only one index file per clear cache would be freed. This fixes that logic to clear everything.

Added unit tests and an integration test to confirm functionality. In addition, modified recall integration tests to get more coverage on the different algo configs. Along with this, had to fix a few things around the computation of recall for non-l2 space types.

Note - the addition of these integration tests will increase the CI time. Its a tradeoff that can be discussed. In future, we can take some items to reduce this time (such as #1540 to decouple CI for lib tests and plugin tests)

Once merged, I will create a PR to merge the feature branch into main so it can be released in 2.13.

Issues Resolved

#1507

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.

codecov · 2024-03-15T17:44:04Z

Codecov Report

Attention: Patch coverage is 95.83333% with 3 lines in your changes are missing coverage. Please review.

Project coverage is 85.32%. Comparing base (8fdeb55) to head (05c4530).

Files	Patch %	Lines
...arch/knn/index/memory/SharedIndexStateManager.java	90.62%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@                         Coverage Diff                          @@
##             feature/remove-redunant-allocs    #1545      +/-   ##
====================================================================
+ Coverage                             85.14%   85.32%   +0.18%     
- Complexity                             1289     1302      +13     
====================================================================
  Files                                   168      169       +1     
  Lines                                  5249     5308      +59     
  Branches                                495      499       +4     
====================================================================
+ Hits                                   4469     4529      +60     
+ Misses                                  573      571       -2     
- Partials                                207      208       +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

navneet1v · 2024-03-16T07:28:51Z

src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java

+     *
+     * @param sharedIndexState to return to the system.
+     */
+    public void release(SharedIndexState sharedIndexState) {


should we wrap this whole function in try and finally to ensure that if any exceptions comes from any of the method we are releasing the write lock?

Yes that makes sense

navneet1v · 2024-03-16T07:34:49Z

src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java

+        if (refCount <= 0) {
+            log.info("Evicting entry from shared index state cache for key {}", sharedIndexState.getModelId());
+            sharedIndexStateCache.remove(sharedIndexState.getModelId());
+            JNIService.freeSharedIndexState(sharedIndexState.getSharedIndexStateAddress(), sharedIndexState.getKnnEngine());
+        }


With this logic it means that if some query is still using the SharedIndexState then we will not be able to release the memory. But if a user has done the clear cache api then in that case we will just skip over this release and return that we have cleared the cache?

The release portion is tied to an IndexAllocation not the query. When the index allocation is closed, we will decrease the reference count to this. Once it goes to 0, we free. If an index is loaded and needs shared state, it will need to obtain read lock and so will be blocked. So, after this completes, it will have to load again.

On clear cache, all IndexAllocations will be closed leading to refcount going down.

Okay, so this line will prevent https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/query/KNNWeight.java#L273 the scenario where query is happening and we cannot release the memory.

yes thats correct

Adds the ability to share index state amongst indices during index load operations into the plugins memory management system. Introduces a manager of the shared state that will properly manage the lifecycle of the shared state. There was a bug in clear cache that had to be fixed to get this change working as well. Previously, only one index file per clear cache would be freed. This fixes that logic to clear everything. Added unit tests and an integration test to confirm functionality. In addition, modified recall integration tests to get more coverage on the different algo configs. Along with this, had to fix a few things around the computation of recall for non-l2 space types. Signed-off-by: John Mazanec <jmazane@amazon.com>

VijayanB · 2024-03-18T02:29:27Z

src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java

+     * @return ShareModelContext
+     */
+    public SharedIndexState get(long indexAddress, String modelId, KNNEngine knnEngine) {
+        this.readWriteLock.readLock().lock();


What happens if we don't add lock?

Its possible that underlying shared index state gets freed in release and this returns a value with an invalid pointer

VijayanB · 2024-03-18T02:36:41Z

src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java

+        this.readWriteLock.writeLock().lock();
+
+        try {
+            if (!sharedIndexStateCache.containsKey(sharedIndexState.getModelId())) {


nit: Why not use get and do null check? This will avoid get call at line no 107

Make sense let me add

VijayanB · 2024-03-18T02:40:42Z

src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java

+         * @return ++referenceCount
+         */
+        private long incRef() {
+            return referenceCount.incrementAndGet();


nit: may be void since you are not using incremented value?

I went back and forth. It is required for decRef() so I figured for symmetry I would make it return long

navneet1v · 2024-03-18T06:17:33Z

src/main/java/org/opensearch/knn/index/IndexUtil.java

+        if (StringUtils.isEmpty(modelId) || KNNEngine.FAISS != knnEngine) {
+            return false;
+        }


checking Engine is faiss or not is already covered in the JNIService.isSharedIndexStateRequired why we are adding it here again? We can avoid that check otherwise if we need to share state for some other engine then we will end up making changes here and also in JNIService.

public static boolean isSharedIndexStateRequired(long indexAddr, KNNEngine knnEngine) { if (KNNEngine.FAISS == knnEngine) { return FaissService.isSharedIndexStateRequired(indexAddr); } return false; }

makes sense. will update

jmazanec15 · 2024-03-18T15:31:12Z

BWC seems unrelated. Created an issue around it: #1556

src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java

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

…#1545) Adds the ability to share index state amongst indices during index load operations into the plugins memory management system. Introduces a manager of the shared state that will properly manage the lifecycle of the shared state. There was a bug in clear cache that had to be fixed to get this change working as well. Previously, only one index file per clear cache would be freed. This fixes that logic to clear everything. Added unit tests and an integration test to confirm functionality. In addition, modified recall integration tests to get more coverage on the different algo configs. Along with this, had to fix a few things around the computation of recall for non-l2 space types. Signed-off-by: John Mazanec <jmazane@amazon.com>

Adds the ability to share index state amongst indices during index load operations into the plugins memory management system. Introduces a manager of the shared state that will properly manage the lifecycle of the shared state. There was a bug in clear cache that had to be fixed to get this change working as well. Previously, only one index file per clear cache would be freed. This fixes that logic to clear everything. Added unit tests and an integration test to confirm functionality. In addition, modified recall integration tests to get more coverage on the different algo configs. Along with this, had to fix a few things around the computation of recall for non-l2 space types. Signed-off-by: John Mazanec <jmazane@amazon.com>

…#1545) Adds the ability to share index state amongst indices during index load operations into the plugins memory management system. Introduces a manager of the shared state that will properly manage the lifecycle of the shared state. There was a bug in clear cache that had to be fixed to get this change working as well. Previously, only one index file per clear cache would be freed. This fixes that logic to clear everything. Added unit tests and an integration test to confirm functionality. In addition, modified recall integration tests to get more coverage on the different algo configs. Along with this, had to fix a few things around the computation of recall for non-l2 space types. Signed-off-by: John Mazanec <jmazane@amazon.com>

Adds the ability to share index state amongst indices during index load operations into the plugins memory management system. Introduces a manager of the shared state that will properly manage the lifecycle of the shared state. There was a bug in clear cache that had to be fixed to get this change working as well. Previously, only one index file per clear cache would be freed. This fixes that logic to clear everything. Added unit tests and an integration test to confirm functionality. In addition, modified recall integration tests to get more coverage on the different algo configs. Along with this, had to fix a few things around the computation of recall for non-l2 space types. Signed-off-by: John Mazanec <jmazane@amazon.com>

jmazanec15 added the skip-changelog label Mar 15, 2024

jmazanec15 requested review from heemin32, navneet1v, VijayanB, vamshin, naveentatikonda, junqiu-lei, martin-gaievski and ryanbogan as code owners March 15, 2024 16:20

jmazanec15 force-pushed the issue-1507 branch from 7f32428 to f1681c7 Compare March 15, 2024 17:20

navneet1v reviewed Mar 16, 2024

View reviewed changes

jmazanec15 force-pushed the issue-1507 branch from f1681c7 to 730381e Compare March 18, 2024 01:51

jmazanec15 requested a review from navneet1v March 18, 2024 01:52

VijayanB reviewed Mar 18, 2024

View reviewed changes

navneet1v reviewed Mar 18, 2024

View reviewed changes

jmazanec15 requested review from navneet1v and VijayanB March 18, 2024 16:01

navneet1v reviewed Mar 18, 2024

View reviewed changes

src/main/java/org/opensearch/knn/index/memory/SharedIndexStateManager.java Outdated Show resolved Hide resolved

Address comments

05c4530

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

jmazanec15 force-pushed the issue-1507 branch from c1c39d2 to 05c4530 Compare March 18, 2024 16:45

jmazanec15 requested a review from navneet1v March 18, 2024 16:46

navneet1v approved these changes Mar 18, 2024

View reviewed changes

VijayanB approved these changes Mar 18, 2024

View reviewed changes

jmazanec15 merged commit da8a936 into opensearch-project:feature/remove-redunant-allocs Mar 18, 2024
37 of 51 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Integrate index state sharing into mem management #1545

Integrate index state sharing into mem management #1545

jmazanec15 commented Mar 15, 2024

codecov bot commented Mar 15, 2024 •

edited

Loading

navneet1v Mar 16, 2024

jmazanec15 Mar 16, 2024

navneet1v Mar 16, 2024

jmazanec15 Mar 16, 2024

navneet1v Mar 17, 2024

jmazanec15 Mar 18, 2024

VijayanB Mar 18, 2024

jmazanec15 Mar 18, 2024

VijayanB Mar 18, 2024 •

edited

Loading

jmazanec15 Mar 18, 2024

VijayanB Mar 18, 2024

jmazanec15 Mar 18, 2024

navneet1v Mar 18, 2024

jmazanec15 Mar 18, 2024

jmazanec15 commented Mar 18, 2024

Integrate index state sharing into mem management #1545

Integrate index state sharing into mem management #1545

Conversation

jmazanec15 commented Mar 15, 2024

Description

Issues Resolved

Check List

codecov bot commented Mar 15, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

VijayanB Mar 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jmazanec15 commented Mar 18, 2024

codecov bot commented Mar 15, 2024 •

edited

Loading

VijayanB Mar 18, 2024 •

edited

Loading