updating bulk update to use sdkclient (#3546) (#3561) (#3562)

* updating bulk update to use sdkclient

* addressed comments

---------

Signed-off-by: Dhrubo Saha <[email protected]>
(cherry picked from commit d8670c5)

Co-authored-by: Dhrubo Saha <[email protected]>

Author: opensearch-trigger-bot[bot]

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

@@ -16,13 +16,13 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.FailedNodeException;
-import org.opensearch.action.bulk.BulkRequest;
+import org.opensearch.action.bulk.BulkItemResponse;
 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;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
@@ -44,6 +44,10 @@
 import org.opensearch.ml.model.MLModelManager;
 import org.opensearch.ml.stats.MLNodeLevelStat;
 import org.opensearch.ml.stats.MLStats;
+import org.opensearch.remote.metadata.client.BulkDataObjectRequest;
+import org.opensearch.remote.metadata.client.SdkClient;
+import org.opensearch.remote.metadata.client.UpdateDataObjectRequest;
+import org.opensearch.remote.metadata.common.SdkClientUtils;
 import org.opensearch.tasks.Task;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
@@ -58,6 +62,7 @@ public class TransportUndeployModelAction extends
     private final MLModelManager mlModelManager;
     private final ClusterService clusterService;
     private final Client client;
+    private final SdkClient sdkClient;
     private final DiscoveryNodeHelper nodeFilter;
     private final MLStats mlStats;
 
@@ -69,6 +74,7 @@ public TransportUndeployModelAction(
         ClusterService clusterService,
         ThreadPool threadPool,
         Client client,
+        SdkClient sdkClient,
         DiscoveryNodeHelper nodeFilter,
         MLStats mlStats
     ) {
@@ -87,19 +93,21 @@ public TransportUndeployModelAction(
 
         this.clusterService = clusterService;
         this.client = client;
+        this.sdkClient = sdkClient;
         this.nodeFilter = nodeFilter;
         this.mlStats = mlStats;
     }
 
     @Override
     protected void doExecute(Task task, MLUndeployModelNodesRequest request, ActionListener<MLUndeployModelNodesResponse> listener) {
         ActionListener<MLUndeployModelNodesResponse> wrappedListener = ActionListener.wrap(undeployModelNodesResponse -> {
-            processUndeployModelResponseAndUpdate(undeployModelNodesResponse, listener);
+            processUndeployModelResponseAndUpdate(request.getTenantId(), undeployModelNodesResponse, listener);
         }, listener::onFailure);
         super.doExecute(task, request, wrappedListener);
     }
 
     void processUndeployModelResponseAndUpdate(
+        String tenantId,
         MLUndeployModelNodesResponse undeployModelNodesResponse,
         ActionListener<MLUndeployModelNodesResponse> listener
     ) {
@@ -145,11 +153,10 @@ void processUndeployModelResponseAndUpdate(
 
         MLSyncUpNodesRequest syncUpRequest = new MLSyncUpNodesRequest(nodeFilter.getAllNodes(), syncUpInput);
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
-            if (actualRemovedNodesMap.size() > 0) {
-                BulkRequest bulkRequest = new BulkRequest();
+            if (!actualRemovedNodesMap.isEmpty()) {
+                BulkDataObjectRequest bulkRequest = BulkDataObjectRequest.builder().globalIndex(ML_MODEL_INDEX).build();
                 Map<String, Boolean> deployToAllNodes = new HashMap<>();
                 for (String modelId : actualRemovedNodesMap.keySet()) {
-                    UpdateRequest updateRequest = new UpdateRequest();
                     List<String> removedNodes = actualRemovedNodesMap.get(modelId);
                     int removedNodeCount = removedNodes.size();
                     /**
@@ -178,7 +185,13 @@ void processUndeployModelResponseAndUpdate(
                         updateDocument.put(MLModel.CURRENT_WORKER_NODE_COUNT_FIELD, newPlanningWorkerNodes.size());
                         deployToAllNodes.put(modelId, false);
                     }
-                    updateRequest.index(ML_MODEL_INDEX).id(modelId).doc(updateDocument);
+
+                    UpdateDataObjectRequest updateRequest = UpdateDataObjectRequest
+                        .builder()
+                        .id(modelId)
+                        .tenantId(tenantId)
+                        .dataObject(updateDocument)
+                        .build();
                     bulkRequest.add(updateRequest).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
                 }
                 syncUpInput.setDeployToAllNodes(deployToAllNodes);
@@ -189,10 +202,33 @@ void processUndeployModelResponseAndUpdate(
                             Arrays.toString(actualRemovedNodesMap.keySet().toArray(new String[0]))
                         );
                 }, e -> { log.error("Failed to update model state as undeployed", e); });
-                client.bulk(bulkRequest, ActionListener.runAfter(actionListener, () -> {
+                ActionListener<BulkResponse> wrappedListener = ActionListener.runAfter(actionListener, () -> {
                     syncUpUndeployedModels(syncUpRequest);
                     listener.onResponse(undeployModelNodesResponse);
-                }));
+                });
+                sdkClient.bulkDataObjectAsync(bulkRequest).whenComplete((r, throwable) -> {
+                    if (throwable != null) {
+                        Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable, OpenSearchStatusException.class);
+                        log.error("Failed to execute BulkDataObject request", cause);
+                        wrappedListener.onFailure(cause);
+                    } else {
+                        try {
+                            BulkResponse bulkResponse = BulkResponse.fromXContent(r.parser());
+                            log
+                                .info(
+                                    "Executed {} bulk operations with {} failures, Took: {}",
+                                    bulkResponse.getItems().length,
+                                    bulkResponse.hasFailures()
+                                        ? Arrays.stream(bulkResponse.getItems()).filter(BulkItemResponse::isFailed).count()
+                                        : 0,
+                                    bulkResponse.getTook()
+                                );
+                            wrappedListener.onResponse(bulkResponse);
+                        } catch (Exception e) {
+                            wrappedListener.onFailure(e);
+                        }
+                    }
+                });
             } else {
                 syncUpUndeployedModels(syncUpRequest);
                 listener.onResponse(undeployModelNodesResponse);