Explorar o código

Inference telemetry (#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 <elasticmachine@users.noreply.github.com>
Jan Kuipers hai 1 ano
pai
achega
a67d5b8986

+ 5 - 0
docs/changelog/102877.yaml

@@ -0,0 +1,5 @@
+pr: 102877
+summary: Add basic telelemetry for the inference feature
+area: Machine Learning
+type: enhancement
+issues: []

+ 5 - 0
docs/reference/rest-api/usage.asciidoc

@@ -197,6 +197,11 @@ GET /_xpack/usage
     },
     "node_count" : 1
   },
+  "inference": {
+    "available" : true,
+    "enabled" : true,
+    "models" : []
+  },
   "logstash" : {
     "available" : true,
     "enabled" : true

+ 1 - 0
server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -184,6 +184,7 @@ public class TransportVersions {
     public static final TransportVersion ESQL_PROFILE = def(8_551_00_0);
     public static final TransportVersion CLUSTER_STATS_RESCORER_USAGE_ADDED = def(8_552_00_0);
     public static final TransportVersion ML_INFERENCE_HF_SERVICE_ADDED = def(8_553_00_0);
+    public static final TransportVersion INFERENCE_USAGE_ADDED = def(8_554_00_0);
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

+ 1 - 0
x-pack/plugin/core/src/main/java/module-info.java

@@ -75,6 +75,7 @@ module org.elasticsearch.xcore {
     exports org.elasticsearch.xpack.core.indexing;
     exports org.elasticsearch.xpack.core.inference.action;
     exports org.elasticsearch.xpack.core.inference.results;
+    exports org.elasticsearch.xpack.core.inference;
     exports org.elasticsearch.xpack.core.logstash;
     exports org.elasticsearch.xpack.core.ml.action;
     exports org.elasticsearch.xpack.core.ml.annotations;

+ 3 - 0
x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/XPackClientPlugin.java

@@ -55,6 +55,7 @@ import org.elasticsearch.xpack.core.ilm.ShrinkAction;
 import org.elasticsearch.xpack.core.ilm.TimeseriesLifecycleType;
 import org.elasticsearch.xpack.core.ilm.UnfollowAction;
 import org.elasticsearch.xpack.core.ilm.WaitForSnapshotAction;
+import org.elasticsearch.xpack.core.inference.InferenceFeatureSetUsage;
 import org.elasticsearch.xpack.core.logstash.LogstashFeatureSetUsage;
 import org.elasticsearch.xpack.core.ml.MachineLearningFeatureSetUsage;
 import org.elasticsearch.xpack.core.ml.MlMetadata;
@@ -133,6 +134,8 @@ public class XPackClientPlugin extends Plugin implements ActionPlugin, NetworkPl
             new NamedWriteableRegistry.Entry(XPackFeatureSet.Usage.class, XPackField.LOGSTASH, LogstashFeatureSetUsage::new),
             // ML
             new NamedWriteableRegistry.Entry(XPackFeatureSet.Usage.class, XPackField.MACHINE_LEARNING, MachineLearningFeatureSetUsage::new),
+            // inference
+            new NamedWriteableRegistry.Entry(XPackFeatureSet.Usage.class, XPackField.INFERENCE, InferenceFeatureSetUsage::new),
             // monitoring
             new NamedWriteableRegistry.Entry(XPackFeatureSet.Usage.class, XPackField.MONITORING, MonitoringFeatureSetUsage::new),
             // security

+ 2 - 0
x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/XPackField.java

@@ -18,6 +18,8 @@ public final class XPackField {
     public static final String GRAPH = "graph";
     /** Name constant for the machine learning feature. */
     public static final String MACHINE_LEARNING = "ml";
+    /** Name constant for the inference feature. */
+    public static final String INFERENCE = "inference";
     /** Name constant for the Logstash feature. */
     public static final String LOGSTASH = "logstash";
     /** Name constant for the Beats feature. */

+ 2 - 0
x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/action/XPackUsageFeatureAction.java

@@ -27,6 +27,7 @@ public class XPackUsageFeatureAction extends ActionType<XPackUsageFeatureRespons
     public static final XPackUsageFeatureAction WATCHER = new XPackUsageFeatureAction(XPackField.WATCHER);
     public static final XPackUsageFeatureAction GRAPH = new XPackUsageFeatureAction(XPackField.GRAPH);
     public static final XPackUsageFeatureAction MACHINE_LEARNING = new XPackUsageFeatureAction(XPackField.MACHINE_LEARNING);
+    public static final XPackUsageFeatureAction INFERENCE = new XPackUsageFeatureAction(XPackField.INFERENCE);
     public static final XPackUsageFeatureAction LOGSTASH = new XPackUsageFeatureAction(XPackField.LOGSTASH);
     public static final XPackUsageFeatureAction EQL = new XPackUsageFeatureAction(XPackField.EQL);
     public static final XPackUsageFeatureAction ESQL = new XPackUsageFeatureAction(XPackField.ESQL);
@@ -64,6 +65,7 @@ public class XPackUsageFeatureAction extends ActionType<XPackUsageFeatureRespons
         FROZEN_INDICES,
         GRAPH,
         INDEX_LIFECYCLE,
+        INFERENCE,
         LOGSTASH,
         MACHINE_LEARNING,
         MONITORING,

+ 116 - 0
x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java

@@ -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 - 0
x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsageTests.java

@@ -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;
+    }
+}

+ 4 - 1
x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -32,6 +32,7 @@ import org.elasticsearch.rest.RestHandler;
 import org.elasticsearch.threadpool.ExecutorBuilder;
 import org.elasticsearch.threadpool.ScalingExecutorBuilder;
 import org.elasticsearch.xpack.core.ClientHelper;
+import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction;
 import org.elasticsearch.xpack.core.inference.action.DeleteInferenceModelAction;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
@@ -39,6 +40,7 @@ import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction;
 import org.elasticsearch.xpack.inference.action.TransportDeleteInferenceModelAction;
 import org.elasticsearch.xpack.inference.action.TransportGetInferenceModelAction;
 import org.elasticsearch.xpack.inference.action.TransportInferenceAction;
+import org.elasticsearch.xpack.inference.action.TransportInferenceUsageAction;
 import org.elasticsearch.xpack.inference.action.TransportPutInferenceModelAction;
 import org.elasticsearch.xpack.inference.external.http.HttpClientManager;
 import org.elasticsearch.xpack.inference.external.http.HttpSettings;
@@ -86,7 +88,8 @@ public class InferencePlugin extends Plugin implements ActionPlugin, ExtensibleP
             new ActionHandler<>(InferenceAction.INSTANCE, TransportInferenceAction.class),
             new ActionHandler<>(GetInferenceModelAction.INSTANCE, TransportGetInferenceModelAction.class),
             new ActionHandler<>(PutInferenceModelAction.INSTANCE, TransportPutInferenceModelAction.class),
-            new ActionHandler<>(DeleteInferenceModelAction.INSTANCE, TransportDeleteInferenceModelAction.class)
+            new ActionHandler<>(DeleteInferenceModelAction.INSTANCE, TransportDeleteInferenceModelAction.class),
+            new ActionHandler<>(XPackUsageFeatureAction.INFERENCE, TransportInferenceUsageAction.class)
         );
     }
 

+ 81 - 0
x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java

@@ -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));
+    }
+}

+ 121 - 0
x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageActionTests.java

@@ -0,0 +1,121 @@
+/*
+ * 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.action.support.PlainActionFuture;
+import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.inference.ServiceSettings;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.protocol.xpack.XPackUsageRequest;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.MockUtils;
+import org.elasticsearch.threadpool.TestThreadPool;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.TransportService;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xpack.core.XPackFeatureSet;
+import org.elasticsearch.xpack.core.XPackField;
+import org.elasticsearch.xpack.core.action.XPackUsageFeatureResponse;
+import org.elasticsearch.xpack.core.inference.InferenceFeatureSetUsage;
+import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
+import org.elasticsearch.xpack.core.watcher.support.xcontent.XContentSource;
+import org.junit.After;
+import org.junit.Before;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.core.Is.is;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class TransportInferenceUsageActionTests extends ESTestCase {
+
+    private Client client;
+    private TransportInferenceUsageAction action;
+
+    @Before
+    public void init() {
+        client = mock(Client.class);
+        ThreadPool threadPool = new TestThreadPool("test");
+        when(client.threadPool()).thenReturn(threadPool);
+
+        TransportService transportService = MockUtils.setupTransportServiceWithThreadpoolExecutor(mock(ThreadPool.class));
+
+        action = new TransportInferenceUsageAction(
+            transportService,
+            mock(ClusterService.class),
+            mock(ThreadPool.class),
+            mock(ActionFilters.class),
+            mock(IndexNameExpressionResolver.class),
+            client
+        );
+    }
+
+    @After
+    public void close() {
+        client.threadPool().shutdown();
+    }
+
+    public void test() throws Exception {
+        doAnswer(invocation -> {
+            @SuppressWarnings("unchecked")
+            var listener = (ActionListener<GetInferenceModelAction.Response>) invocation.getArguments()[2];
+            listener.onResponse(
+                new GetInferenceModelAction.Response(
+                    List.of(
+                        new ModelConfigurations("model-001", TaskType.TEXT_EMBEDDING, "openai", mock(ServiceSettings.class)),
+                        new ModelConfigurations("model-002", TaskType.TEXT_EMBEDDING, "openai", mock(ServiceSettings.class)),
+                        new ModelConfigurations("model-003", TaskType.SPARSE_EMBEDDING, "hugging_face_elser", mock(ServiceSettings.class)),
+                        new ModelConfigurations("model-004", TaskType.TEXT_EMBEDDING, "openai", mock(ServiceSettings.class)),
+                        new ModelConfigurations("model-005", TaskType.SPARSE_EMBEDDING, "openai", mock(ServiceSettings.class)),
+                        new ModelConfigurations("model-006", TaskType.SPARSE_EMBEDDING, "hugging_face_elser", mock(ServiceSettings.class))
+                    )
+                )
+            );
+            return Void.TYPE;
+        }).when(client).execute(any(GetInferenceModelAction.class), any(), any());
+
+        PlainActionFuture<XPackUsageFeatureResponse> future = new PlainActionFuture<>();
+        action.masterOperation(mock(Task.class), mock(XPackUsageRequest.class), mock(ClusterState.class), future);
+
+        BytesStreamOutput out = new BytesStreamOutput();
+        future.get().getUsage().writeTo(out);
+        XPackFeatureSet.Usage usage = new InferenceFeatureSetUsage(out.bytes().streamInput());
+
+        assertThat(usage.name(), is(XPackField.INFERENCE));
+        assertTrue(usage.enabled());
+        assertTrue(usage.available());
+
+        XContentBuilder builder = XContentFactory.jsonBuilder();
+        usage.toXContent(builder, ToXContent.EMPTY_PARAMS);
+        XContentSource source = new XContentSource(builder);
+        assertThat(source.getValue("models"), hasSize(3));
+        assertThat(source.getValue("models.0.service"), is("hugging_face_elser"));
+        assertThat(source.getValue("models.0.task_type"), is("SPARSE_EMBEDDING"));
+        assertThat(source.getValue("models.0.count"), is(2));
+        assertThat(source.getValue("models.1.service"), is("openai"));
+        assertThat(source.getValue("models.1.task_type"), is("SPARSE_EMBEDDING"));
+        assertThat(source.getValue("models.1.count"), is(1));
+        assertThat(source.getValue("models.2.service"), is("openai"));
+        assertThat(source.getValue("models.2.task_type"), is("TEXT_EMBEDDING"));
+        assertThat(source.getValue("models.2.count"), is(3));
+    }
+}

+ 1 - 0
x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java

@@ -416,6 +416,7 @@ public class Constants {
         "cluster:monitor/xpack/usage/graph",
         "cluster:monitor/xpack/usage/health_api",
         "cluster:monitor/xpack/usage/ilm",
+        "cluster:monitor/xpack/usage/inference",
         "cluster:monitor/xpack/usage/logstash",
         "cluster:monitor/xpack/usage/ml",
         "cluster:monitor/xpack/usage/monitoring",