Refine the delete model code

Signed-off-by: zane-neo <[email protected]>
zane-neo · Jan 17, 2024 · 3725de3 · 3725de3
1 parent 4814901
commit 3725de3
Showing 1 changed file with 12 additions and 18 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
@@ -15,7 +15,6 @@
 import static org.opensearch.ml.utils.RestActionUtils.getFetchSourceContext;
 
 import org.opensearch.OpenSearchStatusException;
-import org.opensearch.ResourceNotFoundException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.DocWriteResponse;
 import org.opensearch.action.delete.DeleteRequest;
@@ -39,6 +38,7 @@
 import org.opensearch.index.reindex.BulkByScrollResponse;
 import org.opensearch.index.reindex.DeleteByQueryAction;
 import org.opensearch.index.reindex.DeleteByQueryRequest;
+import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLModel;
 import org.opensearch.ml.common.model.MLModelState;
 import org.opensearch.ml.common.transport.model.MLModelDeleteAction;
@@ -113,6 +113,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                         MLModel mlModel = MLModel.parse(parser, algorithmName);
                         Boolean isHidden = (Boolean) r.getSource().get(IS_HIDDEN_FIELD);
                         MLModelState mlModelState = mlModel.getModelState();
+                        FunctionName functionName = mlModel.getAlgorithm();
                         if (isHidden != null && isHidden) {
                             if (!isSuperAdmin) {
                                 wrappedListener
@@ -124,7 +125,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                                     );
                             } else {
                                 if (isModelNotDeployed(mlModelState)) {
-                                    deleteModel(modelId, actionListener);
+                                    deleteModel(modelId, functionName, actionListener);
                                 } else {
                                     wrappedListener
                                         .onFailure(
@@ -146,7 +147,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                                                 )
                                             );
                                     } else if (isModelNotDeployed(mlModelState)) {
-                                        deleteModel(modelId, actionListener);
+                                        deleteModel(modelId, functionName, actionListener);
                                     } else {
                                         wrappedListener
                                             .onFailure(
@@ -167,7 +168,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                 } else {
                     wrappedListener.onFailure(new OpenSearchStatusException("Failed to find model", RestStatus.NOT_FOUND));
                 }
-            }, e -> { wrappedListener.onFailure(e); }));
+            }, wrappedListener::onFailure));
         } catch (Exception e) {
             log.error("Failed to delete ML model " + modelId, e);
             actionListener.onFailure(e);
@@ -183,13 +184,7 @@ void deleteModelChunks(String modelId, DeleteResponse deleteResponse, ActionList
             if ((r.getBulkFailures() == null || r.getBulkFailures().size() == 0)
                 && (r.getSearchFailures() == null || r.getSearchFailures().size() == 0)) {
                 log.debug("All model chunks are deleted for model {}", modelId);
-                if (deleteResponse != null) {
-                    // If model metaData not found and deleteResponse is null, do not return here.
-                    // ResourceNotFound is returned to notify that this model was deleted.
-                    // This is a walk around to avoid cleaning up model leftovers. Will revisit if
-                    // necessary.
-                    actionListener.onResponse(deleteResponse);
-                }
+                actionListener.onResponse(deleteResponse);
             } else {
                 returnFailure(r, modelId, actionListener);
             }
@@ -212,23 +207,22 @@ private void returnFailure(BulkByScrollResponse response, String modelId, Action
         actionListener.onFailure(new OpenSearchStatusException(errorMessage, RestStatus.INTERNAL_SERVER_ERROR));
     }
 
-    private void deleteModel(String modelId, ActionListener<DeleteResponse> actionListener) {
+    private void deleteModel(String modelId, FunctionName functionName, ActionListener<DeleteResponse> actionListener) {
         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);
+                if (FunctionName.REMOTE != functionName) {
+                    deleteModelChunks(modelId, deleteResponse, actionListener);
+                } else {
+                    actionListener.onResponse(deleteResponse);
+                }
                 deleteModelController(modelId);
             }
 
             @Override
             public void onFailure(Exception e) {
                 log.error("Failed to delete model meta data for model: " + modelId, 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);
             }
         });