Skip to content

Commit 0dddbb5

Browse files
authored
[ML] Parse and index inference model (#48016) (#48152)
This adds parsing an inference model as a possible result of the analytics process. When we do parse such a model we persist a `TrainedModelConfig` into the inference index that contains additional metadata derived from the running job.
1 parent 74812f7 commit 0dddbb5

File tree

23 files changed

+420
-228
lines changed

23 files changed

+420
-228
lines changed

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/TrainedModelConfig.java

Lines changed: 45 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -30,21 +30,21 @@
3030
import java.io.IOException;
3131
import java.time.Instant;
3232
import java.util.Collections;
33+
import java.util.List;
3334
import java.util.Map;
3435
import java.util.Objects;
3536

3637
public class TrainedModelConfig implements ToXContentObject {
3738

38-
public static final String NAME = "trained_model_doc";
39+
public static final String NAME = "trained_model_config";
3940

4041
public static final ParseField MODEL_ID = new ParseField("model_id");
4142
public static final ParseField CREATED_BY = new ParseField("created_by");
4243
public static final ParseField VERSION = new ParseField("version");
4344
public static final ParseField DESCRIPTION = new ParseField("description");
44-
public static final ParseField CREATED_TIME = new ParseField("created_time");
45-
public static final ParseField MODEL_VERSION = new ParseField("model_version");
45+
public static final ParseField CREATE_TIME = new ParseField("create_time");
4646
public static final ParseField DEFINITION = new ParseField("definition");
47-
public static final ParseField MODEL_TYPE = new ParseField("model_type");
47+
public static final ParseField TAGS = new ParseField("tags");
4848
public static final ParseField METADATA = new ParseField("metadata");
4949

5050
public static final ObjectParser<Builder, Void> PARSER = new ObjectParser<>(NAME,
@@ -55,16 +55,15 @@ public class TrainedModelConfig implements ToXContentObject {
5555
PARSER.declareString(TrainedModelConfig.Builder::setCreatedBy, CREATED_BY);
5656
PARSER.declareString(TrainedModelConfig.Builder::setVersion, VERSION);
5757
PARSER.declareString(TrainedModelConfig.Builder::setDescription, DESCRIPTION);
58-
PARSER.declareField(TrainedModelConfig.Builder::setCreatedTime,
59-
(p, c) -> TimeUtil.parseTimeFieldToInstant(p, CREATED_TIME.getPreferredName()),
60-
CREATED_TIME,
58+
PARSER.declareField(TrainedModelConfig.Builder::setCreateTime,
59+
(p, c) -> TimeUtil.parseTimeFieldToInstant(p, CREATE_TIME.getPreferredName()),
60+
CREATE_TIME,
6161
ObjectParser.ValueType.VALUE);
62-
PARSER.declareLong(TrainedModelConfig.Builder::setModelVersion, MODEL_VERSION);
63-
PARSER.declareString(TrainedModelConfig.Builder::setModelType, MODEL_TYPE);
64-
PARSER.declareObject(TrainedModelConfig.Builder::setMetadata, (p, c) -> p.map(), METADATA);
6562
PARSER.declareObject(TrainedModelConfig.Builder::setDefinition,
6663
(p, c) -> TrainedModelDefinition.fromXContent(p),
6764
DEFINITION);
65+
PARSER.declareStringArray(TrainedModelConfig.Builder::setTags, TAGS);
66+
PARSER.declareObject(TrainedModelConfig.Builder::setMetadata, (p, c) -> p.map(), METADATA);
6867
}
6968

7069
public static TrainedModelConfig.Builder fromXContent(XContentParser parser) throws IOException {
@@ -75,30 +74,27 @@ public static TrainedModelConfig.Builder fromXContent(XContentParser parser) thr
7574
private final String createdBy;
7675
private final Version version;
7776
private final String description;
78-
private final Instant createdTime;
79-
private final Long modelVersion;
80-
private final String modelType;
81-
private final Map<String, Object> metadata;
77+
private final Instant createTime;
8278
private final TrainedModelDefinition definition;
79+
private final List<String> tags;
80+
private final Map<String, Object> metadata;
8381

8482
TrainedModelConfig(String modelId,
8583
String createdBy,
8684
Version version,
8785
String description,
88-
Instant createdTime,
89-
Long modelVersion,
90-
String modelType,
86+
Instant createTime,
9187
TrainedModelDefinition definition,
88+
List<String> tags,
9289
Map<String, Object> metadata) {
9390
this.modelId = modelId;
9491
this.createdBy = createdBy;
9592
this.version = version;
96-
this.createdTime = Instant.ofEpochMilli(createdTime.toEpochMilli());
97-
this.modelType = modelType;
93+
this.createTime = Instant.ofEpochMilli(createTime.toEpochMilli());
9894
this.definition = definition;
9995
this.description = description;
96+
this.tags = tags == null ? null : Collections.unmodifiableList(tags);
10097
this.metadata = metadata == null ? null : Collections.unmodifiableMap(metadata);
101-
this.modelVersion = modelVersion;
10298
}
10399

104100
public String getModelId() {
@@ -117,16 +113,12 @@ public String getDescription() {
117113
return description;
118114
}
119115

120-
public Instant getCreatedTime() {
121-
return createdTime;
116+
public Instant getCreateTime() {
117+
return createTime;
122118
}
123119

124-
public Long getModelVersion() {
125-
return modelVersion;
126-
}
127-
128-
public String getModelType() {
129-
return modelType;
120+
public List<String> getTags() {
121+
return tags;
130122
}
131123

132124
public Map<String, Object> getMetadata() {
@@ -156,18 +148,15 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
156148
if (description != null) {
157149
builder.field(DESCRIPTION.getPreferredName(), description);
158150
}
159-
if (createdTime != null) {
160-
builder.timeField(CREATED_TIME.getPreferredName(), CREATED_TIME.getPreferredName() + "_string", createdTime.toEpochMilli());
161-
}
162-
if (modelVersion != null) {
163-
builder.field(MODEL_VERSION.getPreferredName(), modelVersion);
164-
}
165-
if (modelType != null) {
166-
builder.field(MODEL_TYPE.getPreferredName(), modelType);
151+
if (createTime != null) {
152+
builder.timeField(CREATE_TIME.getPreferredName(), CREATE_TIME.getPreferredName() + "_string", createTime.toEpochMilli());
167153
}
168154
if (definition != null) {
169155
builder.field(DEFINITION.getPreferredName(), definition);
170156
}
157+
if (tags != null) {
158+
builder.field(TAGS.getPreferredName(), tags);
159+
}
171160
if (metadata != null) {
172161
builder.field(METADATA.getPreferredName(), metadata);
173162
}
@@ -189,10 +178,9 @@ public boolean equals(Object o) {
189178
Objects.equals(createdBy, that.createdBy) &&
190179
Objects.equals(version, that.version) &&
191180
Objects.equals(description, that.description) &&
192-
Objects.equals(createdTime, that.createdTime) &&
193-
Objects.equals(modelVersion, that.modelVersion) &&
194-
Objects.equals(modelType, that.modelType) &&
181+
Objects.equals(createTime, that.createTime) &&
195182
Objects.equals(definition, that.definition) &&
183+
Objects.equals(tags, that.tags) &&
196184
Objects.equals(metadata, that.metadata);
197185
}
198186

@@ -201,12 +189,11 @@ public int hashCode() {
201189
return Objects.hash(modelId,
202190
createdBy,
203191
version,
204-
createdTime,
205-
modelType,
192+
createTime,
206193
definition,
207194
description,
208-
metadata,
209-
modelVersion);
195+
tags,
196+
metadata);
210197
}
211198

212199

@@ -216,11 +203,10 @@ public static class Builder {
216203
private String createdBy;
217204
private Version version;
218205
private String description;
219-
private Instant createdTime;
220-
private Long modelVersion;
221-
private String modelType;
206+
private Instant createTime;
222207
private Map<String, Object> metadata;
223-
private TrainedModelDefinition.Builder definition;
208+
private List<String> tags;
209+
private TrainedModelDefinition definition;
224210

225211
public Builder setModelId(String modelId) {
226212
this.modelId = modelId;
@@ -246,18 +232,13 @@ public Builder setDescription(String description) {
246232
return this;
247233
}
248234

249-
private Builder setCreatedTime(Instant createdTime) {
250-
this.createdTime = createdTime;
251-
return this;
252-
}
253-
254-
public Builder setModelVersion(Long modelVersion) {
255-
this.modelVersion = modelVersion;
235+
private Builder setCreateTime(Instant createTime) {
236+
this.createTime = createTime;
256237
return this;
257238
}
258239

259-
public Builder setModelType(String modelType) {
260-
this.modelType = modelType;
240+
public Builder setTags(List<String> tags) {
241+
this.tags = tags;
261242
return this;
262243
}
263244

@@ -267,6 +248,11 @@ public Builder setMetadata(Map<String, Object> metadata) {
267248
}
268249

269250
public Builder setDefinition(TrainedModelDefinition.Builder definition) {
251+
this.definition = definition == null ? null : definition.build();
252+
return this;
253+
}
254+
255+
public Builder setDefinition(TrainedModelDefinition definition) {
270256
this.definition = definition;
271257
return this;
272258
}
@@ -277,10 +263,9 @@ public TrainedModelConfig build() {
277263
createdBy,
278264
version,
279265
description,
280-
createdTime,
281-
modelVersion,
282-
modelType,
283-
definition == null ? null : definition.build(),
266+
createTime,
267+
definition,
268+
tags,
284269
metadata);
285270
}
286271
}

client/rest-high-level/src/test/java/org/elasticsearch/client/ml/inference/TrainedModelConfigTests.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
import java.util.Collections;
3232
import java.util.List;
3333
import java.util.function.Predicate;
34+
import java.util.stream.Collectors;
35+
import java.util.stream.Stream;
3436

3537

3638
public class TrainedModelConfigTests extends AbstractXContentTestCase<TrainedModelConfig> {
@@ -58,9 +60,9 @@ protected TrainedModelConfig createTestInstance() {
5860
Version.CURRENT,
5961
randomBoolean() ? null : randomAlphaOfLength(100),
6062
Instant.ofEpochMilli(randomNonNegativeLong()),
61-
randomBoolean() ? null : randomNonNegativeLong(),
62-
randomAlphaOfLength(10),
6363
randomBoolean() ? null : TrainedModelDefinitionTests.createRandomBuilder().build(),
64+
randomBoolean() ? null :
65+
Stream.generate(() -> randomAlphaOfLength(10)).limit(randomIntBetween(0, 5)).collect(Collectors.toList()),
6466
randomBoolean() ? null : Collections.singletonMap(randomAlphaOfLength(10), randomAlphaOfLength(10)));
6567
}
6668

0 commit comments

Comments
 (0)