Enhance batch job task management by adding default action types (#3080) (#3085)

* enhance batch job task management by adding default action types

Signed-off-by: Bhavana Ramaram <[email protected]>
(cherry picked from commit ff6fe67)

Author: rbhavna

common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java

@@ -125,6 +125,11 @@ public Optional<ConnectorAction> findAction(String action) {
         return Optional.empty();
     }
 
+    @Override
+    public void addAction(ConnectorAction action) {
+        actions.add(action);
+    }
+
     @Override
     public void removeCredential() {
         this.credential = null;

common/src/main/java/org/opensearch/ml/common/connector/Connector.java

@@ -65,6 +65,8 @@ public interface Connector extends ToXContentObject, Writeable {
 
     List<ConnectorAction> getActions();
 
+    void addAction(ConnectorAction action);
+
     ConnectorClientConfig getConnectorClientConfig();
 
     String getActionEndpoint(String action, Map<String, String> parameters);

common/src/main/java/org/opensearch/ml/common/output/MLPredictionOutput.java

@@ -6,6 +6,7 @@
 package org.opensearch.ml.common.output;
 
 import java.io.IOException;
+import java.util.Map;
 
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -30,8 +31,12 @@ public class MLPredictionOutput extends MLOutput {
     public static final String STATUS_FIELD = "status";
     public static final String PREDICTION_RESULT_FIELD = "prediction_result";
 
+    // This field will be created for offline batch prediction tasks containing details of the batch job as outputted by the remote server.
+    public static final String REMOTE_JOB_FIELD = "remote_job";
+
     String taskId;
     String status;
+    Map<String, Object> remoteJob;
 
     @ToString.Exclude
     DataFrame predictionResult;
@@ -44,6 +49,14 @@ public MLPredictionOutput(String taskId, String status, DataFrame predictionResu
         this.predictionResult = predictionResult;
     }
 
+    @Builder
+    public MLPredictionOutput(String taskId, String status, Map<String, Object> remoteJob) {
+        super(OUTPUT_TYPE);
+        this.taskId = taskId;
+        this.status = status;
+        this.remoteJob = remoteJob;
+    }
+
     public MLPredictionOutput(StreamInput in) throws IOException {
         super(OUTPUT_TYPE);
         this.taskId = in.readOptionalString();
@@ -56,6 +69,9 @@ public MLPredictionOutput(StreamInput in) throws IOException {
                     break;
             }
         }
+        if (in.readBoolean()) {
+            this.remoteJob = in.readMap(s -> s.readString(), s -> s.readGenericValue());
+        }
     }
 
     @Override
@@ -69,6 +85,12 @@ public void writeTo(StreamOutput out) throws IOException {
         } else {
             out.writeBoolean(false);
         }
+        if (remoteJob != null) {
+            out.writeBoolean(true);
+            out.writeMap(remoteJob, StreamOutput::writeString, StreamOutput::writeGenericValue);
+        } else {
+            out.writeBoolean(false);
+        }
     }
 
     @Override
@@ -87,6 +109,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.endObject();
         }
 
+        if (remoteJob != null) {
+            builder.field(REMOTE_JOB_FIELD, remoteJob);
+        }
+
         builder.endObject();
         return builder;
     }

common/src/test/java/org/opensearch/ml/common/output/MLPredictionOutputTest.java

@@ -9,7 +9,9 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -29,6 +31,7 @@
 public class MLPredictionOutputTest {
 
     MLPredictionOutput output;
+    MLPredictionOutput outputWithRemoteJob;
 
     @Before
     public void setUp() {
@@ -37,12 +40,17 @@ public void setUp() {
         rows.add(new Row(new ColumnValue[] { new IntValue(1) }));
         rows.add(new Row(new ColumnValue[] { new IntValue(2) }));
         DataFrame dataFrame = new DefaultDataFrame(columnMetas, rows);
+        Map<String, Object> remoteJob = new HashMap<>();
+        remoteJob.put("status", "INPROGRESS");
+        remoteJob.put("job_id", "testJobID");
         output = MLPredictionOutput.builder().taskId("test_task_id").status("test_status").predictionResult(dataFrame).build();
+        outputWithRemoteJob = new MLPredictionOutput("test_task_id", "test_status", remoteJob);
     }
 
     @Test
     public void toXContent() throws IOException {
         XContentBuilder builder = XContentFactory.jsonBuilder();
+        XContentBuilder builderWithRemoteJob = XContentFactory.jsonBuilder();
         output.toXContent(builder, ToXContent.EMPTY_PARAMS);
         String jsonStr = builder.toString();
         assertEquals(
@@ -52,6 +60,12 @@ public void toXContent() throws IOException {
                 + "\"value\":2}]}]}}",
             jsonStr
         );
+        outputWithRemoteJob.toXContent(builderWithRemoteJob, ToXContent.EMPTY_PARAMS);
+        String jsonStr2 = builderWithRemoteJob.toString();
+        assertEquals(
+            "{\"task_id\":\"test_task_id\",\"status\":\"test_status\",\"remote_job\":{\"job_id\":\"testJobID\",\"status\":\"INPROGRESS\"}}",
+            jsonStr2
+        );
     }
 
     @Test

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java

@@ -6,6 +6,7 @@
 package org.opensearch.ml.engine.algorithms.remote;
 
 import static org.apache.commons.text.StringEscapeUtils.escapeJson;
+import static org.opensearch.ml.common.connector.ConnectorAction.ActionType.BATCH_PREDICT;
 import static org.opensearch.ml.common.connector.ConnectorAction.ActionType.CANCEL_BATCH_PREDICT;
 import static org.opensearch.ml.common.connector.HttpConnector.RESPONSE_FILTER_FIELD;
 import static org.opensearch.ml.common.connector.MLPreProcessFunction.CONVERT_INPUT_TO_JSON_STRING;
@@ -19,6 +20,7 @@
 import java.net.URI;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -61,6 +63,9 @@ public class ConnectorUtils {
     private static final Aws4Signer signer;
     public static final String SKIP_VALIDATE_MISSING_PARAMETERS = "skip_validating_missing_parameters";
 
+    public static final List<String> SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES = List
+        .of("sagemaker", "openai", "bedrock", "cohere");
+
     static {
         signer = Aws4Signer.create();
     }
@@ -313,4 +318,63 @@ public static SdkHttpFullRequest buildSdkRequest(
         }
         return builder.build();
     }
+
+    public static ConnectorAction createConnectorAction(Connector connector, ConnectorAction.ActionType actionType) {
+        Optional<ConnectorAction> batchPredictAction = connector.findAction(BATCH_PREDICT.name());
+        String predictEndpoint = batchPredictAction.get().getUrl();
+        Map<String, String> parameters = connector.getParameters() != null
+            ? new HashMap<>(connector.getParameters())
+            : Collections.emptyMap();
+
+        // Apply parameter substitution only if needed
+        if (!parameters.isEmpty()) {
+            StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
+            predictEndpoint = substitutor.replace(predictEndpoint);
+        }
+
+        boolean isCancelAction = actionType == CANCEL_BATCH_PREDICT;
+
+        // Initialize the default method and requestBody
+        String method = "POST";
+        String requestBody = null;
+        String url = "";
+
+        switch (getRemoteServerFromURL(predictEndpoint)) {
+            case "sagemaker":
+                url = isCancelAction
+                    ? predictEndpoint.replace("CreateTransformJob", "StopTransformJob")
+                    : predictEndpoint.replace("CreateTransformJob", "DescribeTransformJob");
+                requestBody = "{ \"TransformJobName\" : \"${parameters.TransformJobName}\"}";
+                break;
+            case "openai":
+            case "cohere":
+                url = isCancelAction ? predictEndpoint + "/${parameters.id}/cancel" : predictEndpoint + "/${parameters.id}";
+                method = isCancelAction ? "POST" : "GET";
+                break;
+            case "bedrock":
+                url = isCancelAction
+                    ? predictEndpoint + "/${parameters.processedJobArn}/stop"
+                    : predictEndpoint + "/${parameters.processedJobArn}";
+                method = isCancelAction ? "POST" : "GET";
+                break;
+            default:
+                String errorMessage = isCancelAction
+                    ? "Please configure the action type to cancel the batch job in the connector"
+                    : "Please configure the action type to get the batch job details in the connector";
+                throw new UnsupportedOperationException(errorMessage);
+        }
+
+        return ConnectorAction
+            .builder()
+            .actionType(actionType)
+            .method(method)
+            .url(url)
+            .requestBody(requestBody)
+            .headers(batchPredictAction.get().getHeaders())
+            .build();
+    }
+
+    public static String getRemoteServerFromURL(String url) {
+        return SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES.stream().filter(url::contains).findFirst().orElse("");
+    }
 }