[ML] Restore analytics state if available (#47128)

This commit restores the model state if available in data
frame analytics jobs.

In addition, this changes the start API so that a stopped job
can be restarted. As we now store the progress in the state index
when the task is stopped, we can use it to determine what state
the job was in when it got stopped.

Note that in order to be able to distinguish between a job
that runs for the first time and another that is restarting,
we ensure reindexing progress is reported to be at least 1
for a running task.

Author: dimitris-athanasiou

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.action.support.master.MasterNodeRequest;
 import org.elasticsearch.client.ElasticsearchClient;
 import org.elasticsearch.cluster.metadata.MetaData;
+import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -29,8 +30,11 @@
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
 import org.elasticsearch.xpack.core.ml.job.messages.Messages;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
+import org.elasticsearch.xpack.core.ml.utils.PhaseProgress;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
 import java.util.Objects;
 
 public class StartDataFrameAnalyticsAction extends ActionType<AcknowledgedResponse> {
@@ -150,12 +154,15 @@ public static class TaskParams implements PersistentTaskParams {
 
         public static final Version VERSION_INTRODUCED = Version.V_7_3_0;
 
+        private static final ParseField PROGRESS_ON_START = new ParseField("progress_on_start");
+
         public static ConstructingObjectParser<TaskParams, Void> PARSER = new ConstructingObjectParser<>(
-            MlTasks.DATA_FRAME_ANALYTICS_TASK_NAME, true, a -> new TaskParams((String) a[0], (String) a[1]));
+            MlTasks.DATA_FRAME_ANALYTICS_TASK_NAME, true, a -> new TaskParams((String) a[0], (String) a[1], (List<PhaseProgress>) a[2]));
 
         static {
             PARSER.declareString(ConstructingObjectParser.constructorArg(), DataFrameAnalyticsConfig.ID);
             PARSER.declareString(ConstructingObjectParser.constructorArg(), DataFrameAnalyticsConfig.VERSION);
+            PARSER.declareObjectArray(ConstructingObjectParser.optionalConstructorArg(), PhaseProgress.PARSER, PROGRESS_ON_START);
         }
 
         public static TaskParams fromXContent(XContentParser parser) {
@@ -164,25 +171,36 @@ public static TaskParams fromXContent(XContentParser parser) {
 
         private final String id;
         private final Version version;
+        private final List<PhaseProgress> progressOnStart;
 
-        public TaskParams(String id, Version version) {
+        public TaskParams(String id, Version version, List<PhaseProgress> progressOnStart) {
             this.id = Objects.requireNonNull(id);
             this.version = Objects.requireNonNull(version);
+            this.progressOnStart = Collections.unmodifiableList(progressOnStart);
         }
 
-        private TaskParams(String id, String version) {
-            this(id, Version.fromString(version));
+        private TaskParams(String id, String version, @Nullable List<PhaseProgress> progressOnStart) {
+            this(id, Version.fromString(version), progressOnStart == null ? Collections.emptyList() : progressOnStart);
         }
 
         public TaskParams(StreamInput in) throws IOException {
             this.id = in.readString();
             this.version = Version.readVersion(in);
+            if (in.getVersion().onOrAfter(Version.V_7_5_0)) {
+                progressOnStart = in.readList(PhaseProgress::new);
+            } else {
+                progressOnStart = Collections.emptyList();
+            }
         }
 
         public String getId() {
             return id;
         }
 
+        public List<PhaseProgress> getProgressOnStart() {
+            return progressOnStart;
+        }
+
         @Override
         public String getWriteableName() {
             return MlTasks.DATA_FRAME_ANALYTICS_TASK_NAME;
@@ -197,20 +215,24 @@ public Version getMinimalSupportedVersion() {
         public void writeTo(StreamOutput out) throws IOException {
             out.writeString(id);
             Version.writeVersion(version, out);
+            if (out.getVersion().onOrAfter(Version.V_7_5_0)) {
+                out.writeList(progressOnStart);
+            }
         }
 
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
             builder.field(DataFrameAnalyticsConfig.ID.getPreferredName(), id);
             builder.field(DataFrameAnalyticsConfig.VERSION.getPreferredName(), version);
+            builder.field(PROGRESS_ON_START.getPreferredName(), progressOnStart);
             builder.endObject();
             return builder;
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(id, version);
+            return Objects.hash(id, version, progressOnStart);
         }
 
         @Override
@@ -219,7 +241,9 @@ public boolean equals(Object o) {
             if (o == null || getClass() != o.getClass()) return false;
 
             TaskParams other = (TaskParams) o;
-            return Objects.equals(id, other.id) && Objects.equals(version, other.version);
+            return Objects.equals(id, other.id)
+                && Objects.equals(version, other.version)
+                && Objects.equals(progressOnStart, other.progressOnStart);
         }
     }
 

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

@@ -37,4 +37,9 @@ public interface DataFrameAnalysis extends ToXContentObject, NamedWriteable {
      * @return {@code true} if this analysis persists state that can later be used to restore from a given point
      */
     boolean persistsState();
+
+    /**
+     * Returns the document id for the analysis state
+     */
+    String getStateDocId(String jobId);
 }

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

@@ -174,6 +174,11 @@ public boolean persistsState() {
         return false;
     }
 
+    @Override
+    public String getStateDocId(String jobId) {
+        throw new UnsupportedOperationException("Outlier detection does not support state");
+    }
+
     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

@@ -215,6 +215,11 @@ public boolean persistsState() {
         return true;
     }
 
+    @Override
+    public String getStateDocId(String jobId) {
+        return jobId + "_regression_state#1";
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(dependentVariable, lambda, gamma, eta, maximumNumberTrees, featureBagFraction, predictionFieldName,

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/job/messages/Messages.java

@@ -66,6 +66,7 @@ public final class Messages {
     public static final String DATA_FRAME_ANALYTICS_AUDIT_REUSING_DEST_INDEX = "Using existing destination index [{0}]";
     public static final String DATA_FRAME_ANALYTICS_AUDIT_FINISHED_REINDEXING = "Finished reindexing to destination index [{0}]";
     public static final String DATA_FRAME_ANALYTICS_AUDIT_FINISHED_ANALYSIS = "Finished analysis";
+    public static final String DATA_FRAME_ANALYTICS_AUDIT_RESTORING_STATE = "Restoring from previous model state";
 
     public static final String FILTER_CANNOT_DELETE = "Cannot delete filter [{0}] currently used by jobs {1}";
     public static final String FILTER_CONTAINS_TOO_MANY_ITEMS = "Filter [{0}] contains too many items; up to [{1}] items are allowed";

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/StartDataFrameAnalyticsActionTaskParamsTests.java

@@ -10,8 +10,11 @@
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.test.AbstractSerializingTestCase;
+import org.elasticsearch.xpack.core.ml.utils.PhaseProgress;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 
 public class StartDataFrameAnalyticsActionTaskParamsTests extends AbstractSerializingTestCase<StartDataFrameAnalyticsAction.TaskParams> {
 
@@ -22,7 +25,12 @@ protected StartDataFrameAnalyticsAction.TaskParams doParseInstance(XContentParse
 
     @Override
     protected StartDataFrameAnalyticsAction.TaskParams createTestInstance() {
-        return new StartDataFrameAnalyticsAction.TaskParams(randomAlphaOfLength(10), Version.CURRENT);
+        int phaseCount = randomIntBetween(0, 5);
+        List<PhaseProgress> progressOnStart = new ArrayList<>(phaseCount);
+        for (int i = 0; i < phaseCount; i++) {
+            progressOnStart.add(new PhaseProgress(randomAlphaOfLength(10), randomIntBetween(0, 100)));
+        }
+        return new StartDataFrameAnalyticsAction.TaskParams(randomAlphaOfLength(10), Version.CURRENT, progressOnStart);
     }
 
     @Override

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/dataframe/analyses/OutlierDetectionTests.java

@@ -56,4 +56,10 @@ public void testGetParams_GivenExplicitValues() {
         assertThat((Double) params.get(OutlierDetection.FEATURE_INFLUENCE_THRESHOLD.getPreferredName()),
             is(closeTo(0.42, 1E-9)));
     }
+
+    public void testGetStateDocId() {
+        OutlierDetection outlierDetection = createRandom();
+        assertThat(outlierDetection.persistsState(), is(false));
+        expectThrows(UnsupportedOperationException.class, () -> outlierDetection.getStateDocId("foo"));
+    }
 }

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/dataframe/analyses/RegressionTests.java

@@ -13,6 +13,7 @@
 import java.io.IOException;
 
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
 
 public class RegressionTests extends AbstractSerializingTestCase<Regression> {
 
@@ -124,4 +125,11 @@ public void testRegression_GivenTrainingPercentIsGreaterThan100() {
 
         assertThat(e.getMessage(), equalTo("[training_percent] must be a double in [1, 100]"));
     }
+
+    public void testGetStateDocId() {
+        Regression regression = createRandom();
+        assertThat(regression.persistsState(), is(true));
+        String randomId = randomAlphaOfLength(10);
+        assertThat(regression.getStateDocId(randomId), equalTo(randomId + "_regression_state#1"));
+    }
 }

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

@@ -199,7 +199,8 @@ protected static void assertThatAuditMessagesMatch(String configId, String... ex
         assertBusy(() -> assertTrue(indexExists(AuditorField.NOTIFICATIONS_INDEX)));
         assertBusy(() -> {
             String[] actualAuditMessages = fetchAllAuditMessages(configId);
-            assertThat(actualAuditMessages.length, equalTo(expectedAuditMessagePrefixes.length));
+            assertThat("Messages: " + Arrays.toString(actualAuditMessages), actualAuditMessages.length,
+                equalTo(expectedAuditMessagePrefixes.length));
             for (int i = 0; i < actualAuditMessages.length; i++) {
                 assertThat(actualAuditMessages[i], startsWith(expectedAuditMessagePrefixes[i]));
             }

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

@@ -23,6 +23,7 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
@@ -258,6 +259,87 @@ public void testWithOnlyTrainingRowsAndTrainingPercentIsFifty() throws Exception
         assertModelStatePersisted(jobId);
     }
 
+    public void testStopAndRestart() throws Exception {
+        String jobId = "regression_stop_and_restart";
+        String sourceIndex = jobId + "_source_index";
+
+        BulkRequestBuilder bulkRequestBuilder = client().prepareBulk();
+        bulkRequestBuilder.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+
+        List<Double> featureValues = Arrays.asList(1.0, 2.0, 3.0);
+        List<Double> dependentVariableValues = Arrays.asList(10.0, 20.0, 30.0);
+
+        for (int i = 0; i < 350; i++) {
+            Double field = featureValues.get(i % 3);
+            Double value = dependentVariableValues.get(i % 3);
+
+            IndexRequest indexRequest = new IndexRequest(sourceIndex);
+            indexRequest.source("feature", field, "variable", value);
+            bulkRequestBuilder.add(indexRequest);
+        }
+        BulkResponse bulkResponse = bulkRequestBuilder.get();
+        if (bulkResponse.hasFailures()) {
+            fail("Failed to index data: " + bulkResponse.buildFailureMessage());
+        }
+
+        String destIndex = sourceIndex + "_results";
+        DataFrameAnalyticsConfig config = buildRegressionAnalytics(jobId, new String[] {sourceIndex}, destIndex, null,
+            new Regression("variable"));
+        registerAnalytics(config);
+        putAnalytics(config);
+
+        assertState(jobId, DataFrameAnalyticsState.STOPPED);
+        assertProgress(jobId, 0, 0, 0, 0);
+
+        startAnalytics(jobId);
+
+        // Wait until state is one of REINDEXING or ANALYZING, or until it is STOPPED.
+        assertBusy(() -> {
+            DataFrameAnalyticsState state = getAnalyticsStats(jobId).get(0).getState();
+            assertThat(state, is(anyOf(equalTo(DataFrameAnalyticsState.REINDEXING), equalTo(DataFrameAnalyticsState.ANALYZING),
+                equalTo(DataFrameAnalyticsState.STOPPED))));
+        });
+        stopAnalytics(jobId);
+        waitUntilAnalyticsIsStopped(jobId);
+
+        // Now let's start it again
+        try {
+            startAnalytics(jobId);
+        } catch (Exception e) {
+            if (e.getMessage().equals("Cannot start because the job has already finished")) {
+                // That means the job had managed to complete
+            } else {
+                throw e;
+            }
+        }
+
+        waitUntilAnalyticsIsStopped(jobId);
+
+        SearchResponse sourceData = client().prepareSearch(sourceIndex).setTrackTotalHits(true).setSize(1000).get();
+        for (SearchHit hit : sourceData.getHits()) {
+            GetResponse destDocGetResponse = client().prepareGet().setIndex(config.getDest().getIndex()).setId(hit.getId()).get();
+            assertThat(destDocGetResponse.isExists(), is(true));
+            Map<String, Object> sourceDoc = hit.getSourceAsMap();
+            Map<String, Object> destDoc = destDocGetResponse.getSource();
+            for (String field : sourceDoc.keySet()) {
+                assertThat(destDoc.containsKey(field), is(true));
+                assertThat(destDoc.get(field), equalTo(sourceDoc.get(field)));
+            }
+            assertThat(destDoc.containsKey("ml"), is(true));
+
+            @SuppressWarnings("unchecked")
+            Map<String, Object> resultsObject = (Map<String, Object>) destDoc.get("ml");
+
+            assertThat(resultsObject.containsKey("variable_prediction"), is(true));
+            assertThat(resultsObject.containsKey("is_training"), is(true));
+            assertThat(resultsObject.get("is_training"), is(true));
+        }
+
+        assertProgress(jobId, 100, 100, 100, 100);
+        assertThat(searchStoredProgress(jobId).getHits().getTotalHits().value, equalTo(1L));
+        assertModelStatePersisted(jobId);
+    }
+
     private void assertModelStatePersisted(String jobId) {
         String docId = jobId + "_regression_state#1";
         SearchResponse searchResponse = client().prepareSearch(AnomalyDetectorsIndex.jobStateIndexPattern())

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

@@ -31,6 +31,7 @@
 import java.util.Map;
 
 import static org.hamcrest.Matchers.allOf;
+import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
@@ -476,4 +477,70 @@ public void testModelMemoryLimitLowerThanEstimatedMemoryUsage() throws Exception
             "Created analytics with analysis type [outlier_detection]",
             "Estimated memory usage for this analytics to be");
     }
+
+    public void testOutlierDetectionStopAndRestart() throws Exception {
+        String sourceIndex = "test-outlier-detection-stop-and-restart";
+
+        client().admin().indices().prepareCreate(sourceIndex)
+            .addMapping("_doc", "numeric_1", "type=double", "numeric_2", "type=float", "categorical_1", "type=keyword")
+            .get();
+
+        BulkRequestBuilder bulkRequestBuilder = client().prepareBulk();
+        bulkRequestBuilder.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+
+        int docCount = randomIntBetween(1024, 2048);
+        for (int i = 0; i < docCount; i++) {
+            IndexRequest indexRequest = new IndexRequest(sourceIndex);
+            indexRequest.source("numeric_1", randomDouble(), "numeric_2", randomFloat(), "categorical_1", randomAlphaOfLength(10));
+            bulkRequestBuilder.add(indexRequest);
+        }
+        BulkResponse bulkResponse = bulkRequestBuilder.get();
+        if (bulkResponse.hasFailures()) {
+            fail("Failed to index data: " + bulkResponse.buildFailureMessage());
+        }
+
+        String id = "test_outlier_detection_stop_and_restart";
+        DataFrameAnalyticsConfig config = buildOutlierDetectionAnalytics(
+            id, new String[] {sourceIndex}, sourceIndex + "-results", "custom_ml");
+        registerAnalytics(config);
+        putAnalytics(config);
+
+        assertState(id, DataFrameAnalyticsState.STOPPED);
+        startAnalytics(id);
+
+        // Wait until state is one of REINDEXING or ANALYZING, or until it is STOPPED.
+        assertBusy(() -> {
+            DataFrameAnalyticsState state = getAnalyticsStats(id).get(0).getState();
+            assertThat(state, is(anyOf(equalTo(DataFrameAnalyticsState.REINDEXING), equalTo(DataFrameAnalyticsState.ANALYZING),
+                equalTo(DataFrameAnalyticsState.STOPPED))));
+        });
+        stopAnalytics(id);
+        waitUntilAnalyticsIsStopped(id);
+
+        // Now let's start it again
+        try {
+            startAnalytics(id);
+        } catch (Exception e) {
+            if (e.getMessage().equals("Cannot start because the job has already finished")) {
+                // That means the job had managed to complete
+            } else {
+                throw e;
+            }
+        }
+
+        waitUntilAnalyticsIsStopped(id);
+
+        // Check we've got all docs
+        SearchResponse searchResponse = client().prepareSearch(config.getDest().getIndex()).setTrackTotalHits(true).get();
+        assertThat(searchResponse.getHits().getTotalHits().value, equalTo((long) docCount));
+
+        // Check they all have an outlier_score
+        searchResponse = client().prepareSearch(config.getDest().getIndex())
+            .setTrackTotalHits(true)
+            .setQuery(QueryBuilders.existsQuery("custom_ml.outlier_score")).get();
+        assertThat(searchResponse.getHits().getTotalHits().value, equalTo((long) docCount));
+
+        assertProgress(id, 100, 100, 100, 100);
+        assertThat(searchStoredProgress(id).getHits().getTotalHits().value, equalTo(1L));
+    }
 }