Skip to content

Commit

Permalink
Refactor structure of stats module (#736)
Browse files Browse the repository at this point in the history 
Removes valid stats arg from KNNStatsRequest. Valid stats is a global
constant so there is not a good reason to make it configurable. It is
still written in and out of the stream. This is to prevent breakage in
any kind of backwards compatibility.

Removes KNNStatsConfig class and moves logic directly into the KNNStats
class. KNNStatsConfig just hosted a single map of the stats, so it felt
unneccessary to have the whole class.

Signed-off-by: John Mazanec <[email protected]>
  • Loading branch information
@jmazanec15
jmazanec15 authored Jan 23, 2023
1 parent 8e2ad45 commit 185b54d
Show file tree
Hide file tree
Showing 9 changed files with 136 additions and 152 deletions.
11 changes: 8 additions & 3 deletions 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 @@ -14,10 +15,14 @@
import java.util.List;
import java.util.Map;

import static org.opensearch.knn.plugin.stats.KNNStatsConfig.KNN_STATS;

public class StatsIT extends AbstractRollingUpgradeTestCase {
private KNNStats knnStats = new KNNStats(KNN_STATS);
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
3 changes: 1 addition & 2 deletions src/main/java/org/opensearch/knn/index/KNNCircuitBreaker.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
package org.opensearch.knn.index;

import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
import org.opensearch.knn.plugin.stats.KNNStatsConfig;
import org.opensearch.knn.plugin.stats.StatNames;
import org.opensearch.knn.plugin.transport.KNNStatsAction;
import org.opensearch.knn.plugin.transport.KNNStatsNodeResponse;
Expand Down Expand Up @@ -73,7 +72,7 @@ public void initialize(ThreadPool threadPool, ClusterService clusterService, Cli

// Leader node untriggers CB if all nodes have not reached their max capacity
if (KNNSettings.isCircuitBreakerTriggered() && clusterService.state().nodes().isLocalNodeElectedClusterManager()) {
KNNStatsRequest knnStatsRequest = new KNNStatsRequest(KNNStatsConfig.KNN_STATS.keySet());
KNNStatsRequest knnStatsRequest = new KNNStatsRequest();
knnStatsRequest.addStat(StatNames.CACHE_CAPACITY_REACHED.getName());
knnStatsRequest.timeout(new TimeValue(1000 * 10)); // 10 second timeout

Expand Down
5 changes: 2 additions & 3 deletions src/main/java/org/opensearch/knn/plugin/KNNPlugin.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@
import org.opensearch.index.IndexModule;
import org.opensearch.index.IndexSettings;
import org.opensearch.index.mapper.Mapper;
import org.opensearch.knn.plugin.stats.KNNStatsConfig;
import org.opensearch.knn.plugin.transport.RemoveModelFromCacheAction;
import org.opensearch.knn.plugin.transport.RemoveModelFromCacheTransportAction;
import org.opensearch.knn.plugin.transport.SearchModelAction;
Expand Down Expand Up @@ -202,7 +201,7 @@ public Collection<Object> createComponents(
KNNQueryBuilder.initialize(ModelDao.OpenSearchKNNModelDao.getInstance());
KNNWeight.initialize(ModelDao.OpenSearchKNNModelDao.getInstance());
TrainingModelRequest.initialize(ModelDao.OpenSearchKNNModelDao.getInstance(), clusterService);
knnStats = new KNNStats(KNNStatsConfig.KNN_STATS);
knnStats = new KNNStats();
return ImmutableList.of(knnStats);
}

Expand All @@ -221,7 +220,7 @@ public List<RestHandler> getRestHandlers(
Supplier<DiscoveryNodes> nodesInCluster
) {

RestKNNStatsHandler restKNNStatsHandler = new RestKNNStatsHandler(settings, restController, knnStats);
RestKNNStatsHandler restKNNStatsHandler = new RestKNNStatsHandler();
RestKNNWarmupHandler restKNNWarmupHandler = new RestKNNWarmupHandler(
settings,
restController,
Expand Down
23 changes: 3 additions & 20 deletions src/main/java/org/opensearch/knn/plugin/rest/RestKNNStatsHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,18 +5,14 @@

package org.opensearch.knn.plugin.rest;

import lombok.AllArgsConstructor;
import org.opensearch.knn.plugin.KNNPlugin;
import org.opensearch.knn.plugin.stats.KNNStats;
import org.opensearch.knn.plugin.transport.KNNStatsAction;
import org.opensearch.knn.plugin.transport.KNNStatsRequest;
import com.google.common.collect.ImmutableList;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.client.node.NodeClient;
import org.opensearch.common.Strings;
import org.opensearch.common.settings.Settings;
import org.opensearch.rest.BaseRestHandler;
import org.opensearch.rest.RestController;
import org.opensearch.rest.RestRequest;
import org.opensearch.rest.action.RestActions;

Expand All @@ -33,22 +29,9 @@
* Resthandler for stats api endpoint. The user has the ability to get all stats from
* all nodes or select stats from specific nodes.
*/
@AllArgsConstructor
public class RestKNNStatsHandler extends BaseRestHandler {

private static final Logger LOG = LogManager.getLogger(RestKNNStatsHandler.class);
private static final String NAME = "knn_stats_action";
private KNNStats knnStats;

/**
* Constructor
*
* @param settings Settings
* @param controller Rest Controller
* @param knnStats KNNStats
*/
public RestKNNStatsHandler(Settings settings, RestController controller, KNNStats knnStats) {
this.knnStats = knnStats;
}

@Override
public String getName() {
Expand Down Expand Up @@ -104,7 +87,7 @@ private KNNStatsRequest getRequest(RestRequest request) {
nodeIdsArr = nodesIdsStr.split(",");
}

KNNStatsRequest knnStatsRequest = new KNNStatsRequest(knnStats.getStats().keySet(), nodeIdsArr);
KNNStatsRequest knnStatsRequest = new KNNStatsRequest(nodeIdsArr);
knnStatsRequest.timeout(request.param("timeout"));

// parse the stats the customer wants to see
Expand Down
132 changes: 113 additions & 19 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 @@ -5,6 +5,23 @@

package org.opensearch.knn.plugin.stats;

import com.google.common.cache.CacheStats;
import com.google.common.collect.ImmutableMap;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
import org.opensearch.knn.index.util.KNNEngine;
import org.opensearch.knn.indices.ModelCache;
import org.opensearch.knn.indices.ModelDao;
import org.opensearch.knn.plugin.stats.suppliers.EventOccurredWithinThresholdSupplier;
import org.opensearch.knn.plugin.stats.suppliers.KNNCircuitBreakerSupplier;
import org.opensearch.knn.plugin.stats.suppliers.KNNCounterSupplier;
import org.opensearch.knn.plugin.stats.suppliers.KNNInnerCacheStatsSupplier;
import org.opensearch.knn.plugin.stats.suppliers.LibraryInitializedSupplier;
import org.opensearch.knn.plugin.stats.suppliers.ModelIndexStatusSupplier;
import org.opensearch.knn.plugin.stats.suppliers.ModelIndexingDegradingSupplier;
import org.opensearch.knn.plugin.stats.suppliers.NativeMemoryCacheManagerSupplier;

import java.time.temporal.ChronoUnit;
import java.util.HashMap;
import java.util.Map;

Expand All @@ -13,15 +30,13 @@
*/
public class KNNStats {

private Map<String, KNNStat<?>> knnStats;
private final Map<String, KNNStat<?>> knnStats;

/**
* Constructor
*
* @param knnStats Map that maps name of stat to KNNStat object
*/
public KNNStats(Map<String, KNNStat<?>> knnStats) {
this.knnStats = knnStats;
public KNNStats() {
this.knnStats = buildStatsMap();
}

/**
Expand All @@ -33,20 +48,6 @@ public Map<String, KNNStat<?>> getStats() {
return knnStats;
}

/**
* Get individual stat by stat name
*
* @param key Name of stat
* @return ADStat
* @throws IllegalArgumentException thrown on illegal statName
*/
public KNNStat<?> getStat(String key) throws IllegalArgumentException {
if (!knnStats.keySet().contains(key)) {
throw new IllegalArgumentException("Stat=\"" + key + "\" does not exist");
}
return knnStats.get(key);
}

/**
* Get a map of the stats that are kept at the node level
*
Expand Down Expand Up @@ -75,4 +76,97 @@ private Map<String, KNNStat<?>> getClusterOrNodeStats(Boolean getClusterStats) {
}
return statsMap;
}

private Map<String, KNNStat<?>> buildStatsMap() {
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.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))
)
.put(
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))
)
.put(
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()));
}

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))
)
.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)));
}

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.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))
)
.put(
StatNames.TRAINING_MEMORY_USAGE_PERCENTAGE.getName(),
new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeAsPercentage))
);
}
}
94 changes: 0 additions & 94 deletions src/main/java/org/opensearch/knn/plugin/stats/KNNStatsConfig.java
View file

This file was deleted.

Loading

0 comments on commit 185b54d

Please sign in to comment.