[ML] Handle new actual_memory_usage_bytes field in model size stats.

The `actual_memory_usage_bytes` field represents the real, physical memory allocated to the `autodetect` process as reported by the OS. Reporting this value in the model size stats associated with an AD job is useful, especially in OOM situations.

Author: edsavage

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

@@ -213,6 +213,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion REMOTE_EXCEPTION = def(9_044_0_00);
     public static final TransportVersion ESQL_REMOVE_AGGREGATE_TYPE = def(9_045_0_00);
     public static final TransportVersion ADD_PROJECT_ID_TO_DSL_ERROR_INFO = def(9_046_0_00);
+    public static final TransportVersion ML_AD_ACTUAL_MEMORY_USAGE = def(9_047_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -153,12 +153,13 @@ private static void checkUniqueness(int id, String uniqueId) {
     // V_11 is used in ELSER v2 package configs
     public static final MlConfigVersion V_11 = registerMlConfigVersion(11_00_0_0_99, "79CB2950-57C7-11EE-AE5D-0800200C9A66");
     public static final MlConfigVersion V_12 = registerMlConfigVersion(12_00_0_0_99, "Trained model config prefix strings added");
+    public static final MlConfigVersion V_13 = registerMlConfigVersion(13_00_0_0_99, "Anomaly Detection reports actual memory usage");
 
     /**
      * Reference to the most recent Ml config version.
      * This should be the Ml config version with the highest id.
      */
-    public static final MlConfigVersion CURRENT = V_12;
+    public static final MlConfigVersion CURRENT = V_13;
 
     /**
      * Reference to the first MlConfigVersion that is detached from the

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/job/process/autodetect/state/ModelSizeStats.java

@@ -41,6 +41,7 @@ public class ModelSizeStats implements ToXContentObject, Writeable {
      */
     public static final ParseField MODEL_BYTES_FIELD = new ParseField("model_bytes");
     public static final ParseField PEAK_MODEL_BYTES_FIELD = new ParseField("peak_model_bytes");
+    public static final ParseField ACTUAL_MEMORY_USAGE_BYTES = new ParseField("actual_memory_usage_bytes");
     public static final ParseField MODEL_BYTES_EXCEEDED_FIELD = new ParseField("model_bytes_exceeded");
     public static final ParseField MODEL_BYTES_MEMORY_LIMIT_FIELD = new ParseField("model_bytes_memory_limit");
     public static final ParseField TOTAL_BY_FIELD_COUNT_FIELD = new ParseField("total_by_field_count");
@@ -74,6 +75,7 @@ private static ConstructingObjectParser<Builder, Void> createParser(boolean igno
         parser.declareString((modelSizeStat, s) -> {}, Result.RESULT_TYPE);
         parser.declareLong(Builder::setModelBytes, MODEL_BYTES_FIELD);
         parser.declareLong(Builder::setPeakModelBytes, PEAK_MODEL_BYTES_FIELD);
+        parser.declareLong(Builder::setActualMemoryUsageBytes, ACTUAL_MEMORY_USAGE_BYTES);
         parser.declareLong(Builder::setModelBytesExceeded, MODEL_BYTES_EXCEEDED_FIELD);
         parser.declareLong(Builder::setModelBytesMemoryLimit, MODEL_BYTES_MEMORY_LIMIT_FIELD);
         parser.declareLong(Builder::setBucketAllocationFailuresCount, BUCKET_ALLOCATION_FAILURES_COUNT_FIELD);
@@ -152,14 +154,16 @@ public String toString() {
      * 1. The job's model_memory_limit
      * 2. The current model memory, i.e. what's reported in model_bytes of this object
      * 3. The peak model memory, i.e. what's reported in peak_model_bytes of this object
+     * 4. The actual memory usage, i.e. what's reported in actual_memory_usage_bytes of this object
      * The field storing this enum can also be <code>null</code>, which means the
      * assignment code will decide on the fly - this was the old behaviour prior
      * to 7.11.
      */
     public enum AssignmentMemoryBasis implements Writeable {
         MODEL_MEMORY_LIMIT,
         CURRENT_MODEL_BYTES,
-        PEAK_MODEL_BYTES;
+        PEAK_MODEL_BYTES,
+        ACTUAL_MEMORY_USAGE_BYTES;
 
         public static AssignmentMemoryBasis fromString(String statusName) {
             return valueOf(statusName.trim().toUpperCase(Locale.ROOT));
@@ -183,6 +187,7 @@ public String toString() {
     private final String jobId;
     private final long modelBytes;
     private final Long peakModelBytes;
+    private final Long actualMemoryUsageBytes;
     private final Long modelBytesExceeded;
     private final Long modelBytesMemoryLimit;
     private final long totalByFieldCount;
@@ -206,6 +211,7 @@ private ModelSizeStats(
         String jobId,
         long modelBytes,
         Long peakModelBytes,
+        Long actualMemoryUsageBytes,
         Long modelBytesExceeded,
         Long modelBytesMemoryLimit,
         long totalByFieldCount,
@@ -228,6 +234,7 @@ private ModelSizeStats(
         this.jobId = jobId;
         this.modelBytes = modelBytes;
         this.peakModelBytes = peakModelBytes;
+        this.actualMemoryUsageBytes = actualMemoryUsageBytes;
         this.modelBytesExceeded = modelBytesExceeded;
         this.modelBytesMemoryLimit = modelBytesMemoryLimit;
         this.totalByFieldCount = totalByFieldCount;
@@ -252,6 +259,11 @@ public ModelSizeStats(StreamInput in) throws IOException {
         jobId = in.readString();
         modelBytes = in.readVLong();
         peakModelBytes = in.readOptionalLong();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.ML_AD_ACTUAL_MEMORY_USAGE)) {
+            actualMemoryUsageBytes = in.readOptionalLong();
+        } else {
+            actualMemoryUsageBytes = null;
+        }
         modelBytesExceeded = in.readOptionalLong();
         modelBytesMemoryLimit = in.readOptionalLong();
         totalByFieldCount = in.readVLong();
@@ -293,6 +305,9 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeString(jobId);
         out.writeVLong(modelBytes);
         out.writeOptionalLong(peakModelBytes);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.ML_AD_ACTUAL_MEMORY_USAGE)) {
+            out.writeOptionalLong(actualMemoryUsageBytes);
+        }
         out.writeOptionalLong(modelBytesExceeded);
         out.writeOptionalLong(modelBytesMemoryLimit);
         out.writeVLong(totalByFieldCount);
@@ -339,6 +354,9 @@ public XContentBuilder doXContentBody(XContentBuilder builder) throws IOExceptio
         if (peakModelBytes != null) {
             builder.field(PEAK_MODEL_BYTES_FIELD.getPreferredName(), peakModelBytes);
         }
+        if (actualMemoryUsageBytes != null) {
+            builder.field(ACTUAL_MEMORY_USAGE_BYTES.getPreferredName(), actualMemoryUsageBytes);
+        }
         if (modelBytesExceeded != null) {
             builder.field(MODEL_BYTES_EXCEEDED_FIELD.getPreferredName(), modelBytesExceeded);
         }
@@ -391,6 +409,10 @@ public Long getPeakModelBytes() {
         return peakModelBytes;
     }
 
+    public Long getActualMemoryUsageBytes() {
+        return actualMemoryUsageBytes;
+    }
+
     public Long getModelBytesExceeded() {
         return modelBytesExceeded;
     }
@@ -479,6 +501,7 @@ public int hashCode() {
             jobId,
             modelBytes,
             peakModelBytes,
+            actualMemoryUsageBytes,
             modelBytesExceeded,
             modelBytesMemoryLimit,
             totalByFieldCount,
@@ -517,6 +540,7 @@ public boolean equals(Object other) {
 
         return this.modelBytes == that.modelBytes
             && Objects.equals(this.peakModelBytes, that.peakModelBytes)
+            && this.actualMemoryUsageBytes == that.actualMemoryUsageBytes
             && Objects.equals(this.modelBytesExceeded, that.modelBytesExceeded)
             && Objects.equals(this.modelBytesMemoryLimit, that.modelBytesMemoryLimit)
             && this.totalByFieldCount == that.totalByFieldCount
@@ -543,6 +567,7 @@ public static class Builder {
         private final String jobId;
         private long modelBytes;
         private Long peakModelBytes;
+        private Long actualMemoryUsageBytes;
         private Long modelBytesExceeded;
         private Long modelBytesMemoryLimit;
         private long totalByFieldCount;
@@ -573,6 +598,7 @@ public Builder(ModelSizeStats modelSizeStats) {
             this.jobId = modelSizeStats.jobId;
             this.modelBytes = modelSizeStats.modelBytes;
             this.peakModelBytes = modelSizeStats.peakModelBytes;
+            this.actualMemoryUsageBytes = modelSizeStats.actualMemoryUsageBytes;
             this.modelBytesExceeded = modelSizeStats.modelBytesExceeded;
             this.modelBytesMemoryLimit = modelSizeStats.modelBytesMemoryLimit;
             this.totalByFieldCount = modelSizeStats.totalByFieldCount;
@@ -603,6 +629,12 @@ public Builder setPeakModelBytes(long peakModelBytes) {
             return this;
         }
 
+        public Builder setActualMemoryUsageBytes(long actualMemoryUsageBytes)
+        {
+            this.actualMemoryUsageBytes = actualMemoryUsageBytes;
+            return this;
+        }
+
         public Builder setModelBytesExceeded(long modelBytesExceeded) {
             this.modelBytesExceeded = modelBytesExceeded;
             return this;
@@ -700,6 +732,7 @@ public ModelSizeStats build() {
                 jobId,
                 modelBytes,
                 peakModelBytes,
+                actualMemoryUsageBytes,
                 modelBytesExceeded,
                 modelBytesMemoryLimit,
                 totalByFieldCount,

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/persistence/JobResultsProvider.java

@@ -1590,6 +1590,13 @@ void calculateEstablishedMemoryUsage(
                     handler.accept((storedPeak != null) ? storedPeak : latestModelSizeStats.getModelBytes());
                     return;
                 }
+                case ACTUAL_MEMORY_USAGE_BYTES -> {
+                    Long storedActualMemoryUsageBytes = latestModelSizeStats.getActualMemoryUsageBytes();
+                    handler.accept(
+                        (storedActualMemoryUsageBytes != null) ? storedActualMemoryUsageBytes : latestModelSizeStats.getModelBytes()
+                    );
+                    return;
+                }
             }
         }
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/process/autodetect/AutodetectProcessManager.java

@@ -1076,6 +1076,8 @@ public ByteSizeValue getOpenProcessMemoryUsage() {
                     case MODEL_MEMORY_LIMIT -> Optional.ofNullable(modelSizeStats.getModelBytesMemoryLimit()).orElse(0L);
                     case CURRENT_MODEL_BYTES -> modelSizeStats.getModelBytes();
                     case PEAK_MODEL_BYTES -> Optional.ofNullable(modelSizeStats.getPeakModelBytes()).orElse(modelSizeStats.getModelBytes());
+                    case ACTUAL_MEMORY_USAGE_BYTES -> Optional.ofNullable(modelSizeStats.getActualMemoryUsageBytes())
+                        .orElse(modelSizeStats.getModelBytes());
                 };
                 memoryUsedBytes += Job.PROCESS_MEMORY_OVERHEAD.getBytes();
             }

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/job/process/autodetect/AutodetectProcessManagerTests.java

@@ -834,10 +834,12 @@ public void testGetOpenProcessMemoryUsage() {
         long modelMemoryLimitBytes = ByteSizeValue.ofMb(randomIntBetween(10, 1000)).getBytes();
         long peakModelBytes = randomLongBetween(100000, modelMemoryLimitBytes - 1);
         long modelBytes = randomLongBetween(1, peakModelBytes - 1);
+        long actualMemoryUsageBytes = randomLongBetween(262144, peakModelBytes - 1);
         AssignmentMemoryBasis assignmentMemoryBasis = randomFrom(AssignmentMemoryBasis.values());
         modelSizeStats = new ModelSizeStats.Builder("foo").setModelBytesMemoryLimit(modelMemoryLimitBytes)
             .setPeakModelBytes(peakModelBytes)
             .setModelBytes(modelBytes)
+            .setActualMemoryUsageBytes(actualMemoryUsageBytes)
             .setAssignmentMemoryBasis(assignmentMemoryBasis)
             .build();
         when(autodetectCommunicator.getModelSizeStats()).thenReturn(modelSizeStats);
@@ -850,6 +852,7 @@ public void testGetOpenProcessMemoryUsage() {
             case MODEL_MEMORY_LIMIT -> modelMemoryLimitBytes;
             case CURRENT_MODEL_BYTES -> modelBytes;
             case PEAK_MODEL_BYTES -> peakModelBytes;
+            case ACTUAL_MEMORY_USAGE_BYTES -> actualMemoryUsageBytes;
         };
         assertThat(manager.getOpenProcessMemoryUsage(), equalTo(ByteSizeValue.ofBytes(expectedSizeBytes)));
     }