Backport multiple PRs to main from 2.x (#1652)

* fix parameter name in preprocess function; fix remote model function … (#1362) * fix parameter name in preprocess function; fix remote model function name Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix failed unit test Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> * throw exception when model group not found during update request (#1447) Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * add status code to model tensor (#1443) (#1453) Signed-off-by: Yaliang Wu <ylwu@amazon.com> * register new versions to a model group based on the name provided (#1452) Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * fixing metrics correlation algorithm (#1448) * fixing metrics correlation algorithm Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> * if model version fails to register, update model group accordingly (#1463) * if model version fails to register, update model group accordingly Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * Update Model API (#1350) * Update Model API POC Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Using GetRequest to get model Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Finalize model update API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compile Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compileTest Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add Unit Test Cases for Update Model API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Tune back test coverage thereshold Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add more unit tests on Update model API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add unit test for TransportUpdateModelAction class Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix a test error Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Change exception thrown to failure response Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Move the function judgement to the outter block Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Check if model is undeployed before update model Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add more unit test for update model API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix unit test due to blocking java 11 CI workflow Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Enabling auto bumping model version during registering to a new model group and address reviewers' other concern Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Autobump new model groups' latest version when register to a new model Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Change the REST API method from POST to PUT Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Change the update REST API endpoint Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add a setting to control the update connector API (#1465) * Add a setting to control the update connector API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Enabling the update connnector setting in unit test Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Enabling the update connnector setting in corresponding unit test Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * fix update connector API (#1484) * fix update connector API Signed-off-by: Yaliang Wu <ylwu@amazon.com> * Performance enhacement for predict action by caching model info (#1472) (#1508) * Performance enhacement for predict action by caching model info Signed-off-by: zane-neo <zaniu@amazon.com> * Add context.restore() to avoid missing info Signed-off-by: zane-neo <zaniu@amazon.com> --------- Signed-off-by: zane-neo <zaniu@amazon.com> (cherry picked from commit a985f6e) Co-authored-by: zane-neo <zaniu@amazon.com> * fix failed ut from PR 1472 (#1479) (#1510) * fix failed ut from PR 1472 Signed-off-by: Yaliang Wu <ylwu@amazon.com> * exclude class for low coverage Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit da5d829) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * [Backport to 2.11] throw exception if remote model doesn't return 2xx status code; fix p… (#1477) (#1509) * throw exception if remote model doesn't return 2xx status code; fix p… (#1473) * throw exception if remote model doesn't return 2xx status code; fix predict runner Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix kmeans model deploy bug Signed-off-by: Yaliang Wu <ylwu@amazon.com> * support multiple docs for remote embedding model Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix ut Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix wrong class Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 201c8a8) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * fix no worker node exception for remote embedding model (#1482) (#1511) * fix no worker node exception for remote embedding model Signed-off-by: Yaliang Wu <ylwu@amazon.com> * only add model info to cache if model cache exist Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 6f83b9f) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * fix for delete model group API throwing incorrect error when model index not created (#1485) (#1486) (#1512) * fix for delete model group API throwing incorrect error when model index not created Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> (cherry picked from commit 60ef0fd) Co-authored-by: Bhavana Ramaram <rbhavna@amazon.com> (cherry picked from commit 5544681) Co-authored-by: opensearch-trigger-bot[bot] <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> * fix no worker node error on multi-node cluster (#1487) (#1513) Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit cea1cd6) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * add prefix to show the error is from remote service (#1499) (#1515) Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 3897ad1) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * fix multiple docs support (#1516) Signed-off-by: Yaliang Wu <ylwu@amazon.com> * adding another fix issue to the release note (#1498) (#1514) Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> (cherry picked from commit 440155c) Co-authored-by: Dhrubo Saha <dhrubo@amazon.com> * add bedrockURL to trusted connector regex list (#1461) Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * return parsing exception 400 for parsing errors Signed-off-by: Xun Zhang <xunzh@amazon.com> * add more ut in restupdateconnector Signed-off-by: Xun Zhang <xunzh@amazon.com> * fix format violations Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * Fix model/connector update API to address security concern (#1595) * Fix model/connector update API to address appsec concern Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compile and build failure Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Improve unit test coverage Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix spotless Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Merge update connector feature flag to remote inference feature flag Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compile Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix exception status Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Keep fixing exception status Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Spotless fix Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add UT on parsing exception Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * change XContentFactory to MediaTypeRegistry builder in MLRegisterModelInputTest class Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> Signed-off-by: Sicheng Song <sicheng.song@outlook.com> Signed-off-by: Xun Zhang <xunzh@amazon.com> Co-authored-by: Yaliang Wu <ylwu@amazon.com> Co-authored-by: Dhrubo Saha <dhrubo@amazon.com> Co-authored-by: Sicheng Song <sicheng.song@outlook.com> Co-authored-by: opensearch-trigger-bot[bot] <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Co-authored-by: zane-neo <zaniu@amazon.com> Co-authored-by: Xun Zhang <xunzh@amazon.com>
opensearch-project · Nov 17, 2023 · 9a7075d · 9a7075d
1 parent 5759bf2
commit 9a7075d
Show file tree

Hide file tree

Showing 64 changed files with 3,595 additions and 356 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/CommonValue.java b/common/src/main/java/org/opensearch/ml/common/CommonValue.java
@@ -18,6 +18,7 @@
 public class CommonValue {
 
     public static Integer NO_SCHEMA_VERSION = 0;
+    public static final String REMOTE_SERVICE_ERROR = "Error from remote service: ";
     public static final String USER = "user";
     public static final String META = "_meta";
     public static final String SCHEMA_VERSION_FIELD = "schema_version";

diff --git a/common/src/main/java/org/opensearch/ml/common/MLModel.java b/common/src/main/java/org/opensearch/ml/common/MLModel.java
@@ -191,7 +191,11 @@ public MLModel(StreamInput input) throws IOException{
             modelContentSizeInBytes = input.readOptionalLong();
             modelContentHash = input.readOptionalString();
             if (input.readBoolean()) {
-                modelConfig = new TextEmbeddingModelConfig(input);
+                if (algorithm.equals(FunctionName.METRICS_CORRELATION)) {
+                    modelConfig = new MetricsCorrelationModelConfig(input);
+                } else {
+                    modelConfig = new TextEmbeddingModelConfig(input);
+                }
             }
             createdTime = input.readOptionalInstant();
             lastUpdateTime = input.readOptionalInstant();

diff --git a/common/src/main/java/org/opensearch/ml/common/connector/Connector.java b/common/src/main/java/org/opensearch/ml/common/connector/Connector.java
@@ -30,6 +30,7 @@
 import org.opensearch.ml.common.AccessMode;
 import org.opensearch.ml.common.MLCommonsClassLoader;
 import org.opensearch.ml.common.output.model.ModelTensor;
+import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.ml.common.utils.StringUtils.gson;
@@ -69,6 +70,7 @@ public interface Connector extends ToXContentObject, Writeable {
 
     void writeTo(StreamOutput out) throws IOException;
 
+    void update(MLCreateConnectorInput updateContent, Function<String, String> function);
 
     <T> void parseResponse(T orElse, List<ModelTensor> modelTensors, boolean b) throws IOException;
 

diff --git a/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java b/common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java
@@ -34,6 +34,7 @@
 import static org.opensearch.ml.common.connector.ConnectorProtocols.validateProtocol;
 import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
 import static org.opensearch.ml.common.utils.StringUtils.isJson;
+import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
 
 @Log4j2
 @NoArgsConstructor
@@ -248,6 +249,38 @@ public void writeTo(StreamOutput out) throws IOException {
         }
     }
 
+    @Override
+    public void update(MLCreateConnectorInput updateContent, Function<String, String> function) {
+        if (updateContent.getName() != null) {
+            this.name = updateContent.getName();
+        }
+        if (updateContent.getDescription() != null) {
+            this.description = updateContent.getDescription();
+        }
+        if (updateContent.getVersion() != null) {
+            this.version = updateContent.getVersion();
+        }
+        if (updateContent.getProtocol() != null) {
+            this.protocol = updateContent.getProtocol();
+        }
+        if (updateContent.getParameters() != null && updateContent.getParameters().size() > 0) {
+            this.parameters = updateContent.getParameters();
+        }
+        if (updateContent.getCredential() != null && updateContent.getCredential().size() > 0) {
+            this.credential = updateContent.getCredential();
+            encrypt(function);
+        }
+        if (updateContent.getActions() != null) {
+            this.actions = updateContent.getActions();
+        }
+        if (updateContent.getBackendRoles() != null) {
+            this.backendRoles = updateContent.getBackendRoles();
+        }
+        if (updateContent.getAccess() != null) {
+            this.access = updateContent.getAccess();
+        }
+    }
+
     @Override
     public  <T> T createPredictPayload(Map<String, String> parameters) {
         Optional<ConnectorAction> predictAction = findPredictAction();

diff --git a/common/src/main/java/org/opensearch/ml/common/model/MetricsCorrelationModelConfig.java b/common/src/main/java/org/opensearch/ml/common/model/MetricsCorrelationModelConfig.java
@@ -8,6 +8,7 @@
 import lombok.Builder;
 import lombok.Getter;
 import lombok.Setter;
+import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.ml.common.FunctionName;
@@ -28,6 +29,10 @@ public MetricsCorrelationModelConfig(String modelType, String allConfig) {
         super(modelType, allConfig);
     }
 
+    public MetricsCorrelationModelConfig(StreamInput in) throws IOException{
+        super(in);
+    }
+
     @Override
     public String getWriteableName() {
         return PARSE_FIELD_NAME;

diff --git a/common/src/main/java/org/opensearch/ml/common/output/model/ModelTensors.java b/common/src/main/java/org/opensearch/ml/common/output/model/ModelTensors.java
@@ -7,6 +7,7 @@
 
 import lombok.Builder;
 import lombok.Getter;
+import lombok.Setter;
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.common.io.stream.StreamInput;
@@ -24,7 +25,10 @@
 @Getter
 public class ModelTensors implements Writeable, ToXContentObject {
     public static final String OUTPUT_FIELD = "output";
+    public static final String STATUS_CODE_FIELD = "status_code";
     private List<ModelTensor> mlModelTensors;
+    @Setter
+    private Integer statusCode;
 
     @Builder
     public ModelTensors(List<ModelTensor> mlModelTensors) {
@@ -41,6 +45,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             }
             builder.endArray();
         }
+        if (statusCode != null) {
+            builder.field(STATUS_CODE_FIELD, statusCode);
+        }
         builder.endObject();
         return builder;
     }
@@ -53,6 +60,7 @@ public ModelTensors(StreamInput in) throws IOException {
                 mlModelTensors.add(new ModelTensor(in));
             }
         }
+        statusCode = in.readOptionalInt();
     }
 
     @Override
@@ -66,6 +74,7 @@ public void writeTo(StreamOutput out) throws IOException {
         } else {
             out.writeBoolean(false);
         }
+        out.writeOptionalInt(statusCode);
     }
 
     public void filter(ModelResultFilter resultFilter) {

diff --git a/...on/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java b/...on/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java
@@ -56,6 +56,7 @@ public class MLCreateConnectorInput implements ToXContentObject, Writeable {
     private Boolean addAllBackendRoles;
     private AccessMode access;
     private boolean dryRun = false;
+    private boolean updateConnector = false;
 
     @Builder(toBuilder = true)
     public MLCreateConnectorInput(String name,
@@ -68,9 +69,10 @@ public MLCreateConnectorInput(String name,
                                   List<String> backendRoles,
                                   Boolean addAllBackendRoles,
                                   AccessMode access,
-                                  boolean dryRun
+                                  boolean dryRun,
+                                  boolean updateConnector
     ) {
-        if (!dryRun) {
+        if (!dryRun && !updateConnector) {
             if (name == null) {
                 throw new IllegalArgumentException("Connector name is null");
             }
@@ -92,9 +94,14 @@ public MLCreateConnectorInput(String name,
         this.addAllBackendRoles = addAllBackendRoles;
         this.access = access;
         this.dryRun = dryRun;
+        this.updateConnector = updateConnector;
     }
 
     public static MLCreateConnectorInput parse(XContentParser parser) throws IOException {
+        return parse(parser, false);
+    }
+
+    public static MLCreateConnectorInput parse(XContentParser parser, boolean updateConnector) throws IOException {
         String name = null;
         String description = null;
         String version = null;
@@ -159,7 +166,7 @@ public static MLCreateConnectorInput parse(XContentParser parser) throws IOExcep
                     break;
             }
         }
-        return new MLCreateConnectorInput(name, description, version, protocol, parameters, credential, actions, backendRoles, addAllBackendRoles, access, dryRun);
+        return new MLCreateConnectorInput(name, description, version, protocol, parameters, credential, actions, backendRoles, addAllBackendRoles, access, dryRun, updateConnector);
     }
 
     @Override
@@ -201,10 +208,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
     @Override
     public void writeTo(StreamOutput output) throws IOException {
-        output.writeString(name);
+        output.writeOptionalString(name);
         output.writeOptionalString(description);
-        output.writeString(version);
-        output.writeString(protocol);
+        output.writeOptionalString(version);
+        output.writeOptionalString(protocol);
         if (parameters != null) {
             output.writeBoolean(true);
             output.writeMap(parameters, StreamOutput::writeString, StreamOutput::writeString);
@@ -240,13 +247,14 @@ public void writeTo(StreamOutput output) throws IOException {
             output.writeBoolean(false);
         }
         output.writeBoolean(dryRun);
+        output.writeBoolean(updateConnector);
     }
 
     public MLCreateConnectorInput(StreamInput input) throws IOException {
-        name = input.readString();
+        name = input.readOptionalString();
         description = input.readOptionalString();
-        version = input.readString();
-        protocol = input.readString();
+        version = input.readOptionalString();
+        protocol = input.readOptionalString();
         if (input.readBoolean()) {
             parameters = input.readMap(s -> s.readString(), s -> s.readString());
         }
@@ -268,5 +276,6 @@ public MLCreateConnectorInput(StreamInput input) throws IOException {
             this.access = input.readEnum(AccessMode.class);
         }
         dryRun = input.readBoolean();
+        updateConnector = input.readBoolean();
     }
 }
diff --git a/.../src/main/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequest.java b/.../src/main/java/org/opensearch/ml/common/transport/connector/MLUpdateConnectorRequest.java
@@ -19,32 +19,31 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
-import java.util.Map;
 
 import static org.opensearch.action.ValidateActions.addValidationError;
 
 @Getter
 public class MLUpdateConnectorRequest extends ActionRequest {
     String connectorId;
-    Map<String, Object> updateContent;
+    MLCreateConnectorInput updateContent;
 
     @Builder
-    public MLUpdateConnectorRequest(String connectorId, Map<String, Object> updateContent) {
+    public MLUpdateConnectorRequest(String connectorId, MLCreateConnectorInput updateContent) {
         this.connectorId = connectorId;
         this.updateContent = updateContent;
     }
 
     public MLUpdateConnectorRequest(StreamInput in) throws IOException {
         super(in);
         this.connectorId = in.readString();
-        this.updateContent = in.readMap();
+        this.updateContent = new MLCreateConnectorInput(in);
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeString(this.connectorId);
-        out.writeMap(this.getUpdateContent());
+        this.updateContent.writeTo(out);
     }
 
     @Override
@@ -55,14 +54,17 @@ public ActionRequestValidationException validate() {
             exception = addValidationError("ML connector id can't be null", exception);
         }
 
+        if (updateContent == null) {
+            exception = addValidationError("Update connector content can't be null", exception);
+        }
+
         return exception;
     }
 
     public static MLUpdateConnectorRequest parse(XContentParser parser, String connectorId) throws IOException {
-        Map<String, Object> dataAsMap = null;
-        dataAsMap = parser.map();
+        MLCreateConnectorInput updateContent = MLCreateConnectorInput.parse(parser, true);
 
-        return MLUpdateConnectorRequest.builder().connectorId(connectorId).updateContent(dataAsMap).build();
+        return MLUpdateConnectorRequest.builder().connectorId(connectorId).updateContent(updateContent).build();
     }
 
     public static MLUpdateConnectorRequest fromActionRequest(ActionRequest actionRequest) {

diff --git a/common/src/main/java/org/opensearch/ml/common/transport/model/MLUpdateModelAction.java b/common/src/main/java/org/opensearch/ml/common/transport/model/MLUpdateModelAction.java
@@ -0,0 +1,18 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.model;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.action.update.UpdateResponse;
+
+public class MLUpdateModelAction extends ActionType<UpdateResponse> {
+    public static MLUpdateModelAction INSTANCE = new MLUpdateModelAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/models/update";
+
+    private MLUpdateModelAction() {
+        super(NAME, UpdateResponse::new);
+    }
+}