Skip to content

Commit c5c4e61

Browse files
opensearch-trigger-bot[bot]b4sjoo
opensearch-trigger-bot[bot]
and
b4sjoo
authored
Add a flag to control auto-deploy behavior (opensearch-project#2276) (opensearch-project#2287)
* Add a flag to control auto-deploy behavior Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compilation Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Bump back schema version due to already bumped in guardrail pr Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> (cherry picked from commit 22fcaf0) Co-authored-by: Sicheng Song <sicheng.song@outlook.com>
1 parent 4dbeb28 commit c5c4e61

File tree

10 files changed

+345
-28
lines changed

10 files changed

+345
-28
lines changed

‎common/src/main/java/org/opensearch/ml/common/CommonValue.java

+3
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,9 @@ public class CommonValue {
232232
+ MODEL_MAX_LENGTH_FIELD + "\":{\"type\":\"integer\"},\""
233233
+ ALL_CONFIG_FIELD + "\":{\"type\":\"text\"}}},\n"
234234
+ " \""
235+
+ MLModel.DEPLOY_SETTING_FIELD
236+
+ "\" : {\"type\": \"flat_object\"},\n"
237+
+ " \""
235238
+ MLModel.IS_ENABLED_FIELD
236239
+ "\" : {\"type\": \"boolean\"},\n"
237240
+ " \""

‎common/src/main/java/org/opensearch/ml/common/MLModel.java

+22
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.opensearch.core.xcontent.XContentParser;
1818
import org.opensearch.ml.common.connector.Connector;
1919
import org.opensearch.ml.common.model.Guardrails;
20+
import org.opensearch.ml.common.model.MLDeploySetting;
2021
import org.opensearch.ml.common.model.MLModelConfig;
2122
import org.opensearch.ml.common.controller.MLRateLimiter;
2223
import org.opensearch.ml.common.model.MLModelFormat;
@@ -61,6 +62,7 @@ public class MLModel implements ToXContentObject {
6162
public static final String RATE_LIMITER_FIELD = "rate_limiter";
6263
public static final String IS_CONTROLLER_ENABLED_FIELD = "is_controller_enabled";
6364
public static final String MODEL_CONFIG_FIELD = "model_config";
65+
public static final String DEPLOY_SETTING_FIELD = "deploy_setting"; // optional
6466
public static final String CREATED_TIME_FIELD = "created_time";
6567
public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
6668
@Deprecated
@@ -102,6 +104,7 @@ public class MLModel implements ToXContentObject {
102104
private Long modelContentSizeInBytes;
103105
private String modelContentHash;
104106
private MLModelConfig modelConfig;
107+
private MLDeploySetting deploySetting;
105108
private Boolean isEnabled;
106109
private Boolean isControllerEnabled;
107110
private MLRateLimiter rateLimiter;
@@ -147,6 +150,7 @@ public MLModel(String name,
147150
Boolean isControllerEnabled,
148151
MLRateLimiter rateLimiter,
149152
MLModelConfig modelConfig,
153+
MLDeploySetting deploySetting,
150154
Instant createdTime,
151155
Instant lastUpdateTime,
152156
Instant lastRegisteredTime,
@@ -178,6 +182,7 @@ public MLModel(String name,
178182
this.isControllerEnabled = isControllerEnabled;
179183
this.rateLimiter = rateLimiter;
180184
this.modelConfig = modelConfig;
185+
this.deploySetting = deploySetting;
181186
this.createdTime = createdTime;
182187
this.lastUpdateTime = lastUpdateTime;
183188
this.lastRegisteredTime = lastRegisteredTime;
@@ -226,6 +231,9 @@ public MLModel(StreamInput input) throws IOException {
226231
modelConfig = new TextEmbeddingModelConfig(input);
227232
}
228233
}
234+
if (input.readBoolean()) {
235+
this.deploySetting = new MLDeploySetting(input);
236+
}
229237
isEnabled = input.readOptionalBoolean();
230238
isControllerEnabled = input.readOptionalBoolean();
231239
if (input.readBoolean()) {
@@ -288,6 +296,12 @@ public void writeTo(StreamOutput out) throws IOException {
288296
} else {
289297
out.writeBoolean(false);
290298
}
299+
if (deploySetting != null) {
300+
out.writeBoolean(true);
301+
deploySetting.writeTo(out);
302+
} else {
303+
out.writeBoolean(false);
304+
}
291305
out.writeOptionalBoolean(isEnabled);
292306
out.writeOptionalBoolean(isControllerEnabled);
293307
if (rateLimiter != null) {
@@ -365,6 +379,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
365379
if (modelConfig != null) {
366380
builder.field(MODEL_CONFIG_FIELD, modelConfig);
367381
}
382+
if (deploySetting != null) {
383+
builder.field(DEPLOY_SETTING_FIELD, deploySetting);
384+
}
368385
if (isEnabled != null) {
369386
builder.field(IS_ENABLED_FIELD, isEnabled);
370387
}
@@ -445,6 +462,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
445462
Long modelContentSizeInBytes = null;
446463
String modelContentHash = null;
447464
MLModelConfig modelConfig = null;
465+
MLDeploySetting deploySetting = null;
448466
Boolean isEnabled = null;
449467
Boolean isControllerEnabled = null;
450468
MLRateLimiter rateLimiter = null;
@@ -536,6 +554,9 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
536554
modelConfig = TextEmbeddingModelConfig.parse(parser);
537555
}
538556
break;
557+
case DEPLOY_SETTING_FIELD:
558+
deploySetting = MLDeploySetting.parse(parser);
559+
break;
539560
case IS_ENABLED_FIELD:
540561
isEnabled = parser.booleanValue();
541562
break;
@@ -614,6 +635,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
614635
.modelContentSizeInBytes(modelContentSizeInBytes)
615636
.modelContentHash(modelContentHash)
616637
.modelConfig(modelConfig)
638+
.deploySetting(deploySetting)
617639
.isEnabled(isEnabled)
618640
.isControllerEnabled(isControllerEnabled)
619641
.rateLimiter(rateLimiter)

‎common/src/main/java/org/opensearch/ml/common/model/MLDeploySetting.java

+71
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.model;
7+
8+
import lombok.Builder;
9+
import lombok.Getter;
10+
import lombok.Setter;
11+
import org.opensearch.core.common.io.stream.Writeable;
12+
import org.opensearch.core.common.io.stream.StreamInput;
13+
import org.opensearch.core.common.io.stream.StreamOutput;
14+
import org.opensearch.core.xcontent.ToXContent;
15+
import org.opensearch.core.xcontent.ToXContentObject;
16+
import org.opensearch.core.xcontent.XContentBuilder;
17+
import org.opensearch.core.xcontent.XContentParser;
18+
19+
import java.io.IOException;
20+
21+
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
22+
23+
@Setter
24+
@Getter
25+
public class MLDeploySetting implements ToXContentObject, Writeable {
26+
public static final String IS_AUTO_DEPLOY_ENABLED_FIELD = "is_auto_deploy_enabled";
27+
28+
private Boolean isAutoDeployEnabled;
29+
30+
@Builder(toBuilder = true)
31+
public MLDeploySetting(Boolean isAutoDeployEnabled) {
32+
this.isAutoDeployEnabled = isAutoDeployEnabled;
33+
}
34+
35+
public MLDeploySetting(StreamInput in) throws IOException {
36+
this.isAutoDeployEnabled = in.readOptionalBoolean();
37+
}
38+
39+
@Override
40+
public void writeTo(StreamOutput out) throws IOException {
41+
out.writeOptionalBoolean(isAutoDeployEnabled);
42+
}
43+
44+
public static MLDeploySetting parse(XContentParser parser) throws IOException {
45+
Boolean isAutoDeployEnabled = null;
46+
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
47+
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
48+
String fieldName = parser.currentName();
49+
parser.nextToken();
50+
switch (fieldName) {
51+
case IS_AUTO_DEPLOY_ENABLED_FIELD:
52+
isAutoDeployEnabled = parser.booleanValue();
53+
break;
54+
default:
55+
parser.skipChildren();
56+
break;
57+
}
58+
}
59+
return new MLDeploySetting(isAutoDeployEnabled);
60+
}
61+
62+
@Override
63+
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
64+
builder.startObject();
65+
if (isAutoDeployEnabled != null) {
66+
builder.field(IS_AUTO_DEPLOY_ENABLED_FIELD, isAutoDeployEnabled);
67+
}
68+
builder.endObject();
69+
return builder;
70+
}
71+
}

‎common/src/main/java/org/opensearch/ml/common/transport/model/MLUpdateModelInput.java

+28-6
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,12 @@
1717
import org.opensearch.core.xcontent.XContentParser;
1818
import org.opensearch.ml.common.connector.Connector;
1919
import org.opensearch.ml.common.model.Guardrails;
20+
import org.opensearch.ml.common.model.MLDeploySetting;
2021
import org.opensearch.ml.common.model.MLModelConfig;
2122
import org.opensearch.ml.common.controller.MLRateLimiter;
2223
import org.opensearch.ml.common.model.TextEmbeddingModelConfig;
2324
import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
25+
import org.opensearch.ml.common.transport.register.MLRegisterModelInput;
2426

2527
import java.io.IOException;
2628
import java.time.Instant;
@@ -39,6 +41,7 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
3941
public static final String IS_ENABLED_FIELD = "is_enabled"; // optional
4042
public static final String RATE_LIMITER_FIELD = "rate_limiter"; // optional
4143
public static final String MODEL_CONFIG_FIELD = "model_config"; // optional
44+
public static final String DEPLOY_SETTING_FIELD = "deploy_setting"; // optional
4245
public static final String UPDATED_CONNECTOR_FIELD = "updated_connector"; // passively set when updating the
4346
// internal connector
4447
public static final String CONNECTOR_ID_FIELD = "connector_id"; // optional
@@ -47,8 +50,6 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
4750
// request
4851
public static final String GUARDRAILS_FIELD = "guardrails";
4952

50-
private static final Version MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS = Version.V_2_13_0;
51-
5253
@Getter
5354
private String modelId;
5455
private String description;
@@ -58,6 +59,7 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
5859
private Boolean isEnabled;
5960
private MLRateLimiter rateLimiter;
6061
private MLModelConfig modelConfig;
62+
private MLDeploySetting deploySetting;
6163
private Connector updatedConnector;
6264
private String connectorId;
6365
private MLCreateConnectorInput connector;
@@ -66,7 +68,7 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
6668

6769
@Builder(toBuilder = true)
6870
public MLUpdateModelInput(String modelId, String description, String version, String name, String modelGroupId,
69-
Boolean isEnabled, MLRateLimiter rateLimiter, MLModelConfig modelConfig,
71+
Boolean isEnabled, MLRateLimiter rateLimiter, MLModelConfig modelConfig, MLDeploySetting deploySetting,
7072
Connector updatedConnector, String connectorId, MLCreateConnectorInput connector, Instant lastUpdateTime, Guardrails guardrails) {
7173
this.modelId = modelId;
7274
this.description = description;
@@ -76,6 +78,7 @@ public MLUpdateModelInput(String modelId, String description, String version, St
7678
this.isEnabled = isEnabled;
7779
this.rateLimiter = rateLimiter;
7880
this.modelConfig = modelConfig;
81+
this.deploySetting = deploySetting;
7982
this.updatedConnector = updatedConnector;
8083
this.connectorId = connectorId;
8184
this.connector = connector;
@@ -105,10 +108,13 @@ public MLUpdateModelInput(StreamInput in) throws IOException {
105108
connector = new MLCreateConnectorInput(in);
106109
}
107110
lastUpdateTime = in.readOptionalInstant();
108-
if (streamInputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS)) {
111+
if (streamInputVersion.onOrAfter(MLRegisterModelInput.MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS_AND_AUTO_DEPLOY)) {
109112
if (in.readBoolean()) {
110113
this.guardrails = new Guardrails(in);
111114
}
115+
if (in.readBoolean()) {
116+
this.deploySetting = new MLDeploySetting(in);
117+
}
112118
}
113119
}
114120

@@ -137,6 +143,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
137143
if (modelConfig != null) {
138144
builder.field(MODEL_CONFIG_FIELD, modelConfig);
139145
}
146+
if (deploySetting != null) {
147+
builder.field(DEPLOY_SETTING_FIELD, deploySetting);
148+
}
140149
if (updatedConnector != null) {
141150
builder.field(UPDATED_CONNECTOR_FIELD, updatedConnector);
142151
}
@@ -180,6 +189,9 @@ public XContentBuilder toXContentForUpdateRequestDoc(XContentBuilder builder, Pa
180189
if (modelConfig != null) {
181190
builder.field(MODEL_CONFIG_FIELD, modelConfig);
182191
}
192+
if (deploySetting != null) {
193+
builder.field(DEPLOY_SETTING_FIELD, deploySetting);
194+
}
183195
// Notice that we serialize the updatedConnector to the connector field, in order to be compatible with original internal connector field format.
184196
if (updatedConnector != null) {
185197
builder.field(CONNECTOR_FIELD, updatedConnector);
@@ -232,13 +244,19 @@ public void writeTo(StreamOutput out) throws IOException {
232244
out.writeBoolean(false);
233245
}
234246
out.writeOptionalInstant(lastUpdateTime);
235-
if (streamOutputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS)) {
247+
if (streamOutputVersion.onOrAfter(MLRegisterModelInput.MINIMAL_SUPPORTED_VERSION_FOR_GUARDRAILS_AND_AUTO_DEPLOY)) {
236248
if (guardrails != null) {
237249
out.writeBoolean(true);
238250
guardrails.writeTo(out);
239251
} else {
240252
out.writeBoolean(false);
241253
}
254+
if (deploySetting != null) {
255+
out.writeBoolean(true);
256+
deploySetting.writeTo(out);
257+
} else {
258+
out.writeBoolean(false);
259+
}
242260
}
243261
}
244262

@@ -251,6 +269,7 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
251269
Boolean isEnabled = null;
252270
MLRateLimiter rateLimiter = null;
253271
MLModelConfig modelConfig = null;
272+
MLDeploySetting deploySetting = null;
254273
Connector updatedConnector = null;
255274
String connectorId = null;
256275
MLCreateConnectorInput connector = null;
@@ -280,6 +299,9 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
280299
case MODEL_CONFIG_FIELD:
281300
modelConfig = TextEmbeddingModelConfig.parse(parser);
282301
break;
302+
case DEPLOY_SETTING_FIELD:
303+
deploySetting = MLDeploySetting.parse(parser);
304+
break;
283305
case CONNECTOR_ID_FIELD:
284306
connectorId = parser.text();
285307
break;
@@ -297,6 +319,6 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
297319
// Model ID can only be set through RestRequest. Model version can only be set
298320
// automatically.
299321
return new MLUpdateModelInput(modelId, description, version, name, modelGroupId, isEnabled, rateLimiter,
300-
modelConfig, updatedConnector, connectorId, connector, lastUpdateTime, guardrails);
322+
modelConfig, deploySetting, updatedConnector, connectorId, connector, lastUpdateTime, guardrails);
301323
}
302324
}

0 commit comments

Comments
 (0)