Merge pull request #5367 from breakanalysis/initialize-volumes-earlier

breakanalysis · web-flow · commit c2639bc8f25e · 2022-05-19T16:55:04.000+02:00
Introduce logSteps for apriori unkown volumes
diff --git a/core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/LeafTask.java b/core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/LeafTask.java
@@ -36,7 +36,12 @@ public class LeafTask extends Task {
     @Override
     public void finish() {
         super.finish();
-        setVolume(currentProgress.longValue());
+
+        // This task should now be considered to have 100% progress.
+        if (volume == UNKNOWN_VOLUME) {
+            volume = currentProgress.longValue();
+        }
+        currentProgress.add(volume - currentProgress.longValue());
     }
 
     @Override
diff --git a/core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/ProgressTracker.java b/core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/ProgressTracker.java
@@ -51,6 +51,9 @@ default void logProgress() {
 
     void logProgress(long value, String messageTemplate);
 
+    // prefer setting volume via factory method for leaves
+    // to make root progress available from the start
+    @Deprecated
     void setVolume(long volume);
 
     void logDebug(String message);
@@ -61,6 +64,10 @@ default void logProgress() {
 
     void release();
 
+    void setSteps(long steps);
+
+    void logSteps(long steps);
+
     class EmptyProgressTracker implements ProgressTracker {
 
         @Override
@@ -108,6 +115,16 @@ public void logProgress(long value, String messageTemplate) {
         public void setVolume(long volume) {
         }
 
+        @Override
+        public void setSteps(long steps) {
+
+        }
+
+        @Override
+        public void logSteps(long steps) {
+
+        }
+
         @Override
         public void logDebug(String message) {
 
diff --git a/core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/TaskProgressTracker.java b/core/src/main/java/org/neo4j/gds/core/utils/progress/tasks/TaskProgressTracker.java
@@ -37,12 +37,16 @@
 
 public class TaskProgressTracker implements ProgressTracker {
 
+    private static final long UNKNOWN_STEPS = -1;
+
     private final Task baseTask;
     private final TaskRegistry taskRegistry;
     private final UserLogRegistry userLogRegistry;
     private final TaskProgressLogger taskProgressLogger;
     private final Stack<Task> nestedTasks;
     protected Optional<Task> currentTask;
+    private long currentTotalSteps;
+    private double progressLeftOvers;
 
     public TaskProgressTracker(Task baseTask, Log log, int concurrency, TaskRegistryFactory taskRegistryFactory) {
         this(baseTask, log, concurrency, new JobId(), taskRegistryFactory, EmptyUserLogRegistryFactory.INSTANCE);
@@ -56,6 +60,8 @@ public TaskProgressTracker(
         this.taskRegistry = taskRegistryFactory.newInstance(jobId);
         this.taskProgressLogger = new TaskProgressLogger(log, baseTask, concurrency);
         this.currentTask = Optional.empty();
+        this.currentTotalSteps = UNKNOWN_STEPS;
+        this.progressLeftOvers = 0;
         this.nestedTasks = new Stack<>();
         this.userLogRegistry = userLogRegistryFactory.newInstance();
     }
@@ -76,6 +82,8 @@ public void beginSubTask() {
         nextTask.start();
         taskProgressLogger.logBeginSubTask(nextTask, parentTask());
         currentTask = Optional.of(nextTask);
+        currentTotalSteps = UNKNOWN_STEPS;
+        progressLeftOvers = 0;
     }
 
     @Override
@@ -90,6 +98,26 @@ public void beginSubTask(long taskVolume) {
         setVolume(taskVolume);
     }
 
+    @Override
+    public void setSteps(long steps) {
+        if (steps <= 0) {
+            throw new IllegalStateException(formatWithLocale(
+                "Total steps for task must be at least 1 but was %d",
+                steps
+            ));
+        }
+        currentTotalSteps = steps;
+    }
+
+    @Override
+    public void logSteps(long steps) {
+        long volume = requireCurrentTask().getProgress().volume();
+        double progress = steps * volume / (double) currentTotalSteps + progressLeftOvers;
+        long longProgress = (long) progress;
+        progressLeftOvers = progress - longProgress;
+        logProgress(longProgress);
+    }
+
     @Override
     public void beginSubTask(String expectedTaskDescription, long taskVolume) {
         beginSubTask();
diff --git a/core/src/test/java/org/neo4j/gds/core/utils/progress/tasks/TaskProgressTrackerTest.java b/core/src/test/java/org/neo4j/gds/core/utils/progress/tasks/TaskProgressTrackerTest.java
@@ -208,6 +208,24 @@ void shouldRegisterBaseTaskOnBaseTaskStart() {
         assertThat(taskStore.query("")).containsValue(task);
     }
 
+    @Test
+    void stepsShouldGiveProgress() {
+        var leafTask = Tasks.leaf("leaf", 100);
+        var progressTracker = progressTracker(leafTask);
+        
+        progressTracker.beginSubTask();
+        progressTracker.setSteps(13);
+        progressTracker.logProgress(3);
+        progressTracker.logSteps(1);
+        double expectedDoubleProgressFromFirstStep = 100.0 * 1.0 / 13.0;
+        long progressAfterFirstStep = leafTask.getProgress().progress();
+        assertThat(progressAfterFirstStep).isEqualTo((long) expectedDoubleProgressFromFirstStep + 3);
+
+        progressTracker.logProgress(1);
+        progressTracker.logSteps(4);
+        assertThat(leafTask.getProgress().progress()).isEqualTo(3 + 1 + (long) (100.0 * 5.0 / 13));
+    }
+
     private TaskProgressTracker progressTracker(Task task, Log log) {
         return new TaskProgressTracker(task, log, 1, EmptyTaskRegistryFactory.INSTANCE);
     }
diff --git a/core/src/test/java/org/neo4j/gds/core/utils/progress/tasks/TaskTest.java b/core/src/test/java/org/neo4j/gds/core/utils/progress/tasks/TaskTest.java
@@ -216,22 +216,24 @@ void shouldSetVolumeLate() {
     }
 
     @Test
-    void shouldSetVolumeWhenFinishingTask() {
+    void shouldSetProgressWhenFinishingTask() {
         var task = Tasks.iterativeOpen("root", () -> List.of(Tasks.leaf("leaf")));
         task.start();
         var leaf1 = task.nextSubtask();
         leaf1.start();
         leaf1.logProgress(22L);
         leaf1.finish();
 
+        assertThat(leaf1.getProgress().progress()).isEqualTo(22L);
         assertThat(leaf1.getProgress().volume()).isEqualTo(22L);
         assertThat(task.getProgress().volume()).isEqualTo(Task.UNKNOWN_VOLUME);
 
         var leaf2 = task.nextSubtask();
         leaf2.start();
-        leaf2.logProgress(20L);
+        leaf2.setVolume(20L);
         leaf2.finish();
 
+        assertThat(leaf2.getProgress().progress()).isEqualTo(20L);
         assertThat(leaf2.getProgress().volume()).isEqualTo(20L);
         assertThat(task.getProgress().volume()).isEqualTo(Task.UNKNOWN_VOLUME);
 
diff --git a/ml/ml-algo/src/main/java/org/neo4j/gds/ml/metrics/SignedProbabilities.java b/ml/ml-algo/src/main/java/org/neo4j/gds/ml/metrics/SignedProbabilities.java
@@ -119,7 +119,7 @@ public static SignedProbabilities computeFromLabeledData(
         TerminationFlag terminationFlag,
         ProgressTracker progressTracker
     ) {
-        progressTracker.setVolume(features.size());
+        progressTracker.setSteps(features.size());
 
         var signedProbabilities = SignedProbabilities.create(evaluationQueue.totalSize());
 
@@ -134,7 +134,7 @@ public static SignedProbabilities computeFromLabeledData(
 
                     signedProbabilities.add(probabilityOfPositiveEdge, isEdge);
                 }
-                progressTracker.logProgress(batch.size());
+                progressTracker.logSteps(batch.size());
             },
             terminationFlag
         );
diff --git a/ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/ClassifierTrainer.java b/ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/ClassifierTrainer.java
@@ -28,6 +28,10 @@ public interface ClassifierTrainer {
 
     Classifier train(Features features, HugeLongArray labels, ReadOnlyHugeLongArray trainSet);
 
+    static Task progressTask(String taskName, long volume) {
+        return Tasks.leaf(taskName, volume);
+    }
+
     static Task progressTask(String taskName) {
         return Tasks.leaf(taskName);
     }
diff --git a/pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/BatchLinkFeatureExtractor.java b/pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/BatchLinkFeatureExtractor.java
@@ -61,7 +61,7 @@ public void run() {
             }));
         });
 
-        progressTracker.logProgress(partition.totalDegree());
+        progressTracker.logSteps(partition.totalDegree());
     }
 }
 
diff --git a/pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkFeaturesAndLabelsExtractor.java b/pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkFeaturesAndLabelsExtractor.java
@@ -79,7 +79,7 @@ static FeaturesAndLabels extractFeaturesAndLabels(
         ProgressTracker progressTracker,
         TerminationFlag terminationFlag
     ) {
-        progressTracker.setVolume(graph.relationshipCount() * 2);
+        progressTracker.setSteps(graph.relationshipCount() * 2);
         var features = LinkFeatureExtractor.extractFeatures(
             graph,
             featureSteps,
@@ -125,7 +125,7 @@ private static HugeLongArray extractLabels(
                     }
                     return true;
                 }));
-                progressTracker.logProgress(partition.totalDegree());
+                progressTracker.logSteps(partition.totalDegree());
             }
             );
             relationshipOffset.add(partition.totalDegree());
diff --git a/pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrain.java b/pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrain.java
@@ -95,20 +95,25 @@ public LinkPredictionTrain(
         this.classIdMap = makeClassIdMap();
     }
 
-    public static List<Task> progressTasks(int validationFolds, int numberOfModelSelectionTrials) {
+    public static List<Task> progressTasks(
+        long relationshipCount,
+        LinkPredictionSplitConfig splitConfig,
+        int numberOfModelSelectionTrials
+    ) {
+        var sizes = splitConfig.expectedSetSizes(relationshipCount);
         return List.of(
-            Tasks.leaf("Extract train features"),
+            Tasks.leaf("Extract train features", sizes.trainSize() * 3),
             Tasks.iterativeFixed(
                 "Select best model",
-                () -> List.of(Tasks.leaf("Trial", validationFolds)),
+                () -> List.of(Tasks.leaf("Trial", splitConfig.validationFolds() * sizes.trainSize() * 5)),
                 numberOfModelSelectionTrials
             ),
-            ClassifierTrainer.progressTask("Train best model"),
-            Tasks.leaf("Compute train metrics"),
+            ClassifierTrainer.progressTask("Train best model", sizes.trainSize() * 5),
+            Tasks.leaf("Compute train metrics", sizes.trainSize()),
             Tasks.task(
                 "Evaluate on test data",
-                Tasks.leaf("Extract test features"),
-                Tasks.leaf("Compute test metrics")
+                Tasks.leaf("Extract test features", sizes.testSize() * 3),
+                Tasks.leaf("Compute test metrics", sizes.testSize())
             )
         );
     }
@@ -208,6 +213,7 @@ private void modelSelect(
         int trial = 0;
         while (hyperParameterOptimizer.hasNext()) {
             progressTracker.beginSubTask();
+            progressTracker.setSteps(pipeline.splitConfig().validationFolds());
             var modelParams = hyperParameterOptimizer.next();
             progressTracker.logMessage(formatWithLocale("Method: %s, Parameters: %s", modelParams.method(), modelParams.toMap()));
             var trainStatsBuilder = new ModelStatsBuilder(pipeline.splitConfig().validationFolds());
@@ -242,7 +248,7 @@ private void modelSelect(
                     ProgressTracker.NULL_TRACKER
                 );
 
-                progressTracker.logProgress();
+                progressTracker.logSteps(1);
             }
 
             // insert the candidates' metrics into trainStats and validationStats
diff --git a/pipeline/src/test/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrainTest.java b/pipeline/src/test/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrainTest.java
@@ -461,7 +461,8 @@ void logProgressRF() {
         var log = Neo4jProxy.testLog();
         var progressTracker = new TestProgressTracker(
             progressTask(
-                pipeline.splitConfig().validationFolds(),
+                trainGraph.relationshipCount(),
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ),
             log,
@@ -552,7 +553,8 @@ void logProgressLR() {
         var log = Neo4jProxy.testLog();
         var progressTracker = new TestProgressTracker(
             progressTask(
-                pipeline.splitConfig().validationFolds(),
+                trainGraph.relationshipCount(),
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ),
             log,
@@ -614,7 +616,8 @@ void logProgressLRWithRange() {
         var log = Neo4jProxy.testLog();
         var progressTracker = new TestProgressTracker(
             progressTask(
-                pipeline.splitConfig().validationFolds(),
+                2 * trainGraph.relationshipCount(),
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ),
             log,
@@ -713,8 +716,11 @@ void logProgressLRWithRange() {
             );
     }
 
-    static Task progressTask(int validationFolds, int numberOfModelSelectionTrials) {
-        return Tasks.task("MY TEST TASK", LinkPredictionTrain.progressTasks(validationFolds, numberOfModelSelectionTrials));
+    static Task progressTask(long relationshipCount, LinkPredictionSplitConfig splitConfig, int numberOfModelSelectionTrials) {
+        return Tasks.task(
+            "MY TEST TASK",
+            LinkPredictionTrain.progressTasks(relationshipCount, splitConfig, numberOfModelSelectionTrials)
+        );
     }
 
     private LinkPredictionTrainingPipeline linkPredictionPipeline() {
diff --git a/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/ExhaustiveLinkPrediction.java b/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/ExhaustiveLinkPrediction.java
@@ -76,6 +76,8 @@ ExhaustiveLinkPredictionResult predictLinks(
         Graph graph,
         LinkPredictionSimilarityComputer linkPredictionSimilarityComputer
     ) {
+        progressTracker.setSteps(graph.nodeCount());
+
         var predictionQueue = BoundedLongLongPriorityQueue.max(topN);
 
         var tasks = PartitionUtils.rangePartition(
@@ -139,7 +141,7 @@ public void run() {
                 );
             });
 
-            progressTracker.logProgress(partition.nodeCount());
+            progressTracker.logSteps(partition.nodeCount());
         }
 
         private LongHashSet largerNeighbors(long sourceId) {
diff --git a/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/LinkPredictionPredictPipelineExecutor.java b/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/LinkPredictionPredictPipelineExecutor.java
@@ -102,15 +102,15 @@ public static Task progressTask(
             taskName,
             Tasks.iterativeFixed(
                 "Execute node property steps",
-                () -> List.of(Tasks.leaf("Step")),
+                () -> List.of(Tasks.leaf("Step", 10 * graphStore.getUnion().relationshipCount())),
                 pipeline.nodePropertySteps().size()
             ),
             config.isApproximateStrategy()
                 ? Tasks.task(
                 "Approximate link prediction",
                 KnnFactory.knnTaskTree(graphStore.getUnion(), config.approximateConfig())
             )
-                : Tasks.leaf("Exhaustive link prediction", graphStore.nodeCount())
+                : Tasks.leaf("Exhaustive link prediction", graphStore.getUnion().nodeCount() * graphStore.getUnion().nodeCount() / 2)
         );
     }
 
diff --git a/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineAlgorithmFactory.java b/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineAlgorithmFactory.java
@@ -70,9 +70,15 @@ public String taskName() {
 
     @Override
     public Task progressTask(GraphStore graphStore, LinkPredictionTrainConfig config) {
+        var relationshipCount = config
+            .internalRelationshipTypes(graphStore)
+            .stream()
+            .mapToLong(graphStore::relationshipCount)
+            .sum();
         return LinkPredictionTrainPipelineExecutor.progressTask(
             taskName(),
-            PipelineCatalog.getTyped(config.username(), config.pipeline(), LinkPredictionTrainingPipeline.class)
+            PipelineCatalog.getTyped(config.username(), config.pipeline(), LinkPredictionTrainingPipeline.class),
+            relationshipCount
         );
     }
 
diff --git a/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineExecutor.java b/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineExecutor.java
@@ -81,16 +81,21 @@ public LinkPredictionTrainPipelineExecutor(
         );
     }
 
-    public static Task progressTask(String taskName, LinkPredictionTrainingPipeline pipeline) {
+    public static Task progressTask(String taskName, LinkPredictionTrainingPipeline pipeline, long relationshipCount) {
+        var sizes = pipeline.splitConfig().expectedSetSizes(relationshipCount);
         return Tasks.task(taskName, new ArrayList<>() {{
-            add(Tasks.leaf("Split relationships"));
+            add(Tasks.leaf(
+                "Split relationships",
+                sizes.trainSize() + sizes.featureInputSize() + sizes.testSize() + sizes.testComplementSize()
+            ));
             add(Tasks.iterativeFixed(
                 "Execute node property steps",
-                () -> List.of(Tasks.leaf("Step")),
+                () -> List.of(Tasks.leaf("Step", 10 * sizes.featureInputSize())),
                 pipeline.nodePropertySteps().size()
             ));
             addAll(LinkPredictionTrain.progressTasks(
-                pipeline.splitConfig().validationFolds(),
+                relationshipCount,
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ));
         }});
diff --git a/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/RelationshipSplitter.java b/proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/RelationshipSplitter.java
diff --git a/proc/machine-learning/src/test/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineExecutorTest.java b/proc/machine-learning/src/test/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineExecutorTest.java

Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,10 @@ public interface ClassifierTrainer {`
`28`	`28`
`29`	`29`	`Classifier train(Features features, HugeLongArray labels, ReadOnlyHugeLongArray trainSet);`
`30`	`30`
	`31`	`+ static Task progressTask(String taskName, long volume) {`
	`32`	`+ return Tasks.leaf(taskName, volume);`
	`33`	`+ }`
	`34`	`+`
`31`	`35`	`static Task progressTask(String taskName) {`
`32`	`36`	`return Tasks.leaf(taskName);`
`33`	`37`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ public void run() {`
`61`	`61`	`}));`
`62`	`62`	`});`
`63`	`63`
`64`		`- progressTracker.logProgress(partition.totalDegree());`
	`64`	`+ progressTracker.logSteps(partition.totalDegree());`
`65`	`65`	`}`
`66`	`66`	`}`
`67`	`67`