Make Node2Vec respect TerminationFlag during Training

Co-authored-by: Ioannis Panagiotas <ioannis.panagiotas@neotechnology.com>

Author: vnickolov

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/Node2Vec.java

@@ -110,7 +110,8 @@ public Node2VecResult compute() {
             maybeRandomSeed,
             walks,
             probabilitiesBuilder.build(),
-            progressTracker
+            progressTracker,
+            terminationFlag
         );
 
         var result = node2VecModel.train();

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/Node2VecModel.java

@@ -27,6 +27,7 @@
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.mem.BitUtil;
 import org.neo4j.gds.ml.core.tensor.FloatVector;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -54,6 +55,7 @@ public class Node2VecModel {
     private final RandomWalkProbabilities randomWalkProbabilities;
     private final ProgressTracker progressTracker;
     private final long randomSeed;
+    private final TerminationFlag terminationFlag;
 
     static final double EPSILON = 1e-10;
 
@@ -65,7 +67,8 @@ public class Node2VecModel {
         Optional<Long> maybeRandomSeed,
         CompressedRandomWalks walks,
         RandomWalkProbabilities randomWalkProbabilities,
-        ProgressTracker progressTracker
+        ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
     ) {
         this(
             toOriginalId,
@@ -81,7 +84,8 @@ public class Node2VecModel {
             maybeRandomSeed,
             walks,
             randomWalkProbabilities,
-            progressTracker
+            progressTracker,
+            terminationFlag
         );
     }
 
@@ -99,7 +103,8 @@ public class Node2VecModel {
         Optional<Long> maybeRandomSeed,
         CompressedRandomWalks walks,
         RandomWalkProbabilities randomWalkProbabilities,
-        ProgressTracker progressTracker
+        ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
     ) {
         this.initialLearningRate = initialLearningRate;
         this.minLearningRate = minLearningRate;
@@ -113,6 +118,7 @@ public class Node2VecModel {
         this.randomWalkProbabilities = randomWalkProbabilities;
         this.progressTracker = progressTracker;
         this.randomSeed = maybeRandomSeed.orElseGet(() -> new SplittableRandom().nextLong());
+        this.terminationFlag = terminationFlag;
 
         var random = new Random();
         centerEmbeddings = initializeEmbeddings(toOriginalId, nodeCount, embeddingDimension, random);
@@ -136,6 +142,7 @@ Node2VecResult train() {
 
             RunWithConcurrency.builder()
                 .concurrency(concurrency)
+                .terminationFlag(terminationFlag)
                 .tasks(tasks)
                 .run();
 

algo/src/test/java/org/neo4j/gds/embeddings/node2vec/Node2VecModelTest.java

@@ -31,6 +31,8 @@
 import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.utils.Intersections;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.termination.TerminatedException;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Optional;
 import java.util.Random;
@@ -39,6 +41,7 @@
 import java.util.stream.LongStream;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -92,7 +95,8 @@ void testModel() {
             Optional.empty(),
             walks,
             probabilitiesBuilder.build(),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         );
 
         var trainResult = node2VecModel.train();
@@ -194,7 +198,8 @@ void twoRunsSingleThreadedWithTheSameRandomSeed(int iterations) {
             Optional.of(1337L),
             walks,
             probabilitiesBuilder.build(),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         ).train().embeddings();
 
         var secondRunEmbedding = new Node2VecModel(
@@ -205,7 +210,8 @@ void twoRunsSingleThreadedWithTheSameRandomSeed(int iterations) {
             Optional.of(1337L),
             walks,
             probabilitiesBuilder.build(),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         ).train().embeddings();
 
         for (long node = 0; node < nodeCount; node++) {
@@ -239,7 +245,8 @@ void shouldCreateTrainingTasksWithCorrectRandomSeed() {
                 Optional.of(1L), // Random Seed
                 randomWalksMock,
                 randomWalkProbabilitiesMock,
-                ProgressTracker.NULL_TRACKER
+                ProgressTracker.NULL_TRACKER,
+                TerminationFlag.RUNNING_TRUE
             )
         );
 
@@ -311,7 +318,8 @@ void shouldHaveCorrectLearningRate(){
             Optional.empty(),
             null,
             null,
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            TerminationFlag.RUNNING_TRUE
         );
 
         assertThat(node2VecModel.learningRate(0)).isEqualTo(10f);
@@ -322,4 +330,57 @@ void shouldHaveCorrectLearningRate(){
         assertThat(node2VecModel.learningRate(10000)).isEqualTo(5f);
 
     }
+
+    @Test
+    void shouldRespectTerminationFlag() {
+        var random = new Random(42);
+        int numberOfClusters = 2;
+        int clusterSize = 5;
+        int numberOfWalks = 2;
+        int walkLength = 5;
+
+        var probabilitiesBuilder = new RandomWalkProbabilitiesBuilder(
+            numberOfClusters * clusterSize,
+            new Concurrency(1),
+            0.001,
+            0.75
+        );
+
+        var walks = generateRandomWalks(
+            probabilitiesBuilder,
+            numberOfClusters,
+            clusterSize,
+            numberOfWalks,
+            walkLength,
+            random
+        );
+
+        var trainParameters = new TrainParameters(0.05, 0.0001, 10, 2, 1, 2, EmbeddingInitializer.NORMALIZED);
+
+        var terminationFlag = new TerminationFlag() {
+            private int callCount = 0;
+            @Override
+            public boolean running() {
+                ++callCount;
+                return callCount == 2;
+            }
+        };
+
+        var node2VecModel = new Node2VecModel(
+            nodeId -> nodeId,
+            1000,
+            trainParameters,
+            new Concurrency(4),
+            Optional.of(19L),
+            walks,
+            probabilitiesBuilder.build(),
+            ProgressTracker.NULL_TRACKER,
+            terminationFlag
+        );
+
+        assertThatExceptionOfType(TerminatedException.class)
+            .isThrownBy(node2VecModel::train);
+
+    }
+
 }