Tune LP train task volumes and include splitting and node property steps

Co-Authored-By: Adam Schill Collberg<adam.schill.collberg@protonmail.com>

Author: breakanalysis

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrain.java

@@ -102,17 +102,17 @@ public static List<Task> progressTasks(
     ) {
         var sizes = splitConfig.expectedSetSizes(relationshipCount);
         return List.of(
-            Tasks.leaf("Extract train features", sizes.trainSize()),
+            Tasks.leaf("Extract train features", sizes.trainSize() * 3),
             Tasks.iterativeFixed(
                 "Select best model",
-                () -> List.of(Tasks.leaf("Trial", splitConfig.validationFolds() * sizes.trainSize())),
+                () -> List.of(Tasks.leaf("Trial", splitConfig.validationFolds() * sizes.trainSize() * 5)),
                 numberOfModelSelectionTrials
             ),
-            ClassifierTrainer.progressTask("Train best model", sizes.trainSize()),
+            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", sizes.testSize()),
+                Tasks.leaf("Extract test features", sizes.testSize() * 3),
                 Tasks.leaf("Compute test metrics", sizes.testSize())
             )
         );

proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineExecutor.java

@@ -82,11 +82,15 @@ public LinkPredictionTrainPipelineExecutor(
     }
 
     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(

proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/RelationshipSplitter.java

@@ -71,7 +71,7 @@ public void splitRelationships(
         Optional<Long> randomSeed,
         Optional<String> relationshipWeightProperty
     ) {
-        progressTracker.beginSubTask();
+        progressTracker.beginSubTask("Split relationships");
 
         splitConfig.validateAgainstGraphStore(graphStore);
 
@@ -90,7 +90,7 @@ public void splitRelationships(
 
         graphStore.deleteRelationships(RelationshipType.of(testComplementRelationshipType));
 
-        progressTracker.endSubTask();
+        progressTracker.endSubTask("Split relationships");
     }
 
     private void validateTestSplit(GraphStore graphStore) {