Fix long pending issue when deleting model

Signed-off-by: zane-neo <[email protected]>
zane-neo · Jan 17, 2024 · 4814901 · 4814901
1 parent 433d81b
commit 4814901
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java
@@ -24,6 +24,7 @@
 import org.opensearch.action.get.GetResponse;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.action.support.WriteRequest;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
@@ -212,8 +213,8 @@ private void returnFailure(BulkByScrollResponse response, String modelId, Action
     }
 
     private void deleteModel(String modelId, ActionListener<DeleteResponse> actionListener) {
-        DeleteRequest deleteRequest = new DeleteRequest(ML_MODEL_INDEX, modelId);
-        client.delete(deleteRequest, new ActionListener<DeleteResponse>() {
+        DeleteRequest deleteRequest = new DeleteRequest(ML_MODEL_INDEX, modelId).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+        client.delete(deleteRequest, new ActionListener<>() {
             @Override
             public void onResponse(DeleteResponse deleteResponse) {
                 deleteModelChunks(modelId, deleteResponse, actionListener);
@@ -226,6 +227,7 @@ public void onFailure(Exception e) {
                 if (e instanceof ResourceNotFoundException) {
                     deleteModelChunks(modelId, null, actionListener);
                     deleteModelController(modelId);
+                    actionListener.onFailure(new OpenSearchStatusException("MODEL_NOT_FOUND" + modelId, RestStatus.NOT_FOUND));
                 }
                 actionListener.onFailure(e);
             }

diff --git a/plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelAction.java b/plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelAction.java
@@ -20,6 +20,7 @@
 import org.opensearch.action.bulk.BulkRequest;
 import org.opensearch.action.bulk.BulkResponse;
 import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.WriteRequest;
 import org.opensearch.action.support.nodes.TransportNodesAction;
 import org.opensearch.action.update.UpdateRequest;
 import org.opensearch.client.Client;
@@ -174,7 +175,7 @@ protected MLUndeployModelNodesResponse newResponse(
                             deployToAllNodes.put(modelId, false);
                         }
                         updateRequest.index(ML_MODEL_INDEX).id(modelId).doc(updateDocument);
-                        bulkRequest.add(updateRequest);
+                        bulkRequest.add(updateRequest).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
                     }
                     syncUpInput.setDeployToAllNodes(deployToAllNodes);
                     ActionListener<BulkResponse> actionListener = ActionListener.wrap(r -> {