Skip to content

Commit

Permalink
Add a flag to control auto-deploy behavior (#2276) (#2287)
Browse files Browse the repository at this point in the history
* Add a flag to control auto-deploy behavior

Signed-off-by: Sicheng Song <[email protected]>

* Fix compilation

Signed-off-by: Sicheng Song <[email protected]>

* Bump back schema version due to already bumped in guardrail pr

Signed-off-by: Sicheng Song <[email protected]>

---------

Signed-off-by: Sicheng Song <[email protected]>
(cherry picked from commit 22fcaf0)

Co-authored-by: Sicheng Song <[email protected]>
  • Loading branch information
opensearch-trigger-bot[bot] and b4sjoo authored Mar 27, 2024
1 parent 4dbeb28 commit c5c4e61
Show file tree
Hide file tree
Showing 10 changed files with 345 additions and 28 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,9 @@ public class CommonValue {
+ MODEL_MAX_LENGTH_FIELD + "\":{\"type\":\"integer\"},\""
+ ALL_CONFIG_FIELD + "\":{\"type\":\"text\"}}},\n"
+ " \""
+ MLModel.DEPLOY_SETTING_FIELD
+ "\" : {\"type\": \"flat_object\"},\n"
+ " \""
+ MLModel.IS_ENABLED_FIELD
+ "\" : {\"type\": \"boolean\"},\n"
+ " \""
Expand Down
22 changes: 22 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/MLModel.java
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.ml.common.connector.Connector;
import org.opensearch.ml.common.model.Guardrails;
import org.opensearch.ml.common.model.MLDeploySetting;
import org.opensearch.ml.common.model.MLModelConfig;
import org.opensearch.ml.common.controller.MLRateLimiter;
import org.opensearch.ml.common.model.MLModelFormat;
Expand Down Expand Up @@ -61,6 +62,7 @@ public class MLModel implements ToXContentObject {
public static final String RATE_LIMITER_FIELD = "rate_limiter";
public static final String IS_CONTROLLER_ENABLED_FIELD = "is_controller_enabled";
public static final String MODEL_CONFIG_FIELD = "model_config";
public static final String DEPLOY_SETTING_FIELD = "deploy_setting"; // optional
public static final String CREATED_TIME_FIELD = "created_time";
public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
@Deprecated
Expand Down Expand Up @@ -102,6 +104,7 @@ public class MLModel implements ToXContentObject {
private Long modelContentSizeInBytes;
private String modelContentHash;
private MLModelConfig modelConfig;
private MLDeploySetting deploySetting;
private Boolean isEnabled;
private Boolean isControllerEnabled;
private MLRateLimiter rateLimiter;
Expand Down Expand Up @@ -147,6 +150,7 @@ public MLModel(String name,
Boolean isControllerEnabled,
MLRateLimiter rateLimiter,
MLModelConfig modelConfig,
MLDeploySetting deploySetting,
Instant createdTime,
Instant lastUpdateTime,
Instant lastRegisteredTime,
Expand Down Expand Up @@ -178,6 +182,7 @@ public MLModel(String name,
this.isControllerEnabled = isControllerEnabled;
this.rateLimiter = rateLimiter;
this.modelConfig = modelConfig;
this.deploySetting = deploySetting;
this.createdTime = createdTime;
this.lastUpdateTime = lastUpdateTime;
this.lastRegisteredTime = lastRegisteredTime;
Expand Down Expand Up @@ -226,6 +231,9 @@ public MLModel(StreamInput input) throws IOException {
modelConfig = new TextEmbeddingModelConfig(input);
}
}
if (input.readBoolean()) {
this.deploySetting = new MLDeploySetting(input);
}
isEnabled = input.readOptionalBoolean();
isControllerEnabled = input.readOptionalBoolean();
if (input.readBoolean()) {
Expand Down Expand Up @@ -288,6 +296,12 @@ public void writeTo(StreamOutput out) throws IOException {
} else {
out.writeBoolean(false);
}
if (deploySetting != null) {
out.writeBoolean(true);
deploySetting.writeTo(out);
} else {
out.writeBoolean(false);
}
out.writeOptionalBoolean(isEnabled);
out.writeOptionalBoolean(isControllerEnabled);
if (rateLimiter != null) {
Expand Down Expand Up @@ -365,6 +379,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
if (modelConfig != null) {
builder.field(MODEL_CONFIG_FIELD, modelConfig);
}
if (deploySetting != null) {
builder.field(DEPLOY_SETTING_FIELD, deploySetting);
}
if (isEnabled != null) {
builder.field(IS_ENABLED_FIELD, isEnabled);
}
Expand Down Expand Up @@ -445,6 +462,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
Long modelContentSizeInBytes = null;
String modelContentHash = null;
MLModelConfig modelConfig = null;
MLDeploySetting deploySetting = null;
Boolean isEnabled = null;
Boolean isControllerEnabled = null;
MLRateLimiter rateLimiter = null;
Expand Down Expand Up @@ -536,6 +554,9 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
modelConfig = TextEmbeddingModelConfig.parse(parser);
}
break;
case DEPLOY_SETTING_FIELD:
deploySetting = MLDeploySetting.parse(parser);
break;
case IS_ENABLED_FIELD:
isEnabled = parser.booleanValue();
break;
Expand Down Expand Up @@ -614,6 +635,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
.modelContentSizeInBytes(modelContentSizeInBytes)
.modelContentHash(modelContentHash)
.modelConfig(modelConfig)
.deploySetting(deploySetting)
.isEnabled(isEnabled)
.isControllerEnabled(isControllerEnabled)
.rateLimiter(rateLimiter)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common.model;

import lombok.Builder;
import lombok.Getter;
import lombok.Setter;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.core.xcontent.ToXContentObject;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;

import java.io.IOException;

import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;

@Setter
@Getter
public class MLDeploySetting implements ToXContentObject, Writeable {
public static final String IS_AUTO_DEPLOY_ENABLED_FIELD = "is_auto_deploy_enabled";

private Boolean isAutoDeployEnabled;

@Builder(toBuilder = true)
public MLDeploySetting(Boolean isAutoDeployEnabled) {
this.isAutoDeployEnabled = isAutoDeployEnabled;
}

public MLDeploySetting(StreamInput in) throws IOException {
this.isAutoDeployEnabled = in.readOptionalBoolean();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeOptionalBoolean(isAutoDeployEnabled);
}

public static MLDeploySetting parse(XContentParser parser) throws IOException {
Boolean isAutoDeployEnabled = 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 IS_AUTO_DEPLOY_ENABLED_FIELD:
isAutoDeployEnabled = parser.booleanValue();
break;
default:
parser.skipChildren();
break;
}
}
return new MLDeploySetting(isAutoDeployEnabled);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
builder.startObject();
if (isAutoDeployEnabled != null) {
builder.field(IS_AUTO_DEPLOY_ENABLED_FIELD, isAutoDeployEnabled);
}
builder.endObject();
return builder;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,12 @@
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.ml.common.connector.Connector;
import org.opensearch.ml.common.model.Guardrails;
import org.opensearch.ml.common.model.MLDeploySetting;
import org.opensearch.ml.common.model.MLModelConfig;
import org.opensearch.ml.common.controller.MLRateLimiter;
import org.opensearch.ml.common.model.TextEmbeddingModelConfig;
import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
import org.opensearch.ml.common.transport.register.MLRegisterModelInput;

import java.io.IOException;
import java.time.Instant;
Expand All @@ -39,6 +41,7 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
public static final String IS_ENABLED_FIELD = "is_enabled"; // optional
public static final String RATE_LIMITER_FIELD = "rate_limiter"; // optional
public static final String MODEL_CONFIG_FIELD = "model_config"; // optional
public static final String DEPLOY_SETTING_FIELD = "deploy_setting"; // optional
public static final String UPDATED_CONNECTOR_FIELD = "updated_connector"; // passively set when updating the
// internal connector
public static final String CONNECTOR_ID_FIELD = "connector_id"; // optional
Expand All @@ -47,8 +50,6 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
// request
public static final String GUARDRAILS_FIELD = "guardrails";

private static final Version MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS = Version.V_2_13_0;

@Getter
private String modelId;
private String description;
Expand All @@ -58,6 +59,7 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
private Boolean isEnabled;
private MLRateLimiter rateLimiter;
private MLModelConfig modelConfig;
private MLDeploySetting deploySetting;
private Connector updatedConnector;
private String connectorId;
private MLCreateConnectorInput connector;
Expand All @@ -66,7 +68,7 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {

@Builder(toBuilder = true)
public MLUpdateModelInput(String modelId, String description, String version, String name, String modelGroupId,
Boolean isEnabled, MLRateLimiter rateLimiter, MLModelConfig modelConfig,
Boolean isEnabled, MLRateLimiter rateLimiter, MLModelConfig modelConfig, MLDeploySetting deploySetting,
Connector updatedConnector, String connectorId, MLCreateConnectorInput connector, Instant lastUpdateTime, Guardrails guardrails) {
this.modelId = modelId;
this.description = description;
Expand All @@ -76,6 +78,7 @@ public MLUpdateModelInput(String modelId, String description, String version, St
this.isEnabled = isEnabled;
this.rateLimiter = rateLimiter;
this.modelConfig = modelConfig;
this.deploySetting = deploySetting;
this.updatedConnector = updatedConnector;
this.connectorId = connectorId;
this.connector = connector;
Expand Down Expand Up @@ -105,10 +108,13 @@ public MLUpdateModelInput(StreamInput in) throws IOException {
connector = new MLCreateConnectorInput(in);
}
lastUpdateTime = in.readOptionalInstant();
if (streamInputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS)) {
if (streamInputVersion.onOrAfter(MLRegisterModelInput.MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS_AND_AUTO_DEPLOY)) {
if (in.readBoolean()) {
this.guardrails = new Guardrails(in);
}
if (in.readBoolean()) {
this.deploySetting = new MLDeploySetting(in);
}
}
}

Expand Down Expand Up @@ -137,6 +143,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
if (modelConfig != null) {
builder.field(MODEL_CONFIG_FIELD, modelConfig);
}
if (deploySetting != null) {
builder.field(DEPLOY_SETTING_FIELD, deploySetting);
}
if (updatedConnector != null) {
builder.field(UPDATED_CONNECTOR_FIELD, updatedConnector);
}
Expand Down Expand Up @@ -180,6 +189,9 @@ public XContentBuilder toXContentForUpdateRequestDoc(XContentBuilder builder, Pa
if (modelConfig != null) {
builder.field(MODEL_CONFIG_FIELD, modelConfig);
}
if (deploySetting != null) {
builder.field(DEPLOY_SETTING_FIELD, deploySetting);
}
// Notice that we serialize the updatedConnector to the connector field, in order to be compatible with original internal connector field format.
if (updatedConnector != null) {
builder.field(CONNECTOR_FIELD, updatedConnector);
Expand Down Expand Up @@ -232,13 +244,19 @@ public void writeTo(StreamOutput out) throws IOException {
out.writeBoolean(false);
}
out.writeOptionalInstant(lastUpdateTime);
if (streamOutputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS)) {
if (streamOutputVersion.onOrAfter(MLRegisterModelInput.MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS_AND_AUTO_DEPLOY)) {
if (guardrails != null) {
out.writeBoolean(true);
guardrails.writeTo(out);
} else {
out.writeBoolean(false);
}
if (deploySetting != null) {
out.writeBoolean(true);
deploySetting.writeTo(out);
} else {
out.writeBoolean(false);
}
}
}

Expand All @@ -251,6 +269,7 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
Boolean isEnabled = null;
MLRateLimiter rateLimiter = null;
MLModelConfig modelConfig = null;
MLDeploySetting deploySetting = null;
Connector updatedConnector = null;
String connectorId = null;
MLCreateConnectorInput connector = null;
Expand Down Expand Up @@ -280,6 +299,9 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
case MODEL_CONFIG_FIELD:
modelConfig = TextEmbeddingModelConfig.parse(parser);
break;
case DEPLOY_SETTING_FIELD:
deploySetting = MLDeploySetting.parse(parser);
break;
case CONNECTOR_ID_FIELD:
connectorId = parser.text();
break;
Expand All @@ -297,6 +319,6 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
// Model ID can only be set through RestRequest. Model version can only be set
// automatically.
return new MLUpdateModelInput(modelId, description, version, name, modelGroupId, isEnabled, rateLimiter,
modelConfig, updatedConnector, connectorId, connector, lastUpdateTime, guardrails);
modelConfig, deploySetting, updatedConnector, connectorId, connector, lastUpdateTime, guardrails);
}
}
Loading

0 comments on commit c5c4e61

Please sign in to comment.