Merge pull request #10439 from brs96/213-splittable-random-fix

Node2Vec replace ThreadId with taskId

Author: IoannisPanagiotas

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

@@ -27,7 +27,7 @@ public class NegativeSampleProducer {
 
     private final HugeLongArray contextNodeDistribution;
     private final long cumulativeProbability;
-    private SplittableRandom splittableRandom;
+    private final SplittableRandom splittableRandom;
 
     public NegativeSampleProducer(
         HugeLongArray contextNodeDistribution,

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

@@ -30,9 +30,11 @@
 import org.neo4j.gds.ml.core.tensor.FloatVector;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.Random;
 import java.util.SplittableRandom;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.LongUnaryOperator;
 
 import static org.neo4j.gds.ml.core.tensor.operations.FloatVectorOperations.addInPlace;
@@ -41,7 +43,6 @@
 
 public class Node2VecModel {
 
-
     private final HugeObjectArray<FloatVector> centerEmbeddings;
     private final HugeObjectArray<FloatVector> contextEmbeddings;
     private final double initialLearningRate;
@@ -136,31 +137,7 @@ Node2VecResult train() {
                 initialLearningRate - iteration * learningRateAlpha
             );
 
-            var tasks = PartitionUtils.degreePartitionWithBatchSize(
-                walks.size(),
-                walks::walkLength,
-                BitUtil.ceilDiv(randomWalkProbabilities.sampleCount(), concurrency.value()),
-                partition -> {
-                    var positiveSampleProducer = new PositiveSampleProducer(
-                        walks.iterator(partition.startNode(), partition.nodeCount()),
-                        randomWalkProbabilities.positiveSamplingProbabilities(),
-                        windowSize,
-                        Optional.of(randomSeed)
-                    );
-
-                    return new TrainingTask(
-                        centerEmbeddings,
-                        contextEmbeddings,
-                        positiveSampleProducer,
-                        randomWalkProbabilities.negativeSamplingDistribution(),
-                        learningRate,
-                        negativeSamplingRate,
-                        embeddingDimension,
-                        progressTracker,
-                         randomSeed
-                    );
-                }
-            );
+            var tasks = createTrainingTasks(learningRate);
 
             RunWithConcurrency.builder()
                 .concurrency(concurrency)
@@ -232,12 +209,13 @@ private TrainingTask(
             int negativeSamplingRate,
             int embeddingDimensions,
             ProgressTracker progressTracker,
-            long randomSeed
+            long randomSeed,
+            int taskId
         ) {
             this.centerEmbeddings = centerEmbeddings;
             this.contextEmbeddings = contextEmbeddings;
             this.positiveSampleProducer = positiveSampleProducer;
-            this.negativeSampleProducer = new NegativeSampleProducer(negativeSamples, randomSeed + Thread.currentThread().getId());
+            this.negativeSampleProducer = new NegativeSampleProducer(negativeSamples, randomSeed + taskId);
             this.learningRate = learningRate;
             this.negativeSamplingRate = negativeSamplingRate;
 
@@ -310,4 +288,38 @@ void addAll(FloatConsumer other) {
         }
     }
 
+    List<TrainingTask> createTrainingTasks(float learningRate){
+        AtomicInteger taskIndex = new AtomicInteger(0);
+        return PartitionUtils.degreePartitionWithBatchSize(
+            walks.size(),
+            walks::walkLength,
+            BitUtil.ceilDiv(randomWalkProbabilities.sampleCount(), concurrency.value()),
+            partition -> {
+
+                var taskId = taskIndex.getAndIncrement();
+                var positiveSampleProducer = new PositiveSampleProducer(
+                    walks.iterator(partition.startNode(), partition.nodeCount()),
+                    randomWalkProbabilities.positiveSamplingProbabilities(),
+                    windowSize,
+                    Optional.of(randomSeed),
+                    taskId
+                );
+
+                return new TrainingTask(
+                    centerEmbeddings,
+                    contextEmbeddings,
+                    positiveSampleProducer,
+                    randomWalkProbabilities.negativeSamplingDistribution(),
+                    learningRate,
+                    negativeSamplingRate,
+                    embeddingDimension,
+                    progressTracker,
+                    randomSeed,
+                    taskId
+                );
+            }
+        );
+
+    }
+
 }

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

@@ -48,7 +48,8 @@ public class PositiveSampleProducer {
         Iterator<long[]> walks,
         HugeDoubleArray samplingProbabilities,
         int windowSize,
-        Optional<Long> maybeRandomSeed
+        Optional<Long> maybeRandomSeed,
+        int taskId
     ) {
         this.walks = walks;
         this.samplingProbabilities = samplingProbabilities;
@@ -60,7 +61,7 @@ public class PositiveSampleProducer {
         this.centerWordIndex = -1;
         this.contextWordIndex = 1;
         probabilitySupplier = maybeRandomSeed
-            .map(seed -> new SplittableRandom(Thread.currentThread().getId() + seed))
+            .map(seed -> new SplittableRandom(taskId + seed))
             .orElseGet(() -> new SplittableRandom(ThreadLocalRandom.current().nextLong()));
 
     }

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

@@ -61,7 +61,8 @@ void doesNotCauseStackOverflow() {
             walks.iterator(0, nbrOfWalks),
             HugeDoubleArray.of(LongStream.range(0, nbrOfWalks).mapToDouble((l) -> 1.0).toArray()),
             10,
-            Optional.empty()
+            Optional.empty(),
+            0
         );
 
         var counter = 0L;
@@ -89,7 +90,8 @@ void doesNotCauseStackOverflowDueToBadLuck() {
             walks.iterator(0, nbrOfWalks),
             probabilities,
             10,
-            Optional.empty()
+            Optional.empty(),
+            0
         );
         // does not overflow the stack = passes test
 
@@ -113,7 +115,8 @@ void doesNotAttemptToFetchOutsideBatch() {
             walks.iterator(0, nbrOfWalks / 2),
             HugeDoubleArray.of(LongStream.range(0, nbrOfWalks).mapToDouble((l) -> 1.0).toArray()),
             10,
-            Optional.empty()
+            Optional.empty(),
+            0
         );
 
         var counter = 0L;
@@ -138,7 +141,8 @@ void shouldProducePairsWith(
             walks.iterator(0, walks.size()),
             centerNodeProbabilities,
             windowSize,
-            Optional.empty()
+            Optional.empty(),
+            0
         );
         while (producer.next(buffer)) {
             actualPairs.add(Pair.of(buffer[0], buffer[1]));
@@ -161,7 +165,8 @@ void shouldProducePairsWithBounds() {
             walks.iterator(0, 2),
             centerNodeProbabilities,
             3,
-            Optional.empty()
+            Optional.empty(),
+            0
         );
         while (producer.next(buffer)) {
             actualPairs.add(Pair.of(buffer[0], buffer[1]));
@@ -207,7 +212,8 @@ void shouldRemoveDownsampledWordFromWalk() {
             walks.iterator(0, walks.size()),
             centerNodeProbabilities,
             3,
-            Optional.empty()
+            Optional.empty(),
+            0
         );
 
         while (producer.next(buffer)) {