forked from elastic/elasticsearch
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Inference telemetry (elastic#102877)
* Empty infenrece usage wiring. * Add fake data * Fix NPE for secretSettings == null * Real inference model stats * New transport version * Code polish * Lint fixes * Update docs/changelog/102877.yaml * Update 102877.yaml * Add inference to yamlRestTest * Declare inference usage action as non-operator * TransportInferenceUsageActionTests * Lint fixes * Replace map by ToXContentObject/Writeable * Polish code * AbstractWireSerializingTestCase<InferenceFeatureSetUsage.ModelStats> --------- Co-authored-by: Elastic Machine <[email protected]>
- Loading branch information
1 parent
84dad02
commit a67d5b8
Showing
13 changed files
with
383 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| pr: 102877 | ||
| summary: Add basic telelemetry for the inference feature | ||
| area: Machine Learning | ||
| type: enhancement | ||
| issues: [] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
116 changes: 116 additions & 0 deletions
116
...n/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.core.inference; | ||
|
|
||
| import org.elasticsearch.TransportVersion; | ||
| import org.elasticsearch.TransportVersions; | ||
| import org.elasticsearch.common.io.stream.StreamInput; | ||
| import org.elasticsearch.common.io.stream.StreamOutput; | ||
| import org.elasticsearch.common.io.stream.Writeable; | ||
| import org.elasticsearch.inference.TaskType; | ||
| import org.elasticsearch.xcontent.ToXContentObject; | ||
| import org.elasticsearch.xcontent.XContentBuilder; | ||
| import org.elasticsearch.xpack.core.XPackFeatureSet; | ||
| import org.elasticsearch.xpack.core.XPackField; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Collection; | ||
| import java.util.Objects; | ||
|
|
||
| public class InferenceFeatureSetUsage extends XPackFeatureSet.Usage { | ||
|
|
||
| public static class ModelStats implements ToXContentObject, Writeable { | ||
|
|
||
| private final String service; | ||
| private final TaskType taskType; | ||
| private long count; | ||
|
|
||
| public ModelStats(String service, TaskType taskType) { | ||
| this(service, taskType, 0L); | ||
| } | ||
|
|
||
| public ModelStats(String service, TaskType taskType, long count) { | ||
| this.service = service; | ||
| this.taskType = taskType; | ||
| this.count = count; | ||
| } | ||
|
|
||
| public ModelStats(ModelStats stats) { | ||
| this(stats.service, stats.taskType, stats.count); | ||
| } | ||
|
|
||
| public ModelStats(StreamInput in) throws IOException { | ||
| this.service = in.readString(); | ||
| this.taskType = in.readEnum(TaskType.class); | ||
| this.count = in.readLong(); | ||
| } | ||
|
|
||
| public void add() { | ||
| count++; | ||
| } | ||
|
|
||
| @Override | ||
| public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { | ||
| builder.startObject(); | ||
| builder.field("service", service); | ||
| builder.field("task_type", taskType.name()); | ||
| builder.field("count", count); | ||
| builder.endObject(); | ||
| return builder; | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(StreamOutput out) throws IOException { | ||
| out.writeString(service); | ||
| out.writeEnum(taskType); | ||
| out.writeLong(count); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean equals(Object o) { | ||
| if (this == o) return true; | ||
| if (o == null || getClass() != o.getClass()) return false; | ||
| ModelStats that = (ModelStats) o; | ||
| return count == that.count && Objects.equals(service, that.service) && taskType == that.taskType; | ||
| } | ||
|
|
||
| @Override | ||
| public int hashCode() { | ||
| return Objects.hash(service, taskType, count); | ||
| } | ||
| } | ||
|
|
||
| private final Collection<ModelStats> modelStats; | ||
|
|
||
| public InferenceFeatureSetUsage(Collection<ModelStats> modelStats) { | ||
| super(XPackField.INFERENCE, true, true); | ||
| this.modelStats = modelStats; | ||
| } | ||
|
|
||
| public InferenceFeatureSetUsage(StreamInput in) throws IOException { | ||
| super(in); | ||
| this.modelStats = in.readCollectionAsList(ModelStats::new); | ||
| } | ||
|
|
||
| @Override | ||
| protected void innerXContent(XContentBuilder builder, Params params) throws IOException { | ||
| super.innerXContent(builder, params); | ||
| builder.xContentList("models", modelStats); | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(StreamOutput out) throws IOException { | ||
| super.writeTo(out); | ||
| out.writeCollection(modelStats); | ||
| } | ||
|
|
||
| @Override | ||
| public TransportVersion getMinimalSupportedVersion() { | ||
| return TransportVersions.INFERENCE_USAGE_ADDED; | ||
| } | ||
| } |
41 changes: 41 additions & 0 deletions
41
...e/src/test/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsageTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,41 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.core.inference; | ||
|
|
||
| import com.carrotsearch.randomizedtesting.generators.RandomStrings; | ||
|
|
||
| import org.elasticsearch.common.io.stream.Writeable; | ||
| import org.elasticsearch.inference.TaskType; | ||
| import org.elasticsearch.test.AbstractWireSerializingTestCase; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| public class InferenceFeatureSetUsageTests extends AbstractWireSerializingTestCase<InferenceFeatureSetUsage.ModelStats> { | ||
|
|
||
| @Override | ||
| protected Writeable.Reader<InferenceFeatureSetUsage.ModelStats> instanceReader() { | ||
| return InferenceFeatureSetUsage.ModelStats::new; | ||
| } | ||
|
|
||
| @Override | ||
| protected InferenceFeatureSetUsage.ModelStats createTestInstance() { | ||
| RandomStrings.randomAsciiLettersOfLength(random(), 10); | ||
| return new InferenceFeatureSetUsage.ModelStats( | ||
| randomIdentifier(), | ||
| TaskType.values()[randomInt(TaskType.values().length - 1)], | ||
| randomInt(10) | ||
| ); | ||
| } | ||
|
|
||
| @Override | ||
| protected InferenceFeatureSetUsage.ModelStats mutateInstance(InferenceFeatureSetUsage.ModelStats modelStats) throws IOException { | ||
| InferenceFeatureSetUsage.ModelStats newModelStats = new InferenceFeatureSetUsage.ModelStats(modelStats); | ||
| newModelStats.add(); | ||
| return newModelStats; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
81 changes: 81 additions & 0 deletions
81
...src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.inference.action; | ||
|
|
||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.action.support.ActionFilters; | ||
| import org.elasticsearch.client.internal.Client; | ||
| import org.elasticsearch.client.internal.OriginSettingClient; | ||
| import org.elasticsearch.cluster.ClusterState; | ||
| import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver; | ||
| import org.elasticsearch.cluster.service.ClusterService; | ||
| import org.elasticsearch.common.inject.Inject; | ||
| import org.elasticsearch.inference.ModelConfigurations; | ||
| import org.elasticsearch.inference.TaskType; | ||
| import org.elasticsearch.protocol.xpack.XPackUsageRequest; | ||
| import org.elasticsearch.tasks.Task; | ||
| import org.elasticsearch.threadpool.ThreadPool; | ||
| import org.elasticsearch.transport.TransportService; | ||
| import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction; | ||
| import org.elasticsearch.xpack.core.action.XPackUsageFeatureResponse; | ||
| import org.elasticsearch.xpack.core.action.XPackUsageFeatureTransportAction; | ||
| import org.elasticsearch.xpack.core.inference.InferenceFeatureSetUsage; | ||
| import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction; | ||
|
|
||
| import java.util.Map; | ||
| import java.util.TreeMap; | ||
|
|
||
| import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN; | ||
|
|
||
| public class TransportInferenceUsageAction extends XPackUsageFeatureTransportAction { | ||
|
|
||
| private final Client client; | ||
|
|
||
| @Inject | ||
| public TransportInferenceUsageAction( | ||
| TransportService transportService, | ||
| ClusterService clusterService, | ||
| ThreadPool threadPool, | ||
| ActionFilters actionFilters, | ||
| IndexNameExpressionResolver indexNameExpressionResolver, | ||
| Client client | ||
| ) { | ||
| super( | ||
| XPackUsageFeatureAction.INFERENCE.name(), | ||
| transportService, | ||
| clusterService, | ||
| threadPool, | ||
| actionFilters, | ||
| indexNameExpressionResolver | ||
| ); | ||
| this.client = new OriginSettingClient(client, ML_ORIGIN); | ||
| } | ||
|
|
||
| @Override | ||
| protected void masterOperation( | ||
| Task task, | ||
| XPackUsageRequest request, | ||
| ClusterState state, | ||
| ActionListener<XPackUsageFeatureResponse> listener | ||
| ) throws Exception { | ||
| GetInferenceModelAction.Request getInferenceModelAction = new GetInferenceModelAction.Request("_all", TaskType.ANY); | ||
| client.execute(GetInferenceModelAction.INSTANCE, getInferenceModelAction, ActionListener.wrap(response -> { | ||
| Map<String, InferenceFeatureSetUsage.ModelStats> stats = new TreeMap<>(); | ||
| for (ModelConfigurations model : response.getModels()) { | ||
| String statKey = model.getService() + ":" + model.getTaskType().name(); | ||
| InferenceFeatureSetUsage.ModelStats stat = stats.computeIfAbsent( | ||
| statKey, | ||
| key -> new InferenceFeatureSetUsage.ModelStats(model.getService(), model.getTaskType()) | ||
| ); | ||
| stat.add(); | ||
| } | ||
| InferenceFeatureSetUsage usage = new InferenceFeatureSetUsage(stats.values()); | ||
| listener.onResponse(new XPackUsageFeatureResponse(usage)); | ||
| }, listener::onFailure)); | ||
| } | ||
| } |
Oops, something went wrong.