Add model auto redeploy feature (#852) (#859)

* Add model auto reload feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Add model auto redeploy feature



* Addressed comments and reverted the getEligibleNodes logic



* Change method regarding clear model retry times in model index and addressed naming convention comments



* Fix import * issue



* Rebase upstream code and fix UT failure issue



---------


(cherry picked from commit 1e6b6c4)

Signed-off-by: Zan Niu <[email protected]>
Signed-off-by: Yaliang Wu <[email protected]>
Co-authored-by: zane-neo <[email protected]>

Author: opensearch-trigger-bot[bot]

common/src/main/java/org/opensearch/ml/common/CommonValue.java

@@ -110,6 +110,9 @@ public class CommonValue {
             + MLModel.MODEL_CONTENT_HASH_VALUE_FIELD
             + "\" : {\"type\": \"keyword\"},\n"
             + "      \""
+            + MLModel.AUTO_REDEPLOY_RETRY_TIMES_FIELD
+            + "\" : {\"type\": \"integer\"},\n"
+            + "      \""
             + MLModel.CREATED_TIME_FIELD
             + "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
             + "      \""

common/src/main/java/org/opensearch/ml/common/MLModel.java

@@ -46,6 +46,7 @@ public class MLModel implements ToXContentObject {
     public static final String MODEL_CONTENT_SIZE_IN_BYTES_FIELD = "model_content_size_in_bytes";
     //SHA256 hash value of model content.
     public static final String MODEL_CONTENT_HASH_VALUE_FIELD = "model_content_hash_value";
+
     public static final String MODEL_CONFIG_FIELD = "model_config";
     public static final String CREATED_TIME_FIELD = "created_time";
     public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
@@ -60,6 +61,8 @@ public class MLModel implements ToXContentObject {
     public static final String LAST_UNDEPLOYED_TIME_FIELD = "last_undeployed_time";
 
     public static final String MODEL_ID_FIELD = "model_id";
+    // auto redploy retry times for this model.
+    public static final String AUTO_REDEPLOY_RETRY_TIMES_FIELD = "auto_redeploy_retry_times";
     public static final String CHUNK_NUMBER_FIELD = "chunk_number";
     public static final String TOTAL_CHUNKS_FIELD = "total_chunks";
     public static final String PLANNING_WORKER_NODE_COUNT_FIELD = "planning_worker_node_count";
@@ -88,6 +91,8 @@ public class MLModel implements ToXContentObject {
 
     @Setter
     private String modelId; // model chunk doc only
+
+    private Integer autoRedeployRetryTimes;
     private Integer chunkNumber; // model chunk doc only
     private Integer totalChunks; // model chunk doc only
     private Integer planningWorkerNodeCount; // plan to deploy model to how many nodes
@@ -112,6 +117,7 @@ public MLModel(String name,
                    Instant lastRegisteredTime,
                    Instant lastDeployedTime,
                    Instant lastUndeployedTime,
+                   Integer autoRedeployRetryTimes,
                    String modelId, Integer chunkNumber,
                    Integer totalChunks,
                    Integer planningWorkerNodeCount,
@@ -135,6 +141,7 @@ public MLModel(String name,
         this.lastDeployedTime = lastDeployedTime;
         this.lastUndeployedTime = lastUndeployedTime;
         this.modelId = modelId;
+        this.autoRedeployRetryTimes = autoRedeployRetryTimes;
         this.chunkNumber = chunkNumber;
         this.totalChunks = totalChunks;
         this.planningWorkerNodeCount = planningWorkerNodeCount;
@@ -172,6 +179,7 @@ public MLModel(StreamInput input) throws IOException{
             lastDeployedTime = input.readOptionalInstant();
             lastUndeployedTime = input.readOptionalInstant();
             modelId = input.readOptionalString();
+            autoRedeployRetryTimes = input.readOptionalInt();
             chunkNumber = input.readOptionalInt();
             totalChunks = input.readOptionalInt();
             planningWorkerNodeCount = input.readOptionalInt();
@@ -219,6 +227,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalInstant(lastDeployedTime);
         out.writeOptionalInstant(lastUndeployedTime);
         out.writeOptionalString(modelId);
+        out.writeOptionalInt(autoRedeployRetryTimes);
         out.writeOptionalInt(chunkNumber);
         out.writeOptionalInt(totalChunks);
         out.writeOptionalInt(planningWorkerNodeCount);
@@ -281,6 +290,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
         if (modelId != null) {
             builder.field(MODEL_ID_FIELD, modelId);
         }
+        if (autoRedeployRetryTimes != null) {
+            builder.field(AUTO_REDEPLOY_RETRY_TIMES_FIELD, autoRedeployRetryTimes);
+        }
         if (chunkNumber != null) {
             builder.field(CHUNK_NUMBER_FIELD, chunkNumber);
         }
@@ -327,6 +339,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
         Instant lastDeployedTime = null;
         Instant lastUndeployedTime = null;
         String modelId = null;
+        Integer autoRedeployRetryTimes = null;
         Integer chunkNumber = null;
         Integer totalChunks = null;
         Integer planningWorkerNodeCount = null;
@@ -370,6 +383,9 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
                 case MODEL_ID_FIELD:
                     modelId = parser.text();
                     break;
+                case AUTO_REDEPLOY_RETRY_TIMES_FIELD:
+                    autoRedeployRetryTimes = parser.intValue();
+                    break;
                 case DESCRIPTION_FIELD:
                     description = parser.text();
                     break;
@@ -454,6 +470,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
                 .lastDeployedTime(lastDeployedTime == null? lastLoadedTime : lastDeployedTime)
                 .lastUndeployedTime(lastUndeployedTime == null? lastUnloadedTime : lastUndeployedTime)
                 .modelId(modelId)
+                .autoRedeployRetryTimes(autoRedeployRetryTimes)
                 .chunkNumber(chunkNumber)
                 .totalChunks(totalChunks)
                 .planningWorkerNodeCount(planningWorkerNodeCount)

common/src/main/java/org/opensearch/ml/common/transport/deploy/MLDeployModelInput.java

@@ -21,6 +21,7 @@ public class MLDeployModelInput implements Writeable {
     private String modelContentHash;
     private Integer nodeCount;
     private String coordinatingNodeId;
+    private Boolean isDeployToAllNodes;
     private MLTask mlTask;
 
     public MLDeployModelInput(StreamInput in) throws IOException {
@@ -29,16 +30,18 @@ public MLDeployModelInput(StreamInput in) throws IOException {
         this.modelContentHash = in.readOptionalString();
         this.nodeCount = in.readInt();
         this.coordinatingNodeId = in.readString();
+        this.isDeployToAllNodes = in.readOptionalBoolean();
         this.mlTask = new MLTask(in);
     }
 
     @Builder
-    public MLDeployModelInput(String modelId, String taskId, String modelContentHash, Integer nodeCount, String coordinatingNodeId, MLTask mlTask) {
+    public MLDeployModelInput(String modelId, String taskId, String modelContentHash, Integer nodeCount, String coordinatingNodeId, Boolean isDeployToAllNodes, MLTask mlTask) {
         this.modelId = modelId;
         this.taskId = taskId;
         this.modelContentHash = modelContentHash;
         this.nodeCount = nodeCount;
         this.coordinatingNodeId = coordinatingNodeId;
+        this.isDeployToAllNodes = isDeployToAllNodes;
         this.mlTask = mlTask;
     }
 
@@ -52,6 +55,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(modelContentHash);
         out.writeInt(nodeCount);
         out.writeString(coordinatingNodeId);
+        out.writeOptionalBoolean(isDeployToAllNodes);
         mlTask.writeTo(out);
     }
 

common/src/main/java/org/opensearch/ml/common/transport/sync/MLSyncUpInput.java

@@ -31,19 +31,25 @@ public class MLSyncUpInput implements Writeable {
     // sync running deploy model tasks
     private boolean syncRunningDeployModelTasks;
 
+    // deployToAll flag for models, when deploy/undeploy a model, this will passed to each node to update cache value to make sure
+    // profile API has consistent data with model index.
+    private Map<String, Boolean> deployToAllNodes;
+
     @Builder
     public MLSyncUpInput(boolean getDeployedModels,
                          Map<String, String[]> addedWorkerNodes,
                          Map<String, String[]> removedWorkerNodes,
                          Map<String, Set<String>> modelRoutingTable,
                          Map<String, Set<String>> runningDeployModelTasks,
+                         Map<String, Boolean> deployToAllNodes,
                          boolean clearRoutingTable,
                          boolean syncRunningDeployModelTasks) {
         this.getDeployedModels = getDeployedModels;
         this.addedWorkerNodes = addedWorkerNodes;
         this.removedWorkerNodes = removedWorkerNodes;
         this.modelRoutingTable = modelRoutingTable;
         this.runningDeployModelTasks = runningDeployModelTasks;
+        this.deployToAllNodes = deployToAllNodes;
         this.clearRoutingTable = clearRoutingTable;
         this.syncRunningDeployModelTasks = syncRunningDeployModelTasks;
     }
@@ -64,6 +70,9 @@ public MLSyncUpInput(StreamInput in) throws IOException {
         if (in.readBoolean()) {
             runningDeployModelTasks = in.readMap(StreamInput::readString, s -> s.readSet(StreamInput::readString));
         }
+        if (in.readBoolean()) {
+            deployToAllNodes = in.readMap(StreamInput::readString, StreamInput::readOptionalBoolean);
+        }
         this.clearRoutingTable = in.readBoolean();
         this.syncRunningDeployModelTasks = in.readBoolean();
     }
@@ -95,6 +104,12 @@ public void writeTo(StreamOutput out) throws IOException {
         } else {
             out.writeBoolean(false);
         }
+        if (deployToAllNodes != null && deployToAllNodes.size() > 0) {
+            out.writeBoolean(true);
+            out.writeMap(deployToAllNodes, StreamOutput::writeString, StreamOutput::writeOptionalBoolean);
+        } else {
+            out.writeBoolean(false);
+        }
         out.writeBoolean(clearRoutingTable);
         out.writeBoolean(syncRunningDeployModelTasks);
     }

common/src/main/java/org/opensearch/ml/common/transport/undeploy/MLUndeployModelNodeResponse.java

@@ -20,15 +20,20 @@
 @Getter
 public class MLUndeployModelNodeResponse extends BaseNodeResponse implements ToXContentFragment {
 
+    // Model undeploy status, if there's model running and successfully undeployed, the status is undeployed, if model is not
+    // running on current node, status is not_found
     private Map<String, String> modelUndeployStatus;
-    private Map<String, Integer> modelWorkerNodeCounts;
+    // This is to record before undeploy the model, which nodes are working nodes.
+    private Map<String, String[]> modelWorkerNodeBeforeRemoval;
 
     public MLUndeployModelNodeResponse(DiscoveryNode node,
                                        Map<String, String> modelUndeployStatus,
-                                       Map<String, Integer> modelWorkerNodeCounts) {
+                                       Map<String, String[]> modelWorkerNodeBeforeRemoval
+    ) {
         super(node);
         this.modelUndeployStatus = modelUndeployStatus;
-        this.modelWorkerNodeCounts = modelWorkerNodeCounts;
+        //
+        this.modelWorkerNodeBeforeRemoval = modelWorkerNodeBeforeRemoval;
     }
 
     public MLUndeployModelNodeResponse(StreamInput in) throws IOException {
@@ -37,7 +42,7 @@ public MLUndeployModelNodeResponse(StreamInput in) throws IOException {
             this.modelUndeployStatus = in.readMap(s -> s.readString(), s-> s.readString());
         }
         if (in.readBoolean()) {
-            this.modelWorkerNodeCounts = in.readMap(s -> s.readString(), s-> s.readInt());
+            this.modelWorkerNodeBeforeRemoval = in.readMap(s -> s.readString(), s-> s.readOptionalStringArray());
         }
     }
 
@@ -55,9 +60,9 @@ public void writeTo(StreamOutput out) throws IOException {
         } else {
             out.writeBoolean(false);
         }
-        if (modelWorkerNodeCounts != null) {
+        if (modelWorkerNodeBeforeRemoval != null) {
             out.writeBoolean(true);
-            out.writeMap(modelWorkerNodeCounts, StreamOutput::writeString, StreamOutput::writeInt);
+            out.writeMap(modelWorkerNodeBeforeRemoval, StreamOutput::writeString, StreamOutput::writeOptionalStringArray);
         } else {
             out.writeBoolean(false);
         }

common/src/test/java/org/opensearch/ml/common/transport/deploy/MLDeployModelNodesRequestTest.java

@@ -84,7 +84,7 @@ public void setUp() throws Exception {
     @Test
     public void testConstructorSerialization1() throws IOException {
         String [] nodeIds = {"id1", "id2", "id3"};
-        MLDeployModelInput deployModelInput = new MLDeployModelInput("modelId", "taskId", "modelContentHash", 3, "coordinatingNodeId", mlTask);
+        MLDeployModelInput deployModelInput = new MLDeployModelInput("modelId", "taskId", "modelContentHash", 3, "coordinatingNodeId", true, mlTask);
         MLDeployModelNodeRequest MLDeployModelNodeRequest = new MLDeployModelNodeRequest(
                 new MLDeployModelNodesRequest(nodeIds, deployModelInput)
         );
@@ -104,7 +104,7 @@ public void testConstructorSerialization1() throws IOException {
     @Test
     public void testConstructorSerialization2() throws IOException {
         DiscoveryNode [] nodeIds = {localNode1, localNode2, localNode3};
-        MLDeployModelInput deployModelInput = new MLDeployModelInput("modelId", "taskId", "modelContentHash", 3, "coordinatingNodeId", mlTask);
+        MLDeployModelInput deployModelInput = new MLDeployModelInput("modelId", "taskId", "modelContentHash", 3, "coordinatingNodeId", true, mlTask);
         MLDeployModelNodeRequest MLDeployModelNodeRequest = new MLDeployModelNodeRequest(
                 new MLDeployModelNodesRequest(nodeIds, deployModelInput)
         );
@@ -140,7 +140,7 @@ public void testConstructorSerialization3() throws IOException {
     @Test
     public void testConstructorFromInputStream() throws IOException {
         String [] nodeIds = {"id1", "id2", "id3"};
-        MLDeployModelInput deployModelInput = new MLDeployModelInput("modelId", "taskId", "modelContentHash", 3, "coordinatingNodeId", mlTask);
+        MLDeployModelInput deployModelInput = new MLDeployModelInput("modelId", "taskId", "modelContentHash", 3, "coordinatingNodeId", true, mlTask);
         MLDeployModelNodeRequest MLDeployModelNodeRequest = new MLDeployModelNodeRequest(
                 new MLDeployModelNodesRequest(nodeIds, deployModelInput)
         );

common/src/test/java/org/opensearch/ml/common/transport/undeploy/MLUndeployModelNodeResponseTest.java

@@ -25,7 +25,7 @@ public class MLUndeployModelNodeResponseTest {
     @Mock
     private DiscoveryNode localNode;
 
-    private Map<String, Integer> modelWorkerNodeCounts;
+    private Map<String, String[]> modelWorkerNodeCounts;
 
     @Before
     public void setUp() throws Exception {
@@ -38,7 +38,7 @@ public void setUp() throws Exception {
                 Version.CURRENT
         );
         modelWorkerNodeCounts = new HashMap<>();
-        modelWorkerNodeCounts.put("modelId1", 1);
+        modelWorkerNodeCounts.put("modelId1", new String[]{"node"});
     }
 
     @Test

common/src/test/java/org/opensearch/ml/common/transport/undeploy/MLUndeployModelNodesResponseTest.java

@@ -77,14 +77,14 @@ public void testToXContent() throws IOException {
 
         Map<String, String> modelToUndeployStatus1 = new HashMap<>();
         modelToUndeployStatus1.put("modelId1", "response");
-        Map<String, Integer> modelWorkerNodeCounts1 = new HashMap<>();
-        modelWorkerNodeCounts1.put("modelId1", 1);
+        Map<String, String[]> modelWorkerNodeCounts1 = new HashMap<>();
+        modelWorkerNodeCounts1.put("modelId1", new String[]{"mockNode1"});
         nodes.add(new MLUndeployModelNodeResponse(node1, modelToUndeployStatus1, modelWorkerNodeCounts1));
 
         Map<String, String> modelToUndeployStatus2 = new HashMap<>();
         modelToUndeployStatus2.put("modelId2", "response");
-        Map<String, Integer> modelWorkerNodeCounts2 = new HashMap<>();
-        modelWorkerNodeCounts2.put("modelId2", 2);
+        Map<String, String[]> modelWorkerNodeCounts2 = new HashMap<>();
+        modelWorkerNodeCounts2.put("modelId2", new String[]{"mockNode2"});
         nodes.add(new MLUndeployModelNodeResponse(node2, modelToUndeployStatus2, modelWorkerNodeCounts2));
 
         List<FailedNodeException> failures = new ArrayList<>();

plugin/build.gradle

@@ -277,7 +277,8 @@ List<String> jacocoExclusions = [
         'org.opensearch.ml.action.profile.MLProfileTransportAction',
         'org.opensearch.ml.action.models.DeleteModelTransportAction.1',
         'org.opensearch.ml.rest.RestMLPredictionAction',
-        'org.opensearch.ml.breaker.DiskCircuitBreaker'
+        'org.opensearch.ml.breaker.DiskCircuitBreaker',
+        'org.opensearch.ml.autoredeploy.MLModelAutoReDeployer.SearchRequestBuilderFactory'
 ]
 
 jacocoTestCoverageVerification {

plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelAction.java

@@ -248,6 +248,7 @@ void updateModelDeployStatusAndTriggerOnNodesAction(
             mlModel.getModelContentHash(),
             eligibleNodes.size(),
             localNodeId,
+            deployToAllNodes,
             mlTask
         );
         MLDeployModelNodesRequest deployModelRequest = new MLDeployModelNodesRequest(