Inference autoscaling (#109667)

* Python dev tool for inference autoscaling simulation.

Squashed commit of the following:

commit d98bd3d39d833329ab83a8274885473db41ed08a
Author: Jan Kuipers <[email protected]>
Date:   Mon May 13 17:27:38 2024 +0200

    Increase measurement interval to 10secs

commit e808ae5be52c5ea4d5ff8ccb881a4a80de0254f9
Author: Jan Kuipers <[email protected]>
Date:   Mon May 13 17:09:33 2024 +0200

    jump -> jumps

commit c38cbdebfcec43e6982bb8bd1670519293161154
Author: Jan Kuipers <[email protected]>
Date:   Mon May 13 14:32:42 2024 +0200

    Remove unused estimator

commit 16101f32b539481cd4d648ebb5637a3309853552
Author: Jan Kuipers <[email protected]>
Date:   Mon May 13 14:31:30 2024 +0200

    Measure latency periodically + documentation

commit bc73bf29fde1d772701f0b71a7c8a0908669eb0f
Author: Jan Kuipers <[email protected]>
Date:   Mon May 13 12:53:19 2024 +0200

    Init variance to None

commit 0e73fa836fa9deec6ba55ef1161cc0dd71f35044
Author: Jan Kuipers <[email protected]>
Date:   Mon May 13 11:18:21 2024 +0200

    No autodetection of dynamics changes for latency

commit 75924a744d26a72835529598a6df1a2d22bdaddc
Author: Jan Kuipers <[email protected]>
Date:   Mon May 13 10:10:34 2024 +0200

    Move autoscaling code to own class

commit 23553bb8cccd6ed80ac667b12ec38a6d5562dd29
Author: Jan Kuipers <[email protected]>
Date:   Wed May 8 18:01:55 2024 +0200

    Improved autoscaling simulation

commit 2db606b2bba69d741fa231f369c633ea793294d5
Author: Tom Veasey <[email protected]>
Date:   Tue Apr 30 15:01:40 2024 +0100

    Correct the dependency on allocations

commit 0e45cfbaf901cf9d440efa9b404058a67d000653
Author: Tom Veasey <[email protected]>
Date:   Tue Apr 30 11:11:05 2024 +0100

    Tweak

commit a0f23a4a05875cd5df3863e5ad067b46a67c8cda
Author: Tom Veasey <[email protected]>
Date:   Tue Apr 30 11:09:30 2024 +0100

    Correction

commit f9cdb140d298bd99c64c79f020c058d60bfba134
Author: Tom Veasey <[email protected]>
Date:   Tue Apr 30 09:57:59 2024 +0100

    Allow extrapolation

commit 57eb1a661a2b97412f479606c23c54dfb7887f52
Author: Tom Veasey <[email protected]>
Date:   Tue Apr 30 09:55:17 2024 +0100

    Simplify and estimate average duration rather than rate

commit 36dff17194f2bcf816013b112cf07d70c9eec161
Author: Tom Veasey <[email protected]>
Date:   Mon Apr 29 21:42:25 2024 +0100

    Kalman filter for simple state model for average inference duration as a function of time and allocation count

commit a1b85bd0deeabd5162f2ccd5a28672299025cee5
Author: Jan Kuipers <[email protected]>
Date:   Mon Apr 29 12:15:59 2024 +0200

    Improvements

commit 51040655fcfbfd221f2446542a955fb0f19fb145
Author: Jan Kuipers <[email protected]>
Date:   Mon Apr 29 09:33:10 2024 +0200

    Account for virtual cores / hyperthreading

commit 7a93407ecae6b6044108299a1d05f72cdf0d752a
Author: Jan Kuipers <[email protected]>
Date:   Fri Apr 26 16:58:25 2024 +0200

    Simulator for inference autoscaling.

* Better process variance upon dynamics changes, and propagate dynamics changes to the next iteration.

* Inference autoscaling (WIP)

* Inference autoscaling test scripts

* Debug logs

* Inference autoscaling API

* Update Autoscalers upon cluster changes

* Polish code / fix bugs

* Use correct string formatter

* More fixes

* Autoscaling tests

* spotless

* Remove scripts (moved to ml-data)

* Rebrand to "adaptive allocations".

* Move serialized field to end

* Rebranding leftover

* Improve adaptive allocation timing

* SystemAuditor for scaling messages

* Fix test

* Add documentation

* Update docs/changelog/109667.yaml

* Cooldown of 5mins after scaleup

* Polish code

* High-variance adaptive allocations test

* Fix AdaptiveAllocationsScalerServiceTests

* Fix typo in package name

* Wire adaptive allocations setting into put inference API

* Checkstyle

* Fix serialization of ElserInternalServiceSettings.

* Propagate adaptive allocations settings from put inference request to create trained model request

* Fix CustomElandInternalTextEmbeddingServiceSettingsTests

* Javadocs

* Improvements / fixes

* Disallow setting num_allocations when adaptive allocations is enabled

* Fix AdaptiveAllocationsScalerServiceTests

* spotless

* NPE fixes

* spotless

* Allow autoscaler to update num allocations

* Fix AdaptiveAllocationsScalerServiceTests.

* Fix bug in inference stats api

* Fix PyTorchResultProcessorTests

Author: jan-elastic

docs/changelog/109667.yaml

@@ -0,0 +1,5 @@
+pr: 109667
+summary: Inference autoscaling
+area: Machine Learning
+type: feature
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -210,6 +210,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion VERSIONED_MASTER_NODE_REQUESTS = def(8_701_00_0);
     public static final TransportVersion ML_INFERENCE_AMAZON_BEDROCK_ADDED = def(8_702_00_0);
     public static final TransportVersion ML_INFERENCE_DONT_DELETE_WHEN_SEMANTIC_TEXT_EXISTS = def(8_703_00_0);
+    public static final TransportVersion INFERENCE_ADAPTIVE_ALLOCATIONS = def(8_704_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.core.ml.action;
 
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.action.ActionType;
@@ -18,6 +19,7 @@
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignment;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 
@@ -34,15 +36,22 @@ private CreateTrainedModelAssignmentAction() {
 
     public static class Request extends MasterNodeRequest<Request> {
         private final StartTrainedModelDeploymentAction.TaskParams taskParams;
+        private final AdaptiveAllocationsSettings adaptiveAllocationsSettings;
 
-        public Request(StartTrainedModelDeploymentAction.TaskParams taskParams) {
+        public Request(StartTrainedModelDeploymentAction.TaskParams taskParams, AdaptiveAllocationsSettings adaptiveAllocationsSettings) {
             super(TRAPPY_IMPLICIT_DEFAULT_MASTER_NODE_TIMEOUT);
             this.taskParams = ExceptionsHelper.requireNonNull(taskParams, "taskParams");
+            this.adaptiveAllocationsSettings = adaptiveAllocationsSettings;
         }
 
         public Request(StreamInput in) throws IOException {
             super(in);
             this.taskParams = new StartTrainedModelDeploymentAction.TaskParams(in);
+            if (in.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_ADAPTIVE_ALLOCATIONS)) {
+                this.adaptiveAllocationsSettings = in.readOptionalWriteable(AdaptiveAllocationsSettings::new);
+            } else {
+                this.adaptiveAllocationsSettings = null;
+            }
         }
 
         @Override
@@ -54,24 +63,32 @@ public ActionRequestValidationException validate() {
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
             taskParams.writeTo(out);
+            if (out.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_ADAPTIVE_ALLOCATIONS)) {
+                out.writeOptionalWriteable(adaptiveAllocationsSettings);
+            }
         }
 
         @Override
         public boolean equals(Object o) {
             if (this == o) return true;
             if (o == null || getClass() != o.getClass()) return false;
             Request request = (Request) o;
-            return Objects.equals(taskParams, request.taskParams);
+            return Objects.equals(taskParams, request.taskParams)
+                && Objects.equals(adaptiveAllocationsSettings, request.adaptiveAllocationsSettings);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(taskParams);
+            return Objects.hash(taskParams, adaptiveAllocationsSettings);
         }
 
         public StartTrainedModelDeploymentAction.TaskParams getTaskParams() {
             return taskParams;
         }
+
+        public AdaptiveAllocationsSettings getAdaptiveAllocationsSettings() {
+            return adaptiveAllocationsSettings;
+        }
     }
 
     public static class Response extends ActionResponse implements ToXContentObject {

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

@@ -29,8 +29,10 @@
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.MlConfigVersion;
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
+import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AllocationStatus;
 import org.elasticsearch.xpack.core.ml.inference.assignment.Priority;
+import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignment;
 import org.elasticsearch.xpack.core.ml.job.messages.Messages;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 import org.elasticsearch.xpack.core.ml.utils.MlTaskParams;
@@ -40,7 +42,6 @@
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
-import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
 import static org.elasticsearch.xpack.core.ml.MlTasks.trainedModelAssignmentTaskDescription;
 
 public class StartTrainedModelDeploymentAction extends ActionType<CreateTrainedModelAssignmentAction.Response> {
@@ -99,6 +100,7 @@ public static class Request extends MasterNodeRequest<Request> implements ToXCon
         public static final ParseField QUEUE_CAPACITY = TaskParams.QUEUE_CAPACITY;
         public static final ParseField CACHE_SIZE = TaskParams.CACHE_SIZE;
         public static final ParseField PRIORITY = TaskParams.PRIORITY;
+        public static final ParseField ADAPTIVE_ALLOCATIONS = TrainedModelAssignment.ADAPTIVE_ALLOCATIONS;
 
         public static final ObjectParser<Request, Void> PARSER = new ObjectParser<>(NAME, Request::new);
 
@@ -117,6 +119,12 @@ public static class Request extends MasterNodeRequest<Request> implements ToXCon
                 ObjectParser.ValueType.VALUE
             );
             PARSER.declareString(Request::setPriority, PRIORITY);
+            PARSER.declareObjectOrNull(
+                Request::setAdaptiveAllocationsSettings,
+                (p, c) -> AdaptiveAllocationsSettings.PARSER.parse(p, c).build(),
+                null,
+                ADAPTIVE_ALLOCATIONS
+            );
         }
 
         public static Request parseRequest(String modelId, String deploymentId, XContentParser parser) {
@@ -140,7 +148,8 @@ public static Request parseRequest(String modelId, String deploymentId, XContent
         private TimeValue timeout = DEFAULT_TIMEOUT;
         private AllocationStatus.State waitForState = DEFAULT_WAITFOR_STATE;
         private ByteSizeValue cacheSize;
-        private int numberOfAllocations = DEFAULT_NUM_ALLOCATIONS;
+        private Integer numberOfAllocations;
+        private AdaptiveAllocationsSettings adaptiveAllocationsSettings = null;
         private int threadsPerAllocation = DEFAULT_NUM_THREADS;
         private int queueCapacity = DEFAULT_QUEUE_CAPACITY;
         private Priority priority = DEFAULT_PRIORITY;
@@ -160,7 +169,11 @@ public Request(StreamInput in) throws IOException {
             modelId = in.readString();
             timeout = in.readTimeValue();
             waitForState = in.readEnum(AllocationStatus.State.class);
-            numberOfAllocations = in.readVInt();
+            if (in.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_ADAPTIVE_ALLOCATIONS)) {
+                numberOfAllocations = in.readOptionalVInt();
+            } else {
+                numberOfAllocations = in.readVInt();
+            }
             threadsPerAllocation = in.readVInt();
             queueCapacity = in.readVInt();
             if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_4_0)) {
@@ -171,12 +184,16 @@ public Request(StreamInput in) throws IOException {
             } else {
                 this.priority = Priority.NORMAL;
             }
-
             if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_8_0)) {
                 this.deploymentId = in.readString();
             } else {
                 this.deploymentId = modelId;
             }
+            if (in.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_ADAPTIVE_ALLOCATIONS)) {
+                this.adaptiveAllocationsSettings = in.readOptionalWriteable(AdaptiveAllocationsSettings::new);
+            } else {
+                this.adaptiveAllocationsSettings = null;
+            }
         }
 
         public final void setModelId(String modelId) {
@@ -212,14 +229,34 @@ public Request setWaitForState(AllocationStatus.State waitForState) {
             return this;
         }
 
-        public int getNumberOfAllocations() {
+        public Integer getNumberOfAllocations() {
             return numberOfAllocations;
         }
 
-        public void setNumberOfAllocations(int numberOfAllocations) {
+        public int computeNumberOfAllocations() {
+            if (numberOfAllocations != null) {
+                return numberOfAllocations;
+            } else {
+                if (adaptiveAllocationsSettings == null || adaptiveAllocationsSettings.getMinNumberOfAllocations() == null) {
+                    return DEFAULT_NUM_ALLOCATIONS;
+                } else {
+                    return adaptiveAllocationsSettings.getMinNumberOfAllocations();
+                }
+            }
+        }
+
+        public void setNumberOfAllocations(Integer numberOfAllocations) {
             this.numberOfAllocations = numberOfAllocations;
         }
 
+        public AdaptiveAllocationsSettings getAdaptiveAllocationsSettings() {
+            return adaptiveAllocationsSettings;
+        }
+
+        public void setAdaptiveAllocationsSettings(AdaptiveAllocationsSettings adaptiveAllocationsSettings) {
+            this.adaptiveAllocationsSettings = adaptiveAllocationsSettings;
+        }
+
         public int getThreadsPerAllocation() {
             return threadsPerAllocation;
         }
@@ -258,7 +295,11 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeString(modelId);
             out.writeTimeValue(timeout);
             out.writeEnum(waitForState);
-            out.writeVInt(numberOfAllocations);
+            if (out.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_ADAPTIVE_ALLOCATIONS)) {
+                out.writeOptionalVInt(numberOfAllocations);
+            } else {
+                out.writeVInt(numberOfAllocations);
+            }
             out.writeVInt(threadsPerAllocation);
             out.writeVInt(queueCapacity);
             if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_4_0)) {
@@ -270,6 +311,9 @@ public void writeTo(StreamOutput out) throws IOException {
             if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_8_0)) {
                 out.writeString(deploymentId);
             }
+            if (out.getTransportVersion().onOrAfter(TransportVersions.INFERENCE_ADAPTIVE_ALLOCATIONS)) {
+                out.writeOptionalWriteable(adaptiveAllocationsSettings);
+            }
         }
 
         @Override
@@ -279,7 +323,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(DEPLOYMENT_ID.getPreferredName(), deploymentId);
             builder.field(TIMEOUT.getPreferredName(), timeout.getStringRep());
             builder.field(WAIT_FOR.getPreferredName(), waitForState);
-            builder.field(NUMBER_OF_ALLOCATIONS.getPreferredName(), numberOfAllocations);
+            if (numberOfAllocations != null) {
+                builder.field(NUMBER_OF_ALLOCATIONS.getPreferredName(), numberOfAllocations);
+            }
+            if (adaptiveAllocationsSettings != null) {
+                builder.field(ADAPTIVE_ALLOCATIONS.getPreferredName(), adaptiveAllocationsSettings);
+            }
             builder.field(THREADS_PER_ALLOCATION.getPreferredName(), threadsPerAllocation);
             builder.field(QUEUE_CAPACITY.getPreferredName(), queueCapacity);
             if (cacheSize != null) {
@@ -301,12 +350,25 @@ public ActionRequestValidationException validate() {
                         + Strings.arrayToCommaDelimitedString(VALID_WAIT_STATES)
                 );
             }
-            if (numberOfAllocations < 1) {
-                validationException.addValidationError("[" + NUMBER_OF_ALLOCATIONS + "] must be a positive integer");
+            if (numberOfAllocations != null) {
+                if (numberOfAllocations < 1) {
+                    validationException.addValidationError("[" + NUMBER_OF_ALLOCATIONS + "] must be a positive integer");
+                }
+                if (adaptiveAllocationsSettings != null && adaptiveAllocationsSettings.getEnabled()) {
+                    validationException.addValidationError(
+                        "[" + NUMBER_OF_ALLOCATIONS + "] cannot be set if adaptive allocations is enabled"
+                    );
+                }
             }
             if (threadsPerAllocation < 1) {
                 validationException.addValidationError("[" + THREADS_PER_ALLOCATION + "] must be a positive integer");
             }
+            ActionRequestValidationException autoscaleException = adaptiveAllocationsSettings == null
+                ? null
+                : adaptiveAllocationsSettings.validate();
+            if (autoscaleException != null) {
+                validationException.addValidationErrors(autoscaleException.validationErrors());
+            }
             if (threadsPerAllocation > MAX_THREADS_PER_ALLOCATION || isPowerOf2(threadsPerAllocation) == false) {
                 validationException.addValidationError(
                     "[" + THREADS_PER_ALLOCATION + "] must be a power of 2 less than or equal to " + MAX_THREADS_PER_ALLOCATION
@@ -322,7 +384,7 @@ public ActionRequestValidationException validate() {
                 validationException.addValidationError("[" + TIMEOUT + "] must be positive");
             }
             if (priority == Priority.LOW) {
-                if (numberOfAllocations > 1) {
+                if (numberOfAllocations != null && numberOfAllocations > 1) {
                     validationException.addValidationError("[" + NUMBER_OF_ALLOCATIONS + "] must be 1 when [" + PRIORITY + "] is low");
                 }
                 if (threadsPerAllocation > 1) {
@@ -344,6 +406,7 @@ public int hashCode() {
                 timeout,
                 waitForState,
                 numberOfAllocations,
+                adaptiveAllocationsSettings,
                 threadsPerAllocation,
                 queueCapacity,
                 cacheSize,
@@ -365,7 +428,8 @@ public boolean equals(Object obj) {
                 && Objects.equals(timeout, other.timeout)
                 && Objects.equals(waitForState, other.waitForState)
                 && Objects.equals(cacheSize, other.cacheSize)
-                && numberOfAllocations == other.numberOfAllocations
+                && Objects.equals(numberOfAllocations, other.numberOfAllocations)
+                && Objects.equals(adaptiveAllocationsSettings, other.adaptiveAllocationsSettings)
                 && threadsPerAllocation == other.threadsPerAllocation
                 && queueCapacity == other.queueCapacity
                 && priority == other.priority;
@@ -430,7 +494,7 @@ public static boolean mayAssignToNode(@Nullable DiscoveryNode node) {
             PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), THREADS_PER_ALLOCATION);
             PARSER.declareInt(ConstructingObjectParser.constructorArg(), QUEUE_CAPACITY);
             PARSER.declareField(
-                optionalConstructorArg(),
+                ConstructingObjectParser.optionalConstructorArg(),
                 (p, c) -> ByteSizeValue.parseBytesSizeValue(p.text(), CACHE_SIZE.getPreferredName()),
                 CACHE_SIZE,
                 ObjectParser.ValueType.VALUE