[ML] Add backend for getting ML autoscaling resources (#97362)

add transport action to retrieve autoscaling metrics for ML autoscaling

Author: Hendrik Muhs

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

@@ -71,6 +71,7 @@
     exports org.elasticsearch.xpack.core.logstash;
     exports org.elasticsearch.xpack.core.ml.action;
     exports org.elasticsearch.xpack.core.ml.annotations;
+    exports org.elasticsearch.xpack.core.ml.autoscaling;
     exports org.elasticsearch.xpack.core.ml.calendars;
     exports org.elasticsearch.xpack.core.ml.datafeed.extractor;
     exports org.elasticsearch.xpack.core.ml.datafeed;

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetMlAutoscalingStats.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.ml.action;
+
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.ActionType;
+import org.elasticsearch.action.support.master.AcknowledgedRequest;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.xpack.core.ml.action.GetMlAutoscalingStats.Response;
+import org.elasticsearch.xpack.core.ml.autoscaling.MlAutoscalingStats;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Internal (no-REST) action to retrieve metrics for serverless autoscaling.
+ */
+public class GetMlAutoscalingStats extends ActionType<Response> {
+
+    public static final GetMlAutoscalingStats INSTANCE = new GetMlAutoscalingStats();
+    public static final String NAME = "cluster:monitor/xpack/ml/autoscaling/stats/get";
+
+    public GetMlAutoscalingStats() {
+        super(NAME, Response::new);
+    }
+
+    public static class Request extends AcknowledgedRequest<Request> {
+
+        public Request(TimeValue timeout) {
+            super(timeout);
+        }
+
+        public Request(StreamInput in) throws IOException {
+            super(in);
+        }
+
+        @Override
+        public ActionRequestValidationException validate() {
+            return null;
+        }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, "get_ml_autoscaling_resources", parentTaskId, headers);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(timeout);
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj == null) {
+                return false;
+            }
+            if (getClass() != obj.getClass()) {
+                return false;
+            }
+            GetMlAutoscalingStats.Request other = (GetMlAutoscalingStats.Request) obj;
+            return Objects.equals(timeout, other.timeout);
+        }
+    }
+
+    public static class Response extends ActionResponse {
+
+        private final MlAutoscalingStats autoscalingResources;
+
+        public Response(final MlAutoscalingStats autoscalingResources) {
+            this.autoscalingResources = autoscalingResources;
+        }
+
+        public Response(final StreamInput in) throws IOException {
+            super(in);
+            this.autoscalingResources = new MlAutoscalingStats(in);
+        }
+
+        public MlAutoscalingStats getAutoscalingResources() {
+            return autoscalingResources;
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            autoscalingResources.writeTo(out);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(autoscalingResources);
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj == null) {
+                return false;
+            }
+            if (getClass() != obj.getClass()) {
+                return false;
+            }
+            GetMlAutoscalingStats.Response other = (GetMlAutoscalingStats.Response) obj;
+            return Objects.equals(autoscalingResources, other.autoscalingResources);
+        }
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/autoscaling/MlAutoscalingStats.java

@@ -0,0 +1,57 @@
+/*
+ * 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.autoscaling;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+
+import java.io.IOException;
+
+public record MlAutoscalingStats(
+    int nodes,
+    long memoryInBytesSum,
+    long modelMemoryInBytesSum,
+    int minNodes,
+    long extraSingleNodeModelMemoryInBytes,
+    int extraSingleNodeProcessors,
+    long extraModelMemoryInBytes,
+    int extraProcessors,
+    long removeNodeMemoryInBytes,
+    long perNodeMemoryOverheadInBytes
+) implements Writeable {
+
+    public MlAutoscalingStats(StreamInput in) throws IOException {
+        this(
+            in.readVInt(), // nodes
+            in.readVLong(),  // memoryInBytesSum
+            in.readVLong(), // modelMemoryInBytes
+            in.readVInt(), // minNodes
+            in.readVLong(), // extraSingleNodeModelMemoryInBytes
+            in.readVInt(), // extraSingleNodeProcessors
+            in.readVLong(), // extraModelMemoryInBytes
+            in.readVInt(), // extraProcessors
+            in.readVLong(), // removeNodeMemoryInBytes
+            in.readVLong() // perNodeMemoryOverheadInBytes
+        );
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeVInt(nodes);
+        out.writeVLong(memoryInBytesSum);
+        out.writeVLong(modelMemoryInBytesSum);
+        out.writeVInt(minNodes);
+        out.writeVLong(extraSingleNodeModelMemoryInBytes);
+        out.writeVInt(extraSingleNodeProcessors);
+        out.writeVLong(extraModelMemoryInBytes);
+        out.writeVInt(extraProcessors);
+        out.writeVLong(removeNodeMemoryInBytes);
+        out.writeVLong(perNodeMemoryOverheadInBytes);
+    }
+}

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

@@ -0,0 +1,33 @@
+/*
+ * 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.action;
+
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.test.AbstractWireSerializingTestCase;
+import org.elasticsearch.xpack.core.ml.action.GetMlAutoscalingStats.Request;
+
+import java.io.IOException;
+
+public class GetMlAutoscalingStatsRequestTests extends AbstractWireSerializingTestCase<Request> {
+
+    @Override
+    protected Writeable.Reader<Request> instanceReader() {
+        return Request::new;
+    }
+
+    @Override
+    protected Request createTestInstance() {
+        return new Request(TimeValue.parseTimeValue(randomTimeValue(0, 10_000), "timeout"));
+    }
+
+    @Override
+    protected Request mutateInstance(Request instance) throws IOException {
+        return new Request(TimeValue.timeValueMillis(instance.timeout().millis() + randomIntBetween(1, 1000)));
+    }
+}

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

@@ -0,0 +1,34 @@
+/*
+ * 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.action;
+
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.test.AbstractWireSerializingTestCase;
+import org.elasticsearch.xpack.core.ml.action.GetMlAutoscalingStats.Response;
+
+import java.io.IOException;
+
+import static org.elasticsearch.xpack.core.ml.autoscaling.MlAutoscalingStatsTests.randomAutoscalingResources;
+
+public class GetMlAutoscalingStatsResponseTests extends AbstractWireSerializingTestCase<Response> {
+
+    @Override
+    protected Writeable.Reader<Response> instanceReader() {
+        return Response::new;
+    }
+
+    @Override
+    protected Response createTestInstance() {
+        return new Response(randomAutoscalingResources());
+    }
+
+    @Override
+    protected Response mutateInstance(Response instance) throws IOException {
+        return null; // TODO
+    }
+}

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/autoscaling/MlAutoscalingStatsTests.java

@@ -0,0 +1,46 @@
+/*
+ * 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.autoscaling;
+
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.test.AbstractWireSerializingTestCase;
+
+import java.io.IOException;
+
+public class MlAutoscalingStatsTests extends AbstractWireSerializingTestCase<MlAutoscalingStats> {
+
+    public static MlAutoscalingStats randomAutoscalingResources() {
+        return new MlAutoscalingStats(
+            randomIntBetween(0, 100), // nodes
+            randomNonNegativeLong(), // memoryInBytesSum
+            randomNonNegativeLong(), // modelMemoryInBytes
+            randomIntBetween(0, 100), // minNodes
+            randomNonNegativeLong(), // extraSingleNodeModelMemoryInBytes
+            randomIntBetween(0, 100), // extraSingleNodeProcessors
+            randomNonNegativeLong(), // extraModelMemoryInBytes
+            randomIntBetween(0, 100), // extraProcessors
+            randomNonNegativeLong(), // removeNodeMemoryInBytes
+            randomNonNegativeLong() // perNodeMemoryOverheadInBytes
+        );
+    }
+
+    @Override
+    protected Writeable.Reader<MlAutoscalingStats> instanceReader() {
+        return MlAutoscalingStats::new;
+    }
+
+    @Override
+    protected MlAutoscalingStats createTestInstance() {
+        return randomAutoscalingResources();
+    }
+
+    @Override
+    protected MlAutoscalingStats mutateInstance(MlAutoscalingStats instance) throws IOException {
+        return null; // TODO
+    }
+}

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearning.java

@@ -135,6 +135,7 @@
 import org.elasticsearch.xpack.core.ml.action.GetJobModelSnapshotsUpgradeStatsAction;
 import org.elasticsearch.xpack.core.ml.action.GetJobsAction;
 import org.elasticsearch.xpack.core.ml.action.GetJobsStatsAction;
+import org.elasticsearch.xpack.core.ml.action.GetMlAutoscalingStats;
 import org.elasticsearch.xpack.core.ml.action.GetModelSnapshotsAction;
 import org.elasticsearch.xpack.core.ml.action.GetOverallBucketsAction;
 import org.elasticsearch.xpack.core.ml.action.GetRecordsAction;
@@ -237,6 +238,7 @@
 import org.elasticsearch.xpack.ml.action.TransportGetJobModelSnapshotsUpgradeStatsAction;
 import org.elasticsearch.xpack.ml.action.TransportGetJobsAction;
 import org.elasticsearch.xpack.ml.action.TransportGetJobsStatsAction;
+import org.elasticsearch.xpack.ml.action.TransportGetMlAutoscalingStats;
 import org.elasticsearch.xpack.ml.action.TransportGetModelSnapshotsAction;
 import org.elasticsearch.xpack.ml.action.TransportGetOverallBucketsAction;
 import org.elasticsearch.xpack.ml.action.TransportGetRecordsAction;
@@ -1415,6 +1417,7 @@ public List<RestHandler> getRestHandlers(
         actionHandlers.add(new ActionHandler<>(SetResetModeAction.INSTANCE, TransportSetResetModeAction.class));
         // Included in this section as it's used by MlMemoryAction
         actionHandlers.add(new ActionHandler<>(TrainedModelCacheInfoAction.INSTANCE, TransportTrainedModelCacheInfoAction.class));
+        actionHandlers.add(new ActionHandler<>(GetMlAutoscalingStats.INSTANCE, TransportGetMlAutoscalingStats.class));
         if (machineLearningExtension.get().isAnomalyDetectionEnabled()) {
             actionHandlers.add(new ActionHandler<>(GetJobsAction.INSTANCE, TransportGetJobsAction.class));
             actionHandlers.add(new ActionHandler<>(GetJobsStatsAction.INSTANCE, TransportGetJobsStatsAction.class));