Serialize all models into cluster metadata

CHANGELOG.md

@@ -20,6 +20,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Support script score when doc value is disabled [#1573](https://github.com/opensearch-project/k-NN/pull/1573)
 * Implemented the Streaming Feature to stream vectors from Java to JNI layer to enable creation of larger segments for vector indices [#1604](https://github.com/opensearch-project/k-NN/pull/1604)
 * Remove unnecessary toString conversion of vector field and added some minor optimization in KNNCodec [1613](https://github.com/opensearch-project/k-NN/pull/1613)
+* Serialize all models into cluster metadata [#1499](https://github.com/opensearch-project/k-NN/pull/1499)
 ### Bug Fixes
 ### Infrastructure
 * Add micro-benchmark module in k-NN plugin for benchmark streaming vectors to JNI layer functionality. [#1583](https://github.com/opensearch-project/k-NN/pull/1583)

src/main/java/org/opensearch/knn/index/IndexUtil.java

@@ -23,6 +23,7 @@
 import org.opensearch.knn.index.util.KNNEngine;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
+import org.opensearch.knn.indices.ModelUtil;
 import org.opensearch.knn.jni.JNIService;
 
 import java.io.File;
@@ -197,7 +198,7 @@ public static ValidationException validateKnnField(
             }
 
             ModelMetadata modelMetadata = modelDao.getMetadata(modelId);
-            if (modelMetadata == null) {
+            if (!ModelUtil.isModelPresent(modelMetadata) || !ModelUtil.isModelCreated(modelMetadata)) {
                 exception.addValidationError(String.format("Model \"%s\" for field \"%s\" does not exist.", modelId, field));
                 return exception;
             }
@@ -284,4 +285,5 @@ public static boolean isSharedIndexStateRequired(KNNEngine knnEngine, String mod
         }
         return JNIService.isSharedIndexStateRequired(indexAddr, knnEngine);
     }
+
 }

src/main/java/org/opensearch/knn/index/mapper/ModelFieldMapper.java

@@ -11,6 +11,7 @@
 import org.opensearch.index.mapper.ParseContext;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
+import org.opensearch.knn.indices.ModelUtil;
 
 import java.io.IOException;
 
@@ -50,7 +51,7 @@ protected void parseCreateField(ParseContext context) throws IOException {
         // model when ingestion starts.
         ModelMetadata modelMetadata = this.modelDao.getMetadata(modelId);
 
-        if (modelMetadata == null) {
+        if (!ModelUtil.isModelPresent(modelMetadata) || !ModelUtil.isModelCreated(modelMetadata)) {
             throw new IllegalStateException(
                 String.format(
                     "Model \"%s\" from %s's mapping does not exist. Because the \"%s\" parameter is not updatable, this index will need to be recreated with a valid model.",

src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java

@@ -24,6 +24,7 @@
 import org.opensearch.knn.index.util.KNNEngine;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
+import org.opensearch.knn.indices.ModelUtil;
 import org.opensearch.knn.plugin.stats.KNNCounter;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
@@ -517,7 +518,7 @@ private ModelMetadata getModelMetadataForField(KNNVectorFieldMapper.KNNVectorFie
         }
 
         ModelMetadata modelMetadata = modelDao.getMetadata(modelId);
-        if (modelMetadata == null) {
+        if (!ModelUtil.isModelPresent(modelMetadata) || !ModelUtil.isModelCreated(modelMetadata)) {
             throw new IllegalArgumentException(String.format("Model ID '%s' does not exist.", modelId));
         }
         return modelMetadata;

src/main/java/org/opensearch/knn/index/query/KNNWeight.java

@@ -40,6 +40,7 @@
 import org.opensearch.knn.index.util.KNNEngine;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
+import org.opensearch.knn.indices.ModelUtil;
 import org.opensearch.knn.jni.JNIService;
 import org.opensearch.knn.plugin.stats.KNNCounter;
 
@@ -213,7 +214,7 @@ private Map<Integer, Float> doANNSearch(final LeafReaderContext context, final B
         String modelId = fieldInfo.getAttribute(MODEL_ID);
         if (modelId != null) {
             ModelMetadata modelMetadata = modelDao.getMetadata(modelId);
-            if (modelMetadata == null) {
+            if (!ModelUtil.isModelPresent(modelMetadata) || !ModelUtil.isModelCreated(modelMetadata)) {
                 throw new RuntimeException("Model \"" + modelId + "\" does not exist.");
             }
 

src/main/java/org/opensearch/knn/indices/ModelDao.java

@@ -341,13 +341,7 @@ private void putInternal(Model model, ActionListener<IndexResponse> listener, Do
                 );
             }, listener::onFailure);
 
-            // After the model is indexed, update metadata only if the model is in CREATED state
-            ActionListener<IndexResponse> onIndexListener;
-            if (ModelState.CREATED.equals(model.getModelMetadata().getState())) {
-                onIndexListener = getUpdateModelMetadataListener(model.getModelMetadata(), onMetaListener);
-            } else {
-                onIndexListener = onMetaListener;
-            }
+            ActionListener<IndexResponse> onIndexListener = getUpdateModelMetadataListener(model.getModelMetadata(), onMetaListener);
 
             // Create the model index if it does not already exist
             Runnable indexModelRunnable = () -> indexRequestBuilder.execute(onIndexListener);

src/main/java/org/opensearch/knn/indices/ModelUtil.java

@@ -0,0 +1,24 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.knn.indices;
+
+public class ModelUtil {
+
+    public static boolean isModelPresent(ModelMetadata modelMetadata) {
+        return modelMetadata != null;
+    }
+
+    public static boolean isModelCreated(ModelMetadata modelMetadata) {
+        return modelMetadata.getState().equals(ModelState.CREATED);
+    }
+
+}

src/main/java/org/opensearch/knn/training/TrainingJobClusterStateListener.java

@@ -112,7 +112,7 @@ protected void updateModelsNewCluster() throws IOException, InterruptedException
                 Model model = modelDao.get(modelId);
                 ModelMetadata modelMetadata = model.getModelMetadata();
                 if (modelMetadata.getState().equals(ModelState.TRAINING)) {
-                    updateModelStateAsFailed(model, "Training failed to complete as cluster crashed");
+                    updateModelStateAsFailed(modelId, modelMetadata, "Training failed to complete as cluster crashed");
                 }
             }
         }
@@ -123,11 +123,10 @@ protected void updateModelsNodesRemoved(List<DiscoveryNode> removedNodes) throws
             List<String> modelIds = searchModelIds();
             for (DiscoveryNode removedNode : removedNodes) {
                 for (String modelId : modelIds) {
-                    Model model = modelDao.get(modelId);
-                    ModelMetadata modelMetadata = model.getModelMetadata();
+                    ModelMetadata modelMetadata = modelDao.getMetadata(modelId);
                     if (modelMetadata.getNodeAssignment().equals(removedNode.getEphemeralId())
                         && modelMetadata.getState().equals(ModelState.TRAINING)) {
-                        updateModelStateAsFailed(model, "Training failed to complete as node dropped");
+                        updateModelStateAsFailed(modelId, modelMetadata, "Training failed to complete as node dropped");
                     }
                 }
             }
@@ -158,9 +157,11 @@ public void onFailure(Exception e) {
         return modelIds;
     }
 
-    private void updateModelStateAsFailed(Model model, String msg) throws IOException {
-        model.getModelMetadata().setState(ModelState.FAILED);
-        model.getModelMetadata().setError(msg);
+    private void updateModelStateAsFailed(String modelId, ModelMetadata modelMetadata, String msg) throws IOException, ExecutionException,
+        InterruptedException {
+        modelMetadata.setState(ModelState.FAILED);
+        modelMetadata.setError(msg);
+        Model model = new Model(modelMetadata, null, modelId);
         modelDao.update(model, new ActionListener<IndexResponse>() {
             @Override
             public void onResponse(IndexResponse indexResponse) {

src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java

@@ -39,6 +39,7 @@
 import org.opensearch.knn.index.util.KNNEngine;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
+import org.opensearch.knn.indices.ModelState;
 import org.opensearch.plugins.SearchPlugin;
 
 import java.io.IOException;
@@ -661,6 +662,7 @@ public void testDoToQuery_FromModel() {
         when(modelMetadata.getDimension()).thenReturn(4);
         when(modelMetadata.getKnnEngine()).thenReturn(KNNEngine.FAISS);
         when(modelMetadata.getSpaceType()).thenReturn(SpaceType.COSINESIMIL);
+        when(modelMetadata.getState()).thenReturn(ModelState.CREATED);
         ModelDao modelDao = mock(ModelDao.class);
         when(modelDao.getMetadata(modelId)).thenReturn(modelMetadata);
         KNNQueryBuilder.initialize(modelDao);
@@ -690,6 +692,7 @@ public void testDoToQuery_whenFromModel_whenDoRadiusSearch_whenDistanceThreshold
         when(modelMetadata.getDimension()).thenReturn(4);
         when(modelMetadata.getKnnEngine()).thenReturn(KNNEngine.FAISS);
         when(modelMetadata.getSpaceType()).thenReturn(SpaceType.L2);
+        when(modelMetadata.getState()).thenReturn(ModelState.CREATED);
         ModelDao modelDao = mock(ModelDao.class);
         when(modelDao.getMetadata(modelId)).thenReturn(modelMetadata);
         KNNQueryBuilder.initialize(modelDao);
@@ -722,6 +725,7 @@ public void testDoToQuery_whenFromModel_whenDoRadiusSearch_whenScoreThreshold_th
         when(modelMetadata.getDimension()).thenReturn(4);
         when(modelMetadata.getKnnEngine()).thenReturn(KNNEngine.FAISS);
         when(modelMetadata.getSpaceType()).thenReturn(SpaceType.L2);
+        when(modelMetadata.getState()).thenReturn(ModelState.CREATED);
         ModelDao modelDao = mock(ModelDao.class);
         when(modelDao.getMetadata(modelId)).thenReturn(modelMetadata);
         KNNQueryBuilder.initialize(modelDao);

src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java

@@ -45,6 +45,7 @@
 import org.opensearch.knn.index.util.KNNEngine;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
+import org.opensearch.knn.indices.ModelState;
 import org.opensearch.knn.jni.JNIService;
 
 import java.io.IOException;
@@ -167,6 +168,7 @@ public void testQueryScoreForFaissWithModel() {
         ModelMetadata modelMetadata = mock(ModelMetadata.class);
         when(modelMetadata.getKnnEngine()).thenReturn(KNNEngine.FAISS);
         when(modelMetadata.getSpaceType()).thenReturn(spaceType);
+        when(modelMetadata.getState()).thenReturn(ModelState.CREATED);
         when(modelDao.getMetadata(eq("modelId"))).thenReturn(modelMetadata);
 
         KNNWeight.initialize(modelDao);

src/test/java/org/opensearch/knn/plugin/transport/TrainingModelRequestTests.java

@@ -661,6 +661,7 @@ public void testValidation_valid_trainingIndexBuiltFromModel() {
         // Mock the model dao to return metadata for modelId to recognize it is a duplicate
         ModelMetadata trainingFieldModelMetadata = mock(ModelMetadata.class);
         when(trainingFieldModelMetadata.getDimension()).thenReturn(dimension);
+        when(trainingFieldModelMetadata.getState()).thenReturn(ModelState.CREATED);
 
         ModelDao modelDao = mock(ModelDao.class);
         when(modelDao.getMetadata(modelId)).thenReturn(null);

src/test/java/org/opensearch/knn/training/TrainingJobClusterStateListenerTests.java

@@ -101,6 +101,7 @@ public void testUpdateModelsNewCluster() throws IOException, InterruptedExceptio
         ModelDao modelDao = mock(ModelDao.class);
         when(modelDao.isCreated()).thenReturn(true);
         when(modelDao.get(modelId)).thenReturn(model);
+        when(modelDao.getMetadata(modelId)).thenReturn(modelMetadata);
         doAnswer(invocationOnMock -> {
             SearchResponse searchResponse = mock(SearchResponse.class);
             SearchHits searchHits = mock(SearchHits.class);
@@ -144,6 +145,7 @@ public void testUpdateModelsNodesRemoved() throws IOException, InterruptedExcept
         ModelDao modelDao = mock(ModelDao.class);
         when(modelDao.isCreated()).thenReturn(true);
         when(modelDao.get(modelId)).thenReturn(model);
+        when(modelDao.getMetadata(modelId)).thenReturn(modelMetadata);
         DiscoveryNode node1 = mock(DiscoveryNode.class);
         when(node1.getEphemeralId()).thenReturn("test-node-model-match");
         DiscoveryNode node2 = mock(DiscoveryNode.class);