migrate link prediction train estimate

Author: lassewesth

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

@@ -25,6 +25,7 @@
 import org.neo4j.gds.compat.GdsVersionInfoProvider;
 import org.neo4j.gds.core.model.CatalogModelContainer;
 import org.neo4j.gds.core.model.Model;
+import org.neo4j.gds.core.model.ModelCatalog;
 import org.neo4j.gds.mem.MemoryEstimation;
 import org.neo4j.gds.mem.MemoryEstimations;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
@@ -40,6 +41,7 @@
 import org.neo4j.gds.ml.pipeline.linkPipeline.LinkPredictionPredictPipeline;
 import org.neo4j.gds.ml.pipeline.linkPipeline.LinkPredictionTrainingPipeline;
 import org.neo4j.gds.ml.training.TrainingStatistics;
+import org.neo4j.gds.procedures.algorithms.AlgorithmsProcedureFacade;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -103,21 +105,23 @@ public static Task progressTask(String taskName, LinkPredictionTrainingPipeline
     }
 
     public static MemoryEstimation estimate(
-        ExecutionContext executionContext,
         LinkPredictionTrainingPipeline pipeline,
-        LinkPredictionTrainConfig configuration
+        LinkPredictionTrainConfig configuration,
+        ModelCatalog modelCatalog,
+        AlgorithmsProcedureFacade algorithmsProcedureFacade,
+        String username
     ) {
         pipeline.validateTrainingParameterSpace();
 
         var splitEstimations = splitEstimation(
             pipeline.splitConfig(),
             configuration.targetRelationshipType(),
-            pipeline.relationshipWeightProperty(executionContext.modelCatalog(), executionContext.username())
+            pipeline.relationshipWeightProperty(modelCatalog, username)
         );
 
         MemoryEstimation maxOverNodePropertySteps = NodePropertyStepExecutor.estimateNodePropertySteps(
-            executionContext.algorithmsProcedureFacade(),
-            executionContext.modelCatalog(),
+            algorithmsProcedureFacade,
+            modelCatalog,
             configuration.username(),
             pipeline.nodePropertySteps(),
             configuration.nodeLabels(),

pipeline/src/test/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrainPipelineExecutorTest.java

@@ -501,7 +501,13 @@ void estimateWithDifferentNodePropertySteps(
             );
 
             var actualRange = LinkPredictionTrainPipelineExecutor
-                .estimate(ImmutableExecutionContext.EMPTY, pipeline, config)
+                .estimate(
+                    pipeline,
+                    config,
+                    ImmutableExecutionContext.EMPTY.modelCatalog(),
+                    ImmutableExecutionContext.EMPTY.algorithmsProcedureFacade(),
+                    ImmutableExecutionContext.EMPTY.username()
+                )
                 .estimate(graphDimensions, config.concurrency())
                 .memoryUsage();
 
@@ -523,9 +529,11 @@ void failEstimateOnEmptyParameterSpace() {
             LinkPredictionTrainingPipeline pipeline = new LinkPredictionTrainingPipeline();
 
             assertThatThrownBy(() -> LinkPredictionTrainPipelineExecutor.estimate(
-                ExecutionContext.EMPTY,
                 pipeline,
-                config
+                config,
+                ExecutionContext.EMPTY.modelCatalog(),
+                ExecutionContext.EMPTY.algorithmsProcedureFacade(),
+                ExecutionContext.EMPTY.username()
             ))
                 .hasMessage("Need at least one model candidate for training.");
         }

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

@@ -22,10 +22,10 @@
 import org.neo4j.gds.BaseProc;
 import org.neo4j.gds.core.model.ModelCatalog;
 import org.neo4j.gds.executor.ExecutionContext;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.executor.ProcedureExecutor;
 import org.neo4j.gds.ml.pipeline.PipelineCompanion;
 import org.neo4j.gds.applications.algorithms.machinery.MemoryEstimateResult;
+import org.neo4j.gds.procedures.GraphDataScienceProcedures;
 import org.neo4j.procedure.Context;
 import org.neo4j.procedure.Description;
 import org.neo4j.procedure.Mode;
@@ -38,6 +38,8 @@
 import static org.neo4j.procedure.Mode.READ;
 
 public class LinkPredictionPipelineTrainProc extends BaseProc {
+    @Context
+    public GraphDataScienceProcedures facade;
 
     @Context
     public ModelCatalog modelCatalog;
@@ -61,12 +63,7 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
         @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
     ) {
-        PipelineCompanion.preparePipelineConfig(graphNameOrConfiguration, algoConfiguration);
-        return new MemoryEstimationExecutor<>(
-            new LinkPredictionPipelineTrainSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphNameOrConfiguration, algoConfiguration);
+        return facade.pipelines().linkPrediction().trainEstimate(graphNameOrConfiguration, algoConfiguration);
     }
 
     @Override

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

@@ -93,7 +93,13 @@ public MemoryEstimation memoryEstimation(LinkPredictionTrainConfig configuration
             LinkPredictionTrainingPipeline.class
         );
 
-        return LinkPredictionTrainPipelineExecutor.estimate(executionContext, pipeline, configuration);
+        return LinkPredictionTrainPipelineExecutor.estimate(
+            pipeline,
+            configuration,
+            executionContext.modelCatalog(),
+            executionContext.algorithmsProcedureFacade(),
+            executionContext.username()
+        );
     }
 
     @Override

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/DimensionTransformerForLinkPredictionTrain.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.procedures.pipelines;
+
+import org.neo4j.gds.api.User;
+import org.neo4j.gds.applications.algorithms.machinery.DimensionTransformer;
+import org.neo4j.gds.core.GraphDimensions;
+import org.neo4j.gds.ml.pipeline.linkPipeline.train.LinkPredictionTrainConfig;
+
+class DimensionTransformerForLinkPredictionTrain implements DimensionTransformer {
+    private final PipelineRepository pipelineRepository;
+    private final LinkPredictionTrainConfig configuration;
+
+    DimensionTransformerForLinkPredictionTrain(
+        PipelineRepository pipelineRepository,
+        LinkPredictionTrainConfig configuration
+    ) {
+        this.pipelineRepository = pipelineRepository;
+        this.configuration = configuration;
+    }
+
+    @Override
+    public GraphDimensions transform(GraphDimensions graphDimensions) {
+        // inject expected relationship set sizes which are used in the estimation of the TrainPipelineExecutor
+        // this allows to compute the MemoryTree over a single graphDimension
+        var user = new User(configuration.username(), false);
+        var pipelineName = PipelineName.parse(configuration.pipeline());
+
+        var pipeline = pipelineRepository.getLinkPredictionTrainingPipeline(user, pipelineName);
+
+        var splitConfig = pipeline.splitConfig();
+        var targetRelationshipType = configuration.targetRelationshipType();
+
+        return splitConfig.expectedGraphDimensions(graphDimensions, targetRelationshipType);
+    }
+}

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/LinkPredictionFacade.java

@@ -196,4 +196,20 @@ public Stream<MemoryEstimateResult> streamEstimate(
 
         return Stream.of(result);
     }
+
+    public Stream<MemoryEstimateResult> trainEstimate(
+        Object graphNameOrConfiguration,
+        Map<String, Object> rawConfiguration
+    ) {
+        PipelineCompanion.preparePipelineConfig(graphNameOrConfiguration, rawConfiguration);
+
+        var configuration = pipelineConfigurationParser.parseLinkPredictionTrainConfig(rawConfiguration);
+
+        var result = pipelineApplications.linkPredictionTrainEstimate(
+            graphNameOrConfiguration,
+            configuration
+        );
+
+        return Stream.of(result);
+    }
 }

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/PipelineApplications.java

@@ -52,6 +52,8 @@
 import org.neo4j.gds.ml.pipeline.linkPipeline.LinkFeatureStepFactory;
 import org.neo4j.gds.ml.pipeline.linkPipeline.LinkPredictionTrainingPipeline;
 import org.neo4j.gds.ml.pipeline.linkPipeline.linkfunctions.LinkFeatureStepConfiguration;
+import org.neo4j.gds.ml.pipeline.linkPipeline.train.LinkPredictionTrainConfig;
+import org.neo4j.gds.ml.pipeline.linkPipeline.train.LinkPredictionTrainPipelineExecutor;
 import org.neo4j.gds.ml.pipeline.nodePipeline.NodeFeatureStep;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.NodeClassificationTrainingPipeline;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.train.NodeClassificationModelResult;
@@ -407,6 +409,26 @@ Stream<StreamResult> linkPredictionStream(GraphName graphName, Map<String, Objec
         );
     }
 
+    MemoryEstimateResult linkPredictionTrainEstimate(
+        Object graphNameOrConfiguration,
+        LinkPredictionTrainConfig configuration
+    ) {
+        var estimate = linkPredictionTrainMemoryEstimation(configuration);
+
+        var memoryEstimation = MemoryEstimations.builder("Link Prediction Pipeline Executor")
+            .add("Pipeline executor", estimate)
+            .build();
+
+        var dimensionTransformer = new DimensionTransformerForLinkPredictionTrain(pipelineRepository, configuration);
+
+        return algorithmEstimationTemplate.estimate(
+            configuration,
+            graphNameOrConfiguration,
+            memoryEstimation,
+            dimensionTransformer
+        );
+    }
+
     MemoryEstimateResult nodeClassificationPredictEstimate(
         Object graphNameOrConfiguration,
         NodeClassificationPredictPipelineBaseConfig configuration
@@ -626,6 +648,28 @@ private MemoryEstimation linkPredictionMemoryEstimation(LinkPredictionPredictPip
         return linkPredictionPipelineEstimator.estimate(model, configuration);
     }
 
+    private MemoryEstimation linkPredictionTrainMemoryEstimation(LinkPredictionTrainConfig configuration) {
+        var specifiedUser = new User(configuration.username(), false);
+        var pipelineName = PipelineName.parse(configuration.pipeline());
+
+        var pipeline = pipelineRepository.getLinkPredictionTrainingPipeline(
+            specifiedUser,
+            pipelineName
+        );
+
+        var estimate = LinkPredictionTrainPipelineExecutor.estimate(
+            pipeline,
+            configuration,
+            modelCatalog,
+            algorithmsProcedureFacade,
+            user.getUsername()
+        );
+
+        return MemoryEstimations.builder("LinkPredictionPipelineTrain")
+            .add(estimate)
+            .build();
+    }
+
     private MemoryEstimation nodeClassificationPredictMemoryEstimation(NodeClassificationPredictPipelineBaseConfig configuration) {
         var modelName = configuration.modelName();
         var username = configuration.username();

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/PipelineConfigurationParser.java

@@ -31,6 +31,7 @@
 import org.neo4j.gds.ml.pipeline.linkPipeline.LinkPredictionSplitConfig;
 import org.neo4j.gds.ml.pipeline.linkPipeline.linkfunctions.LinkFeatureStepConfiguration;
 import org.neo4j.gds.ml.pipeline.linkPipeline.linkfunctions.LinkFeatureStepConfigurationImpl;
+import org.neo4j.gds.ml.pipeline.linkPipeline.train.LinkPredictionTrainConfig;
 import org.neo4j.gds.ml.pipeline.nodePipeline.NodePropertyPredictionSplitConfig;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.train.NodeClassificationPipelineTrainConfig;
 
@@ -76,6 +77,10 @@ LinkPredictionPredictPipelineStreamConfig parseLinkPredictionPredictPipelineStre
         return parseConfiguration(LinkPredictionPredictPipelineStreamConfig::of, configuration);
     }
 
+    LinkPredictionTrainConfig parseLinkPredictionTrainConfig(Map<String, Object> configuration) {
+        return parseConfiguration(LinkPredictionTrainConfig::of, configuration);
+    }
+
     TunableTrainerConfig parseLogisticRegressionTrainerConfig(Map<String, Object> configuration) {
         return parseTrainerConfiguration(
             configuration,