add more stats: connector count, connector/config index status; fix model count bug (#1181)

ylwu-amzn · web-flow · commit 89dab70084a2 · 2023-08-16T15:31:44.000-07:00
Signed-off-by: Yaliang Wu &lt;ylwu@amazon.com&gt;
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java
@@ -34,6 +34,7 @@
 @Log4j2
 public class EncryptorImpl implements Encryptor {
 
+    public static final String MASTER_KEY_NOT_READY_ERROR = "The ML encryption master key has not been initialized yet. Please retry after waiting for 10 seconds.";
     private ClusterService clusterService;
     private Client client;
     private volatile String masterKey;
@@ -114,15 +115,15 @@ private void initMasterKey() {
                         String masterKey = (String) r.getSourceAsMap().get(MASTER_KEY);
                         this.masterKey = masterKey;
                     } else {
-                        exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet"));
+                        exceptionRef.set(new ResourceNotFoundException(MASTER_KEY_NOT_READY_ERROR));
                     }
                 }, e -> {
                     log.error("Failed to get ML encryption master key", e);
                     exceptionRef.set(e);
                 }), latch));
             }
         } else {
-            exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet"));
+            exceptionRef.set(new ResourceNotFoundException(MASTER_KEY_NOT_READY_ERROR));
             latch.countDown();
         }
 
@@ -141,7 +142,7 @@ private void initMasterKey() {
             }
         }
         if (masterKey == null) {
-            throw new ResourceNotFoundException("ML encryption master key not initialized yet");
+            throw new ResourceNotFoundException(MASTER_KEY_NOT_READY_ERROR);
         }
     }
 }
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/utils/ScriptUtils.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/utils/ScriptUtils.java
@@ -37,6 +37,7 @@ public static Optional<String> executePreprocessFunction(ScriptService scriptSer
         }
         return Optional.empty();
     }
+
     public static Optional<String> executePostprocessFunction(ScriptService scriptService,
                                                               String postProcessFunction,
                                                               String resultJson) {
diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java
@@ -31,6 +31,7 @@
 import static org.opensearch.ml.common.CommonValue.CREATE_TIME_FIELD;
 import static org.opensearch.ml.common.CommonValue.MASTER_KEY;
 import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;
+import static org.opensearch.ml.engine.encryptor.EncryptorImpl.MASTER_KEY_NOT_READY_ERROR;
 
 public class EncryptorImplTest {
     @Rule
@@ -121,7 +122,7 @@ public void decrypt() {
     @Test
     public void encrypt_NullMasterKey_NullMasterKey_MasterKeyNotExistInIndex() {
         exceptionRule.expect(ResourceNotFoundException.class);
-        exceptionRule.expectMessage("ML encryption master key not initialized yet");
+        exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
 
         doAnswer(invocation -> {
             ActionListener<GetResponse> listener = invocation.getArgument(1);
@@ -155,7 +156,7 @@ public void decrypt_NullMasterKey_GetMasterKey_Exception() {
     @Test
     public void decrypt_MLConfigIndexNotFound() {
         exceptionRule.expect(ResourceNotFoundException.class);
-        exceptionRule.expectMessage("ML encryption master key not initialized yet");
+        exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
 
         Metadata metadata = new Metadata.Builder().indices(ImmutableMap.of()).build();
         when(clusterState.metadata()).thenReturn(metadata);
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelGroupManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelGroupManager.java
@@ -77,7 +77,9 @@ public void createModelGroup(MLRegisterModelGroupInput input, ActionListener<Str
                                     new IllegalArgumentException(
                                         "The name you provided is already being used by another model with ID: "
                                             + id
-                                            + ". Please provide a different name"
+                                            + ". Please provide a different name or add \"model_group_id\": \""
+                                            + id
+                                            + "\" to request body"
                                     )
                                 );
                         }
diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -5,6 +5,8 @@
 
 package org.opensearch.ml.plugin;
 
+import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_TASK_INDEX;
 
@@ -293,8 +295,12 @@ public Collection<Object> createComponents(
         Map<Enum, MLStat<?>> stats = new ConcurrentHashMap<>();
         // cluster level stats
         stats.put(MLClusterLevelStat.ML_MODEL_INDEX_STATUS, new MLStat<>(true, new IndexStatusSupplier(indexUtils, ML_MODEL_INDEX)));
+        stats
+            .put(MLClusterLevelStat.ML_CONNECTOR_INDEX_STATUS, new MLStat<>(true, new IndexStatusSupplier(indexUtils, ML_CONNECTOR_INDEX)));
+        stats.put(MLClusterLevelStat.ML_CONFIG_INDEX_STATUS, new MLStat<>(true, new IndexStatusSupplier(indexUtils, ML_CONFIG_INDEX)));
         stats.put(MLClusterLevelStat.ML_TASK_INDEX_STATUS, new MLStat<>(true, new IndexStatusSupplier(indexUtils, ML_TASK_INDEX)));
         stats.put(MLClusterLevelStat.ML_MODEL_COUNT, new MLStat<>(true, new CounterSupplier()));
+        stats.put(MLClusterLevelStat.ML_CONNECTOR_COUNT, new MLStat<>(true, new CounterSupplier()));
         // node level stats
         stats.put(MLNodeLevelStat.ML_NODE_EXECUTING_TASK_COUNT, new MLStat<>(false, new CounterSupplier()));
         stats.put(MLNodeLevelStat.ML_NODE_TOTAL_REQUEST_COUNT, new MLStat<>(false, new CounterSupplier()));
@@ -452,7 +458,7 @@ public List<RestHandler> getRestHandlers(
         IndexNameExpressionResolver indexNameExpressionResolver,
         Supplier<DiscoveryNodes> nodesInCluster
     ) {
-        RestMLStatsAction restMLStatsAction = new RestMLStatsAction(mlStats, clusterService, indexUtils);
+        RestMLStatsAction restMLStatsAction = new RestMLStatsAction(mlStats, clusterService, indexUtils, xContentRegistry);
         RestMLTrainingAction restMLTrainingAction = new RestMLTrainingAction();
         RestMLTrainAndPredictAction restMLTrainAndPredictAction = new RestMLTrainAndPredictAction();
         RestMLPredictionAction restMLPredictionAction = new RestMLPredictionAction(mlModelManager);
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLStatsAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLStatsAction.java
@@ -6,6 +6,7 @@
 package org.opensearch.ml.rest;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
 import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
 import static org.opensearch.ml.utils.RestActionUtils.splitCommaSeparatedParam;
@@ -27,6 +28,7 @@
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
@@ -54,17 +56,26 @@ public class RestMLStatsAction extends BaseRestHandler {
     private MLStats mlStats;
     private ClusterService clusterService;
     private IndexUtils indexUtils;
+    private NamedXContentRegistry xContentRegistry;
+    private static final String QUERY_ALL_MODEL_META_DOC =
+        "{\"query\":{\"bool\":{\"must_not\":{\"exists\":{\"field\":\"chunk_number\"}}}}}";
 
     /**
      * Constructor
      * @param mlStats MLStats object
      * @param clusterService cluster service
      * @param indexUtils index util
      */
-    public RestMLStatsAction(MLStats mlStats, ClusterService clusterService, IndexUtils indexUtils) {
+    public RestMLStatsAction(
+        MLStats mlStats,
+        ClusterService clusterService,
+        IndexUtils indexUtils,
+        NamedXContentRegistry xContentRegistry
+    ) {
         this.mlStats = mlStats;
         this.clusterService = clusterService;
         this.indexUtils = indexUtils;
+        this.xContentRegistry = xContentRegistry;
     }
 
     @Override
@@ -109,14 +120,27 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             if (finalMlStatsInput.getTargetStatLevels().contains(MLStatLevel.CLUSTER)
                 && (finalMlStatsInput.retrieveAllClusterLevelStats()
                     || finalMlStatsInput.getClusterLevelStats().contains(MLClusterLevelStat.ML_MODEL_COUNT))) {
-                indexUtils.getNumberOfDocumentsInIndex(ML_MODEL_INDEX, ActionListener.wrap(count -> {
-                    clusterStatsMap.put(MLClusterLevelStat.ML_MODEL_COUNT, count);
-                    getNodeStats(finalMlStatsInput, clusterStatsMap, client, mlStatsNodesRequest, channel);
-                }, e -> {
-                    String errorMessage = "Failed to get ML model count";
-                    log.error(errorMessage, e);
-                    onFailure(channel, RestStatus.INTERNAL_SERVER_ERROR, errorMessage, e);
-                }));
+                indexUtils
+                    .getNumberOfDocumentsInIndex(
+                        ML_MODEL_INDEX,
+                        QUERY_ALL_MODEL_META_DOC,
+                        xContentRegistry,
+                        ActionListener.wrap(modelCount -> {
+                            clusterStatsMap.put(MLClusterLevelStat.ML_MODEL_COUNT, modelCount);
+                            indexUtils.getNumberOfDocumentsInIndex(ML_CONNECTOR_INDEX, ActionListener.wrap(connectorCount -> {
+                                clusterStatsMap.put(MLClusterLevelStat.ML_CONNECTOR_COUNT, connectorCount);
+                                getNodeStats(finalMlStatsInput, clusterStatsMap, client, mlStatsNodesRequest, channel);
+                            }, e -> {
+                                String errorMessage = "Failed to get ML model count";
+                                log.error(errorMessage, e);
+                                onFailure(channel, RestStatus.INTERNAL_SERVER_ERROR, errorMessage, e);
+                            }));
+                        }, e -> {
+                            String errorMessage = "Failed to get ML model count";
+                            log.error(errorMessage, e);
+                            onFailure(channel, RestStatus.INTERNAL_SERVER_ERROR, errorMessage, e);
+                        })
+                    );
             } else {
                 getNodeStats(finalMlStatsInput, clusterStatsMap, client, mlStatsNodesRequest, channel);
             }
diff --git a/plugin/src/main/java/org/opensearch/ml/stats/MLClusterLevelStat.java b/plugin/src/main/java/org/opensearch/ml/stats/MLClusterLevelStat.java
@@ -11,8 +11,11 @@
  */
 public enum MLClusterLevelStat {
     ML_MODEL_INDEX_STATUS,
+    ML_CONNECTOR_INDEX_STATUS,
+    ML_CONFIG_INDEX_STATUS,
     ML_TASK_INDEX_STATUS,
-    ML_MODEL_COUNT;
+    ML_MODEL_COUNT,
+    ML_CONNECTOR_COUNT;
 
     public static MLClusterLevelStat from(String value) {
         try {
diff --git a/plugin/src/main/java/org/opensearch/ml/utils/IndexUtils.java b/plugin/src/main/java/org/opensearch/ml/utils/IndexUtils.java
@@ -8,13 +8,22 @@
 import java.util.List;
 import java.util.Locale;
 
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.admin.indices.stats.IndicesStatsRequest;
+import org.opensearch.action.search.SearchRequest;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.health.ClusterIndexHealth;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.common.xcontent.LoggingDeprecationHandler;
+import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
+import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.search.builder.SearchSourceBuilder;
 
 public class IndexUtils {
     /**
@@ -96,4 +105,33 @@ public void getNumberOfDocumentsInIndex(String indexName, ActionListener<Long> l
         }
     }
 
+    // TODO: add connector count stats
+    public void getNumberOfDocumentsInIndex(
+        String indexName,
+        String searchQuery,
+        NamedXContentRegistry xContentRegistry,
+        ActionListener<Long> listener
+    ) {
+        if (clusterService.state().getRoutingTable().hasIndex(indexName)) {
+            try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+                SearchRequest searchRequest = new SearchRequest();
+                XContentParser parser = XContentType.JSON
+                    .xContent()
+                    .createParser(xContentRegistry, LoggingDeprecationHandler.INSTANCE, searchQuery);
+                SearchSourceBuilder builder = SearchSourceBuilder.fromXContent(parser);
+                builder.fetchSource(false);
+                searchRequest.source(builder).indices(indexName);
+
+                client.search(searchRequest, ActionListener.runBefore(ActionListener.wrap(r -> {
+                    long count = r.getHits().getTotalHits().value;
+                    listener.onResponse(count);
+                }, e -> { listener.onFailure(e); }), () -> context.restore()));
+            } catch (Exception e) {
+                throw new OpenSearchStatusException("Failed to search index " + indexName, RestStatus.BAD_REQUEST);
+            }
+        } else {
+            listener.onResponse(0L);
+        }
+    }
+
 }
diff --git a/plugin/src/test/java/org/opensearch/ml/model/MLModelGroupManagerTests.java b/plugin/src/test/java/org/opensearch/ml/model/MLModelGroupManagerTests.java
@@ -129,7 +129,7 @@ public void test_SuccessAddAllBackendRolesTrue() {
         verify(actionListener).onResponse(argumentCaptor.capture());
     }
 
-    public void test_ModelGroupNameNotUnique() throws IOException {
+    public void test_ModelGroupNameNotUnique() throws IOException {//
         SearchResponse searchResponse = createModelGroupSearchResponse(1);
         doAnswer(invocation -> {
             ActionListener<SearchResponse> listener = invocation.getArgument(1);
@@ -143,7 +143,7 @@ public void test_ModelGroupNameNotUnique() throws IOException {
         ArgumentCaptor<Exception> argumentCaptor = ArgumentCaptor.forClass(Exception.class);
         verify(actionListener).onFailure(argumentCaptor.capture());
         assertEquals(
-            "The name you provided is already being used by another model with ID: model_group_ID. Please provide a different name",
+            "The name you provided is already being used by another model with ID: model_group_ID. Please provide a different name or add \"model_group_id\": \"model_group_ID\" to request body",
             argumentCaptor.getValue().getMessage()
         );
 
diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLStatsActionTests.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLStatsActionTests.java

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,7 @@`
`34`	`34`	`@Log4j2`
`35`	`35`	`public class EncryptorImpl implements Encryptor {`
`36`	`36`
	`37`	`+ public static final String MASTER_KEY_NOT_READY_ERROR = "The ML encryption master key has not been initialized yet. Please retry after waiting for 10 seconds.";`
`37`	`38`	`private ClusterService clusterService;`
`38`	`39`	`private Client client;`
`39`	`40`	`private volatile String masterKey;`
`@@ -114,15 +115,15 @@ private void initMasterKey() {`
`114`	`115`	`String masterKey = (String) r.getSourceAsMap().get(MASTER_KEY);`
`115`	`116`	`this.masterKey = masterKey;`
`116`	`117`	`} else {`
`117`		`- exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet"));`
	`118`	`+ exceptionRef.set(new ResourceNotFoundException(MASTER_KEY_NOT_READY_ERROR));`
`118`	`119`	`}`
`119`	`120`	`}, e -> {`
`120`	`121`	`log.error("Failed to get ML encryption master key", e);`
`121`	`122`	`exceptionRef.set(e);`
`122`	`123`	`}), latch));`
`123`	`124`	`}`
`124`	`125`	`} else {`
`125`		`- exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet"));`
	`126`	`+ exceptionRef.set(new ResourceNotFoundException(MASTER_KEY_NOT_READY_ERROR));`
`126`	`127`	`latch.countDown();`
`127`	`128`	`}`
`128`	`129`
`@@ -141,7 +142,7 @@ private void initMasterKey() {`
`141`	`142`	`}`
`142`	`143`	`}`
`143`	`144`	`if (masterKey == null) {`
`144`		`- throw new ResourceNotFoundException("ML encryption master key not initialized yet");`
	`145`	`+ throw new ResourceNotFoundException(MASTER_KEY_NOT_READY_ERROR);`
`145`	`146`	`}`
`146`	`147`	`}`
`147`	`148`	`}`
Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ public static Optional<String> executePreprocessFunction(ScriptService scriptSer`
`37`	`37`	`}`
`38`	`38`	`return Optional.empty();`
`39`	`39`	`}`
	`40`	`+`
`40`	`41`	`public static Optional<String> executePostprocessFunction(ScriptService scriptService,`
`41`	`42`	`String postProcessFunction,`
`42`	`43`	`String resultJson) {`
Original file line number	Diff line number	Diff line change
`@@ -77,7 +77,9 @@ public void createModelGroup(MLRegisterModelGroupInput input, ActionListener<Str`
`77`	`77`	`new IllegalArgumentException(`
`78`	`78`	`"The name you provided is already being used by another model with ID: "`
`79`	`79`	`+ id`
`80`		`- + ". Please provide a different name"`
	`80`	`+ + ". Please provide a different name or add \"model_group_id\": \""`
	`81`	`+ + id`
	`82`	`+ + "\" to request body"`
`81`	`83`	`)`
`82`	`84`	`);`
`83`	`85`	`}`