Separate QuantizationStateCache teardown from KNNTestCase

[owenhalpert]

Description

QuantizationStateCacheManager.getInstance() will create an instance of a QuantizationStateCache if one is not already present. Since the cache is not created by default (like the NativeMemoryCache is), the call to QuantizationStateCacheManager.getInstance().close() in KNNTestCase will wastefully create instances in many cases. There are very few test classes in which a QuantizationStateCache is created, so it will be better to move this logic somewhere more specific.

Importantly, a side effect of the current flow is KNNWeightTests mocks KNNSettings but doesn't mock the QS cache settings, so when the cache is created due to the call above, a NullPointerException is thrown and all the tests in the class fail when running in isolation:
java.lang.NullPointerException: Cannot invoke "org.opensearch.core.common.unit.ByteSizeValue.getKb()" because the return value of "org.opensearch.knn.index.KNNSettings.getSettingValue(String)" is null

This error doesn't show up when running the full suite of tests, presumably because the default setting is applied elsewhere.

Check List

• New functionality includes testing.
• New functionality has been documented.
• [NA] API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
• [NA] Public documentation issue/PR created.

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.

[owenhalpert]

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.

[kotwanikunal]

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.

[Gankris96]

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

[kotwanikunal]

nit: I don't think we need this explicitly.

[Gankris96]

I guess the close of NativeMemoryCacheManager should have this. Because that is where the executor.shutdown() is called. The same place should be used to cleanup all tasks. I see the maintenanceTask.cancel() is called. Post that we should clear out the threadpool; no?

Actually, this should not even be required as @kotwanikunal mentioned

[owenhalpert]

@Gankris96 the only consideration there is that the threadpool is set only once in KNNPlugin where we have access to pass in OpenSearch's threadpool. Once close() is called and that is set to null, we would lose that reference unless KNNPlugin.createComponents is called again. I would need to investigate the impact of getting rid of that reference fully.

[Gankris96]

I think the problem is that the threadpool is getting shutdown in the tests before you clear out the QSCache object. That is what causes the threadpool != null check to succeed but the mockNode.shutDown() terminates the threadpool.

For example

k-NN/src/test/java/org/opensearch/knn/index/KNNSettingsTests.java

Line 64 in f5cf255

mockNode.close();

is where the node is closed, releasing all non-active threadpool assignments. In your case it wouldn't be active since QuantizationStateCache has not been initiailzed but you are setting it on a static variable.

[kotwanikunal]

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