Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Separate QuantizationStateCache teardown from KNNTestCase #2406

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Separate QuantizationStateCache teardown from KNNTestCase #2406

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
8b14477
Separate QuantizationStateCache teardown from KNNTestCase
owenhalpert Jan 18, 2025
bec302f
Add QSC cleanup to NativeEngines990KnnVectorsFormatTests
owenhalpert Jan 21, 2025
812c383
Add terminated check to cache maintenance start
owenhalpert Jan 21, 2025
045f4bd
Spotless
owenhalpert Jan 21, 2025
3f2ebb0
Add specific cleanup to NativeEngines990KnnVectorsFormatTests
owenhalpert Jan 21, 2025
39b5e99
Typo
owenhalpert Jan 21, 2025
06fd772
Assure cache threadpools are re-set to null so subsequent tests don't…
owenhalpert Jan 21, 2025
61901fc
Generic setup/teardown names
owenhalpert Jan 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.common.xcontent.XContentFactory;
import org.opensearch.index.IndexService;
import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
import org.opensearch.plugins.Plugin;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.test.OpenSearchSingleNodeTestCase;
Expand Down Expand Up @@ -87,7 +86,6 @@ protected boolean resetNodeAfterTest() {
public void tearDown() throws Exception {
NativeMemoryCacheManager.getInstance().invalidateAll();
NativeMemoryCacheManager.getInstance().close();
QuantizationStateCacheManager.getInstance().close();
NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance().close();
NativeMemoryLoadStrategy.TrainingLoadStrategy.getInstance().close();
NativeMemoryLoadStrategy.AnonymousLoadStrategy.getInstance().close();
Expand Down
5 changes: 1 addition & 4 deletions src/test/java/org/opensearch/knn/KNNTestCase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,8 @@
import org.opensearch.core.common.bytes.BytesReference;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.common.xcontent.XContentHelper;
import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
import org.opensearch.test.OpenSearchTestCase;

import java.io.IOException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
Expand Down Expand Up @@ -75,7 +73,7 @@ protected boolean enableWarningsCheck() {
return false;
}

public void resetState() throws IOException {
public void resetState() {
// Reset all of the counters
for (KNNCounter knnCounter : KNNCounter.values()) {
knnCounter.set(0L);
Expand All @@ -85,7 +83,6 @@ public void resetState() throws IOException {
// Clean up the cache
NativeMemoryCacheManager.getInstance().invalidateAll();
NativeMemoryCacheManager.getInstance().close();
QuantizationStateCacheManager.getInstance().close();
}

private void initKNNSettings() {
Expand Down
3 changes: 2 additions & 1 deletion ...org/opensearch/knn/quantization/models/quantizationState/QuantizationStateCacheTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,9 @@ public void setThreadPool() {
}

@After
public void terminateThreadPool() {
public void stopMaintenance() throws IOException {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: you can use cleanup() or tearDown() as a generic method name for test class cleanup.

terminate(threadPool);
QuantizationStateCache.getInstance().close();
Copy link
Contributor Author

@owenhalpert owenhalpert Jan 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This call to close() (which solely cancels the maintenance) isn't necessary because of how the RescheduledRunnable is implemented. When we terminate the threadpool on line 46, the next task is rejected and the task is canceled anyway.

I thought I'd include it for clarity and to immediately cancel the task, but I wanted to flag it to see if anyone had thoughts.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can have it for code clarity. My concern is the setThreadPool(null) pattern.
I know it exists solely for the tests but the threadpool termination and close() should ideally take care of all the cleanup.

Going one step further - I think the setThreadPool should be on a final variable with a non-null check.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess NativeMemoryCacheManager will always exist because it is created as part of plugin initialization. Thats why the NativeMemoryCacheManager.getInstance().close() works butQuantizationStateCacheManager.getInstance().close(); doesn't

Why cant we just have something like public boolean quantizationCacheEnabled; as part of KNNTestCase.java
and inside the resetState() have a check like -

// Clean up the cache
        NativeMemoryCacheManager.getInstance().invalidateAll();
        NativeMemoryCacheManager.getInstance().close();
        if (quantizationCacheEnabled) {
            QuantizationStateCacheManager.getInstance().close();
        }

Only in the tests where we actually have QuantizationStateCacheManager created, we can go set this value to true

}

@SneakyThrows
Expand Down
Loading