add batch predict job actiontype in connector

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Author: Zhangxunmt

common/src/main/java/org/opensearch/ml/common/PredictMode.java

@@ -0,0 +1,18 @@
+package org.opensearch.ml.common;
+
+import java.util.Locale;
+
+public enum PredictMode {
+    PREDICT,
+    BATCH,
+    ASYNC,
+    STREAMING;
+
+    public static PredictMode from(String value) {
+        try {
+            return PredictMode.valueOf(value.toUpperCase(Locale.ROOT));
+        } catch (Exception e) {
+            throw new IllegalArgumentException("Wrong Predict mode");
+        }
+    }
+}

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

@@ -183,6 +183,7 @@ public static ConnectorAction parse(XContentParser parser) throws IOException {
 
     public enum ActionType {
         PREDICT,
-        EXECUTE
+        EXECUTE,
+        BATCH
     }
 }

common/src/main/java/org/opensearch/ml/common/dataset/remote/RemoteInferenceInputDataSet.java

@@ -10,6 +10,7 @@
 import lombok.Setter;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.ml.common.PredictMode;
 import org.opensearch.ml.common.annotation.InputDataSet;
 import org.opensearch.ml.common.dataset.MLInputDataType;
 import org.opensearch.ml.common.dataset.MLInputDataset;
@@ -23,11 +24,17 @@ public class RemoteInferenceInputDataSet extends MLInputDataset {
 
     @Setter
     private Map<String, String> parameters;
+    private PredictMode predictMode;
 
     @Builder(toBuilder = true)
-    public RemoteInferenceInputDataSet(Map<String, String> parameters) {
+    public RemoteInferenceInputDataSet(Map<String, String> parameters, PredictMode predictMode) {
         super(MLInputDataType.REMOTE);
         this.parameters = parameters;
+        this.predictMode = predictMode;
+    }
+
+    public RemoteInferenceInputDataSet(Map<String, String> parameters) {
+        this(parameters, null);
     }
 
     public RemoteInferenceInputDataSet(StreamInput streamInput) throws IOException {

common/src/main/java/org/opensearch/ml/common/input/MLInput.java

@@ -35,6 +35,7 @@
 import java.util.Map;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.opensearch.ml.common.input.remote.RemoteInferenceMLInput.PREDICT_MODE_FIELD;
 
 /**
  * ML input data: algorithm name, parameters and input data set.
@@ -196,6 +197,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
                     RemoteInferenceInputDataSet remoteInferenceInputDataSet = (RemoteInferenceInputDataSet) this.inputDataset;
                     Map<String, String> parameters = remoteInferenceInputDataSet.getParameters();
                     builder.field(PARAMETERS_FIELD, parameters);
+                    builder.field(PREDICT_MODE_FIELD, remoteInferenceInputDataSet.getPredictMode());
                     break;
                 default:
                     break;

common/src/main/java/org/opensearch/ml/common/input/remote/RemoteInferenceMLInput.java

@@ -9,6 +9,7 @@
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.ml.common.FunctionName;
+import org.opensearch.ml.common.PredictMode;
 import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
 import org.opensearch.ml.common.utils.StringUtils;
@@ -21,6 +22,7 @@
 @org.opensearch.ml.common.annotation.MLInput(functionNames = {FunctionName.REMOTE})
 public class RemoteInferenceMLInput extends MLInput {
     public static final String PARAMETERS_FIELD = "parameters";
+    public static final String PREDICT_MODE_FIELD = "mode";
 
     public RemoteInferenceMLInput(StreamInput in) throws IOException {
         super(in);
@@ -34,21 +36,26 @@ public void writeTo(StreamOutput out) throws IOException {
     public RemoteInferenceMLInput(XContentParser parser, FunctionName functionName) throws IOException {
         super();
         this.algorithm = functionName;
+        Map<String, String> parameters = null;
+        PredictMode predictMode = null;
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
             String fieldName = parser.currentName();
             parser.nextToken();
 
             switch (fieldName) {
                 case PARAMETERS_FIELD:
-                    Map<String, String> parameters = StringUtils.getParameterMap(parser.map());
-                    inputDataset = new RemoteInferenceInputDataSet(parameters);
+                    parameters = StringUtils.getParameterMap(parser.map());
                     break;
+                case PREDICT_MODE_FIELD:
+                    predictMode = PredictMode.from(parser.text());
                 default:
                     parser.skipChildren();
                     break;
             }
         }
+        predictMode = predictMode == null? PredictMode.PREDICT:predictMode;
+        inputDataset = new RemoteInferenceInputDataSet(parameters, predictMode);
     }
 
 }

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

@@ -17,7 +17,9 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLModel;
+import org.opensearch.ml.common.PredictMode;
 import org.opensearch.ml.common.connector.Connector;
+import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
 import org.opensearch.ml.common.exception.MLException;
 import org.opensearch.ml.common.input.MLInput;
 import org.opensearch.ml.common.model.MLGuard;
@@ -70,7 +72,12 @@ public void asyncPredict(MLInput mlInput, ActionListener<MLTaskResponse> actionL
             return;
         }
         try {
-            connectorExecutor.executeAction(PREDICT.name(), mlInput, actionListener);
+            PredictMode predictMode = null;
+            if (mlInput.getInputDataset() instanceof RemoteInferenceInputDataSet) {
+                predictMode = ((RemoteInferenceInputDataSet) mlInput.getInputDataset()).getPredictMode();
+            }
+            predictMode = predictMode == null?PredictMode.PREDICT:predictMode;
+            connectorExecutor.executeAction(predictMode.toString(), mlInput, actionListener);
         } catch (RuntimeException e) {
             log.error("Failed to call remote model.", e);
             actionListener.onFailure(e);

plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java

@@ -132,6 +132,7 @@ private MLCommonsSettings() {}
             ImmutableList
                 .of(
                     "^https://runtime\\.sagemaker\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
+                    "^https://api\\.sagemaker\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
                     "^https://api\\.openai\\.com/.*$",
                     "^https://api\\.cohere\\.ai/.*$",
                     "^https://bedrock-runtime\\..*[a-z0-9-]\\.amazonaws\\.com/.*$"