Improve GCN infer override and weighted payload handling

Author: aotenjou

geaflow/geaflow-common/src/main/java/org/apache/geaflow/common/config/InferConfigHelper.java

@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.geaflow.common.config;
+
+import static org.apache.geaflow.common.config.keys.FrameworkConfigKeys.INFER_ENV_USER_TRANSFORM_CLASSNAME;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public final class InferConfigHelper {
+
+    private InferConfigHelper() {
+    }
+
+    public static boolean hasTransformClassName(Configuration configuration) {
+        return getTransformClassName(configuration) != null;
+    }
+
+    public static String getTransformClassName(Configuration configuration) {
+        if (configuration == null || !configuration.contains(INFER_ENV_USER_TRANSFORM_CLASSNAME)) {
+            return null;
+        }
+        String transformClassName = configuration.getString(
+            INFER_ENV_USER_TRANSFORM_CLASSNAME.getKey(), null);
+        if (transformClassName == null) {
+            return null;
+        }
+        String normalizedClassName = transformClassName.trim();
+        return normalizedClassName.isEmpty() ? null : normalizedClassName;
+    }
+
+    public static Configuration buildInferConfiguration(Configuration baseConfig,
+                                                        String pythonTransformClassName) {
+        if (baseConfig == null || pythonTransformClassName == null) {
+            return baseConfig;
+        }
+        String normalizedClassName = pythonTransformClassName.trim();
+        if (normalizedClassName.isEmpty()) {
+            return baseConfig;
+        }
+        if (normalizedClassName.equals(getTransformClassName(baseConfig))) {
+            return baseConfig;
+        }
+        Map<String, String> configMap = new HashMap<>(baseConfig.getConfigMap());
+        configMap.put(INFER_ENV_USER_TRANSFORM_CLASSNAME.getKey(), normalizedClassName);
+        Configuration configuration = new Configuration(configMap);
+        configuration.setMasterId(baseConfig.getMasterId());
+        return configuration;
+    }
+}

geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/main/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/DynamicGraphVertexCentricComputeOp.java

@@ -20,10 +20,8 @@
 package org.apache.geaflow.operator.impl.graph.compute.dynamic;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import org.apache.geaflow.api.function.iterator.RichIteratorFunction;
 import org.apache.geaflow.api.graph.base.algo.AbstractIncVertexCentricComputeAlgo;
@@ -32,6 +30,7 @@
 import org.apache.geaflow.api.graph.function.vc.base.IncGraphInferContext;
 import org.apache.geaflow.collector.ICollector;
 import org.apache.geaflow.common.config.Configuration;
+import org.apache.geaflow.common.config.InferConfigHelper;
 import org.apache.geaflow.common.config.keys.FrameworkConfigKeys;
 import org.apache.geaflow.common.exception.GeaflowRuntimeException;
 import org.apache.geaflow.infer.InferContext;
@@ -196,14 +195,6 @@ public void close() throws IOException {
     }
 
     static Configuration buildInferConfiguration(Configuration baseConfig, String pythonTransformClassName) {
-        if (pythonTransformClassName == null || pythonTransformClassName.trim().isEmpty()) {
-            return baseConfig;
-        }
-        Map<String, String> configMap = new HashMap<>(baseConfig.getConfigMap());
-        configMap.put(FrameworkConfigKeys.INFER_ENV_USER_TRANSFORM_CLASSNAME.getKey(),
-            pythonTransformClassName);
-        Configuration configuration = new Configuration(configMap);
-        configuration.setMasterId(baseConfig.getMasterId());
-        return configuration;
+        return InferConfigHelper.buildInferConfiguration(baseConfig, pythonTransformClassName);
     }
 }

geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/DynamicVertexRuntimeContext.java

@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.geaflow.dsl.common.algo;
+
+import org.apache.geaflow.dsl.common.data.RowVertex;
+
+public interface DynamicVertexRuntimeContext<K, M> extends AlgorithmRuntimeContext<K, M> {
+
+    void setVertexId(K vertexId);
+
+    RowVertex loadVertex();
+}

geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/org/apache/geaflow/dsl/udf/graph/GCN.java

@@ -19,48 +19,59 @@
 
 package org.apache.geaflow.dsl.udf.graph;
 
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
+import org.apache.geaflow.common.config.ConfigHelper;
+import org.apache.geaflow.common.config.Configuration;
+import org.apache.geaflow.common.config.InferConfigHelper;
+import org.apache.geaflow.common.config.keys.FrameworkConfigKeys;
+import org.apache.geaflow.common.exception.GeaflowRuntimeException;
 import org.apache.geaflow.common.type.Types;
 import org.apache.geaflow.dsl.common.algo.AlgorithmRuntimeContext;
 import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction;
+import org.apache.geaflow.dsl.common.algo.DynamicVertexRuntimeContext;
 import org.apache.geaflow.dsl.common.algo.IncrementalAlgorithmUserFunction;
 import org.apache.geaflow.dsl.common.data.Row;
-import org.apache.geaflow.dsl.common.data.RowEdge;
 import org.apache.geaflow.dsl.common.data.RowVertex;
 import org.apache.geaflow.dsl.common.data.impl.ObjectRow;
 import org.apache.geaflow.dsl.common.function.Description;
 import org.apache.geaflow.dsl.common.types.ArrayType;
 import org.apache.geaflow.dsl.common.types.GraphSchema;
 import org.apache.geaflow.dsl.common.types.StructType;
 import org.apache.geaflow.dsl.common.types.TableField;
+import org.apache.geaflow.dsl.udf.graph.gcn.AbstractGCNSubgraphAdapter;
 import org.apache.geaflow.dsl.udf.graph.gcn.GCNConfig;
 import org.apache.geaflow.dsl.udf.graph.gcn.GCNFeatureCollector;
 import org.apache.geaflow.dsl.udf.graph.gcn.GCNInferPayload;
 import org.apache.geaflow.dsl.udf.graph.gcn.GCNInferResult;
 import org.apache.geaflow.dsl.udf.graph.gcn.GCNResultParser;
 import org.apache.geaflow.dsl.udf.graph.gcn.GCNSubgraphBuilder;
+import org.apache.geaflow.infer.InferContext;
 import org.apache.geaflow.model.graph.edge.EdgeDirection;
+import org.apache.geaflow.model.graph.edge.IEdge;
 
 @Description(name = "gcn", description = "built-in udga for GCN inference")
 public class GCN implements AlgorithmUserFunction<Object, Object>, IncrementalAlgorithmUserFunction {
 
     private AlgorithmRuntimeContext<Object, Object> context;
+    private DynamicVertexRuntimeContext<Object, Object> dynamicContext;
     private GCNConfig config;
     private GCNSubgraphBuilder subgraphBuilder;
+    private InferContext<Object> localInferContext;
     private final GCNFeatureCollector featureCollector = new GCNFeatureCollector();
     private final GCNResultParser resultParser = new GCNResultParser();
 
     @Override
     public void init(AlgorithmRuntimeContext<Object, Object> context, Object[] params) {
         this.context = context;
+        this.dynamicContext = context instanceof DynamicVertexRuntimeContext
+            ? (DynamicVertexRuntimeContext<Object, Object>) context : null;
         this.config = parseConfig(params);
         this.subgraphBuilder = new GCNSubgraphBuilder(config);
+        this.localInferContext = createLocalInferContext(context.getConfig());
     }
 
     @Override
@@ -69,13 +80,17 @@ public void process(RowVertex vertex, Optional<Row> updatedValues, Iterator<Obje
             return;
         }
         GCNInferPayload payload = subgraphBuilder.build(vertex.getId(), new DynamicGraphAdapter(vertex));
-        Object rawResult = context.infer(payload);
+        Object rawResult = infer(payload);
         GCNInferResult result = resultParser.parse(vertex.getId(), rawResult);
         context.take(ObjectRow.create(result.toRowValues()));
     }
 
     @Override
     public void finish(RowVertex graphVertex, Optional<Row> updatedValues) {
+        if (localInferContext != null) {
+            localInferContext.close();
+            localInferContext = null;
+        }
     }
 
     @Override
@@ -102,65 +117,88 @@ private GCNConfig parseConfig(Object[] params) {
         return new GCNConfig(numHops, numSamplesPerHop, className);
     }
 
-    private class DynamicGraphAdapter implements GCNSubgraphBuilder.GraphAdapter {
+    private Object infer(GCNInferPayload payload) {
+        if (localInferContext == null) {
+            return context.infer(payload);
+        }
+        try {
+            return localInferContext.infer(payload);
+        } catch (Exception e) {
+            throw new GeaflowRuntimeException("GCN local infer failed", e);
+        }
+    }
+
+    private InferContext<Object> createLocalInferContext(Configuration baseConfig) {
+        if (!ConfigHelper.getBooleanOrDefault(baseConfig.getConfigMap(),
+            FrameworkConfigKeys.INFER_ENV_ENABLE.getKey(), false)) {
+            return null;
+        }
+        String globalClassName = InferConfigHelper.getTransformClassName(baseConfig);
+        if (Objects.equals(config.getPythonTransformClassName(), globalClassName)) {
+            return null;
+        }
+        return createInferContext(InferConfigHelper.buildInferConfiguration(baseConfig,
+            config.getPythonTransformClassName()));
+    }
+
+    protected InferContext<Object> createInferContext(Configuration configuration) {
+        return new InferContext<>(configuration);
+    }
+
+    private class DynamicGraphAdapter extends AbstractGCNSubgraphAdapter<RowVertex> {
 
         private final GraphSchema graphSchema = context.getGraphSchema();
-        private final RowVertex rootVertex;
+        private final Object rootId;
 
         DynamicGraphAdapter(RowVertex rootVertex) {
-            this.rootVertex = rootVertex;
+            super(rootVertex.getId(), rootVertex, config.getFeatureDimLimit());
+            this.rootId = rootVertex.getId();
         }
 
         @Override
-        public List<Object> loadNeighbors(Object nodeId) {
+        protected List<? extends IEdge<Object, ?>> loadEdges(Object nodeId) {
             RowVertex current = switchVertex(nodeId);
-            if (current == null) {
-                restoreVertex(rootVertex.getId());
-                return Collections.emptyList();
-            }
-            List<RowEdge> edges = context.loadEdges(EdgeDirection.BOTH);
-            List<Object> neighbors = new ArrayList<>(edges.size());
-            for (RowEdge edge : edges) {
-                Object neighborId = nodeId.equals(edge.getSrcId()) ? edge.getTargetId()
-                    : edge.getSrcId();
-                neighbors.add(neighborId);
+            try {
+                if (current == null) {
+                    return Collections.emptyList();
+                }
+                return context.loadEdges(EdgeDirection.BOTH);
+            } finally {
+                restoreVertex();
             }
-            restoreVertex(rootVertex.getId());
-            return neighbors;
         }
 
         @Override
-        public double[] loadFeatures(Object nodeId) {
-            RowVertex vertex = nodeId.equals(rootVertex.getId()) ? rootVertex : switchVertex(nodeId);
+        protected RowVertex loadNode(Object nodeId) {
+            RowVertex vertex = switchVertex(nodeId);
             try {
-                return vertex == null ? new double[config.getFeatureDimLimit()]
-                    : featureCollector.collectFromRowVertex(vertex, graphSchema,
-                        config.getFeatureDimLimit());
+                return vertex;
             } finally {
-                restoreVertex(rootVertex.getId());
+                restoreVertex();
             }
         }
 
+        @Override
+        protected double[] extractFeatures(RowVertex vertex) {
+            return featureCollector.collectFromRowVertex(vertex, graphSchema,
+                config.getFeatureDimLimit());
+        }
+
         private RowVertex switchVertex(Object nodeId) {
-            try {
-                Method setVertexId = context.getClass().getMethod("setVertexId", Object.class);
-                setVertexId.invoke(context, nodeId);
-                Method loadVertex = context.getClass().getMethod("loadVertex");
-                return (RowVertex) loadVertex.invoke(context);
-            } catch (NoSuchMethodException e) {
-                throw new IllegalStateException("GCN requires dynamic runtime context with loadVertex support", e);
-            } catch (IllegalAccessException | InvocationTargetException e) {
-                throw new IllegalStateException("Failed to switch runtime context vertex", e);
+            if (dynamicContext == null) {
+                throw new IllegalStateException(
+                    "GCN requires DynamicVertexRuntimeContext to load sampled vertices");
             }
+            dynamicContext.setVertexId(nodeId);
+            return dynamicContext.loadVertex();
         }
 
-        private void restoreVertex(Object rootId) {
-            try {
-                Method setVertexId = context.getClass().getMethod("setVertexId", Object.class);
-                setVertexId.invoke(context, rootId);
-            } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
-                throw new IllegalStateException("Failed to restore runtime context vertex", e);
+        private void restoreVertex() {
+            if (dynamicContext == null) {
+                throw new IllegalStateException(
+                    "GCN requires DynamicVertexRuntimeContext to restore sampled vertices");
             }
+            dynamicContext.setVertexId(rootId);
         }
     }
 }