Make root progress available and sane for NC predict

Author: adamnsch

ml/ml-algo/src/main/java/org/neo4j/gds/ml/nodeClassification/NodeClassificationPredict.java

@@ -22,7 +22,6 @@
 import org.jetbrains.annotations.Nullable;
 import org.neo4j.gds.Algorithm;
 import org.neo4j.gds.annotation.ValueClass;
-import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.core.utils.TerminationFlag;
 import org.neo4j.gds.core.utils.mem.MemoryEstimation;
 import org.neo4j.gds.core.utils.mem.MemoryEstimations;
@@ -74,8 +73,8 @@ public NodeClassificationPredict(
         );
     }
 
-    public static Task progressTask(Graph graph) {
-        return Tasks.leaf("Node classification predict", graph.nodeCount());
+    public static Task progressTask(long nodeCount) {
+        return Tasks.leaf("Node classification predict", nodeCount);
     }
 
     public static MemoryEstimation memoryEstimation(
@@ -138,6 +137,7 @@ public static MemoryEstimation memoryEstimationWithDerivedBatchSize(
     @Override
     public NodeClassificationResult compute() {
         progressTracker.beginSubTask();
+        progressTracker.setSteps(features.size());
         var predictedProbabilities = initProbabilities();
         var predictedClasses = predictor.predict(predictedProbabilities);
         progressTracker.endSubTask();

ml/ml-algo/src/test/java/org/neo4j/gds/ml/nodeClassification/NodeClassificationPredictTest.java

@@ -222,7 +222,7 @@ void shouldLogProgress() {
 
         var log = Neo4jProxy.testLog();
         var progressTracker = new TaskProgressTracker(
-            NodeClassificationPredict.progressTask(graph),
+            NodeClassificationPredict.progressTask(graph.nodeCount()),
             log,
             1,
             new JobId(),

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineAlgorithmFactory.java

@@ -52,8 +52,7 @@ public Task progressTask(GraphStore graphStore, CONFIG config) {
             modelCatalog,
             config.modelName(),
             config.username()
-        ).customInfo()
-            .pipeline();
+        ).customInfo().pipeline();
 
         return NodeClassificationPredictPipelineExecutor.progressTask(taskName(), trainingPipeline, graphStore);
     }

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineExecutor.java

@@ -71,10 +71,10 @@ public static Task progressTask(String taskName, NodePropertyPredictPipeline pip
             taskName,
             Tasks.iterativeFixed(
                 "Execute node property steps",
-                () -> List.of(Tasks.leaf("Step")),
+                () -> List.of(Tasks.leaf("Step", 10 * graphStore.getUnion().nodeCount())),
                 pipeline.nodePropertySteps().size()
             ),
-            Tasks.leaf("Node classification predict", graphStore.getUnion().nodeCount())
+            NodeClassificationPredict.progressTask(graphStore.getUnion().nodeCount())
         );
     }