[ML] inference performance optimizations and refactor (#57674)

This is a major refactor of the underlying inference logic.

The main refactor is now we are separating the model configuration and 
the inference interfaces. 

This has the following benefits:
 - we can store extra things with the model that are not 
   necessary for inference (i.e. treenode split information gain)
 - we can optimize inference separate from model serialization and storage.
 - The user is oblivious to the optimizations (other than seeing the benefits).

A major part of this commit is removing all inference related methods from the
trained model configurations (ensemble, tree, etc.) and moving them to a new class.

This new class satisfies a new interface that is ONLY for inference.

The optimizations applied currently are:
- feature maps are flattened once
- feature extraction only happens once at the highest level 
  (improves inference + feature importance through put)
- Only storing what we need for inference + feature importance on heap

Author: benwtrent

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

@@ -10,6 +10,8 @@
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.unit.ByteSizeUnit;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.common.xcontent.ToXContentObject;
@@ -37,7 +39,7 @@ public final class InferenceToXContentCompressor {
     // Either 10% of the configured JVM heap, or 1 GB, which ever is smaller
     private static final long MAX_INFLATED_BYTES = Math.min(
         (long)((0.10) * JvmInfo.jvmInfo().getMem().getHeapMax().getBytes()),
-        1_000_000_000); // 1 gb maximum
+        new ByteSizeValue(1, ByteSizeUnit.GB).getBytes());
 
     private InferenceToXContentCompressor() {}
 
@@ -46,9 +48,9 @@ public static <T extends ToXContentObject> String deflate(T objectToCompress) th
         return deflate(reference);
     }
 
-    static <T> T inflate(String compressedString,
-                         CheckedFunction<XContentParser, T, IOException> parserFunction,
-                         NamedXContentRegistry xContentRegistry) throws IOException {
+    public static <T> T inflate(String compressedString,
+                                CheckedFunction<XContentParser, T, IOException> parserFunction,
+                                NamedXContentRegistry xContentRegistry) throws IOException {
         try(XContentParser parser = JsonXContent.jsonXContent.createParser(xContentRegistry,
             LoggingDeprecationHandler.INSTANCE,
             inflate(compressedString, MAX_INFLATED_BYTES))) {

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

@@ -30,6 +30,9 @@
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.ensemble.StrictlyParsedOutputAggregator;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.ensemble.WeightedMode;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.ensemble.WeightedSum;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.inference.EnsembleInferenceModel;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.inference.InferenceModel;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.inference.TreeInferenceModel;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.langident.LangIdentNeuralNetwork;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.tree.Tree;
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.FrequencyEncoding;
@@ -119,6 +122,13 @@ public List<NamedXContentRegistry.Entry> getNamedXContentParsers() {
         namedXContent.add(new NamedXContentRegistry.Entry(InferenceConfigUpdate.class, RegressionConfigUpdate.NAME,
             RegressionConfigUpdate::fromXContentStrict));
 
+        // Inference models
+        namedXContent.add(new NamedXContentRegistry.Entry(InferenceModel.class, Ensemble.NAME, EnsembleInferenceModel::fromXContent));
+        namedXContent.add(new NamedXContentRegistry.Entry(InferenceModel.class, Tree.NAME, TreeInferenceModel::fromXContent));
+        namedXContent.add(new NamedXContentRegistry.Entry(InferenceModel.class,
+            LangIdentNeuralNetwork.NAME,
+            LangIdentNeuralNetwork::fromXContentLenient));
+
         return namedXContent;
     }
 

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

@@ -33,6 +33,7 @@
 
 import java.io.IOException;
 import java.time.Instant;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -607,6 +608,14 @@ public Builder validate(boolean forCreation) {
             if (input != null && input.getFieldNames().isEmpty()) {
                 validationException = addValidationError("[input.field_names] must not be empty", validationException);
             }
+            if (input != null && input.getFieldNames()
+                .stream()
+                .filter(s -> s.contains("."))
+                .flatMap(s -> Arrays.stream(Strings.delimitedListToStringArray(s, ".")))
+                .anyMatch(String::isEmpty)) {
+                validationException = addValidationError("[input.field_names] must only contain valid dot delimited field names",
+                    validationException);
+            }
             if (forCreation) {
                 validationException = checkIllegalSetting(version, VERSION.getPreferredName(), validationException);
                 validationException = checkIllegalSetting(createdBy, CREATED_BY.getPreferredName(), validationException);

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

@@ -20,8 +20,6 @@
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.LenientlyParsedPreProcessor;
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.PreProcessor;
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.StrictlyParsedPreProcessor;
-import org.elasticsearch.xpack.core.ml.inference.results.InferenceResults;
-import org.elasticsearch.xpack.core.ml.inference.trainedmodel.InferenceConfig;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.LenientlyParsedTrainedModel;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.StrictlyParsedTrainedModel;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TrainedModel;
@@ -32,9 +30,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
 
 public class TrainedModelDefinition implements ToXContentObject, Writeable, Accountable {
@@ -73,7 +69,6 @@ public static TrainedModelDefinition.Builder fromXContent(XContentParser parser,
 
     private final TrainedModel trainedModel;
     private final List<PreProcessor> preProcessors;
-    private Map<String, String> decoderMap;
 
     private TrainedModelDefinition(TrainedModel trainedModel, List<PreProcessor> preProcessors) {
         this.trainedModel = ExceptionsHelper.requireNonNull(trainedModel, TRAINED_MODEL);
@@ -116,37 +111,6 @@ public List<PreProcessor> getPreProcessors() {
         return preProcessors;
     }
 
-    void preProcess(Map<String, Object> fields) {
-        preProcessors.forEach(preProcessor -> preProcessor.process(fields));
-    }
-
-    public InferenceResults infer(Map<String, Object> fields, InferenceConfig config) {
-        preProcess(fields);
-        if (config.requestingImportance() && trainedModel.supportsFeatureImportance() == false) {
-            throw ExceptionsHelper.badRequestException(
-                "Feature importance is not supported for the configured model of type [{}]",
-                trainedModel.getName());
-        }
-        return trainedModel.infer(fields,
-            config,
-            config.requestingImportance() ? getDecoderMap() : Collections.emptyMap());
-    }
-
-    private Map<String, String> getDecoderMap() {
-        if (decoderMap != null) {
-            return decoderMap;
-        }
-        synchronized (this) {
-            if (decoderMap != null) {
-                return decoderMap;
-            }
-            this.decoderMap = preProcessors.stream()
-                .map(PreProcessor::reverseLookup)
-                .collect(HashMap::new, Map::putAll, Map::putAll);
-            return decoderMap;
-        }
-    }
-
     @Override
     public String toString() {
         return Strings.toString(this);
@@ -218,14 +182,6 @@ public Builder setTrainedModel(TrainedModel trainedModel) {
             return this;
         }
 
-        private Builder setTrainedModel(List<TrainedModel> trainedModel) {
-            if (trainedModel.size() != 1) {
-                throw ExceptionsHelper.badRequestException("[{}] must have exactly one trained model defined.",
-                    TRAINED_MODEL.getPreferredName());
-            }
-            return setTrainedModel(trainedModel.get(0));
-        }
-
         private void setProcessorsInOrder(boolean value) {
             this.processorsInOrder = value;
         }

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

@@ -14,7 +14,6 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
-import org.elasticsearch.xpack.core.ml.utils.MapHelper;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -109,7 +108,7 @@ public String getName() {
 
     @Override
     public void process(Map<String, Object> fields) {
-        Object value = MapHelper.dig(field, fields);
+        Object value = fields.get(field);
         if (value == null) {
             return;
         }

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

@@ -14,7 +14,6 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
-import org.elasticsearch.xpack.core.ml.utils.MapHelper;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -94,7 +93,7 @@ public String getName() {
 
     @Override
     public void process(Map<String, Object> fields) {
-        Object value = MapHelper.dig(field, fields);
+        Object value = fields.get(field);
         if (value == null) {
             return;
         }

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

@@ -14,7 +14,6 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
-import org.elasticsearch.xpack.core.ml.utils.MapHelper;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -120,7 +119,7 @@ public String getName() {
 
     @Override
     public void process(Map<String, Object> fields) {
-        Object value = MapHelper.dig(field, fields);
+        Object value = fields.get(field);
         if (value == null) {
             return;
         }

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

@@ -10,17 +10,16 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Map;
 import java.util.Objects;
 
 public class RawInferenceResults implements InferenceResults {
 
     public static final String NAME = "raw";
 
     private final double[] value;
-    private final Map<String, double[]> featureImportance;
+    private final double[][] featureImportance;
 
-    public RawInferenceResults(double[] value, Map<String, double[]> featureImportance) {
+    public RawInferenceResults(double[] value, double[][] featureImportance) {
         this.value = value;
         this.featureImportance = featureImportance;
     }
@@ -29,7 +28,7 @@ public double[] getValue() {
         return value;
     }
 
-    public Map<String, double[]> getFeatureImportance() {
+    public double[][] getFeatureImportance() {
         return featureImportance;
     }
 
@@ -44,7 +43,7 @@ public boolean equals(Object object) {
         if (object == null || getClass() != object.getClass()) { return false; }
         RawInferenceResults that = (RawInferenceResults) object;
         return Arrays.equals(value, that.value)
-            && Objects.equals(featureImportance, that.featureImportance);
+            && Arrays.deepEquals(featureImportance, that.featureImportance);
     }
 
     @Override

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

@@ -7,29 +7,12 @@
 
 import org.apache.lucene.util.Accountable;
 import org.elasticsearch.Version;
-import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.io.stream.NamedWriteable;
-import org.elasticsearch.xpack.core.ml.inference.results.InferenceResults;
 import org.elasticsearch.xpack.core.ml.utils.NamedXContentObject;
 
-import java.util.Map;
 
 public interface TrainedModel extends NamedXContentObject, NamedWriteable, Accountable {
 
-    /**
-     * Infer against the provided fields
-     *
-     * NOTE: Must be thread safe
-     *
-     * @param fields The fields and their values to infer against
-     * @param config The configuration options for inference
-     * @param featureDecoderMap A map for decoding feature value names to their originating feature.
-     *                          Necessary for feature influence.
-     * @return The predicted value. For classification this will be discrete values (e.g. 0.0, or 1.0).
-     *                              For regression this is continuous.
-     */
-    InferenceResults infer(Map<String, Object> fields, InferenceConfig config, @Nullable Map<String, String> featureDecoderMap);
-
     /**
      * @return {@link TargetType} for the model.
      */
@@ -49,21 +32,6 @@ public interface TrainedModel extends NamedXContentObject, NamedWriteable, Accou
      */
     long estimatedNumOperations();
 
-    /**
-     * @return Does the model support feature importance
-     */
-    boolean supportsFeatureImportance();
-
-    /**
-     * Calculates the importance of each feature reference by the model for the passed in field values
-     *
-     * NOTE: Must be thread safe
-     * @param fields The fields inferring against
-     * @param featureDecoder A Map translating processed feature names to their original feature names
-     * @return A {@code Map<String, double[]>} mapping each featureName to its importance
-     */
-    Map<String, double[]> featureImportance(Map<String, Object> fields, Map<String, String> featureDecoder);
-
     default Version getMinimalCompatibilityVersion() {
         return Version.V_7_6_0;
     }