Skip to content

Commit c2946d1

Browse files
benwtrentbenwtrent
authored
[ML][Inference] Adding _stats endpoint for inference (#48492)
[ML][Inference] Adding _stats endpoint for inference. Initially only contains ingest stats and pipeline counts.
1 parent 5e503fd commit c2946d1

File tree

12 files changed

+1277
-7
lines changed

12 files changed

+1277
-7
lines changed

‎server/src/main/java/org/elasticsearch/ingest/IngestStats.java‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import java.util.HashMap;
3333
import java.util.List;
3434
import java.util.Map;
35+
import java.util.Objects;
3536
import java.util.concurrent.TimeUnit;
3637

3738
public class IngestStats implements Writeable, ToXContentFragment {
@@ -135,6 +136,21 @@ public Map<String, List<ProcessorStat>> getProcessorStats() {
135136
return processorStats;
136137
}
137138

139+
@Override
140+
public boolean equals(Object o) {
141+
if (this == o) return true;
142+
if (o == null || getClass() != o.getClass()) return false;
143+
IngestStats that = (IngestStats) o;
144+
return Objects.equals(totalStats, that.totalStats)
145+
&& Objects.equals(pipelineStats, that.pipelineStats)
146+
&& Objects.equals(processorStats, that.processorStats);
147+
}
148+
149+
@Override
150+
public int hashCode() {
151+
return Objects.hash(totalStats, pipelineStats, processorStats);
152+
}
153+
138154
public static class Stats implements Writeable, ToXContentFragment {
139155

140156
private final long ingestCount;
@@ -203,6 +219,22 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
203219
builder.field("failed", ingestFailedCount);
204220
return builder;
205221
}
222+
223+
@Override
224+
public boolean equals(Object o) {
225+
if (this == o) return true;
226+
if (o == null || getClass() != o.getClass()) return false;
227+
IngestStats.Stats that = (IngestStats.Stats) o;
228+
return Objects.equals(ingestCount, that.ingestCount)
229+
&& Objects.equals(ingestTimeInMillis, that.ingestTimeInMillis)
230+
&& Objects.equals(ingestFailedCount, that.ingestFailedCount)
231+
&& Objects.equals(ingestCurrent, that.ingestCurrent);
232+
}
233+
234+
@Override
235+
public int hashCode() {
236+
return Objects.hash(ingestCount, ingestTimeInMillis, ingestFailedCount, ingestCurrent);
237+
}
206238
}
207239

208240
/**
@@ -255,6 +287,20 @@ public String getPipelineId() {
255287
public Stats getStats() {
256288
return stats;
257289
}
290+
291+
@Override
292+
public boolean equals(Object o) {
293+
if (this == o) return true;
294+
if (o == null || getClass() != o.getClass()) return false;
295+
IngestStats.PipelineStat that = (IngestStats.PipelineStat) o;
296+
return Objects.equals(pipelineId, that.pipelineId)
297+
&& Objects.equals(stats, that.stats);
298+
}
299+
300+
@Override
301+
public int hashCode() {
302+
return Objects.hash(pipelineId, stats);
303+
}
258304
}
259305

260306
/**
@@ -276,5 +322,20 @@ public String getName() {
276322
public Stats getStats() {
277323
return stats;
278324
}
325+
326+
327+
@Override
328+
public boolean equals(Object o) {
329+
if (this == o) return true;
330+
if (o == null || getClass() != o.getClass()) return false;
331+
IngestStats.ProcessorStat that = (IngestStats.ProcessorStat) o;
332+
return Objects.equals(name, that.name)
333+
&& Objects.equals(stats, that.stats);
334+
}
335+
336+
@Override
337+
public int hashCode() {
338+
return Objects.hash(name, stats);
339+
}
279340
}
280341
}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/XPackClientPlugin.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@
7777
import org.elasticsearch.xpack.core.ml.action.DeleteForecastAction;
7878
import org.elasticsearch.xpack.core.ml.action.DeleteJobAction;
7979
import org.elasticsearch.xpack.core.ml.action.DeleteModelSnapshotAction;
80+
import org.elasticsearch.xpack.core.ml.action.DeleteTrainedModelAction;
8081
import org.elasticsearch.xpack.core.ml.action.EstimateMemoryUsageAction;
8182
import org.elasticsearch.xpack.core.ml.action.EvaluateDataFrameAction;
8283
import org.elasticsearch.xpack.core.ml.action.FinalizeJobExecutionAction;
@@ -98,6 +99,8 @@
9899
import org.elasticsearch.xpack.core.ml.action.GetModelSnapshotsAction;
99100
import org.elasticsearch.xpack.core.ml.action.GetOverallBucketsAction;
100101
import org.elasticsearch.xpack.core.ml.action.GetRecordsAction;
102+
import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsAction;
103+
import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsStatsAction;
101104
import org.elasticsearch.xpack.core.ml.action.InferModelAction;
102105
import org.elasticsearch.xpack.core.ml.action.IsolateDatafeedAction;
103106
import org.elasticsearch.xpack.core.ml.action.KillProcessAction;
@@ -344,6 +347,9 @@ public List<ActionType<? extends ActionResponse>> getClientActions() {
344347
EvaluateDataFrameAction.INSTANCE,
345348
EstimateMemoryUsageAction.INSTANCE,
346349
InferModelAction.INSTANCE,
350+
GetTrainedModelsAction.INSTANCE,
351+
DeleteTrainedModelAction.INSTANCE,
352+
GetTrainedModelsStatsAction.INSTANCE,
347353
// security
348354
ClearRealmCacheAction.INSTANCE,
349355
ClearRolesCacheAction.INSTANCE,

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetTrainedModelsStatsAction.java‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
package org.elasticsearch.xpack.core.ml.action;
7+
8+
import org.elasticsearch.ElasticsearchException;
9+
import org.elasticsearch.action.ActionRequestBuilder;
10+
import org.elasticsearch.action.ActionType;
11+
import org.elasticsearch.client.ElasticsearchClient;
12+
import org.elasticsearch.common.ParseField;
13+
import org.elasticsearch.common.io.stream.StreamInput;
14+
import org.elasticsearch.common.io.stream.StreamOutput;
15+
import org.elasticsearch.common.io.stream.Writeable;
16+
import org.elasticsearch.common.xcontent.ToXContentObject;
17+
import org.elasticsearch.common.xcontent.XContentBuilder;
18+
import org.elasticsearch.ingest.IngestStats;
19+
import org.elasticsearch.xpack.core.action.AbstractGetResourcesRequest;
20+
import org.elasticsearch.xpack.core.action.AbstractGetResourcesResponse;
21+
import org.elasticsearch.xpack.core.action.util.QueryPage;
22+
import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
23+
24+
import java.io.IOException;
25+
import java.util.ArrayList;
26+
import java.util.Collections;
27+
import java.util.List;
28+
import java.util.Map;
29+
import java.util.Objects;
30+
import java.util.Set;
31+
32+
public class GetTrainedModelsStatsAction extends ActionType<GetTrainedModelsStatsAction.Response> {
33+
34+
public static final GetTrainedModelsStatsAction INSTANCE = new GetTrainedModelsStatsAction();
35+
public static final String NAME = "cluster:monitor/xpack/ml/inference/stats/get";
36+
37+
public static final ParseField MODEL_ID = new ParseField("model_id");
38+
public static final ParseField PIPELINE_COUNT = new ParseField("pipeline_count");
39+
40+
private GetTrainedModelsStatsAction() {
41+
super(NAME, GetTrainedModelsStatsAction.Response::new);
42+
}
43+
44+
public static class Request extends AbstractGetResourcesRequest {
45+
46+
public static final ParseField ALLOW_NO_MATCH = new ParseField("allow_no_match");
47+
48+
public Request() {
49+
setAllowNoResources(true);
50+
}
51+
52+
public Request(String id) {
53+
setResourceId(id);
54+
setAllowNoResources(true);
55+
}
56+
57+
public Request(StreamInput in) throws IOException {
58+
super(in);
59+
}
60+
61+
@Override
62+
public String getResourceIdField() {
63+
return TrainedModelConfig.MODEL_ID.getPreferredName();
64+
}
65+
66+
}
67+
68+
public static class RequestBuilder extends ActionRequestBuilder<Request, Response> {
69+
70+
public RequestBuilder(ElasticsearchClient client, GetTrainedModelsStatsAction action) {
71+
super(client, action, new Request());
72+
}
73+
}
74+
75+
public static class Response extends AbstractGetResourcesResponse<Response.TrainedModelStats> {
76+
77+
public static class TrainedModelStats implements ToXContentObject, Writeable {
78+
private final String modelId;
79+
private final IngestStats ingestStats;
80+
private final int pipelineCount;
81+
82+
private static final IngestStats EMPTY_INGEST_STATS = new IngestStats(new IngestStats.Stats(0, 0, 0, 0),
83+
Collections.emptyList(),
84+
Collections.emptyMap());
85+
86+
public TrainedModelStats(String modelId, IngestStats ingestStats, int pipelineCount) {
87+
this.modelId = Objects.requireNonNull(modelId);
88+
this.ingestStats = ingestStats == null ? EMPTY_INGEST_STATS : ingestStats;
89+
if (pipelineCount < 0) {
90+
throw new ElasticsearchException("[{}] must be a greater than or equal to 0", PIPELINE_COUNT.getPreferredName());
91+
}
92+
this.pipelineCount = pipelineCount;
93+
}
94+
95+
public TrainedModelStats(StreamInput in) throws IOException {
96+
modelId = in.readString();
97+
ingestStats = new IngestStats(in);
98+
pipelineCount = in.readVInt();
99+
}
100+
101+
@Override
102+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
103+
builder.startObject();
104+
builder.field(MODEL_ID.getPreferredName(), modelId);
105+
builder.field(PIPELINE_COUNT.getPreferredName(), pipelineCount);
106+
if (pipelineCount > 0) {
107+
// Ingest stats is a fragment
108+
ingestStats.toXContent(builder, params);
109+
}
110+
builder.endObject();
111+
return builder;
112+
}
113+
114+
@Override
115+
public void writeTo(StreamOutput out) throws IOException {
116+
out.writeString(modelId);
117+
ingestStats.writeTo(out);
118+
out.writeVInt(pipelineCount);
119+
}
120+
121+
@Override
122+
public int hashCode() {
123+
return Objects.hash(modelId, ingestStats, pipelineCount);
124+
}
125+
126+
@Override
127+
public boolean equals(Object obj) {
128+
if (obj == null) {
129+
return false;
130+
}
131+
if (getClass() != obj.getClass()) {
132+
return false;
133+
}
134+
TrainedModelStats other = (TrainedModelStats) obj;
135+
return Objects.equals(this.modelId, other.modelId)
136+
&& Objects.equals(this.ingestStats, other.ingestStats)
137+
&& Objects.equals(this.pipelineCount, other.pipelineCount);
138+
}
139+
}
140+
141+
public static final ParseField RESULTS_FIELD = new ParseField("trained_model_stats");
142+
143+
public Response(StreamInput in) throws IOException {
144+
super(in);
145+
}
146+
147+
public Response(QueryPage<Response.TrainedModelStats> trainedModels) {
148+
super(trainedModels);
149+
}
150+
151+
@Override
152+
protected Reader<Response.TrainedModelStats> getReader() {
153+
return Response.TrainedModelStats::new;
154+
}
155+
156+
public static class Builder {
157+
158+
private long totalModelCount;
159+
private Set<String> expandedIds;
160+
private Map<String, IngestStats> ingestStatsMap;
161+
162+
public Builder setTotalModelCount(long totalModelCount) {
163+
this.totalModelCount = totalModelCount;
164+
return this;
165+
}
166+
167+
public Builder setExpandedIds(Set<String> expandedIds) {
168+
this.expandedIds = expandedIds;
169+
return this;
170+
}
171+
172+
public Set<String> getExpandedIds() {
173+
return this.expandedIds;
174+
}
175+
176+
public Builder setIngestStatsByModelId(Map<String, IngestStats> ingestStatsByModelId) {
177+
this.ingestStatsMap = ingestStatsByModelId;
178+
return this;
179+
}
180+
181+
public Response build() {
182+
List<TrainedModelStats> trainedModelStats = new ArrayList<>(expandedIds.size());
183+
expandedIds.forEach(id -> {
184+
IngestStats ingestStats = ingestStatsMap.get(id);
185+
trainedModelStats.add(new TrainedModelStats(id, ingestStats, ingestStats == null ?
186+
0 :
187+
ingestStats.getPipelineStats().size()));
188+
});
189+
return new Response(new QueryPage<>(trainedModelStats, totalModelCount, RESULTS_FIELD));
190+
}
191+
}
192+
}
193+
194+
}

‎x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetTrainedModelsStatsActionResponseTests.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
package org.elasticsearch.xpack.core.ml.action;
7+
8+
import org.elasticsearch.common.io.stream.Writeable;
9+
import org.elasticsearch.ingest.IngestStats;
10+
import org.elasticsearch.test.AbstractWireSerializingTestCase;
11+
import org.elasticsearch.xpack.core.action.util.QueryPage;
12+
import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsStatsAction.Response;
13+
14+
import java.util.List;
15+
import java.util.function.Function;
16+
import java.util.stream.Collectors;
17+
import java.util.stream.Stream;
18+
19+
20+
public class GetTrainedModelsStatsActionResponseTests extends AbstractWireSerializingTestCase<Response> {
21+
22+
@Override
23+
protected Response createTestInstance() {
24+
int listSize = randomInt(10);
25+
List<Response.TrainedModelStats> trainedModelStats = Stream.generate(() -> randomAlphaOfLength(10))
26+
.limit(listSize).map(id ->
27+
new Response.TrainedModelStats(id,
28+
randomBoolean() ? randomIngestStats() : null,
29+
randomIntBetween(0, 10))
30+
)
31+
.collect(Collectors.toList());
32+
return new Response(new QueryPage<>(trainedModelStats, randomLongBetween(listSize, 1000), Response.RESULTS_FIELD));
33+
}
34+
35+
private IngestStats randomIngestStats() {
36+
List<String> pipelineIds = Stream.generate(()-> randomAlphaOfLength(10))
37+
.limit(randomIntBetween(0, 10))
38+
.collect(Collectors.toList());
39+
return new IngestStats(
40+
new IngestStats.Stats(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong()),
41+
pipelineIds.stream().map(id -> new IngestStats.PipelineStat(id, randomStats())).collect(Collectors.toList()),
42+
pipelineIds.stream().collect(Collectors.toMap(Function.identity(), (v) -> randomProcessorStats())));
43+
}
44+
45+
private IngestStats.Stats randomStats(){
46+
return new IngestStats.Stats(randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong());
47+
}
48+
49+
private List<IngestStats.ProcessorStat> randomProcessorStats() {
50+
return Stream.generate(() -> randomAlphaOfLength(10))
51+
.limit(randomIntBetween(0, 10))
52+
.map(name -> new IngestStats.ProcessorStat(name, randomStats()))
53+
.collect(Collectors.toList());
54+
}
55+
56+
@Override
57+
protected Writeable.Reader<Response> instanceReader() {
58+
return Response::new;
59+
}
60+
}

‎x-pack/plugin/ml/qa/ml-with-security/build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,8 @@ integTest.runner {
130130
'ml/inference_crud/Test delete with missing model',
131131
'ml/inference_crud/Test get given missing trained model',
132132
'ml/inference_crud/Test get given expression without matches and allow_no_match is false',
133+
'ml/inference_stats_crud/Test get stats given missing trained model',
134+
'ml/inference_stats_crud/Test get stats given expression without matches and allow_no_match is false',
133135
'ml/jobs_crud/Test cannot create job with existing categorizer state document',
134136
'ml/jobs_crud/Test cannot create job with existing quantiles document',
135137
'ml/jobs_crud/Test cannot create job with existing result document',

0 commit comments

Comments
 (0)