Merge pull request #9848 from IoannisPanagiotas/memory-tracking-reporting

Memory listing procedures

Author: IoannisPanagiotas

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/ComputationService.java

@@ -34,11 +34,13 @@ class ComputationService {
     private final Log log;
     private final MemoryGuard memoryGuard;
     private final AlgorithmMetricsService algorithmMetricsService;
+    private final String username;
 
-    ComputationService(Log log, MemoryGuard memoryGuard, AlgorithmMetricsService algorithmMetricsService) {
+    ComputationService(String username,Log log, MemoryGuard memoryGuard, AlgorithmMetricsService algorithmMetricsService) {
         this.log = log;
         this.memoryGuard = memoryGuard;
         this.algorithmMetricsService = algorithmMetricsService;
+        this.username = username;
     }
 
     <CONFIGURATION extends AlgoBaseConfig, RESULT_FROM_ALGORITHM> RESULT_FROM_ALGORITHM computeAlgorithm(
@@ -49,7 +51,7 @@ <CONFIGURATION extends AlgoBaseConfig, RESULT_FROM_ALGORITHM> RESULT_FROM_ALGORI
         Computation<RESULT_FROM_ALGORITHM> computation,
         DimensionTransformer dimensionTransformer
     ) {
-        memoryGuard.assertAlgorithmCanRun(estimationSupplier, graphResources.graphStore(), configuration, label, dimensionTransformer);
+        memoryGuard.assertAlgorithmCanRun(username,estimationSupplier, graphResources.graphStore(), configuration, label, dimensionTransformer);
 
         return computeWithMetrics(graphResources, label, computation);
     }

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/DefaultAlgorithmProcessingTemplate.java

@@ -57,7 +57,7 @@ public static DefaultAlgorithmProcessingTemplate create(
         MemoryGuard memoryGuard,
         RequestScopedDependencies requestScopedDependencies
     ) {
-        var algorithmComputer = new ComputationService(log, memoryGuard, algorithmMetricsService);
+        var algorithmComputer = new ComputationService(requestScopedDependencies.getUser().getUsername(),log, memoryGuard, algorithmMetricsService);
 
         return new DefaultAlgorithmProcessingTemplate(
             graphStoreCatalogService,

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/DefaultMemoryGuard.java

@@ -61,6 +61,7 @@ public static DefaultMemoryGuard create(
 
     @Override
     public synchronized <CONFIGURATION extends AlgoBaseConfig> void assertAlgorithmCanRun(
+        String username,
         Supplier<MemoryEstimation> estimationFactory,
         GraphStore graphStore,
         CONFIGURATION configuration,
@@ -80,11 +81,11 @@ public synchronized <CONFIGURATION extends AlgoBaseConfig> void assertAlgorithmC
 
             var bytesToReserve = memoryRequirement.requiredMemory();
             if (configuration.sudo()) {
-                memoryTracker.track(configuration.jobId(), bytesToReserve);
+                memoryTracker.track(username,label.asString(), configuration.jobId(), bytesToReserve);
                 return;
             }
 
-            memoryTracker.tryToTrack(configuration.jobId(), bytesToReserve);
+            memoryTracker.tryToTrack(username, label.asString(), configuration.jobId(), bytesToReserve);
 
         } catch (MemoryEstimationNotImplementedException e) {
             log.info("Memory usage estimate not available for " + label + ", skipping guard");

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/MemoryGuard.java

@@ -35,13 +35,13 @@ public interface MemoryGuard {
     MemoryGuard DISABLED = new MemoryGuard() {
         @Override
         public <CONFIGURATION extends AlgoBaseConfig> void assertAlgorithmCanRun(
+            String username,
             Supplier<MemoryEstimation> estimationFactory,
             GraphStore graphStore,
             CONFIGURATION configuration,
             Label label,
             DimensionTransformer dimensionTransformer
         ) {
-            // do nothing
         }
     };
 
@@ -51,6 +51,7 @@ public <CONFIGURATION extends AlgoBaseConfig> void assertAlgorithmCanRun(
      * @throws IllegalStateException when there is not enough memory available to run the algorithm in the given configuration on the given graph
      */
     <CONFIGURATION extends AlgoBaseConfig> void assertAlgorithmCanRun(
+        String username,
         Supplier<MemoryEstimation> estimationFactory,
         GraphStore graphStore,
         CONFIGURATION configuration,

applications/algorithms/machinery/src/test/java/org/neo4j/gds/applications/algorithms/machinery/DefaultMemoryGuardTest.java

@@ -21,6 +21,7 @@
 
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.api.GraphStore;
+import org.neo4j.gds.api.User;
 import org.neo4j.gds.applications.services.GraphDimensionFactory;
 import org.neo4j.gds.core.GraphDimensions;
 import org.neo4j.gds.core.concurrency.Concurrency;
@@ -67,7 +68,8 @@ void shouldAllowExecution() {
 
         // there is enough memory available
         memoryGuard.assertAlgorithmCanRun(
-            () -> memoryEstimation,
+            User.DEFAULT.getUsername(),
+            ()-> memoryEstimation,
             graphStore,
             configuration,
             new StandardLabel("some label"),
@@ -97,6 +99,7 @@ void shouldGuardExecutionUsingMinimumEstimate() {
         // uh oh
         try {
             memoryGuard.assertAlgorithmCanRun(
+                User.DEFAULT.getUsername(),
                 () -> memoryEstimation,
                 graphStore,
                 configuration,
@@ -132,6 +135,7 @@ void shouldGuardExecutionUsingMaximumEstimate() {
         // uh oh
         try {
             memoryGuard.assertAlgorithmCanRun(
+                User.DEFAULT.getUsername(),
                 () -> memoryEstimation,
                 graphStore,
                 configuration,
@@ -166,6 +170,7 @@ void shouldRespectSudoFlag() {
         when(memoryTree.memoryUsage()).thenReturn(MemoryRange.of(43, 99));
 
         assertThatIllegalStateException().isThrownBy(() -> memoryGuard.assertAlgorithmCanRun(
+            User.DEFAULT.getUsername(),
             () -> memoryEstimation,
             graphStore,
             new ExampleConfiguration(false),
@@ -175,6 +180,7 @@ void shouldRespectSudoFlag() {
 
         // now with sudo
         assertDoesNotThrow(() -> memoryGuard.assertAlgorithmCanRun(
+            User.DEFAULT.getUsername(),
             () -> memoryEstimation,
             graphStore,
             new ExampleConfiguration(true),

applications/graph-store-catalog/src/integrationTest/java/org/neo4j/gds/applications/graphstorecatalog/MemoryUsageValidatorTest.java

@@ -55,11 +55,12 @@ void shouldPassOnSufficientMemory() {
         var dimensions = GraphDimensions.builder().nodeCount(1000).build();
         var memoryTree = MemoryTree.empty();
 
-        assertThatNoException().isThrownBy(() -> new MemoryUsageValidator(new MemoryTracker(10000000, Log.noOpLog()),
+        assertThatNoException().isThrownBy(() -> new MemoryUsageValidator("foo",new MemoryTracker(10000000, Log.noOpLog()),
             false,
             Log.noOpLog()
         )
             .tryValidateMemoryUsage(
+                "task",
                 TestConfig.empty(),
                 (config) -> new MemoryTreeWithDimensions(memoryTree, dimensions)
             ));
@@ -70,11 +71,12 @@ void shouldFailOnInsufficientMemory() {
         var dimensions = GraphDimensions.builder().nodeCount(1000).build();
         var memoryTree = new TestTree("test", MemoryRange.of(42));
 
-        assertThatThrownBy(() -> new MemoryUsageValidator(new MemoryTracker(21, Log.noOpLog()),
+        assertThatThrownBy(() -> new MemoryUsageValidator("foo", new MemoryTracker(21, Log.noOpLog()),
             false,
             Log.noOpLog()
         )
             .tryValidateMemoryUsage(
+                "task",
                 TestConfig.empty(),
                 (config) -> new MemoryTreeWithDimensions(memoryTree, dimensions)
         ))
@@ -87,11 +89,12 @@ void shouldNotFailOnInsufficientMemoryIfInSudoMode() {
         var dimensions = GraphDimensions.builder().nodeCount(1000).build();
         var memoryTree = new TestTree("test", MemoryRange.of(42));
 
-        assertThatNoException().isThrownBy(() -> new MemoryUsageValidator(new MemoryTracker(21, Log.noOpLog()),
+        assertThatNoException().isThrownBy(() -> new MemoryUsageValidator("foo", new MemoryTracker(21, Log.noOpLog()),
             false,
             Log.noOpLog()
         )
             .tryValidateMemoryUsage(
+                "task",
                 TestConfig.of(CypherMapWrapper.empty().withBoolean("sudo", true)),
                 (config) -> new MemoryTreeWithDimensions(memoryTree, dimensions)
             ));
@@ -103,11 +106,15 @@ void shouldLogWhenFailing() {
         var dimensions = GraphDimensions.builder().nodeCount(1000).build();
         var memoryTree = new TestTree("test", MemoryRange.of(42));
         var memoryUsageValidator = new MemoryUsageValidator(
-            new MemoryTracker(21, log), false, log
+            "foo",
+            new MemoryTracker(21, log),
+            false,
+            log
         );
 
         assertThatIllegalStateException().isThrownBy(
             () -> memoryUsageValidator.tryValidateMemoryUsage(
+                "task",
                 TestConfig.of(CypherMapWrapper.empty()),
                 (config -> new MemoryTreeWithDimensions(memoryTree, dimensions))
             )
@@ -140,6 +147,7 @@ static Stream<Arguments> input() {
     void doesNotThrow(MemoryEstimation estimation, boolean useMaxMemoryUsage) {
         var memoryTrackerMock = mock(MemoryTracker.class);
         var memoryUsageValidator = new MemoryUsageValidator(
+            "foo",
             memoryTrackerMock,
             false,
             Log.noOpLog()
@@ -148,6 +156,7 @@ void doesNotThrow(MemoryEstimation estimation, boolean useMaxMemoryUsage) {
         var memoryTreeWithDimensions = new MemoryTreeWithDimensions(memoryTree, TEST_DIMENSIONS);
 
         assertDoesNotThrow(() -> memoryUsageValidator.validateMemoryUsage(
+            "task",
             memoryTreeWithDimensions.memoryTree.memoryUsage(), 10_000,
             useMaxMemoryUsage,
             new JobId("foo"), Log.noOpLog()
@@ -158,6 +167,7 @@ void doesNotThrow(MemoryEstimation estimation, boolean useMaxMemoryUsage) {
     @MethodSource("input")
     void throwsOnMinUsageExceeded(MemoryEstimation estimation, boolean ignored) {
         var memoryUsageValidator = new MemoryUsageValidator(
+            "foo",
             null,
             false,
             Log.noOpLog()
@@ -167,6 +177,7 @@ void throwsOnMinUsageExceeded(MemoryEstimation estimation, boolean ignored) {
         var memoryTreeWithDimensions = new MemoryTreeWithDimensions(memoryTree, TEST_DIMENSIONS);
 
         assertThatThrownBy(() -> memoryUsageValidator.validateMemoryUsage(
+            "task",
             memoryTreeWithDimensions.memoryTree.memoryUsage(), 1,
             false,
             new JobId("foo"), Log.noOpLog()
@@ -179,6 +190,7 @@ void throwsOnMinUsageExceeded(MemoryEstimation estimation, boolean ignored) {
     @MethodSource("input")
     void throwsOnMaxUsageExceeded(MemoryEstimation estimation, boolean ignored) {
         var memoryUsageValidator = new MemoryUsageValidator(
+            "foo",
             null,
             false,
             Log.noOpLog()
@@ -188,7 +200,9 @@ void throwsOnMaxUsageExceeded(MemoryEstimation estimation, boolean ignored) {
         var memoryTreeWithDimensions = new MemoryTreeWithDimensions(memoryTree, TEST_DIMENSIONS);
 
         assertThatThrownBy(() -> memoryUsageValidator.validateMemoryUsage(
-            memoryTreeWithDimensions.memoryTree.memoryUsage(), 1,
+            "task",
+            memoryTreeWithDimensions.memoryTree.memoryUsage(),
+            1,
             true,
             new JobId("foo"), Log.noOpLog()
         ))

applications/graph-store-catalog/src/main/java/org/neo4j/gds/applications/graphstorecatalog/GraphProjectMemoryUsageService.java

@@ -39,13 +39,15 @@ public class GraphProjectMemoryUsageService {
     private final Log log;
     private final GraphDatabaseService graphDatabaseService;
     private final MemoryTracker memoryTracker;
+    private final String username;
 
-    public GraphProjectMemoryUsageService(Log log, GraphDatabaseService graphDatabaseService,
+    public GraphProjectMemoryUsageService(String username,Log log, GraphDatabaseService graphDatabaseService,
         MemoryTracker memoryTracker
     ) {
         this.log = log;
         this.graphDatabaseService = graphDatabaseService;
         this.memoryTracker = memoryTracker;
+        this.username = username;
     }
 
     public void validateMemoryUsage(
@@ -57,6 +59,7 @@ public void validateMemoryUsage(
         GraphProjectConfig configuration
     ) {
         memoryUsageValidator().tryValidateMemoryUsage(
+            "Loading",
             configuration,
             graphProjectConfig -> getEstimate(
                 databaseId,
@@ -105,7 +108,7 @@ private MemoryUsageValidator memoryUsageValidator() {
             .resolveDependency(Config.class);
         var useMaxMemoryEstimation = neo4jConfig.get(GdsSettings.validateUsingMaxMemoryEstimation());
 
-        return new MemoryUsageValidator(memoryTracker, useMaxMemoryEstimation, log);
+        return new MemoryUsageValidator(username,memoryTracker, useMaxMemoryEstimation, log);
     }
 
     private GraphLoaderContext graphLoaderContext(

applications/graph-store-catalog/src/main/java/org/neo4j/gds/applications/graphstorecatalog/MemoryUsageValidator.java

@@ -40,14 +40,17 @@ public class MemoryUsageValidator {
     private final Log log;
     private final boolean useMaxMemoryEstimation;
     private final MemoryTracker memoryTracker;
+    private final String username;
 
-    public MemoryUsageValidator(MemoryTracker memoryTracker, boolean useMaxMemoryEstimation, Log log) {
+    public MemoryUsageValidator(String username,MemoryTracker memoryTracker, boolean useMaxMemoryEstimation, Log log) {
         this.log = log;
         this.useMaxMemoryEstimation = useMaxMemoryEstimation;
         this.memoryTracker = memoryTracker;
+        this.username = username;
     }
 
     public <C extends BaseConfig & JobIdConfig> MemoryRange tryValidateMemoryUsage(
+        String taskName,
         C config,
         Function<C, MemoryTreeWithDimensions> runEstimation
     ) {
@@ -58,11 +61,14 @@ public <C extends BaseConfig & JobIdConfig> MemoryRange tryValidateMemoryUsage(
             if (config.sudo()) {
                 log.debug("Sudo mode: Won't check for available memory.");
                 memoryTracker.track(
+                    username,
+                    taskName,
                     config.jobId(),
                     useMaxMemoryEstimation ? estimatedMemoryRange.max : estimatedMemoryRange.min
                 );
             } else {
                 validateMemoryUsage(
+                    username,
                     estimatedMemoryRange,
                     memoryTracker.availableMemory(),
                     useMaxMemoryEstimation,
@@ -78,6 +84,7 @@ public <C extends BaseConfig & JobIdConfig> MemoryRange tryValidateMemoryUsage(
     }
 
     void validateMemoryUsage(
+        String taskName,
         MemoryRange estimatedMemoryRange,
         long availableBytes,
         boolean useMaxMemoryEstimation,
@@ -86,6 +93,7 @@ void validateMemoryUsage(
     ) {
         if (useMaxMemoryEstimation) {
             validateMemoryUsage(
+                taskName,
                 availableBytes,
                 estimatedMemoryRange.max,
                 "maximum",
@@ -97,6 +105,7 @@ void validateMemoryUsage(
             );
         } else {
             validateMemoryUsage(
+                taskName,
                 availableBytes,
                 estimatedMemoryRange.min,
                 "minimum",
@@ -106,6 +115,7 @@ void validateMemoryUsage(
     }
 
     private void validateMemoryUsage(
+        String taskName,
         long availableBytes,
         long requiredBytes,
         String memoryString,
@@ -137,6 +147,6 @@ private void validateMemoryUsage(
             log.info(message);
             throw new IllegalStateException(message);
         }
-        memoryTracker.track(jobId, requiredBytes);
+        memoryTracker.track(username,taskName,jobId, requiredBytes);
     }
 }