[ML][Inference] Adding model memory estimations (#48323)

* [ML][Inference] Adding model memory estimations

* addressing PR comments

Author: benwtrent

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

@@ -174,7 +174,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         }
         builder.timeField(CREATE_TIME.getPreferredName(), CREATE_TIME.getPreferredName() + "_string", createTime.toEpochMilli());
         if (definition != null) {
-            builder.field(DEFINITION.getPreferredName(), definition);
+            builder.field(DEFINITION.getPreferredName(), definition, params);
         }
         builder.field(TAGS.getPreferredName(), tags);
         if (metadata != null) {

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

@@ -5,11 +5,15 @@
  */
 package org.elasticsearch.xpack.core.ml.inference;
 
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.Accountables;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.ToXContentObject;
@@ -25,16 +29,21 @@
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TrainedModel;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 import org.elasticsearch.xpack.core.ml.utils.NamedXContentObjectHelper;
+import org.elasticsearch.xpack.core.ml.utils.ToXContentParams;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
-public class TrainedModelDefinition implements ToXContentObject, Writeable {
+public class TrainedModelDefinition implements ToXContentObject, Writeable, Accountable {
 
+    private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(TrainedModelDefinition.class);
     public static final String NAME = "trained_mode_definition";
+    public static final String HEAP_MEMORY_ESTIMATION = "heap_memory_estimation";
 
     public static final ParseField TRAINED_MODEL = new ParseField("trained_model");
     public static final ParseField PREPROCESSORS = new ParseField("preprocessors");
@@ -105,6 +114,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             PREPROCESSORS.getPreferredName(),
             preProcessors);
         builder.field(INPUT.getPreferredName(), input);
+        if (params.paramAsBoolean(ToXContentParams.FOR_INTERNAL_STORAGE, false) == false) {
+            builder.humanReadableField(HEAP_MEMORY_ESTIMATION + "_bytes",
+                HEAP_MEMORY_ESTIMATION,
+                new ByteSizeValue(ramBytesUsed()));
+        }
         builder.endObject();
         return builder;
     }
@@ -150,6 +164,26 @@ public int hashCode() {
         return Objects.hash(trainedModel, input, preProcessors);
     }
 
+    @Override
+    public long ramBytesUsed() {
+        long size = SHALLOW_SIZE;
+        size += RamUsageEstimator.sizeOf(trainedModel);
+        size += RamUsageEstimator.sizeOf(input);
+        size += RamUsageEstimator.sizeOfCollection(preProcessors);
+        return size;
+    }
+
+    @Override
+    public Collection<Accountable> getChildResources() {
+        List<Accountable> accountables = new ArrayList<>(preProcessors.size() + 2);
+        accountables.add(Accountables.namedAccountable("input", input));
+        accountables.add(Accountables.namedAccountable("trained_model", trainedModel));
+        for(PreProcessor preProcessor : preProcessors) {
+            accountables.add(Accountables.namedAccountable("pre_processor_" + preProcessor.getName(), preProcessor));
+        }
+        return accountables;
+    }
+
     public static class Builder {
 
         private List<PreProcessor> preProcessors;
@@ -204,8 +238,9 @@ public TrainedModelDefinition build() {
         }
     }
 
-    public static class Input implements ToXContentObject, Writeable {
+    public static class Input implements ToXContentObject, Writeable, Accountable {
 
+        private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(Input.class);
         public static final String NAME = "trained_mode_definition_input";
         public static final ParseField FIELD_NAMES = new ParseField("field_names");
 
@@ -265,6 +300,15 @@ public int hashCode() {
             return Objects.hash(fieldNames);
         }
 
+        @Override
+        public long ramBytesUsed() {
+            return SHALLOW_SIZE + RamUsageEstimator.sizeOfCollection(fieldNames);
+        }
+
+        @Override
+        public String toString() {
+            return Strings.toString(this);
+        }
     }
 
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/preprocessing/FrequencyEncoding.java

@@ -5,7 +5,9 @@
  */
 package org.elasticsearch.xpack.core.ml.inference.preprocessing;
 
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
@@ -25,6 +27,8 @@
  */
 public class FrequencyEncoding implements LenientlyParsedPreProcessor, StrictlyParsedPreProcessor {
 
+    private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(FrequencyEncoding.class);
+
     public static final ParseField NAME = new ParseField("frequency_encoding");
     public static final ParseField FIELD = new ParseField("field");
     public static final ParseField FEATURE_NAME = new ParseField("feature_name");
@@ -143,4 +147,17 @@ public int hashCode() {
         return Objects.hash(field, featureName, frequencyMap);
     }
 
+    @Override
+    public long ramBytesUsed() {
+        long size = SHALLOW_SIZE;
+        size += RamUsageEstimator.sizeOf(field);
+        size += RamUsageEstimator.sizeOf(featureName);
+        size += RamUsageEstimator.sizeOfMap(frequencyMap);
+        return size;
+    }
+
+    @Override
+    public String toString() {
+        return Strings.toString(this);
+    }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/preprocessing/OneHotEncoding.java

@@ -5,7 +5,9 @@
  */
 package org.elasticsearch.xpack.core.ml.inference.preprocessing;
 
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
@@ -23,6 +25,7 @@
  */
 public class OneHotEncoding implements LenientlyParsedPreProcessor, StrictlyParsedPreProcessor {
 
+    private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(OneHotEncoding.class);
     public static final ParseField NAME = new ParseField("one_hot_encoding");
     public static final ParseField FIELD = new ParseField("field");
     public static final ParseField HOT_MAP = new ParseField("hot_map");
@@ -127,4 +130,16 @@ public int hashCode() {
         return Objects.hash(field, hotMap);
     }
 
+    @Override
+    public long ramBytesUsed() {
+        long size = SHALLOW_SIZE;
+        size += RamUsageEstimator.sizeOf(field);
+        size += RamUsageEstimator.sizeOfMap(hotMap);
+        return size;
+    }
+
+    @Override
+    public String toString() {
+        return Strings.toString(this);
+    }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/preprocessing/PreProcessor.java

@@ -5,6 +5,7 @@
  */
 package org.elasticsearch.xpack.core.ml.inference.preprocessing;
 
+import org.apache.lucene.util.Accountable;
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.xpack.core.ml.utils.NamedXContentObject;
 
@@ -14,7 +15,7 @@
  * Describes a pre-processor for a defined machine learning model
  * This processor should take a set of fields and return the modified set of fields.
  */
-public interface PreProcessor extends NamedXContentObject, NamedWriteable {
+public interface PreProcessor extends NamedXContentObject, NamedWriteable, Accountable {
 
     /**
      * Process the given fields and their values and return the modified map.

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/preprocessing/TargetMeanEncoding.java

@@ -5,7 +5,9 @@
  */
 package org.elasticsearch.xpack.core.ml.inference.preprocessing;
 
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
@@ -25,6 +27,7 @@
  */
 public class TargetMeanEncoding implements LenientlyParsedPreProcessor, StrictlyParsedPreProcessor {
 
+    private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(TargetMeanEncoding.class);
     public static final ParseField NAME = new ParseField("target_mean_encoding");
     public static final ParseField FIELD = new ParseField("field");
     public static final ParseField FEATURE_NAME = new ParseField("feature_name");
@@ -158,4 +161,17 @@ public int hashCode() {
         return Objects.hash(field, featureName, meanMap, defaultValue);
     }
 
+    @Override
+    public long ramBytesUsed() {
+        long size = SHALLOW_SIZE;
+        size += RamUsageEstimator.sizeOf(field);
+        size += RamUsageEstimator.sizeOf(featureName);
+        size += RamUsageEstimator.sizeOfMap(meanMap);
+        return size;
+    }
+
+    @Override
+    public String toString() {
+        return Strings.toString(this);
+    }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/TrainedModel.java

@@ -5,6 +5,7 @@
  */
 package org.elasticsearch.xpack.core.ml.inference.trainedmodel;
 
+import org.apache.lucene.util.Accountable;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.xpack.core.ml.inference.results.InferenceResults;
@@ -13,7 +14,7 @@
 import java.util.List;
 import java.util.Map;
 
-public interface TrainedModel extends NamedXContentObject, NamedWriteable {
+public interface TrainedModel extends NamedXContentObject, NamedWriteable, Accountable {
 
     /**
      * @return List of featureNames expected by the model. In the order that they are expected

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ensemble/Ensemble.java

@@ -5,6 +5,9 @@
  */
 package org.elasticsearch.xpack.core.ml.inference.trainedmodel.ensemble;
 
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.Accountables;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -29,6 +32,8 @@
 import org.elasticsearch.xpack.core.ml.utils.NamedXContentObjectHelper;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -39,6 +44,7 @@
 
 public class Ensemble implements LenientlyParsedTrainedModel, StrictlyParsedTrainedModel {
 
+    private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(Ensemble.class);
     // TODO should we have regression/classification sub-classes that accept the builder?
     public static final ParseField NAME = new ParseField("ensemble");
     public static final ParseField FEATURE_NAMES = new ParseField("feature_names");
@@ -249,6 +255,26 @@ public static Builder builder() {
         return new Builder();
     }
 
+    @Override
+    public long ramBytesUsed() {
+        long size = SHALLOW_SIZE;
+        size += RamUsageEstimator.sizeOfCollection(featureNames);
+        size += RamUsageEstimator.sizeOfCollection(classificationLabels);
+        size += RamUsageEstimator.sizeOfCollection(models);
+        size += outputAggregator.ramBytesUsed();
+        return size;
+    }
+
+    @Override
+    public Collection<Accountable> getChildResources() {
+        List<Accountable> accountables = new ArrayList<>(models.size() + 1);
+        for (TrainedModel model : models) {
+            accountables.add(Accountables.namedAccountable(model.getName(), model));
+        }
+        accountables.add(Accountables.namedAccountable(outputAggregator.getName(), outputAggregator));
+        return Collections.unmodifiableCollection(accountables);
+    }
+
     public static class Builder {
         private List<String> featureNames;
         private List<TrainedModel> trainedModels;