[ML] DFA result processor should only skip rows and model chunks on cancel (elastic#60113)

When the job is force-closed or shutting down due to a fatal error we clean
up all cancellable job operations. This includes cancelling the results processor.
However, this means that we might not persist objects that are written from the
process like stats, memory usage, etc.

In hindsight, we do not gain from cancelling the results processor in its
entirety. It makes more sense to skip row results and model chunks but keep
stats and instrumentation about the job as the latter may contain useful information
to understand what happened to the job.

Author: dimitris-athanasiou

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

@@ -434,7 +434,6 @@ synchronized void stop() {
             if (inferenceRunner.get() != null) {
                 inferenceRunner.get().cancel();
             }
-            statsPersister.cancel();
             if (process.get() != null) {
                 try {
                     process.get().kill();

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

@@ -56,6 +56,7 @@ public class AnalyticsResultProcessor {
     private final ChunkedTrainedModelPersister chunkedTrainedModelPersister;
     private volatile String failure;
     private volatile boolean isCancelled;
+    private long processedRows;
 
     private volatile String latestModelId;
 
@@ -92,31 +93,17 @@ public void awaitForCompletion() {
 
     public void cancel() {
         dataFrameRowsJoiner.cancel();
-        statsPersister.cancel();
         isCancelled = true;
     }
 
     public void process(AnalyticsProcess<AnalyticsResult> process) {
         long totalRows = process.getConfig().rows();
-        long processedRows = 0;
 
         // TODO When java 9 features can be used, we will not need the local variable here
         try (DataFrameRowsJoiner resultsJoiner = dataFrameRowsJoiner) {
             Iterator<AnalyticsResult> iterator = process.readAnalyticsResults();
             while (iterator.hasNext()) {
-                if (isCancelled) {
-                    break;
-                }
-                AnalyticsResult result = iterator.next();
-                processResult(result, resultsJoiner);
-                if (result.getRowResults() != null) {
-                    if (processedRows == 0) {
-                        LOGGER.info("[{}] Started writing results", analytics.getId());
-                        auditor.info(analytics.getId(), Messages.getMessage(Messages.DATA_FRAME_ANALYTICS_AUDIT_STARTED_WRITING_RESULTS));
-                    }
-                    processedRows++;
-                    updateResultsProgress(processedRows >= totalRows ? 100 : (int) (processedRows * 100.0 / totalRows));
-                }
+                processResult(iterator.next(), resultsJoiner, totalRows);
             }
         } catch (Exception e) {
             if (isCancelled) {
@@ -141,10 +128,10 @@ private void completeResultsProgress() {
         statsHolder.getProgressTracker().updateWritingResultsProgress(100);
     }
 
-    private void processResult(AnalyticsResult result, DataFrameRowsJoiner resultsJoiner) {
+    private void processResult(AnalyticsResult result, DataFrameRowsJoiner resultsJoiner, long totalRows) {
         RowResults rowResults = result.getRowResults();
-        if (rowResults != null) {
-            resultsJoiner.processRowResults(rowResults);
+        if (rowResults != null && isCancelled == false) {
+            processRowResult(resultsJoiner, totalRows, rowResults);
         }
         PhaseProgress phaseProgress = result.getPhaseProgress();
         if (phaseProgress != null) {
@@ -157,7 +144,7 @@ private void processResult(AnalyticsResult result, DataFrameRowsJoiner resultsJo
             latestModelId = chunkedTrainedModelPersister.createAndIndexInferenceModelMetadata(modelSize);
         }
         TrainedModelDefinitionChunk trainedModelDefinitionChunk = result.getTrainedModelDefinitionChunk();
-        if (trainedModelDefinitionChunk != null) {
+        if (trainedModelDefinitionChunk != null && isCancelled == false) {
             chunkedTrainedModelPersister.createAndIndexInferenceModelDoc(trainedModelDefinitionChunk);
         }
         MemoryUsage memoryUsage = result.getMemoryUsage();
@@ -181,6 +168,16 @@ private void processResult(AnalyticsResult result, DataFrameRowsJoiner resultsJo
         }
     }
 
+    private void processRowResult(DataFrameRowsJoiner rowsJoiner, long totalRows, RowResults rowResults) {
+        rowsJoiner.processRowResults(rowResults);
+        if (processedRows == 0) {
+            LOGGER.info("[{}] Started writing results", analytics.getId());
+            auditor.info(analytics.getId(), Messages.getMessage(Messages.DATA_FRAME_ANALYTICS_AUDIT_STARTED_WRITING_RESULTS));
+        }
+        processedRows++;
+        updateResultsProgress(processedRows >= totalRows ? 100 : (int) (processedRows * 100.0 / totalRows));
+    }
+
     private void setAndReportFailure(Exception e) {
         LOGGER.error(new ParameterizedMessage("[{}] Error processing results; ", analytics.getId()), e);
         failure = "error processing results; " + e.getMessage();

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/process/results/AnalyticsResult.java

@@ -66,14 +66,14 @@ public class AnalyticsResult implements ToXContentObject {
     private final ModelSizeInfo modelSizeInfo;
     private final TrainedModelDefinitionChunk trainedModelDefinitionChunk;
 
-    public AnalyticsResult(@Nullable RowResults rowResults,
-                           @Nullable PhaseProgress phaseProgress,
-                           @Nullable MemoryUsage memoryUsage,
-                           @Nullable OutlierDetectionStats outlierDetectionStats,
-                           @Nullable ClassificationStats classificationStats,
-                           @Nullable RegressionStats regressionStats,
-                           @Nullable ModelSizeInfo modelSizeInfo,
-                           @Nullable TrainedModelDefinitionChunk trainedModelDefinitionChunk) {
+    private AnalyticsResult(@Nullable RowResults rowResults,
+                            @Nullable PhaseProgress phaseProgress,
+                            @Nullable MemoryUsage memoryUsage,
+                            @Nullable OutlierDetectionStats outlierDetectionStats,
+                            @Nullable ClassificationStats classificationStats,
+                            @Nullable RegressionStats regressionStats,
+                            @Nullable ModelSizeInfo modelSizeInfo,
+                            @Nullable TrainedModelDefinitionChunk trainedModelDefinitionChunk) {
         this.rowResults = rowResults;
         this.phaseProgress = phaseProgress;
         this.memoryUsage = memoryUsage;
@@ -172,4 +172,75 @@ public int hashCode() {
         return Objects.hash(rowResults, phaseProgress, memoryUsage, outlierDetectionStats, classificationStats,
             regressionStats, modelSizeInfo, trainedModelDefinitionChunk);
     }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    public static class Builder {
+
+        private RowResults rowResults;
+        private PhaseProgress phaseProgress;
+        private MemoryUsage memoryUsage;
+        private OutlierDetectionStats outlierDetectionStats;
+        private ClassificationStats classificationStats;
+        private RegressionStats regressionStats;
+        private ModelSizeInfo modelSizeInfo;
+        private TrainedModelDefinitionChunk trainedModelDefinitionChunk;
+
+        private Builder() {}
+
+        public Builder setRowResults(RowResults rowResults) {
+            this.rowResults = rowResults;
+            return this;
+        }
+
+        public Builder setPhaseProgress(PhaseProgress phaseProgress) {
+            this.phaseProgress = phaseProgress;
+            return this;
+        }
+
+        public Builder setMemoryUsage(MemoryUsage memoryUsage) {
+            this.memoryUsage = memoryUsage;
+            return this;
+        }
+
+        public Builder setOutlierDetectionStats(OutlierDetectionStats outlierDetectionStats) {
+            this.outlierDetectionStats = outlierDetectionStats;
+            return this;
+        }
+
+        public Builder setClassificationStats(ClassificationStats classificationStats) {
+            this.classificationStats = classificationStats;
+            return this;
+        }
+
+        public Builder setRegressionStats(RegressionStats regressionStats) {
+            this.regressionStats = regressionStats;
+            return this;
+        }
+
+        public Builder setModelSizeInfo(ModelSizeInfo modelSizeInfo) {
+            this.modelSizeInfo = modelSizeInfo;
+            return this;
+        }
+
+        public Builder setTrainedModelDefinitionChunk(TrainedModelDefinitionChunk trainedModelDefinitionChunk) {
+            this.trainedModelDefinitionChunk = trainedModelDefinitionChunk;
+            return this;
+        }
+
+        public AnalyticsResult build() {
+            return new AnalyticsResult(
+                rowResults,
+                phaseProgress,
+                memoryUsage,
+                outlierDetectionStats,
+                classificationStats,
+                regressionStats,
+                modelSizeInfo,
+                trainedModelDefinitionChunk
+            );
+        }
+    }
 }

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

@@ -29,7 +29,6 @@ public class StatsPersister {
     private final String jobId;
     private final ResultsPersisterService resultsPersisterService;
     private final DataFrameAnalyticsAuditor auditor;
-    private volatile boolean isCancelled;
 
     public StatsPersister(String jobId, ResultsPersisterService resultsPersisterService, DataFrameAnalyticsAuditor auditor) {
         this.jobId = Objects.requireNonNull(jobId);
@@ -38,18 +37,14 @@ public StatsPersister(String jobId, ResultsPersisterService resultsPersisterServ
     }
 
     public void persistWithRetry(ToXContentObject result, Function<String, String> docIdSupplier) {
-        if (isCancelled) {
-            return;
-        }
-
         try {
             resultsPersisterService.indexWithRetry(jobId,
                 MlStatsIndex.writeAlias(),
                 result,
                 new ToXContent.MapParams(Collections.singletonMap(ToXContentParams.FOR_INTERNAL_STORAGE, "true")),
                 WriteRequest.RefreshPolicy.NONE,
                 docIdSupplier.apply(jobId),
-                () -> isCancelled == false,
+                () -> true,
                 errorMsg -> auditor.error(jobId,
                     "failed to persist result with id [" + docIdSupplier.apply(jobId) + "]; " + errorMsg)
             );
@@ -59,8 +54,4 @@ public void persistWithRetry(ToXContentObject result, Function<String, String> d
             LOGGER.error(() -> new ParameterizedMessage("[{}] Failed indexing stats result", jobId), e);
         }
     }
-
-    public void cancel() {
-        isCancelled = true;
-    }
 }

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/dataframe/process/AnalyticsProcessManagerTests.java

@@ -60,7 +60,7 @@ public class AnalyticsProcessManagerTests extends ESTestCase {
     private static final String CONFIG_ID = "config-id";
     private static final int NUM_ROWS = 100;
     private static final int NUM_COLS = 4;
-    private static final AnalyticsResult PROCESS_RESULT = new AnalyticsResult(null, null, null, null, null, null, null, null);
+    private static final AnalyticsResult PROCESS_RESULT = AnalyticsResult.builder().build();
 
     private Client client;
     private DataFrameAnalyticsAuditor auditor;