WIP

jed326 · Feb 25, 2025 · 9a77801 · 9a77801
1 parent 90aa335
commit 9a77801
Show file tree

Hide file tree

Showing 3 changed files with 137 additions and 1 deletion.
diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNRemoteIndexBuildValue.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNRemoteIndexBuildValue.java
@@ -0,0 +1,89 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.plugin.stats;
+
+import lombok.Getter;
+
+import java.util.concurrent.atomic.LongAdder;
+
+public enum KNNRemoteIndexBuildValue {
+
+    // Repository Stats
+    WRITE_SUCCESS_COUNT("write_success_count"),
+    WRITE_FAILURE_COUNT("write_failure_count"),
+    WRITE_SIZE("write_size_in_bytes"),
+    WRITE_TIME("write_time_in_millis"),
+    READ_SUCCESS_COUNT("read_success_count"),
+    READ_FAILURE_COUNT("read_failure_count"),
+    READ_SIZE("read_size_in_bytes"),
+    READ_TIME("read_time_in_millis"),
+
+    BUILD_REQUEST_SUCCESS_COUNT("build_request_success_count"),
+    BUILD_REQUEST_FAILURE_COUNT("build_request_failure_count"),
+    STATUS_REQUEST_SUCCESS_COUNT("status_request_success_count"),
+    STATUS_REQUEST_FAILURE_COUNT("status_request_failure_count"),
+    INDEX_BUILD_SUCCESS_COUNT("index_build_success_count"),
+    INDEX_BUILD_FAILURE_COUNT("index_build_failure_count"),
+    WAITING_TIME("waiting_time_in_ms"),
+
+    REMOTE_FLUSH_TIME("remote_flush_time_in_millis"),
+    REMOTE_MERGE_TIME("remote_merge_time_in_millis");
+
+    @Getter
+    private final String name;
+    private final LongAdder value;
+
+    /**
+     * Constructor
+     *
+     * @param name name of the graph value
+     */
+    KNNRemoteIndexBuildValue(String name) {
+        this.name = name;
+        this.value = new LongAdder();
+    }
+
+    /**
+     * Get the graph value
+     *
+     * @return value
+     */
+    public Long getValue() {
+        return value.longValue();
+    }
+
+    /**
+     * Increment the graph value
+     */
+    public void increment() {
+        value.increment();
+    }
+
+    /**
+     * Decrement the graph value
+     */
+    public void decrement() {
+        value.decrement();
+    }
+
+    /**
+     * Increment the graph value by a specified amount
+     *
+     * @param delta The amount to increment
+     */
+    public void incrementBy(long delta) {
+        value.add(delta);
+    }
+
+    /**
+     * Decrement the graph value by a specified amount
+     *
+     * @param delta The amount to decrement
+     */
+    public void decrementBy(long delta) {
+        value.add(delta * -1);
+    }
+}
diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java
@@ -8,8 +8,9 @@
 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.common.featureflags.KNNFeatureFlags;
 import org.opensearch.knn.index.engine.KNNEngine;
+import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
 import org.opensearch.knn.indices.ModelCache;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.plugin.stats.suppliers.EventOccurredWithinThresholdSupplier;
@@ -86,6 +87,10 @@ private Map<String, KNNStat<?>> buildStatsMap() {
         addScriptStats(builder);
         addModelStats(builder);
         addGraphStats(builder);
+        if (KNNFeatureFlags.isKNNRemoteVectorBuildEnabled())
+        {
+            addRemoteIndexBuildStats(builder);
+        }
         return builder.build();
     }
 
@@ -210,12 +215,51 @@ private Map<String, Map<String, Object>> createGraphStatsMap() {
         mergeMap.put(KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.getName(), KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.getValue());
         mergeMap.put(KNNGraphValue.MERGE_TOTAL_DOCS.getName(), KNNGraphValue.MERGE_TOTAL_DOCS.getValue());
         mergeMap.put(KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getName(), KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue());
+        if (KNNFeatureFlags.isKNNRemoteVectorBuildEnabled())
+        {
+            mergeMap.put(KNNRemoteIndexBuildValue.REMOTE_MERGE_TIME.getName(), KNNRemoteIndexBuildValue.REMOTE_MERGE_TIME.getValue());
+        }
         Map<String, Object> refreshMap = new HashMap<>();
         refreshMap.put(KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getName(), KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue());
         refreshMap.put(KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getName(), KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getValue());
+        if (KNNFeatureFlags.isKNNRemoteVectorBuildEnabled())
+        {
+            mergeMap.put(KNNRemoteIndexBuildValue.REMOTE_FLUSH_TIME.getName(), KNNRemoteIndexBuildValue.REMOTE_FLUSH_TIME.getValue());
+        }
         Map<String, Map<String, Object>> graphStatsMap = new HashMap<>();
         graphStatsMap.put(StatNames.MERGE.getName(), mergeMap);
         graphStatsMap.put(StatNames.REFRESH.getName(), refreshMap);
         return graphStatsMap;
     }
+
+
+    private void addRemoteIndexBuildStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
+        builder.put(StatNames.REMOTE_VECTOR_INDEX_BUILD_STATS.getName(), new KNNStat<>(false, this::createRemoteIndexStatsMap));
+    }
+
+    private Map<String, Map<String, Object>> createRemoteIndexStatsMap() {
+        Map<String, Object> clientStatsMap = new HashMap<>();
+        clientStatsMap.put(KNNRemoteIndexBuildValue.BUILD_REQUEST_SUCCESS_COUNT.getName(), KNNRemoteIndexBuildValue.BUILD_REQUEST_SUCCESS_COUNT.getValue());
+        clientStatsMap.put(KNNRemoteIndexBuildValue.BUILD_REQUEST_FAILURE_COUNT.getName(), KNNRemoteIndexBuildValue.BUILD_REQUEST_FAILURE_COUNT.getValue());
+        clientStatsMap.put(KNNRemoteIndexBuildValue.STATUS_REQUEST_SUCCESS_COUNT.getName(), KNNRemoteIndexBuildValue.STATUS_REQUEST_SUCCESS_COUNT.getValue());
+        clientStatsMap.put(KNNRemoteIndexBuildValue.STATUS_REQUEST_FAILURE_COUNT.getName(), KNNRemoteIndexBuildValue.STATUS_REQUEST_FAILURE_COUNT.getValue());
+        clientStatsMap.put(KNNRemoteIndexBuildValue.INDEX_BUILD_SUCCESS_COUNT.getName(), KNNRemoteIndexBuildValue.INDEX_BUILD_SUCCESS_COUNT.getValue());
+        clientStatsMap.put(KNNRemoteIndexBuildValue.INDEX_BUILD_FAILURE_COUNT.getName(), KNNRemoteIndexBuildValue.INDEX_BUILD_FAILURE_COUNT.getValue());
+
+        Map<String, Object> repoStatsMap = new HashMap<>();
+        repoStatsMap.put(KNNRemoteIndexBuildValue.WRITE_SUCCESS_COUNT.getName(), KNNRemoteIndexBuildValue.WRITE_SUCCESS_COUNT.getValue());
+        repoStatsMap.put(KNNRemoteIndexBuildValue.WRITE_FAILURE_COUNT.getName(), KNNRemoteIndexBuildValue.WRITE_FAILURE_COUNT.getValue());
+        repoStatsMap.put(KNNRemoteIndexBuildValue.WRITE_SIZE.getName(), KNNRemoteIndexBuildValue.WRITE_SIZE.getValue());
+        repoStatsMap.put(KNNRemoteIndexBuildValue.WRITE_TIME.getName(), KNNRemoteIndexBuildValue.WRITE_TIME.getValue());
+        repoStatsMap.put(KNNRemoteIndexBuildValue.READ_SUCCESS_COUNT.getName(), KNNRemoteIndexBuildValue.READ_SUCCESS_COUNT.getValue());
+        repoStatsMap.put(KNNRemoteIndexBuildValue.READ_FAILURE_COUNT.getName(), KNNRemoteIndexBuildValue.READ_FAILURE_COUNT.getValue());
+        repoStatsMap.put(KNNRemoteIndexBuildValue.READ_SIZE.getName(), KNNRemoteIndexBuildValue.READ_SIZE.getValue());
+        repoStatsMap.put(KNNRemoteIndexBuildValue.READ_TIME.getName(), KNNRemoteIndexBuildValue.READ_TIME.getValue());
+
+        Map<String, Map<String, Object>> remoteIndexBuildStatsMap = new HashMap<>();
+        remoteIndexBuildStatsMap.put(StatNames.REMOTE_INDEX_BUILD_SERVICE_STATS.getName(), clientStatsMap);
+        remoteIndexBuildStatsMap.put(StatNames.REPOSITORY_STATS.getName(), repoStatsMap);
+
+        return remoteIndexBuildStatsMap;
+    }
 }
diff --git a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java
@@ -45,6 +45,9 @@ public enum StatNames {
     GRAPH_STATS("graph_stats"),
     REFRESH("refresh"),
     MERGE("merge"),
+    REMOTE_VECTOR_INDEX_BUILD_STATS("remote_vector_index_build_stats"),
+    REMOTE_INDEX_BUILD_SERVICE_STATS("remote_index_build_service_stats"),
+    REPOSITORY_STATS("repository_stats"),
     MIN_SCORE_QUERY_REQUESTS(KNNCounter.MIN_SCORE_QUERY_REQUESTS.getName()),
     MIN_SCORE_QUERY_WITH_FILTER_REQUESTS(KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS.getName()),
     MAX_DISTANCE_QUERY_REQUESTS(KNNCounter.MAX_DISTANCE_QUERY_REQUESTS.getName()),