Skip to content

Commit

Permalink
Remove latestSettings cache from KNNSettings (#727)
Browse files Browse the repository at this point in the history 
Removes the latestSettings cache from the KNNSettings class.
latestSettings cache gets updated in a consumer when the particular
settings are updated.

KNNSettings.getSettingValue would pull from this cache and then fallback
to the default if it is not present. However, in the case when the
settings are set via opensearch.yml config file, the settings update
consumers never get called, so the config values never get put in the
cache. This leads to getSettingValue to always return the default
instead of the value specified in the config file.

To fix this, this change refactors getSettingValue to pull from the
cluster settings and removes the latestSetting cache. However, because
the dynamicCacheSettings have consumers that rebuild the
NativeMemoryCacheManager cache when they are changed, the logic for
passing the parameters to rebuild this cache had to change as well.

Also, switches configuration of cache manager to use DTO.

Signed-off-by: John Mazanec <[email protected]>
  • Loading branch information
@jmazanec15
jmazanec15 authored Jan 19, 2023
1 parent a79a58b commit 8e2ad45
Show file tree
Hide file tree
Showing 6 changed files with 242 additions and 77 deletions.
102 changes: 38 additions & 64 deletions src/main/java/org/opensearch/knn/index/KNNSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,16 +20,17 @@
import org.opensearch.common.unit.TimeValue;
import org.opensearch.index.IndexModule;
import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
import org.opensearch.knn.index.memory.NativeMemoryCacheManagerDto;
import org.opensearch.monitor.jvm.JvmInfo;
import org.opensearch.monitor.os.OsProbe;

import java.security.InvalidParameterException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.util.stream.Stream;

Expand All @@ -41,16 +42,15 @@

/**
* This class defines
* 1. KNN settings to hold the HNSW algorithm parameters.
* https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md
* 1. KNN settings to hold the <a href="https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md">HNSW algorithm parameters</a>.
* 2. KNN settings to enable/disable plugin, circuit breaker settings
* 3. KNN settings to manage graphs loaded in native memory
*/
public class KNNSettings {

private static Logger logger = LogManager.getLogger(KNNSettings.class);
private static final Logger logger = LogManager.getLogger(KNNSettings.class);
private static KNNSettings INSTANCE;
private static OsProbe osProbe = OsProbe.getInstance();
private static final OsProbe osProbe = OsProbe.getInstance();

private static final int INDEX_THREAD_QTY_MAX = 32;

Expand Down Expand Up @@ -85,6 +85,7 @@ public class KNNSettings {
public static final Integer KNN_DEFAULT_CIRCUIT_BREAKER_UNSET_PERCENTAGE = 75;
public static final Integer KNN_DEFAULT_MODEL_CACHE_SIZE_LIMIT_PERCENTAGE = 10; // By default, set aside 10% of the JVM for the limit
public static final Integer KNN_MAX_MODEL_CACHE_SIZE_LIMIT_PERCENTAGE = 25; // Model cache limit cannot exceed 25% of the JVM heap
public static final String KNN_DEFAULT_MEMORY_CIRCUIT_BREAKER_LIMIT = "50%";

/**
* Settings Definition
Expand Down Expand Up @@ -233,7 +234,13 @@ public class KNNSettings {
put(KNN_MEMORY_CIRCUIT_BREAKER_ENABLED, Setting.boolSetting(KNN_MEMORY_CIRCUIT_BREAKER_ENABLED, true, NodeScope, Dynamic));
put(
KNN_MEMORY_CIRCUIT_BREAKER_LIMIT,
knnMemoryCircuitBreakerSetting(KNN_MEMORY_CIRCUIT_BREAKER_LIMIT, "50%", NodeScope, Dynamic)
new Setting<>(
KNNSettings.KNN_MEMORY_CIRCUIT_BREAKER_LIMIT,
KNNSettings.KNN_DEFAULT_MEMORY_CIRCUIT_BREAKER_LIMIT,
(s) -> parseknnMemoryCircuitBreakerValue(s, KNNSettings.KNN_MEMORY_CIRCUIT_BREAKER_LIMIT),
NodeScope,
Dynamic
)
);

/**
Expand All @@ -247,9 +254,6 @@ public class KNNSettings {
}
};

/** Latest setting value for each registered key. Thread-safe is required. */
private final Map<String, Object> latestSettings = new ConcurrentHashMap<>();

private ClusterService clusterService;
private Client client;

Expand All @@ -262,35 +266,32 @@ public static synchronized KNNSettings state() {
return INSTANCE;
}

public void setSettingsUpdateConsumers() {
for (Setting<?> setting : dynamicCacheSettings.values()) {
clusterService.getClusterSettings().addSettingsUpdateConsumer(setting, newVal -> {
logger.debug("The value of setting [{}] changed to [{}]", setting.getKey(), newVal);
latestSettings.put(setting.getKey(), newVal);

// Rebuild the cache with updated limit
NativeMemoryCacheManager.getInstance().rebuildCache();
});
}
private void setSettingsUpdateConsumers() {
clusterService.getClusterSettings().addSettingsUpdateConsumer(updatedSettings -> {
// When any of the dynamic settings are updated, rebuild the cache with the updated values. Use the current
// cluster settings values as defaults.
NativeMemoryCacheManagerDto.NativeMemoryCacheManagerDtoBuilder builder = NativeMemoryCacheManagerDto.builder();

/**
* We do not have to rebuild the cache for below settings
*/
clusterService.getClusterSettings()
.addSettingsUpdateConsumer(
KNN_CIRCUIT_BREAKER_TRIGGERED_SETTING,
newVal -> { latestSettings.put(KNN_CIRCUIT_BREAKER_TRIGGERED, newVal); }
builder.isWeightLimited(
updatedSettings.getAsBoolean(KNN_MEMORY_CIRCUIT_BREAKER_ENABLED, getSettingValue(KNN_MEMORY_CIRCUIT_BREAKER_ENABLED))
);
clusterService.getClusterSettings()
.addSettingsUpdateConsumer(
KNN_CIRCUIT_BREAKER_UNSET_PERCENTAGE_SETTING,
newVal -> { latestSettings.put(KNN_CIRCUIT_BREAKER_UNSET_PERCENTAGE, newVal); }

builder.maxWeight(((ByteSizeValue) getSettingValue(KNN_MEMORY_CIRCUIT_BREAKER_LIMIT)).getKb());
if (updatedSettings.hasValue(KNN_MEMORY_CIRCUIT_BREAKER_LIMIT)) {
builder.maxWeight(((ByteSizeValue) getSetting(KNN_MEMORY_CIRCUIT_BREAKER_LIMIT).get(updatedSettings)).getKb());
}

builder.isExpirationLimited(
updatedSettings.getAsBoolean(KNN_CACHE_ITEM_EXPIRY_ENABLED, getSettingValue(KNN_CACHE_ITEM_EXPIRY_ENABLED))
);
clusterService.getClusterSettings()
.addSettingsUpdateConsumer(
KNN_ALGO_PARAM_INDEX_THREAD_QTY_SETTING,
newVal -> { latestSettings.put(KNN_ALGO_PARAM_INDEX_THREAD_QTY, newVal); }

builder.expiryTimeInMin(
updatedSettings.getAsTime(KNN_CACHE_ITEM_EXPIRY_TIME_MINUTES, getSettingValue(KNN_CACHE_ITEM_EXPIRY_TIME_MINUTES))
.getMinutes()
);

NativeMemoryCacheManager.getInstance().rebuildCache(builder.build());
}, new ArrayList<>(dynamicCacheSettings.values()));
}

/**
Expand All @@ -302,10 +303,10 @@ public void setSettingsUpdateConsumers() {
*/
@SuppressWarnings("unchecked")
public <T> T getSettingValue(String key) {
return (T) latestSettings.getOrDefault(key, getSetting(key).getDefault(Settings.EMPTY));
return (T) clusterService.getClusterSettings().get(getSetting(key));
}

public Setting<?> getSetting(String key) {
private Setting<?> getSetting(String key) {
if (dynamicCacheSettings.containsKey(key)) {
return dynamicCacheSettings.get(key);
}
Expand Down Expand Up @@ -364,19 +365,6 @@ public void initialize(Client client, ClusterService clusterService) {
setSettingsUpdateConsumers();
}

/**
* Creates a setting which specifies a circuit breaker memory limit. This can either be
* specified as an absolute bytes value or as a percentage.
*
* @param key the key for the setting
* @param defaultValue the default value for this setting
* @param properties properties properties for this setting like scope, filtering...
* @return the setting object
*/
public static Setting<ByteSizeValue> knnMemoryCircuitBreakerSetting(String key, String defaultValue, Setting.Property... properties) {
return new Setting<>(key, defaultValue, (s) -> parseknnMemoryCircuitBreakerValue(s, key), properties);
}

public static ByteSizeValue parseknnMemoryCircuitBreakerValue(String sValue, String settingName) {
settingName = Objects.requireNonNull(settingName);
if (sValue != null && sValue.endsWith("%")) {
Expand Down Expand Up @@ -436,24 +424,11 @@ public void onFailure(Exception e) {
* @return efSearch value
*/
public static int getEfSearchParam(String index) {
return getIndexSettingValue(index, KNN_ALGO_PARAM_EF_SEARCH, 512);
}

/**
*
* @param index Name of the index
* @return spaceType name in KNN plugin
*/
public static String getSpaceType(String index) {
return KNNSettings.state().clusterService.state()
.getMetadata()
.index(index)
.getSettings()
.get(KNN_SPACE_TYPE, SpaceType.DEFAULT.getValue());
}

public static int getIndexSettingValue(String index, String settingName, int defaultValue) {
return KNNSettings.state().clusterService.state().getMetadata().index(index).getSettings().getAsInt(settingName, defaultValue);
.getAsInt(KNNSettings.KNN_ALGO_PARAM_EF_SEARCH, 512);
}

public void setClusterService(ClusterService clusterService) {
Expand All @@ -475,7 +450,6 @@ public void validate(String value) {
public void onIndexModule(IndexModule module) {
module.addSettingsUpdateConsumer(INDEX_KNN_ALGO_PARAM_EF_SEARCH_SETTING, newVal -> {
logger.debug("The value of [KNN] setting [{}] changed to [{}]", KNN_ALGO_PARAM_EF_SEARCH, newVal);
latestSettings.put(KNN_ALGO_PARAM_EF_SEARCH, newVal);
// TODO: replace cache-rebuild with index reload into the cache
NativeMemoryCacheManager.getInstance().rebuildCache();
});
Expand Down
54 changes: 42 additions & 12 deletions src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,11 +40,11 @@ public class NativeMemoryCacheManager implements Closeable {

public static String GRAPH_COUNT = "graph_count";

private static Logger logger = LogManager.getLogger(NativeMemoryCacheManager.class);
private static final Logger logger = LogManager.getLogger(NativeMemoryCacheManager.class);
private static NativeMemoryCacheManager INSTANCE;

private Cache<String, NativeMemoryAllocation> cache;
private ExecutorService executor;
private final ExecutorService executor;
private AtomicBoolean cacheCapacityReached;
private long maxWeight;

Expand All @@ -68,20 +68,31 @@ public static synchronized NativeMemoryCacheManager getInstance() {
}

private void initialize() {
initialize(
NativeMemoryCacheManagerDto.builder()
.isWeightLimited(KNNSettings.state().getSettingValue(KNNSettings.KNN_MEMORY_CIRCUIT_BREAKER_ENABLED))
.maxWeight(KNNSettings.getCircuitBreakerLimit().getKb())
.isExpirationLimited(KNNSettings.state().getSettingValue(KNNSettings.KNN_CACHE_ITEM_EXPIRY_ENABLED))
.expiryTimeInMin(
((TimeValue) KNNSettings.state().getSettingValue(KNNSettings.KNN_CACHE_ITEM_EXPIRY_TIME_MINUTES)).getMinutes()
)
.build()
);
}

private void initialize(NativeMemoryCacheManagerDto nativeMemoryCacheDTO) {
CacheBuilder<String, NativeMemoryAllocation> cacheBuilder = CacheBuilder.newBuilder()
.recordStats()
.concurrencyLevel(1)
.removalListener(this::onRemoval);

if (KNNSettings.state().getSettingValue(KNNSettings.KNN_MEMORY_CIRCUIT_BREAKER_ENABLED)) {
maxWeight = KNNSettings.getCircuitBreakerLimit().getKb();
cacheBuilder.maximumWeight(maxWeight).weigher((k, v) -> v.getSizeInKB());
if (nativeMemoryCacheDTO.isWeightLimited()) {
this.maxWeight = nativeMemoryCacheDTO.getMaxWeight();
cacheBuilder.maximumWeight(this.maxWeight).weigher((k, v) -> v.getSizeInKB());
}

if (KNNSettings.state().getSettingValue(KNNSettings.KNN_CACHE_ITEM_EXPIRY_ENABLED)) {
long expiryTime = ((TimeValue) KNNSettings.state().getSettingValue(KNNSettings.KNN_CACHE_ITEM_EXPIRY_TIME_MINUTES))
.getMinutes();
cacheBuilder.expireAfterAccess(expiryTime, TimeUnit.MINUTES);
if (nativeMemoryCacheDTO.isExpirationLimited()) {
cacheBuilder.expireAfterAccess(nativeMemoryCacheDTO.getExpiryTimeInMin(), TimeUnit.MINUTES);
}

cacheCapacityReached = new AtomicBoolean(false);
Expand All @@ -93,13 +104,32 @@ private void initialize() {
* Evicts all entries from the cache and rebuilds.
*/
public synchronized void rebuildCache() {
rebuildCache(
NativeMemoryCacheManagerDto.builder()
.isWeightLimited(KNNSettings.state().getSettingValue(KNNSettings.KNN_MEMORY_CIRCUIT_BREAKER_ENABLED))
.maxWeight(KNNSettings.getCircuitBreakerLimit().getKb())
.isExpirationLimited(KNNSettings.state().getSettingValue(KNNSettings.KNN_CACHE_ITEM_EXPIRY_ENABLED))
.expiryTimeInMin(
((TimeValue) KNNSettings.state().getSettingValue(KNNSettings.KNN_CACHE_ITEM_EXPIRY_TIME_MINUTES)).getMinutes()
)
.build()
);
}

/**
* Evict all entries from the cache and rebuilds
*
* @param nativeMemoryCacheDTO DTO for cache configuration
*/
public synchronized void rebuildCache(NativeMemoryCacheManagerDto nativeMemoryCacheDTO) {
logger.info("KNN Cache rebuilding.");

// TODO: Does this really need to be executed with an executor? Also, does invalidateAll really need to be
// called?
// TODO: Does this really need to be executed with an executor?
executor.execute(() -> {
// Explicitly invalidate all so that we do not have to wait for garbage collection to be invoked to
// free up native memory
cache.invalidateAll();
initialize();
initialize(nativeMemoryCacheDTO);
});
}

Expand Down
18 changes: 18 additions & 0 deletions src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerDto.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.knn.index.memory;

import lombok.Builder;
import lombok.Value;

@Value
@Builder
public class NativeMemoryCacheManagerDto {
boolean isWeightLimited;
long maxWeight;
boolean isExpirationLimited;
long expiryTimeInMin;
}
37 changes: 36 additions & 1 deletion src/test/java/org/opensearch/knn/KNNTestCase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,31 +5,66 @@

package org.opensearch.knn;

import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.settings.ClusterSettings;
import org.opensearch.common.settings.Setting;
import org.opensearch.common.settings.Settings;
import org.opensearch.knn.index.KNNSettings;
import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
import org.opensearch.knn.plugin.stats.KNNCounter;
import org.opensearch.common.bytes.BytesReference;
import org.opensearch.common.xcontent.XContentBuilder;
import org.opensearch.common.xcontent.XContentHelper;
import org.opensearch.test.OpenSearchTestCase;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import static org.mockito.Mockito.when;

/**
* Base class for integration tests for KNN plugin. Contains several methods for testing KNN ES functionality.
*/
public class KNNTestCase extends OpenSearchTestCase {

@Mock
protected ClusterService clusterService;
private AutoCloseable openMocks;

@Override
public void setUp() throws Exception {
super.setUp();
openMocks = MockitoAnnotations.openMocks(this);
}

@Override
public void tearDown() throws Exception {
super.tearDown();
resetState();
openMocks.close();
}

public static void resetState() {
public void resetState() {
// Reset all of the counters
for (KNNCounter knnCounter : KNNCounter.values()) {
knnCounter.set(0L);
}

Set<Setting<?>> defaultClusterSettings = new HashSet<>(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
defaultClusterSettings.addAll(
KNNSettings.state()
.getSettings()
.stream()
.filter(s -> s.getProperties().contains(Setting.Property.NodeScope))
.collect(Collectors.toList())
);
when(clusterService.getClusterSettings()).thenReturn(new ClusterSettings(Settings.EMPTY, defaultClusterSettings));
KNNSettings.state().setClusterService(clusterService);

// Clean up the cache
NativeMemoryCacheManager.getInstance().invalidateAll();
NativeMemoryCacheManager.getInstance().close();
Expand Down
Loading

0 comments on commit 8e2ad45

Please sign in to comment.