From 11b99b1196ab233d594e01154bbaadce4398973d Mon Sep 17 00:00:00 2001 From: Veselin Nikolov Date: Tue, 14 Nov 2023 09:51:35 +0000 Subject: [PATCH] Hook AlgorithmMetricsService to ProcedureExecutor Co-authored-by: Ioannis Panagiotas --- executor/build.gradle | 1 + .../neo4j/gds/executor/ExecutionContext.java | 9 +++++++++ .../neo4j/gds/executor/ProcedureExecutor.java | 19 ++++++++++++++++--- .../gds/executor/ProcedureExecutorTest.java | 3 +++ proc/common/build.gradle | 1 + .../src/main/java/org/neo4j/gds/BaseProc.java | 5 +++++ procedures/extension/build.gradle | 1 + .../OpenGraphDataScienceExtension.java | 14 +++++++++++++- .../integration/ExtensionBuilder.java | 5 ++++- 9 files changed, 53 insertions(+), 5 deletions(-) diff --git a/executor/build.gradle b/executor/build.gradle index 4f9c3a06fb..aaf31e574f 100644 --- a/executor/build.gradle +++ b/executor/build.gradle @@ -10,6 +10,7 @@ dependencies { annotationProcessor group: 'org.immutables', name: 'builder', version: ver.'immutables' annotationProcessor group: 'org.immutables', name: 'value', version: ver.'immutables' + implementation project(':algorithm-metrics-api') implementation project(':annotations') implementation project(':algo') implementation project(':algo-common') diff --git a/executor/src/main/java/org/neo4j/gds/executor/ExecutionContext.java b/executor/src/main/java/org/neo4j/gds/executor/ExecutionContext.java index ee333204b6..d3759c4f04 100644 --- a/executor/src/main/java/org/neo4j/gds/executor/ExecutionContext.java +++ b/executor/src/main/java/org/neo4j/gds/executor/ExecutionContext.java @@ -21,6 +21,8 @@ import org.jetbrains.annotations.Nullable; import org.neo4j.common.DependencyResolver; +import org.neo4j.gds.algorithms.metrics.AlgorithmMetricsService; +import org.neo4j.gds.algorithms.metrics.PassthroughAlgorithmMetricRegistrar; import org.neo4j.gds.annotation.ValueClass; import org.neo4j.gds.api.AlgorithmMetaDataSetter; import org.neo4j.gds.api.CloseableResourceRegistry; @@ -70,6 +72,8 @@ public interface ExecutionContext { boolean isGdsAdmin(); + AlgorithmMetricsService algorithmMetricsService(); + @Nullable RelationshipStreamExporterBuilder relationshipStreamExporterBuilder(); @@ -173,6 +177,11 @@ public boolean isGdsAdmin() { return false; } + @Override + public AlgorithmMetricsService algorithmMetricsService() { + return new AlgorithmMetricsService(new PassthroughAlgorithmMetricRegistrar()); + } + @Override public @Nullable RelationshipStreamExporterBuilder relationshipStreamExporterBuilder() { return null; diff --git a/executor/src/main/java/org/neo4j/gds/executor/ProcedureExecutor.java b/executor/src/main/java/org/neo4j/gds/executor/ProcedureExecutor.java index 93b0d78462..7611e96771 100644 --- a/executor/src/main/java/org/neo4j/gds/executor/ProcedureExecutor.java +++ b/executor/src/main/java/org/neo4j/gds/executor/ProcedureExecutor.java @@ -23,6 +23,7 @@ import org.neo4j.gds.AlgorithmFactory; import org.neo4j.gds.GraphAlgorithmFactory; import org.neo4j.gds.GraphStoreAlgorithmFactory; +import org.neo4j.gds.algorithms.metrics.AlgorithmMetricsService; import org.neo4j.gds.api.Graph; import org.neo4j.gds.api.GraphStore; import org.neo4j.gds.config.AlgoBaseConfig; @@ -110,7 +111,8 @@ public RESULT compute( algo.getProgressTracker().setEstimatedResourceFootprint(memoryEstimationInBytes, config.concurrency()); - ALGO_RESULT result = executeAlgorithm(builder, algo); + + ALGO_RESULT result = executeAlgorithm(builder, algo, executionContext.algorithmMetricsService()); var computationResult = builder .graph(graph) @@ -125,15 +127,26 @@ public RESULT compute( private ALGO_RESULT executeAlgorithm( ImmutableComputationResult.Builder builder, - ALGO algo + ALGO algo, + AlgorithmMetricsService algorithmMetricsService ) { return runWithExceptionLogging( "Computation failed", () -> { - try (ProgressTimer ignored = ProgressTimer.start(builder::computeMillis)) { + var algorithmMetric = algorithmMetricsService.create( + // we don't want to use `spec.name()` because it's different for the different procedure modes; + // we want to capture the algorithm name as defined by the algorithm factory `taskName()` + algoSpec.algorithmFactory(executionContext).taskName() + ); + try ( + ProgressTimer ignored = ProgressTimer.start(builder::computeMillis); + algorithmMetric; + ) { + algorithmMetric.start(); return algo.compute(); } catch (Throwable e) { algo.getProgressTracker().endSubTaskWithFailure(); + algorithmMetric.failed(); throw e; } finally { if (algoSpec.releaseProgressTask()) { diff --git a/executor/src/test/java/org/neo4j/gds/executor/ProcedureExecutorTest.java b/executor/src/test/java/org/neo4j/gds/executor/ProcedureExecutorTest.java index 147f5beead..5c92e79796 100644 --- a/executor/src/test/java/org/neo4j/gds/executor/ProcedureExecutorTest.java +++ b/executor/src/test/java/org/neo4j/gds/executor/ProcedureExecutorTest.java @@ -23,6 +23,8 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.neo4j.gds.ProcedureCallContextReturnColumns; +import org.neo4j.gds.algorithms.metrics.AlgorithmMetricsService; +import org.neo4j.gds.algorithms.metrics.PassthroughAlgorithmMetricRegistrar; import org.neo4j.gds.api.AlgorithmMetaDataSetter; import org.neo4j.gds.api.CloseableResourceRegistry; import org.neo4j.gds.api.GraphStore; @@ -136,6 +138,7 @@ private ExecutionContext executionContext(TaskStore taskStore) { .algorithmMetaDataSetter(AlgorithmMetaDataSetter.EMPTY) .nodeLookup(NodeLookup.EMPTY) .userLogRegistryFactory(EmptyUserLogRegistryFactory.INSTANCE) + .algorithmMetricsService(new AlgorithmMetricsService(new PassthroughAlgorithmMetricRegistrar())) .build(); } diff --git a/proc/common/build.gradle b/proc/common/build.gradle index 6a95c597a9..5c5c5cbef2 100644 --- a/proc/common/build.gradle +++ b/proc/common/build.gradle @@ -15,6 +15,7 @@ dependencies { implementation project(':annotations') implementation project(':algo-common') + implementation project(':algorithm-metrics-api') implementation project(':config-api') implementation project(':core') implementation project(':core-write') diff --git a/proc/common/src/main/java/org/neo4j/gds/BaseProc.java b/proc/common/src/main/java/org/neo4j/gds/BaseProc.java index 3ccb80ce63..2320155cb2 100644 --- a/proc/common/src/main/java/org/neo4j/gds/BaseProc.java +++ b/proc/common/src/main/java/org/neo4j/gds/BaseProc.java @@ -19,6 +19,7 @@ */ package org.neo4j.gds; +import org.neo4j.gds.algorithms.metrics.AlgorithmMetricsService; import org.neo4j.gds.api.DatabaseId; import org.neo4j.gds.compat.GraphDatabaseApiProxy; import org.neo4j.gds.config.BaseConfig; @@ -75,6 +76,9 @@ public abstract class BaseProc { @Context public Username username = Username.EMPTY_USERNAME; + @Context + public AlgorithmMetricsService algorithmMetricsService; + protected String username() { return username.username(); } @@ -157,6 +161,7 @@ public ExecutionContext executionContext() { .algorithmMetaDataSetter(new TransactionAlgorithmMetaDataSetter(transaction)) .nodeLookup(new TransactionNodeLookup(transaction)) .isGdsAdmin(transactionContext().isGdsAdmin()) + .algorithmMetricsService(algorithmMetricsService) .build(); } diff --git a/procedures/extension/build.gradle b/procedures/extension/build.gradle index e3ed8a1c61..f51938d5e7 100644 --- a/procedures/extension/build.gradle +++ b/procedures/extension/build.gradle @@ -18,6 +18,7 @@ dependencies { compileOnly(group: 'org.neo4j', name: 'neo4j-logging', version: ver.'neo4j') { transitive = false } // the necessary GDS things for the extension to construct the application + implementation project(':algorithm-metrics-api') implementation project(':config-api') implementation project(':core') implementation project(':core-utils') diff --git a/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java b/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java index a6e8cf9508..9d38c84ddd 100644 --- a/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java +++ b/procedures/extension/src/main/java/org/neo4j/gds/extension/OpenGraphDataScienceExtension.java @@ -21,6 +21,8 @@ import org.neo4j.annotations.service.ServiceProvider; import org.neo4j.configuration.Config; +import org.neo4j.gds.algorithms.metrics.AlgorithmMetricsService; +import org.neo4j.gds.algorithms.metrics.PassthroughAlgorithmMetricRegistrar; import org.neo4j.gds.applications.graphstorecatalog.CatalogBusinessFacade; import org.neo4j.gds.core.write.NativeExportBuildersProvider; import org.neo4j.gds.procedures.GraphDataScience; @@ -64,10 +66,20 @@ public Lifecycle newInstance(ExtensionContext extensionContext, Dependencies dep // we have no extra checks to do in OpenGDS Optional> businessFacadeDecorator = Optional.empty(); + var algorithmMetricsService = new AlgorithmMetricsService(new PassthroughAlgorithmMetricRegistrar()); + extensionBuilder .withComponent( GraphDataScience.class, - () -> extensionBuilder.gdsProvider(exporterBuildersProviderService, businessFacadeDecorator) + () -> extensionBuilder.gdsProvider( + exporterBuildersProviderService, + businessFacadeDecorator, + algorithmMetricsService + ) + ) + .withComponent( + AlgorithmMetricsService.class, + () -> ctx -> algorithmMetricsService ) .registerExtension(); diff --git a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/ExtensionBuilder.java b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/ExtensionBuilder.java index 9e288c0752..511560b2c0 100644 --- a/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/ExtensionBuilder.java +++ b/procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/ExtensionBuilder.java @@ -20,6 +20,7 @@ package org.neo4j.gds.procedures.integration; import org.neo4j.function.ThrowingFunction; +import org.neo4j.gds.algorithms.metrics.AlgorithmMetricsService; import org.neo4j.gds.applications.graphstorecatalog.CatalogBusinessFacade; import org.neo4j.gds.core.loading.GraphStoreCatalogService; import org.neo4j.gds.core.utils.progress.ProgressFeatureSettings; @@ -187,10 +188,12 @@ public void registerExtension() { * * @param exporterBuildersProviderService The catalog of writers * @param businessFacadeDecorator Any checks added across requests + * @param algorithmMetricsService */ public ThrowingFunction gdsProvider( ExporterBuildersProviderService exporterBuildersProviderService, - Optional> businessFacadeDecorator + Optional> businessFacadeDecorator, + AlgorithmMetricsService algorithmMetricsService ) { var catalogFacadeProvider = createCatalogFacadeProvider( exporterBuildersProviderService,