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

Remove latestSettings cache from KNNSettings #727

Merged
merged 4 commits into from
Jan 19, 2023
Merged

Remove latestSettings cache from KNNSettings #727

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
728b2a8
Remove latestSettings cache from KNNSettings
jmazanec15 Jan 13, 2023
4769666
Address review comments
jmazanec15 Jan 18, 2023
896816c
Minor test cleanup
jmazanec15 Jan 18, 2023
3a0277d
Address comments
jmazanec15 Jan 19, 2023
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
102 changes: 38 additions & 64 deletions src/main/java/org/opensearch/knn/index/KNNSettings.java
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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