Skip to content

Commit d0895bb

Browse files
b4sjoob4sjoo
authored
Adding UT coverage for in-cache update and fine-tuning throttling feature (opensearch-project#1837)
1 parent a14521d commit d0895bb

File tree

88 files changed

+3232
-2842
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

88 files changed

+3232
-2842
lines changed

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

+422-420
Large diffs are not rendered by default.

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

+58-57
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,8 @@ public class MLModel implements ToXContentObject {
5656

5757
// Model level quota and throttling control
5858
public static final String IS_ENABLED_FIELD = "is_enabled";
59-
public static final String MODEL_RATE_LIMITER_CONFIG_FIELD = "model_rate_limiter_config";
60-
public static final String IS_MODEL_CONTROLLER_ENABLED_FIELD = "is_model_controller_enabled";
59+
public static final String RATE_LIMITER_FIELD = "rate_limiter";
60+
public static final String IS_CONTROLLER_ENABLED_FIELD = "is_controller_enabled";
6161
public static final String MODEL_CONFIG_FIELD = "model_config";
6262
public static final String CREATED_TIME_FIELD = "created_time";
6363
public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
@@ -100,8 +100,8 @@ public class MLModel implements ToXContentObject {
100100
private String modelContentHash;
101101
private MLModelConfig modelConfig;
102102
private Boolean isEnabled;
103-
private Boolean isModelControllerEnabled;
104-
private MLRateLimiter modelRateLimiterConfig;
103+
private Boolean isControllerEnabled;
104+
private MLRateLimiter rateLimiter;
105105
private Instant createdTime;
106106
private Instant lastUpdateTime;
107107
private Instant lastRegisteredTime;
@@ -120,7 +120,8 @@ public class MLModel implements ToXContentObject {
120120
private String[] planningWorkerNodes; // plan to deploy model to these nodes
121121
private boolean deployToAllNodes;
122122

123-
//is domain manager creates any special hidden model in the cluster this status will be true. Otherwise,
123+
// is domain manager creates any special hidden model in the cluster this status
124+
// will be true. Otherwise,
124125
// False by default
125126
private Boolean isHidden;
126127
@Setter
@@ -129,35 +130,35 @@ public class MLModel implements ToXContentObject {
129130

130131
@Builder(toBuilder = true)
131132
public MLModel(String name,
132-
String modelGroupId,
133-
FunctionName algorithm,
134-
String version,
135-
String content,
136-
User user,
137-
String description,
138-
MLModelFormat modelFormat,
139-
MLModelState modelState,
140-
Long modelContentSizeInBytes,
141-
String modelContentHash,
142-
Boolean isEnabled,
143-
Boolean isModelControllerEnabled,
144-
MLRateLimiter modelRateLimiterConfig,
145-
MLModelConfig modelConfig,
146-
Instant createdTime,
147-
Instant lastUpdateTime,
148-
Instant lastRegisteredTime,
149-
Instant lastDeployedTime,
150-
Instant lastUndeployedTime,
151-
Integer autoRedeployRetryTimes,
152-
String modelId, Integer chunkNumber,
153-
Integer totalChunks,
154-
Integer planningWorkerNodeCount,
155-
Integer currentWorkerNodeCount,
156-
String[] planningWorkerNodes,
157-
boolean deployToAllNodes,
158-
Boolean isHidden,
159-
Connector connector,
160-
String connectorId) {
133+
String modelGroupId,
134+
FunctionName algorithm,
135+
String version,
136+
String content,
137+
User user,
138+
String description,
139+
MLModelFormat modelFormat,
140+
MLModelState modelState,
141+
Long modelContentSizeInBytes,
142+
String modelContentHash,
143+
Boolean isEnabled,
144+
Boolean isControllerEnabled,
145+
MLRateLimiter rateLimiter,
146+
MLModelConfig modelConfig,
147+
Instant createdTime,
148+
Instant lastUpdateTime,
149+
Instant lastRegisteredTime,
150+
Instant lastDeployedTime,
151+
Instant lastUndeployedTime,
152+
Integer autoRedeployRetryTimes,
153+
String modelId, Integer chunkNumber,
154+
Integer totalChunks,
155+
Integer planningWorkerNodeCount,
156+
Integer currentWorkerNodeCount,
157+
String[] planningWorkerNodes,
158+
boolean deployToAllNodes,
159+
Boolean isHidden,
160+
Connector connector,
161+
String connectorId) {
161162
this.name = name;
162163
this.modelGroupId = modelGroupId;
163164
this.algorithm = algorithm;
@@ -170,8 +171,8 @@ public MLModel(String name,
170171
this.modelContentSizeInBytes = modelContentSizeInBytes;
171172
this.modelContentHash = modelContentHash;
172173
this.isEnabled = isEnabled;
173-
this.isModelControllerEnabled = isModelControllerEnabled;
174-
this.modelRateLimiterConfig = modelRateLimiterConfig;
174+
this.isControllerEnabled = isControllerEnabled;
175+
this.rateLimiter = rateLimiter;
175176
this.modelConfig = modelConfig;
176177
this.createdTime = createdTime;
177178
this.lastUpdateTime = lastUpdateTime;
@@ -191,7 +192,7 @@ public MLModel(String name,
191192
this.connectorId = connectorId;
192193
}
193194

194-
public MLModel(StreamInput input) throws IOException{
195+
public MLModel(StreamInput input) throws IOException {
195196
name = input.readOptionalString();
196197
algorithm = input.readEnum(FunctionName.class);
197198
version = input.readString();
@@ -219,9 +220,9 @@ public MLModel(StreamInput input) throws IOException{
219220
}
220221
}
221222
isEnabled = input.readOptionalBoolean();
222-
isModelControllerEnabled = input.readOptionalBoolean();
223+
isControllerEnabled = input.readOptionalBoolean();
223224
if (input.readBoolean()) {
224-
modelRateLimiterConfig = new MLRateLimiter(input);
225+
rateLimiter = new MLRateLimiter(input);
225226
}
226227
createdTime = input.readOptionalInstant();
227228
lastUpdateTime = input.readOptionalInstant();
@@ -278,10 +279,10 @@ public void writeTo(StreamOutput out) throws IOException {
278279
out.writeBoolean(false);
279280
}
280281
out.writeOptionalBoolean(isEnabled);
281-
out.writeOptionalBoolean(isModelControllerEnabled);
282-
if (modelRateLimiterConfig != null) {
282+
out.writeOptionalBoolean(isControllerEnabled);
283+
if (rateLimiter != null) {
283284
out.writeBoolean(true);
284-
modelRateLimiterConfig.writeTo(out);
285+
rateLimiter.writeTo(out);
285286
} else {
286287
out.writeBoolean(false);
287288
}
@@ -351,11 +352,11 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
351352
if (isEnabled != null) {
352353
builder.field(IS_ENABLED_FIELD, isEnabled);
353354
}
354-
if (isModelControllerEnabled != null) {
355-
builder.field(IS_MODEL_CONTROLLER_ENABLED_FIELD, isModelControllerEnabled);
355+
if (isControllerEnabled != null) {
356+
builder.field(IS_CONTROLLER_ENABLED_FIELD, isControllerEnabled);
356357
}
357-
if (modelRateLimiterConfig != null) {
358-
builder.field(MODEL_RATE_LIMITER_CONFIG_FIELD, modelRateLimiterConfig);
358+
if (rateLimiter != null) {
359+
builder.field(RATE_LIMITER_FIELD, rateLimiter);
359360
}
360361
if (createdTime != null) {
361362
builder.field(CREATED_TIME_FIELD, createdTime.toEpochMilli());
@@ -426,8 +427,8 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
426427
String modelContentHash = null;
427428
MLModelConfig modelConfig = null;
428429
Boolean isEnabled = null;
429-
Boolean isModelControllerEnabled = null;
430-
MLRateLimiter modelRateLimiterConfig = null;
430+
Boolean isControllerEnabled = null;
431+
MLRateLimiter rateLimiter = null;
431432
Instant createdTime = null;
432433
Instant lastUpdateTime = null;
433434
Instant lastUploadedTime = null;
@@ -516,11 +517,11 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
516517
case IS_ENABLED_FIELD:
517518
isEnabled = parser.booleanValue();
518519
break;
519-
case IS_MODEL_CONTROLLER_ENABLED_FIELD:
520-
isModelControllerEnabled = parser.booleanValue();
520+
case IS_CONTROLLER_ENABLED_FIELD:
521+
isControllerEnabled = parser.booleanValue();
521522
break;
522-
case MODEL_RATE_LIMITER_CONFIG_FIELD:
523-
modelRateLimiterConfig = MLRateLimiter.parse(parser);
523+
case RATE_LIMITER_FIELD:
524+
rateLimiter = MLRateLimiter.parse(parser);
524525
break;
525526
case PLANNING_WORKER_NODE_COUNT_FIELD:
526527
planningWorkerNodeCount = parser.intValue();
@@ -589,13 +590,13 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
589590
.modelContentHash(modelContentHash)
590591
.modelConfig(modelConfig)
591592
.isEnabled(isEnabled)
592-
.isModelControllerEnabled(isModelControllerEnabled)
593-
.modelRateLimiterConfig(modelRateLimiterConfig)
593+
.isControllerEnabled(isControllerEnabled)
594+
.rateLimiter(rateLimiter)
594595
.createdTime(createdTime)
595596
.lastUpdateTime(lastUpdateTime)
596-
.lastRegisteredTime(lastRegisteredTime == null? lastUploadedTime : lastRegisteredTime)
597-
.lastDeployedTime(lastDeployedTime == null? lastLoadedTime : lastDeployedTime)
598-
.lastUndeployedTime(lastUndeployedTime == null? lastUnloadedTime : lastUndeployedTime)
597+
.lastRegisteredTime(lastRegisteredTime == null ? lastUploadedTime : lastRegisteredTime)
598+
.lastDeployedTime(lastDeployedTime == null ? lastLoadedTime : lastDeployedTime)
599+
.lastUndeployedTime(lastUndeployedTime == null ? lastUnloadedTime : lastUndeployedTime)
599600
.modelId(modelId)
600601
.autoRedeployRetryTimes(autoRedeployRetryTimes)
601602
.chunkNumber(chunkNumber)

‎common/src/main/java/org/opensearch/ml/common/controller/MLModelController.java ‎common/src/main/java/org/opensearch/ml/common/controller/MLController.java

+41-38
Original file line numberDiff line numberDiff line change
@@ -29,25 +29,26 @@
2929
import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
3030

3131
@Data
32-
public class MLModelController implements ToXContentObject, Writeable {
32+
public class MLController implements ToXContentObject, Writeable {
3333

3434
public static final String MODEL_ID_FIELD = "model_id"; // mandatory
35-
public static final String USER_RATE_LIMITER_CONFIG = "user_rate_limiter_config";
35+
public static final String USER_RATE_LIMITER = "user_rate_limiter";
3636

3737
@Getter
3838
private String modelId;
39-
// The String is the username field where the MLRateLimiter is its corresponding rate limiter config.
40-
private Map<String, MLRateLimiter> userRateLimiterConfig;
39+
// The String is the username field where the MLRateLimiter is its corresponding
40+
// rate limiter config.
41+
private Map<String, MLRateLimiter> userRateLimiter;
4142

4243
@Builder(toBuilder = true)
43-
public MLModelController(String modelId, Map<String, MLRateLimiter> userRateLimiterConfig) {
44+
public MLController(String modelId, Map<String, MLRateLimiter> userRateLimiter) {
4445
this.modelId = modelId;
45-
this.userRateLimiterConfig = userRateLimiterConfig;
46+
this.userRateLimiter = userRateLimiter;
4647
}
4748

48-
public static MLModelController parse(XContentParser parser) throws IOException {
49+
public static MLController parse(XContentParser parser) throws IOException {
4950
String modelId = null;
50-
Map<String, MLRateLimiter> userRateLimiterConfig = new HashMap<>();
51+
Map<String, MLRateLimiter> userRateLimiter = new HashMap<>();
5152

5253
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
5354
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -58,15 +59,16 @@ public static MLModelController parse(XContentParser parser) throws IOException
5859
case MODEL_ID_FIELD:
5960
modelId = parser.text();
6061
break;
61-
case USER_RATE_LIMITER_CONFIG:
62-
Map<String, String> userRateLimiterConfigStringMap = getParameterMap(parser.map());
63-
userRateLimiterConfigStringMap.forEach((user, rateLimiterString) -> {
62+
case USER_RATE_LIMITER:
63+
Map<String, String> userRateLimiterStringMap = getParameterMap(parser.map());
64+
userRateLimiterStringMap.forEach((user, rateLimiterString) -> {
6465
try {
65-
XContentParser rateLimiterParser = XContentType.JSON.xContent().createParser(NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, rateLimiterString);
66+
XContentParser rateLimiterParser = XContentType.JSON.xContent().createParser(
67+
NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, rateLimiterString);
6668
rateLimiterParser.nextToken();
6769
MLRateLimiter rateLimiter = MLRateLimiter.parse(rateLimiterParser);
6870
if (!rateLimiter.isEmpty()) {
69-
userRateLimiterConfig.put(user, rateLimiter);
71+
userRateLimiter.put(user, rateLimiter);
7072
}
7173
} catch (IOException e) {
7274
throw new RuntimeException(e);
@@ -79,22 +81,23 @@ public static MLModelController parse(XContentParser parser) throws IOException
7981
}
8082
}
8183
// Model ID can only be set through RestRequest.
82-
return new MLModelController(modelId, userRateLimiterConfig);
84+
return new MLController(modelId, userRateLimiter);
8385
}
8486

85-
public MLModelController(StreamInput in) throws IOException{
87+
public MLController(StreamInput in) throws IOException {
8688
modelId = in.readString();
8789
if (in.readBoolean()) {
88-
userRateLimiterConfig = in.readMap(StreamInput::readString, MLRateLimiter::new);
90+
userRateLimiter = in.readMap(StreamInput::readString, MLRateLimiter::new);
8991
}
9092
}
9193

9294
@Override
9395
public void writeTo(StreamOutput out) throws IOException {
9496
out.writeString(modelId);
95-
if (userRateLimiterConfig != null) {
97+
if (userRateLimiter != null) {
9698
out.writeBoolean(true);
97-
out.writeMap(userRateLimiterConfig, StreamOutput::writeString, (streamOutput, rateLimiter) -> rateLimiter.writeTo(streamOutput));
99+
out.writeMap(userRateLimiter, StreamOutput::writeString,
100+
(streamOutput, rateLimiter) -> rateLimiter.writeTo(streamOutput));
98101
} else {
99102
out.writeBoolean(false);
100103
}
@@ -104,28 +107,28 @@ public void writeTo(StreamOutput out) throws IOException {
104107
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
105108
builder.startObject();
106109
builder.field(MODEL_ID_FIELD, modelId);
107-
if (userRateLimiterConfig != null) {
108-
builder.field(USER_RATE_LIMITER_CONFIG, userRateLimiterConfig);
110+
if (userRateLimiter != null) {
111+
builder.field(USER_RATE_LIMITER, userRateLimiter);
109112
}
110113
builder.endObject();
111114
return builder;
112115
}
113116

114-
115117
/**
116-
* Checks if a deployment is required after updating the MLModelController.
118+
* Checks if a deployment is required after updating the MLController.
117119
*
118-
* @param updateContent The updated MLModelController object.
120+
* @param updateContent The updated MLController object.
119121
* @return True if a deployment is required, false otherwise.
120122
*/
121-
public boolean isDeployRequiredAfterUpdate(MLModelController updateContent) {
122-
if (updateContent != null && updateContent.getUserRateLimiterConfig() != null && !updateContent.getUserRateLimiterConfig().isEmpty()) {
123-
Map<String, MLRateLimiter> updateUserRateLimiterConfig = updateContent.getUserRateLimiterConfig();
124-
for (Map.Entry<String, MLRateLimiter> entry : updateUserRateLimiterConfig.entrySet()) {
123+
public boolean isDeployRequiredAfterUpdate(MLController updateContent) {
124+
if (updateContent != null && updateContent.getUserRateLimiter() != null
125+
&& !updateContent.getUserRateLimiter().isEmpty()) {
126+
Map<String, MLRateLimiter> updateUserRateLimiter = updateContent.getUserRateLimiter();
127+
for (Map.Entry<String, MLRateLimiter> entry : updateUserRateLimiter.entrySet()) {
125128
String newUser = entry.getKey();
126129
MLRateLimiter newRateLimiter = entry.getValue();
127-
if (this.userRateLimiterConfig.containsKey(newUser)) {
128-
MLRateLimiter oldRateLimiter = this.userRateLimiterConfig.get(newUser);
130+
if (this.userRateLimiter.containsKey(newUser)) {
131+
MLRateLimiter oldRateLimiter = this.userRateLimiter.get(newUser);
129132
if (MLRateLimiter.isDeployRequiredAfterUpdate(oldRateLimiter, newRateLimiter)) {
130133
return true;
131134
}
@@ -139,16 +142,16 @@ public boolean isDeployRequiredAfterUpdate(MLModelController updateContent) {
139142
return false;
140143
}
141144

142-
public void update(MLModelController updateContent) {
143-
Map<String, MLRateLimiter> updateUserRateLimiterConfig = updateContent.getUserRateLimiterConfig();
144-
if (updateUserRateLimiterConfig != null && !updateUserRateLimiterConfig.isEmpty()) {
145-
updateUserRateLimiterConfig.forEach((user, rateLimiter) -> {
145+
public void update(MLController updateContent) {
146+
Map<String, MLRateLimiter> updateUserRateLimiter = updateContent.getUserRateLimiter();
147+
if (updateUserRateLimiter != null && !updateUserRateLimiter.isEmpty()) {
148+
updateUserRateLimiter.forEach((user, rateLimiter) -> {
146149
// rateLimiter can't be null due to parsing exception
147-
if (this.userRateLimiterConfig.containsKey(user)) {
148-
this.userRateLimiterConfig.get(user).update(rateLimiter);
149-
} else {
150-
this.userRateLimiterConfig.put(user, rateLimiter);
151-
}
150+
if (this.userRateLimiter.containsKey(user)) {
151+
this.userRateLimiter.get(user).update(rateLimiter);
152+
} else {
153+
this.userRateLimiter.put(user, rateLimiter);
154+
}
152155
});
153156
}
154157
}

0 commit comments

Comments
 (0)