Properly designate model state for actively training models when nodes crash or leave cluster

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.11...2.x)
 ### Features
 * Add parent join support for lucene knn [#1182](https://github.com/opensearch-project/k-NN/pull/1182)
+* Properly designate model state for actively training models when nodes crash or leave cluster [#1317](https://github.com/opensearch-project/k-NN/pull/1317)
 ### Enhancements
 ### Bug Fixes
 * Fix use-after-free case on nmslib search path [#1305](https://github.com/opensearch-project/k-NN/pull/1305)

qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ModelIT.java

@@ -257,6 +257,6 @@ public String modelIndexMapping(String fieldName, String modelId) throws IOExcep
     }
 
     private ModelMetadata getModelMetadata() {
-        return new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, 4, ModelState.CREATED, "2021-03-27", "test model", "");
+        return new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, 4, ModelState.CREATED, "2021-03-27", "test model", "", "");
     }
 }

src/main/java/org/opensearch/knn/common/KNNConstants.java

@@ -45,6 +45,7 @@ public class KNNConstants {
     public static final String MODEL_TIMESTAMP = "timestamp";
     public static final String MODEL_DESCRIPTION = "description";
     public static final String MODEL_ERROR = "error";
+    public static final String MODEL_NODE_ASSIGNMENT = "node_assignment";
     public static final String PARAM_SIZE = "size";
     public static final Integer SEARCH_MODEL_MIN_SIZE = 1;
     public static final Integer SEARCH_MODEL_MAX_SIZE = 1000;

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

@@ -35,9 +35,11 @@
 public class IndexUtil {
 
     private static final Version MINIMAL_SUPPORTED_VERSION_FOR_IGNORE_UNMAPPED = Version.V_2_11_0;
+    private static final Version MINIMAL_SUPPORTED_VERSION_FOR_MODEL_NODE_ASSIGNMENT = Version.V_2_12_0;
     private static final Map<String, Version> minimalRequiredVersionMap = new HashMap<String, Version>() {
         {
             put("ignore_unmapped", MINIMAL_SUPPORTED_VERSION_FOR_IGNORE_UNMAPPED);
+            put("model_node_assignment", MINIMAL_SUPPORTED_VERSION_FOR_MODEL_NODE_ASSIGNMENT);
         }
     };
 

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

@@ -287,6 +287,7 @@ private void putInternal(Model model, ActionListener<IndexResponse> listener, Do
                     put(KNNConstants.MODEL_TIMESTAMP, modelMetadata.getTimestamp());
                     put(KNNConstants.MODEL_DESCRIPTION, modelMetadata.getDescription());
                     put(KNNConstants.MODEL_ERROR, modelMetadata.getError());
+                    put(KNNConstants.MODEL_NODE_ASSIGNMENT, modelMetadata.getNodeAssignment());
                 }
             };
 

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

@@ -19,6 +19,7 @@
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.knn.common.KNNConstants;
+import org.opensearch.knn.index.IndexUtil;
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.util.KNNEngine;
 
@@ -34,6 +35,7 @@
 import static org.opensearch.knn.common.KNNConstants.MODEL_ERROR;
 import static org.opensearch.knn.common.KNNConstants.MODEL_STATE;
 import static org.opensearch.knn.common.KNNConstants.MODEL_TIMESTAMP;
+import static org.opensearch.knn.common.KNNConstants.MODEL_NODE_ASSIGNMENT;
 
 public class ModelMetadata implements Writeable, ToXContentObject {
 
@@ -47,6 +49,7 @@ public class ModelMetadata implements Writeable, ToXContentObject {
     final private String timestamp;
     final private String description;
     private String error;
+    private String nodeAssignment;
 
     /**
      * Constructor
@@ -64,6 +67,7 @@ public ModelMetadata(StreamInput in) throws IOException {
         // which is checked in constructor and setters
         this.description = in.readString();
         this.error = in.readString();
+        this.nodeAssignment = in.readOptionalString();
     }
 
     /**
@@ -84,7 +88,8 @@ public ModelMetadata(
         ModelState modelState,
         String timestamp,
         String description,
-        String error
+        String error,
+        String nodeAssignment
     ) {
         this.knnEngine = Objects.requireNonNull(knnEngine, "knnEngine must not be null");
         this.spaceType = Objects.requireNonNull(spaceType, "spaceType must not be null");
@@ -104,6 +109,7 @@ public ModelMetadata(
         this.timestamp = Objects.requireNonNull(timestamp, "timestamp must not be null");
         this.description = Objects.requireNonNull(description, "description must not be null");
         this.error = Objects.requireNonNull(error, "error must not be null");
+        this.nodeAssignment = nodeAssignment;
     }
 
     /**
@@ -169,6 +175,19 @@ public String getError() {
         return error;
     }
 
+    /**
+     * getter for model's node assignment
+     *
+     * @return nodeAssignment
+     */
+    public String getNodeAssignment() {
+        return nodeAssignment;
+    }
+
+    public void setNodeAssignment(String nodeAssignment) {
+        this.nodeAssignment = nodeAssignment;
+    }
+
     /**
      * setter for model's state
      *
@@ -197,7 +216,8 @@ public String toString() {
             getState().toString(),
             timestamp,
             description,
-            error
+            error,
+            nodeAssignment
         );
     }
 
@@ -240,10 +260,10 @@ public int hashCode() {
     public static ModelMetadata fromString(String modelMetadataString) {
         String[] modelMetadataArray = modelMetadataString.split(DELIMITER, -1);
 
-        if (modelMetadataArray.length != 7) {
+        if (modelMetadataArray.length != 8 && modelMetadataArray.length != 7) {
             throw new IllegalArgumentException(
                 "Illegal format for model metadata. Must be of the form "
-                    + "\"<KNNEngine>,<SpaceType>,<Dimension>,<ModelState>,<Timestamp>,<Description>,<Error>\"."
+                    + "\"<KNNEngine>,<SpaceType>,<Dimension>,<ModelState>,<Timestamp>,<Description>,<Error>,<NodeAssignment>\"."
             );
         }
 
@@ -254,8 +274,11 @@ public static ModelMetadata fromString(String modelMetadataString) {
         String timestamp = modelMetadataArray[4];
         String description = modelMetadataArray[5];
         String error = modelMetadataArray[6];
-
-        return new ModelMetadata(knnEngine, spaceType, dimension, modelState, timestamp, description, error);
+        if (modelMetadataArray.length == 8) {
+            String nodeAssignment = modelMetadataArray[7];
+            return new ModelMetadata(knnEngine, spaceType, dimension, modelState, timestamp, description, error, nodeAssignment);
+        }
+        return new ModelMetadata(knnEngine, spaceType, dimension, modelState, timestamp, description, error, "");
     }
 
     private static String objectToString(Object value) {
@@ -278,10 +301,11 @@ public static ModelMetadata getMetadataFromSourceMap(final Map<String, Object> m
         Object engine = modelSourceMap.get(KNNConstants.KNN_ENGINE);
         Object space = modelSourceMap.get(KNNConstants.METHOD_PARAMETER_SPACE_TYPE);
         Object dimension = modelSourceMap.get(KNNConstants.DIMENSION);
-        Object state = modelSourceMap.get(KNNConstants.MODEL_STATE);
-        Object timestamp = modelSourceMap.get(KNNConstants.MODEL_TIMESTAMP);
+        Object state = modelSourceMap.get(MODEL_STATE);
+        Object timestamp = modelSourceMap.get(MODEL_TIMESTAMP);
         Object description = modelSourceMap.get(KNNConstants.MODEL_DESCRIPTION);
         Object error = modelSourceMap.get(KNNConstants.MODEL_ERROR);
+        Object nodeAssignment = modelSourceMap.get(KNNConstants.MODEL_NODE_ASSIGNMENT);
 
         ModelMetadata modelMetadata = new ModelMetadata(
             KNNEngine.getEngine(objectToString(engine)),
@@ -290,7 +314,8 @@ public static ModelMetadata getMetadataFromSourceMap(final Map<String, Object> m
             ModelState.getModelState(objectToString(state)),
             objectToString(timestamp),
             objectToString(description),
-            objectToString(error)
+            objectToString(error),
+            objectToString(nodeAssignment)
         );
         return modelMetadata;
     }
@@ -304,6 +329,9 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeString(getTimestamp());
         out.writeString(getDescription());
         out.writeString(getError());
+        if (IndexUtil.isClusterOnOrAfterMinRequiredVersion("model_node_assignment")) {
+            out.writeString(getNodeAssignment());
+        }
     }
 
     @Override
@@ -316,6 +344,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(METHOD_PARAMETER_SPACE_TYPE, getSpaceType().getValue());
         builder.field(DIMENSION, getDimension());
         builder.field(KNN_ENGINE, getKnnEngine().getName());
+        if (IndexUtil.isClusterOnOrAfterMinRequiredVersion("model_node_assignment")) {
+            builder.field(MODEL_NODE_ASSIGNMENT, getNodeAssignment());
+        }
         return builder;
     }
 }

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

@@ -23,7 +23,8 @@
 public enum ModelState implements Writeable {
     TRAINING("training"),
     CREATED("created"),
-    FAILED("failed");
+    FAILED("failed"),
+    ZOMBIE("zombie");
 
     private final String name;
 
@@ -79,6 +80,10 @@ public static ModelState getModelState(String name) {
             return FAILED;
         }
 
+        if (ZOMBIE.getName().equals(name)) {
+            return ZOMBIE;
+        }
+
         throw new IllegalArgumentException("Unable to find model state: \"" + name + "\"");
     }
 }

src/main/java/org/opensearch/knn/plugin/KNNPlugin.java

@@ -199,7 +199,7 @@ public Collection<Object> createComponents(
         KNNClusterUtil.instance().initialize(clusterService);
         ModelDao.OpenSearchKNNModelDao.initialize(client, clusterService, environment.settings());
         ModelCache.initialize(ModelDao.OpenSearchKNNModelDao.getInstance(), clusterService);
-        TrainingJobRunner.initialize(threadPool, ModelDao.OpenSearchKNNModelDao.getInstance());
+        TrainingJobRunner.initialize(threadPool, ModelDao.OpenSearchKNNModelDao.getInstance(), clusterService);
         KNNCircuitBreaker.getInstance().initialize(threadPool, clusterService, client);
         KNNQueryBuilder.initialize(ModelDao.OpenSearchKNNModelDao.getInstance());
         KNNWeight.initialize(ModelDao.OpenSearchKNNModelDao.getInstance());

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

@@ -81,6 +81,7 @@ public TrainingJob(
                 ModelState.TRAINING,
                 ZonedDateTime.now(ZoneOffset.UTC).toString(),
                 description,
+                "",
                 ""
             ),
             null,