Skip to content

Commit

Permalink
Address review comments
Browse files Browse the repository at this point in the history 
Signed-off-by: John Mazanec <jmazane@amazon.com>
  • Loading branch information
@jmazanec15
jmazanec15 committed Jan 23, 2023
1 parent ad93488 commit c7844ed
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 44 deletions.
9 changes: 8 additions & 1 deletion qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/StatsIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
package org.opensearch.knn.bwc;

import org.apache.hc.core5.http.io.entity.EntityUtils;
import org.junit.Before;
import org.opensearch.client.Response;
import org.opensearch.client.ResponseException;
import org.opensearch.knn.plugin.stats.KNNStats;
Expand All @@ -15,7 +16,13 @@
import java.util.Map;

public class StatsIT extends AbstractRollingUpgradeTestCase {
private KNNStats knnStats = new KNNStats();
private KNNStats knnStats;

@Before
public void setUp() throws Exception {
super.setUp();
this.knnStats = new KNNStats();
}

// Validate if all the KNN Stats metrics from old version are present in new version
public void testAllMetricStatsReturned() throws Exception {
Expand Down
106 changes: 65 additions & 41 deletions src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,71 +78,95 @@ private Map<String, KNNStat<?>> getClusterOrNodeStats(Boolean getClusterStats) {
}

private Map<String, KNNStat<?>> buildStatsMap() {
return ImmutableMap.<String, KNNStat<?>>builder()
.put(StatNames.HIT_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::hitCount)))
ImmutableMap.Builder<String, KNNStat<?>> builder = ImmutableMap.<String, KNNStat<?>>builder();
addQueryStats(builder);
addNativeMemoryStats(builder);
addEngineStats(builder);
addScriptStats(builder);
addModelStats(builder);
return builder.build();
}

private void addQueryStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
builder.put(StatNames.KNN_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_REQUESTS)))
.put(
StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS))
);

}

private void addNativeMemoryStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
builder.put(StatNames.HIT_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::hitCount)))
.put(StatNames.MISS_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::missCount)))
.put(StatNames.LOAD_SUCCESS_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::loadSuccessCount)))
.put(StatNames.LOAD_EXCEPTION_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::loadExceptionCount)))
.put(
StatNames.LOAD_EXCEPTION_COUNT.getName(),
new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::loadExceptionCount))
)
.put(StatNames.TOTAL_LOAD_TIME.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::totalLoadTime)))
.put(StatNames.EVICTION_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::evictionCount)))
.put(
StatNames.GRAPH_MEMORY_USAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeInKilobytes))
StatNames.GRAPH_MEMORY_USAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeInKilobytes))
)
.put(
StatNames.GRAPH_MEMORY_USAGE_PERCENTAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeAsPercentage))
StatNames.GRAPH_MEMORY_USAGE_PERCENTAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeAsPercentage))
)
.put(
StatNames.INDICES_IN_CACHE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesCacheStats))
StatNames.INDICES_IN_CACHE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesCacheStats))
)
.put(
StatNames.CACHE_CAPACITY_REACHED.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::isCacheCapacityReached))
StatNames.CACHE_CAPACITY_REACHED.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::isCacheCapacityReached))
)
.put(StatNames.GRAPH_QUERY_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_ERRORS)))
.put(StatNames.GRAPH_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_REQUESTS)))
.put(StatNames.GRAPH_INDEX_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_ERRORS)))
.put(StatNames.GRAPH_INDEX_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_REQUESTS)))
.put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), new KNNStat<>(true, new KNNCircuitBreakerSupplier()))
.put(StatNames.MODEL_INDEX_STATUS.getName(), new KNNStat<>(true, new ModelIndexStatusSupplier<>(ModelDao::getHealthStatus)))
.put(StatNames.KNN_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_REQUESTS)))
.put(
StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS))
)
.put(StatNames.SCRIPT_COMPILATIONS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATIONS)))
.put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), new KNNStat<>(true, new KNNCircuitBreakerSupplier()));
}

private void addEngineStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
builder.put(StatNames.FAISS_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.FAISS)))
.put(StatNames.NMSLIB_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.NMSLIB)))
.put(StatNames.LUCENE_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.LUCENE)));
}

private void addScriptStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
builder.put(StatNames.SCRIPT_COMPILATIONS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATIONS)))
.put(
StatNames.SCRIPT_COMPILATION_ERRORS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATION_ERRORS))
StatNames.SCRIPT_COMPILATION_ERRORS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATION_ERRORS))
)
.put(StatNames.SCRIPT_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_QUERY_REQUESTS)))
.put(StatNames.SCRIPT_QUERY_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_QUERY_ERRORS)))
.put(
StatNames.INDEXING_FROM_MODEL_DEGRADED.getName(),
new KNNStat<>(
false,
new EventOccurredWithinThresholdSupplier(
new ModelIndexingDegradingSupplier(ModelCache::getEvictedDueToSizeAt),
KNNConstants.MODEL_CACHE_CAPACITY_ATROPHY_THRESHOLD_IN_MINUTES,
ChronoUnit.MINUTES
)
)
.put(StatNames.SCRIPT_QUERY_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_QUERY_ERRORS)));
}

private void addModelStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
builder.put(
StatNames.INDEXING_FROM_MODEL_DEGRADED.getName(),
new KNNStat<>(
false,
new EventOccurredWithinThresholdSupplier(
new ModelIndexingDegradingSupplier(ModelCache::getEvictedDueToSizeAt),
KNNConstants.MODEL_CACHE_CAPACITY_ATROPHY_THRESHOLD_IN_MINUTES,
ChronoUnit.MINUTES
)
)
.put(StatNames.FAISS_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.FAISS)))
.put(StatNames.NMSLIB_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.NMSLIB)))
.put(StatNames.LUCENE_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.LUCENE)))
)
.put(StatNames.MODEL_INDEX_STATUS.getName(), new KNNStat<>(true, new ModelIndexStatusSupplier<>(ModelDao::getHealthStatus)))
.put(StatNames.TRAINING_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.TRAINING_REQUESTS)))
.put(StatNames.TRAINING_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.TRAINING_ERRORS)))
.put(
StatNames.TRAINING_MEMORY_USAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeInKilobytes))
StatNames.TRAINING_MEMORY_USAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeInKilobytes))
)
.put(
StatNames.TRAINING_MEMORY_USAGE_PERCENTAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeAsPercentage))
)
.build();
StatNames.TRAINING_MEMORY_USAGE_PERCENTAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeAsPercentage))
);
}
}
7 changes: 5 additions & 2 deletions src/main/java/org/opensearch/knn/plugin/transport/KNNStatsRequest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,16 @@ public class KNNStatsRequest extends BaseNodesRequest<KNNStatsRequest> {
* Key indicating all stats should be retrieved
*/
public static final String ALL_STATS_KEY = "_all";
private Set<String> validStats = StatNames.getNames();
private Set<String> statsToBeRetrieved;
private final Set<String> validStats;
private final Set<String> statsToBeRetrieved;

/**
* Empty constructor needed for KNNStatsTransportAction
*/
public KNNStatsRequest() {
super((String[]) null);
validStats = StatNames.getNames();
statsToBeRetrieved = new HashSet<>();
}

/**
Expand All @@ -53,6 +55,7 @@ public KNNStatsRequest(StreamInput in) throws IOException {
*/
public KNNStatsRequest(String... nodeIds) {
super(nodeIds);
validStats = StatNames.getNames();
statsToBeRetrieved = new HashSet<>();
}

Expand Down

0 comments on commit c7844ed

Please sign in to comment.