[ML] Add deployment threading details and memory usage to telemetry (#113099) (#113516)

davidkyle · web-flow · commit cc3caa228def · 2024-09-25T22:35:09.000+10:00
Adds deployment threading options and a new memory section reporting
the memory usage for each of the ml features
# Conflicts:
#	server/src/main/java/org/elasticsearch/TransportVersions.java
diff --git a/docs/reference/rest-api/usage.asciidoc b/docs/reference/rest-api/usage.asciidoc
@@ -195,7 +195,13 @@ GET /_xpack/usage
         }
       }
     },
-    "node_count" : 1
+    "node_count" : 1,
+    "memory": {
+      anomaly_detectors_memory_bytes: 0,
+      data_frame_analytics_memory_bytes: 0,
+      pytorch_inference_memory_bytes: 0,
+      total_used_memory_bytes: 0
+    }
   },
   "inference": {
     "available" : true,
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningFeatureSetUsage.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningFeatureSetUsage.java
@@ -31,11 +31,13 @@ public class MachineLearningFeatureSetUsage extends XPackFeatureSet.Usage {
     public static final String NODE_COUNT = "node_count";
     public static final String DATA_FRAME_ANALYTICS_JOBS_FIELD = "data_frame_analytics_jobs";
     public static final String INFERENCE_FIELD = "inference";
+    public static final String MEMORY_FIELD = "memory";
 
     private final Map<String, Object> jobsUsage;
     private final Map<String, Object> datafeedsUsage;
     private final Map<String, Object> analyticsUsage;
     private final Map<String, Object> inferenceUsage;
+    private final Map<String, Object> memoryUsage;
     private final int nodeCount;
 
     public MachineLearningFeatureSetUsage(
@@ -45,13 +47,15 @@ public MachineLearningFeatureSetUsage(
         Map<String, Object> datafeedsUsage,
         Map<String, Object> analyticsUsage,
         Map<String, Object> inferenceUsage,
+        Map<String, Object> memoryUsage,
         int nodeCount
     ) {
         super(XPackField.MACHINE_LEARNING, available, enabled);
         this.jobsUsage = Objects.requireNonNull(jobsUsage);
         this.datafeedsUsage = Objects.requireNonNull(datafeedsUsage);
         this.analyticsUsage = Objects.requireNonNull(analyticsUsage);
         this.inferenceUsage = Objects.requireNonNull(inferenceUsage);
+        this.memoryUsage = Objects.requireNonNull(memoryUsage);
         this.nodeCount = nodeCount;
     }
 
@@ -62,6 +66,11 @@ public MachineLearningFeatureSetUsage(StreamInput in) throws IOException {
         this.analyticsUsage = in.readGenericMap();
         this.inferenceUsage = in.readGenericMap();
         this.nodeCount = in.readInt();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.ML_TELEMETRY_MEMORY_ADDED)) {
+            this.memoryUsage = in.readGenericMap();
+        } else {
+            this.memoryUsage = Map.of();
+        }
     }
 
     @Override
@@ -77,6 +86,9 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeGenericMap(analyticsUsage);
         out.writeGenericMap(inferenceUsage);
         out.writeInt(nodeCount);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ML_TELEMETRY_MEMORY_ADDED)) {
+            out.writeGenericMap(memoryUsage);
+        }
     }
 
     @Override
@@ -86,9 +98,51 @@ protected void innerXContent(XContentBuilder builder, Params params) throws IOEx
         builder.field(DATAFEEDS_FIELD, datafeedsUsage);
         builder.field(DATA_FRAME_ANALYTICS_JOBS_FIELD, analyticsUsage);
         builder.field(INFERENCE_FIELD, inferenceUsage);
+        builder.field(MEMORY_FIELD, memoryUsage);
         if (nodeCount >= 0) {
             builder.field(NODE_COUNT, nodeCount);
         }
     }
 
+    public Map<String, Object> getJobsUsage() {
+        return jobsUsage;
+    }
+
+    public Map<String, Object> getDatafeedsUsage() {
+        return datafeedsUsage;
+    }
+
+    public Map<String, Object> getAnalyticsUsage() {
+        return analyticsUsage;
+    }
+
+    public Map<String, Object> getInferenceUsage() {
+        return inferenceUsage;
+    }
+
+    public Map<String, Object> getMemoryUsage() {
+        return memoryUsage;
+    }
+
+    public int getNodeCount() {
+        return nodeCount;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        MachineLearningFeatureSetUsage that = (MachineLearningFeatureSetUsage) o;
+        return nodeCount == that.nodeCount
+            && Objects.equals(jobsUsage, that.jobsUsage)
+            && Objects.equals(datafeedsUsage, that.datafeedsUsage)
+            && Objects.equals(analyticsUsage, that.analyticsUsage)
+            && Objects.equals(inferenceUsage, that.inferenceUsage)
+            && Objects.equals(memoryUsage, that.memoryUsage);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(jobsUsage, datafeedsUsage, analyticsUsage, inferenceUsage, memoryUsage, nodeCount);
+    }
 }
diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/MachineLearningFeatureSetUsageTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/MachineLearningFeatureSetUsageTests.java
@@ -0,0 +1,75 @@
+/*
+ * 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.ml;
+
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.core.Tuple;
+
+import java.io.IOException;
+import java.util.Collections;
+
+public class MachineLearningFeatureSetUsageTests extends AbstractBWCWireSerializationTestCase<MachineLearningFeatureSetUsage> {
+    @Override
+    protected Writeable.Reader<MachineLearningFeatureSetUsage> instanceReader() {
+        return MachineLearningFeatureSetUsage::new;
+    }
+
+    @Override
+    protected MachineLearningFeatureSetUsage createTestInstance() {
+        boolean enabled = randomBoolean();
+
+        if (enabled == false) {
+            return new MachineLearningFeatureSetUsage(
+                randomBoolean(),
+                enabled,
+                Collections.emptyMap(),
+                Collections.emptyMap(),
+                Collections.emptyMap(),
+                Collections.emptyMap(),
+                Collections.emptyMap(),
+                0
+            );
+        } else {
+            return new MachineLearningFeatureSetUsage(
+                randomBoolean(),
+                enabled,
+                randomMap(0, 4, () -> new Tuple<>(randomAlphaOfLength(4), randomAlphaOfLength(4))),
+                randomMap(0, 4, () -> new Tuple<>(randomAlphaOfLength(4), randomAlphaOfLength(4))),
+                randomMap(0, 4, () -> new Tuple<>(randomAlphaOfLength(4), randomAlphaOfLength(4))),
+                randomMap(0, 4, () -> new Tuple<>(randomAlphaOfLength(4), randomAlphaOfLength(4))),
+                randomMap(0, 4, () -> new Tuple<>(randomAlphaOfLength(4), randomAlphaOfLength(4))),
+                randomIntBetween(1, 10)
+            );
+        }
+    }
+
+    @Override
+    protected MachineLearningFeatureSetUsage mutateInstance(MachineLearningFeatureSetUsage instance) throws IOException {
+        return null;
+    }
+
+    @Override
+    protected MachineLearningFeatureSetUsage mutateInstanceForVersion(MachineLearningFeatureSetUsage instance, TransportVersion version) {
+        if (version.before(TransportVersions.ML_TELEMETRY_MEMORY_ADDED)) {
+            return new MachineLearningFeatureSetUsage(
+                instance.available(),
+                instance.enabled(),
+                instance.getJobsUsage(),
+                instance.getDatafeedsUsage(),
+                instance.getAnalyticsUsage(),
+                instance.getInferenceUsage(),
+                Collections.emptyMap(),
+                instance.getNodeCount()
+            );
+        }
+
+        return instance;
+    }
+}
diff --git a/x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelIT.java b/x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelIT.java
@@ -1120,6 +1120,28 @@ public void testStartMultipleLowPriorityDeployments() throws Exception {
         }
     }
 
+    @SuppressWarnings("unchecked")
+    public void testDeploymentThreadsIncludedInUsage() throws IOException {
+        String modelId = "deployment_threads_in_usage";
+        createPassThroughModel(modelId);
+        putModelDefinition(modelId);
+        putVocabulary(List.of("these", "are", "my", "words"), modelId);
+        startDeployment(modelId);
+
+        Request request = new Request("GET", "/_xpack/usage");
+        var usage = entityAsMap(client().performRequest(request).getEntity());
+
+        var ml = (Map<String, Object>) usage.get("ml");
+        assertNotNull(usage.toString(), ml);
+        var inference = (Map<String, Object>) ml.get("inference");
+        var deployments = (Map<String, Object>) inference.get("deployments");
+        var deploymentStats = (List<Map<String, Object>>) deployments.get("stats_by_model");
+        for (var stat : deploymentStats) {
+            assertThat(stat.toString(), (Integer) stat.get("num_threads"), greaterThanOrEqualTo(1));
+            assertThat(stat.toString(), (Integer) stat.get("num_allocations"), greaterThanOrEqualTo(1));
+        }
+    }
+
     private void putModelDefinition(String modelId) throws IOException {
         putModelDefinition(modelId, BASE_64_ENCODED_MODEL, RAW_MODEL_SIZE);
     }
diff --git a/x-pack/plugin/ml/qa/single-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/MlUsageIT.java b/x-pack/plugin/ml/qa/single-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/MlUsageIT.java
@@ -0,0 +1,35 @@
+/*
+ * 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.ml.integration;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.test.rest.ESRestTestCase;
+
+import java.io.IOException;
+import java.util.Map;
+
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+
+// Test the phone home/telemetry data
+public class MlUsageIT extends ESRestTestCase {
+
+    @SuppressWarnings("unchecked")
+    public void testMLUsage() throws IOException {
+        Request request = new Request("GET", "/_xpack/usage");
+        var usage = entityAsMap(client().performRequest(request).getEntity());
+
+        var ml = (Map<String, Object>) usage.get("ml");
+        assertNotNull(usage.toString(), ml);
+        var memoryUsage = (Map<String, Object>) ml.get("memory");
+        assertNotNull(ml.toString(), memoryUsage);
+        assertThat(memoryUsage.toString(), (Integer) memoryUsage.get("anomaly_detectors_memory_bytes"), greaterThanOrEqualTo(0));
+        assertThat(memoryUsage.toString(), (Integer) memoryUsage.get("data_frame_analytics_memory_bytes"), greaterThanOrEqualTo(0));
+        assertThat(memoryUsage.toString(), (Integer) memoryUsage.get("pytorch_inference_memory_bytes"), greaterThanOrEqualTo(0));
+        assertThat(memoryUsage.toString(), (Integer) memoryUsage.get("total_used_memory_bytes"), greaterThanOrEqualTo(0));
+    }
+}
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearningUsageTransportAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearningUsageTransportAction.java
diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/MachineLearningInfoTransportActionTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/MachineLearningInfoTransportActionTests.java
