Exit early when no nodes service the model

Now when entering this method its guaranteed to write to index first before sending back the MLUndeploy response. And will also send back a exception if the write back fails

Signed-off-by: Brian Flores <[email protected]>

Author: brianf-aws

plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelsAction.java

@@ -40,6 +40,7 @@
 import org.opensearch.ml.common.transport.deploy.MLDeployModelRequest;
 import org.opensearch.ml.common.transport.undeploy.MLUndeployModelAction;
 import org.opensearch.ml.common.transport.undeploy.MLUndeployModelNodesRequest;
+import org.opensearch.ml.common.transport.undeploy.MLUndeployModelNodesResponse;
 import org.opensearch.ml.common.transport.undeploy.MLUndeployModelsAction;
 import org.opensearch.ml.common.transport.undeploy.MLUndeployModelsRequest;
 import org.opensearch.ml.common.transport.undeploy.MLUndeployModelsResponse;
@@ -162,15 +163,20 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUnde
     private void undeployModels(String[] targetNodeIds, String[] modelIds, ActionListener<MLUndeployModelsResponse> listener) {
         MLUndeployModelNodesRequest mlUndeployModelNodesRequest = new MLUndeployModelNodesRequest(targetNodeIds, modelIds);
 
-        client.execute(MLUndeployModelAction.INSTANCE, mlUndeployModelNodesRequest, ActionListener.wrap(r -> {
-            if (r.getNodes().isEmpty()) {
-                bulkSetModelIndexToUndeploy(modelIds);
+        client.execute(MLUndeployModelAction.INSTANCE, mlUndeployModelNodesRequest, ActionListener.wrap(response -> {
+            if (response.getNodes().isEmpty()) {
+                bulkSetModelIndexToUndeploy(modelIds, listener, response);
+                return;
             }
-            listener.onResponse(new MLUndeployModelsResponse(r));
+            listener.onResponse(new MLUndeployModelsResponse(response));
         }, listener::onFailure));
     }
 
-    private void bulkSetModelIndexToUndeploy(String[] modelIds) {
+    private void bulkSetModelIndexToUndeploy(
+        String[] modelIds,
+        ActionListener<MLUndeployModelsResponse> listener,
+        MLUndeployModelNodesResponse response
+    ) {
         BulkRequest bulkUpdateRequest = new BulkRequest();
         for (String modelId : modelIds) {
             UpdateRequest updateRequest = new UpdateRequest();
@@ -186,10 +192,16 @@ private void bulkSetModelIndexToUndeploy(String[] modelIds) {
             updateRequest.index(ML_MODEL_INDEX).id(modelId).doc(builder.build());
             bulkUpdateRequest.add(updateRequest);
         }
+
         bulkUpdateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
-        log.info("No models service: {}", modelIds.toString());
-        client.bulk(bulkUpdateRequest, ActionListener.wrap(br -> { log.debug("Successfully set modelIds to UNDEPLOY in index"); }, e -> {
+        log.info("No nodes service: {}", modelIds.toString());
+
+        client.bulk(bulkUpdateRequest, ActionListener.wrap(br -> {
+            log.debug("Successfully set modelIds to UNDEPLOY in index");
+            listener.onResponse(new MLUndeployModelsResponse(response));
+        }, e -> {
             log.error("Failed to set modelIds to UNDEPLOY in index", e);
+            listener.onFailure(e);
         }));
     }