[7.x][ML] Perform test inference on java (#58877) (#59298)

Since we are able to load the inference model
and perform inference in java, we no longer need
to rely on the analytics process to be performing
test inference on the docs that were not used for
training. The benefit is that we do not need to
send test docs and fit them in memory of the c++
process.

Backport of #58877

Co-authored-by: Dimitris Athanasiou <[email protected]>

Co-authored-by: Benjamin Trent <[email protected]>

Author: dimitris-athanasiou

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/analyses/Classification.java

@@ -370,6 +370,11 @@ public InferenceConfig inferenceConfig(FieldInfo fieldInfo) {
             .build();
     }
 
+    @Override
+    public boolean supportsInference() {
+        return true;
+    }
+
     public static String extractJobIdFromStateDoc(String stateDocId) {
         int suffixIndex = stateDocId.lastIndexOf(STATE_DOC_ID_SUFFIX);
         return suffixIndex <= 0 ? null : stateDocId.substring(0, suffixIndex);

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/analyses/DataFrameAnalysis.java

@@ -78,6 +78,11 @@ public interface DataFrameAnalysis extends ToXContentObject, NamedWriteable {
     @Nullable
     InferenceConfig inferenceConfig(FieldInfo fieldInfo);
 
+    /**
+     * @return {@code true} if this analysis trains a model that can be used for inference
+     */
+    boolean supportsInference();
+
     /**
      * Summarizes information about the fields that is necessary for analysis to generate
      * the parameters needed for the process configuration.

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/analyses/OutlierDetection.java

@@ -257,6 +257,11 @@ public InferenceConfig inferenceConfig(FieldInfo fieldInfo) {
         return null;
     }
 
+    @Override
+    public boolean supportsInference() {
+        return false;
+    }
+
     public enum Method {
         LOF, LDOF, DISTANCE_KTH_NN, DISTANCE_KNN;
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/dataframe/analyses/Regression.java

@@ -283,6 +283,11 @@ public InferenceConfig inferenceConfig(FieldInfo fieldInfo) {
             .build();
     }
 
+    @Override
+    public boolean supportsInference() {
+        return true;
+    }
+
     public static String extractJobIdFromStateDoc(String stateDocId) {
         int suffixIndex = stateDocId.lastIndexOf(STATE_DOC_ID_SUFFIX);
         return suffixIndex <= 0 ? null : stateDocId.substring(0, suffixIndex);

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

@@ -17,6 +17,8 @@
 
 import java.io.IOException;
 import java.util.Collections;
+import java.util.Map;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Objects;
 import java.util.stream.Collectors;
@@ -137,18 +139,20 @@ public Object predictedValue() {
     public void writeResult(IngestDocument document, String parentResultField) {
         ExceptionsHelper.requireNonNull(document, "document");
         ExceptionsHelper.requireNonNull(parentResultField, "parentResultField");
-        document.setFieldValue(parentResultField + "." + this.resultsField,
-            predictionFieldType.transformPredictedValue(value(), valueAsString()));
-        if (topClasses.size() > 0) {
-            document.setFieldValue(parentResultField + "." + topNumClassesField,
-                topClasses.stream().map(TopClassEntry::asValueMap).collect(Collectors.toList()));
+        document.setFieldValue(parentResultField, asMap());
+    }
+
+    @Override
+    public Map<String, Object> asMap() {
+        Map<String, Object> map = new LinkedHashMap<>();
+        map.put(resultsField, predictionFieldType.transformPredictedValue(value(), valueAsString()));
+        if (topClasses.isEmpty() == false) {
+            map.put(topNumClassesField, topClasses.stream().map(TopClassEntry::asValueMap).collect(Collectors.toList()));
         }
-        if (getFeatureImportance().size() > 0) {
-            document.setFieldValue(parentResultField + "." + FEATURE_IMPORTANCE, getFeatureImportance()
-                .stream()
-                .map(FeatureImportance::toMap)
-                .collect(Collectors.toList()));
+        if (getFeatureImportance().isEmpty() == false) {
+            map.put(FEATURE_IMPORTANCE, getFeatureImportance().stream().map(FeatureImportance::toMap).collect(Collectors.toList()));
         }
+        return map;
     }
 
     @Override

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

@@ -9,9 +9,13 @@
 import org.elasticsearch.common.xcontent.ToXContentFragment;
 import org.elasticsearch.ingest.IngestDocument;
 
+import java.util.Map;
+
 public interface InferenceResults extends NamedWriteable, ToXContentFragment {
 
     void writeResult(IngestDocument document, String parentResultField);
 
+    Map<String, Object> asMap();
+
     Object predictedValue();
 }

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

@@ -11,6 +11,7 @@
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.Objects;
 
 public class RawInferenceResults implements InferenceResults {
@@ -57,6 +58,10 @@ public void writeResult(IngestDocument document, String parentResultField) {
         throw new UnsupportedOperationException("[raw] does not support writing inference results");
     }
 
+    @Override
+    public Map<String, Object> asMap() {
+        throw new UnsupportedOperationException("[raw] does not support map conversion");
+    }
     @Override
     public Object predictedValue() {
         return null;

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

@@ -15,7 +15,9 @@
 
 import java.io.IOException;
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Collectors;
 
@@ -85,13 +87,17 @@ public Object predictedValue() {
     public void writeResult(IngestDocument document, String parentResultField) {
         ExceptionsHelper.requireNonNull(document, "document");
         ExceptionsHelper.requireNonNull(parentResultField, "parentResultField");
-        document.setFieldValue(parentResultField + "." + this.resultsField, value());
-        if (getFeatureImportance().size() > 0) {
-            document.setFieldValue(parentResultField + ".feature_importance", getFeatureImportance()
-                .stream()
-                .map(FeatureImportance::toMap)
-                .collect(Collectors.toList()));
+        document.setFieldValue(parentResultField, asMap());
+    }
+
+    @Override
+    public Map<String, Object> asMap() {
+        Map<String, Object> map = new LinkedHashMap<>();
+        map.put(resultsField, value());
+        if (getFeatureImportance().isEmpty() == false) {
+            map.put(FEATURE_IMPORTANCE, getFeatureImportance().stream().map(FeatureImportance::toMap).collect(Collectors.toList()));
         }
+        return map;
     }
 
     @Override

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

@@ -13,6 +13,8 @@
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 
 import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.Objects;
 
 public class WarningInferenceResults implements InferenceResults {
@@ -56,7 +58,14 @@ public int hashCode() {
     public void writeResult(IngestDocument document, String parentResultField) {
         ExceptionsHelper.requireNonNull(document, "document");
         ExceptionsHelper.requireNonNull(parentResultField, "resultField");
-        document.setFieldValue(parentResultField + "." + NAME, warning);
+        document.setFieldValue(parentResultField, asMap());
+    }
+
+    @Override
+    public Map<String, Object> asMap() {
+        Map<String, Object> asMap = new LinkedHashMap<>();
+        asMap.put(NAME, warning);
+        return asMap;
     }
 
     @Override

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/ClassificationEvaluationIT.java

@@ -649,4 +649,9 @@ private static void indexDistinctAnimals(String indexName, int distinctAnimalCou
             fail("Failed to index data: " + bulkResponse.buildFailureMessage());
         }
     }
+
+    @Override
+    boolean supportsInference() {
+        return true;
+    }
 }

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/ClassificationIT.java

@@ -893,4 +893,9 @@ private String stateDocId() {
     private String expectedDestIndexAuditMessage() {
         return (analysisUsesExistingDestIndex ? "Using existing" : "Creating") + " destination index [" + destIndex + "]";
     }
+
+    @Override
+    boolean supportsInference() {
+        return true;
+    }
 }

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/ExplainDataFrameAnalyticsIT.java

@@ -123,7 +123,12 @@ public void testTrainingPercentageIsApplied() throws IOException {
 
         explainResponse = explainDataFrame(config);
 
-        assertThat(explainResponse.getMemoryEstimation().getExpectedMemoryWithoutDisk(), 
+        assertThat(explainResponse.getMemoryEstimation().getExpectedMemoryWithoutDisk(),
                    lessThanOrEqualTo(allDataUsedForTraining));
     }
+
+    @Override
+    boolean supportsInference() {
+        return false;
+    }
 }

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeDataFrameAnalyticsIntegTestCase.java

@@ -219,6 +219,8 @@ protected void assertProgressComplete(String id) {
             progress.stream().allMatch(phaseProgress -> phaseProgress.getProgressPercent() == 100), is(true));
     }
 
+    abstract boolean supportsInference();
+
     private List<PhaseProgress> getProgress(String id) {
         GetDataFrameAnalyticsStatsAction.Response.Stats stats = getAnalyticsStats(id);
         assertThat(stats.getId(), equalTo(id));
@@ -227,7 +229,12 @@ private List<PhaseProgress> getProgress(String id) {
         assertThat(progress.size(), greaterThanOrEqualTo(4));
         assertThat(progress.get(0).getPhase(), equalTo("reindexing"));
         assertThat(progress.get(1).getPhase(), equalTo("loading_data"));
-        assertThat(progress.get(progress.size() - 1).getPhase(), equalTo("writing_results"));
+        if (supportsInference()) {
+            assertThat(progress.get(progress.size() - 2).getPhase(), equalTo("writing_results"));
+            assertThat(progress.get(progress.size() - 1).getPhase(), equalTo("inference"));
+        } else {
+            assertThat(progress.get(progress.size() - 1).getPhase(), equalTo("writing_results"));
+        }
         return progress;
     }
 

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/OutlierDetectionWithMissingFieldsIT.java

@@ -111,4 +111,9 @@ public void testMissingFields() throws Exception {
         assertProgressComplete(id);
         assertThat(searchStoredProgress(id).getHits().getTotalHits().value, equalTo(1L));
     }
+
+    @Override
+    boolean supportsInference() {
+        return false;
+    }
 }

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/RegressionEvaluationIT.java

@@ -155,4 +155,9 @@ private static void indexHousesData(String indexName) {
             fail("Failed to index data: " + bulkResponse.buildFailureMessage());
         }
     }
+
+    @Override
+    boolean supportsInference() {
+        return true;
+    }
 }

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/RegressionIT.java

@@ -540,4 +540,9 @@ private static Map<String, Object> getMlResultsObjectFromDestDoc(Map<String, Obj
     protected String stateDocId() {
         return jobId + "_regression_state#1";
     }
+
+    @Override
+    boolean supportsInference() {
+        return true;
+    }
 }

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/RunDataFrameAnalyticsIT.java

@@ -745,4 +745,9 @@ public void testOutlierDetectionWithCustomParams() throws Exception {
             "Started writing results",
             "Finished analysis");
     }
+
+    @Override
+    boolean supportsInference() {
+        return false;
+    }
 }

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

@@ -693,8 +693,13 @@ public Collection<Object> createComponents(Client client, ClusterService cluster
         this.modelLoadingService.set(modelLoadingService);
 
         // Data frame analytics components
-        AnalyticsProcessManager analyticsProcessManager = new AnalyticsProcessManager(client, threadPool, analyticsProcessFactory,
-            dataFrameAnalyticsAuditor, trainedModelProvider, resultsPersisterService);
+        AnalyticsProcessManager analyticsProcessManager = new AnalyticsProcessManager(client,
+            threadPool,
+            analyticsProcessFactory,
+            dataFrameAnalyticsAuditor,
+            trainedModelProvider,
+            modelLoadingService,
+            resultsPersisterService);
         MemoryUsageEstimationProcessManager memoryEstimationProcessManager =
             new MemoryUsageEstimationProcessManager(
                 threadPool.generic(), threadPool.executor(MachineLearning.JOB_COMMS_THREAD_POOL_NAME), memoryEstimationProcessFactory);

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

@@ -195,7 +195,7 @@ private void searchStats(DataFrameAnalyticsConfig config, ActionListener<Stats>
         logger.debug("[{}] Gathering stats for stopped task", config.getId());
 
         RetrievedStatsHolder retrievedStatsHolder = new RetrievedStatsHolder(
-            ProgressTracker.fromZeroes(config.getAnalysis().getProgressPhases()).report());
+            ProgressTracker.fromZeroes(config.getAnalysis().getProgressPhases(), config.getAnalysis().supportsInference()).report());
 
         MultiSearchRequest multiSearchRequest = new MultiSearchRequest();
         multiSearchRequest.add(buildStoredProgressSearch(config.getId()));

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

@@ -84,7 +84,8 @@ public void execute(DataFrameAnalyticsTask task, DataFrameAnalyticsState current
                 // At this point we have the config at hand and we can reset the progress tracker
                 // to use the analyses phases. We preserve reindexing progress as if reindexing was
                 // finished it will not be reset.
-                task.getStatsHolder().resetProgressTrackerPreservingReindexingProgress(config.getAnalysis().getProgressPhases());
+                task.getStatsHolder().resetProgressTrackerPreservingReindexingProgress(config.getAnalysis().getProgressPhases(),
+                    config.getAnalysis().supportsInference());
 
                 switch(currentState) {
                     // If we are STARTED, it means the job was started because the start API was called.

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

@@ -49,6 +49,11 @@ public final class DestinationIndex {
 
     public static final String ID_COPY = "ml__id_copy";
 
+    /**
+     * The field that indicates whether a doc was used for training or not
+     */
+    public static final String IS_TRAINING = "is_training";
+
     // Metadata fields
     static final String CREATION_DATE_MILLIS = "creation_date_in_millis";
     static final String VERSION = "version";