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 <[email protected]> * fix failed unit test Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> * throw exception when model group not found during update request (#1447) Signed-off-by: Bhavana Ramaram <[email protected]> * add status code to model tensor (#1443) (#1453) Signed-off-by: Yaliang Wu <[email protected]> * register new versions to a model group based on the name provided (#1452) Signed-off-by: Bhavana Ramaram <[email protected]> * fixing metrics correlation algorithm (#1448) * fixing metrics correlation algorithm Signed-off-by: Dhrubo Saha <[email protected]> * 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 <[email protected]> * Update Model API (#1350) * Update Model API POC Signed-off-by: Sicheng Song <[email protected]> * Using GetRequest to get model Signed-off-by: Sicheng Song <[email protected]> * Finalize model update API Signed-off-by: Sicheng Song <[email protected]> * Fix compile Signed-off-by: Sicheng Song <[email protected]> * Fix compileTest Signed-off-by: Sicheng Song <[email protected]> * Add Unit Test Cases for Update Model API Signed-off-by: Sicheng Song <[email protected]> * Tune back test coverage thereshold Signed-off-by: Sicheng Song <[email protected]> * Add more unit tests on Update model API Signed-off-by: Sicheng Song <[email protected]> * Add unit test for TransportUpdateModelAction class Signed-off-by: Sicheng Song <[email protected]> * Fix a test error Signed-off-by: Sicheng Song <[email protected]> * Change exception thrown to failure response Signed-off-by: Sicheng Song <[email protected]> * Move the function judgement to the outter block Signed-off-by: Sicheng Song <[email protected]> * Check if model is undeployed before update model Signed-off-by: Sicheng Song <[email protected]> * Add more unit test for update model API Signed-off-by: Sicheng Song <[email protected]> * Fix unit test due to blocking java 11 CI workflow Signed-off-by: Sicheng Song <[email protected]> * Enabling auto bumping model version during registering to a new model group and address reviewers' other concern Signed-off-by: Sicheng Song <[email protected]> * Autobump new model groups' latest version when register to a new model Signed-off-by: Sicheng Song <[email protected]> * Change the REST API method from POST to PUT Signed-off-by: Sicheng Song <[email protected]> * Change the update REST API endpoint Signed-off-by: Sicheng Song <[email protected]> --------- Signed-off-by: Sicheng Song <[email protected]> * Add a setting to control the update connector API (#1465) * Add a setting to control the update connector API Signed-off-by: Sicheng Song <[email protected]> * Enabling the update connnector setting in unit test Signed-off-by: Sicheng Song <[email protected]> * Enabling the update connnector setting in corresponding unit test Signed-off-by: Sicheng Song <[email protected]> --------- Signed-off-by: Sicheng Song <[email protected]> * fix update connector API (#1484) * fix update connector API Signed-off-by: Yaliang Wu <[email protected]> * 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 <[email protected]> * Add context.restore() to avoid missing info Signed-off-by: zane-neo <[email protected]> --------- Signed-off-by: zane-neo <[email protected]> (cherry picked from commit a985f6e) Co-authored-by: zane-neo <[email protected]> * fix failed ut from PR 1472 (#1479) (#1510) * fix failed ut from PR 1472 Signed-off-by: Yaliang Wu <[email protected]> * exclude class for low coverage Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> (cherry picked from commit da5d829) Co-authored-by: Yaliang Wu <[email protected]> * [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 <[email protected]> * fix kmeans model deploy bug Signed-off-by: Yaliang Wu <[email protected]> * support multiple docs for remote embedding model Signed-off-by: Yaliang Wu <[email protected]> * fix ut Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> * fix wrong class Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> (cherry picked from commit 201c8a8) Co-authored-by: Yaliang Wu <[email protected]> * 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 <[email protected]> * only add model info to cache if model cache exist Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> (cherry picked from commit 6f83b9f) Co-authored-by: Yaliang Wu <[email protected]> * 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 <[email protected]> (cherry picked from commit 60ef0fd) Co-authored-by: Bhavana Ramaram <[email protected]> (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 <[email protected]> (cherry picked from commit cea1cd6) Co-authored-by: Yaliang Wu <[email protected]> * add prefix to show the error is from remote service (#1499) (#1515) Signed-off-by: Yaliang Wu <[email protected]> (cherry picked from commit 3897ad1) Co-authored-by: Yaliang Wu <[email protected]> * fix multiple docs support (#1516) Signed-off-by: Yaliang Wu <[email protected]> * adding another fix issue to the release note (#1498) (#1514) Signed-off-by: Dhrubo Saha <[email protected]> (cherry picked from commit 440155c) Co-authored-by: Dhrubo Saha <[email protected]> * add bedrockURL to trusted connector regex list (#1461) Signed-off-by: Bhavana Ramaram <[email protected]> * return parsing exception 400 for parsing errors Signed-off-by: Xun Zhang <[email protected]> * add more ut in restupdateconnector Signed-off-by: Xun Zhang <[email protected]> * fix format violations Signed-off-by: Bhavana Ramaram <[email protected]> * Fix model/connector update API to address security concern (#1595) * Fix model/connector update API to address appsec concern Signed-off-by: Sicheng Song <[email protected]> * Fix compile and build failure Signed-off-by: Sicheng Song <[email protected]> * Improve unit test coverage Signed-off-by: Sicheng Song <[email protected]> * Fix spotless Signed-off-by: Sicheng Song <[email protected]> * Merge update connector feature flag to remote inference feature flag Signed-off-by: Sicheng Song <[email protected]> * Fix compile Signed-off-by: Sicheng Song <[email protected]> * Fix exception status Signed-off-by: Sicheng Song <[email protected]> * Keep fixing exception status Signed-off-by: Sicheng Song <[email protected]> * Spotless fix Signed-off-by: Sicheng Song <[email protected]> * Add UT on parsing exception Signed-off-by: Sicheng Song <[email protected]> --------- Signed-off-by: Sicheng Song <[email protected]> * change XContentFactory to MediaTypeRegistry builder in MLRegisterModelInputTest class Signed-off-by: Bhavana Ramaram <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> Signed-off-by: Bhavana Ramaram <[email protected]> Signed-off-by: Dhrubo Saha <[email protected]> Signed-off-by: Sicheng Song <[email protected]> Signed-off-by: Xun Zhang <[email protected]> Co-authored-by: Yaliang Wu <[email protected]> Co-authored-by: Dhrubo Saha <[email protected]> Co-authored-by: Sicheng Song <[email protected]> Co-authored-by: opensearch-trigger-bot[bot] <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Co-authored-by: zane-neo <[email protected]> Co-authored-by: Xun Zhang <[email protected]>
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);
+ }
+}