| /* |
| * 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.beam.runners.dataflow.worker.graph; |
| |
| import com.google.api.services.dataflow.model.FlattenInstruction; |
| import com.google.api.services.dataflow.model.MultiOutputInfo; |
| import com.google.api.services.dataflow.model.ParallelInstruction; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Set; |
| import java.util.function.Function; |
| import java.util.function.Supplier; |
| import org.apache.beam.runners.dataflow.worker.graph.Edges.Edge; |
| import org.apache.beam.runners.dataflow.worker.graph.Edges.HappensBeforeEdge; |
| import org.apache.beam.runners.dataflow.worker.graph.Edges.MultiOutputInfoEdge; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode; |
| import org.apache.beam.sdk.fn.IdGenerator; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.Graphs; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.ImmutableNetwork; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.MutableNetwork; |
| |
| /** |
| * Splits the instruction graph into SDK and runner harness portions replacing the SDK sub-graphs |
| * with operations linked by gRPC ports. |
| * |
| * <p>This transform expects that all {@link FlattenInstruction}s have been removed. {@link |
| * FlattenInstruction}s are not required because Beam Fn API multiplexers handle multiple inputs. |
| * Also, {@link FlattenInstruction}s complicate the implementation as to what fused section of the |
| * graph executes in the Runner harness versus what executes in the SDK harness since you need to |
| * walk the instruction graph across flattens to find whether the neighboring instruction executes |
| * in the Runner or in the SDK harness. |
| * |
| * <p>Note that this implementation minimizes the amount of data that will be transferred across a |
| * port by creating new edges which represent a smaller portion of each PCollection. For example: |
| * |
| * <pre><code> |
| * RunnerSource --\ /--> RunnerParDo |
| * out |
| * CustomSource --/ \--> SdkParDo |
| * </code></pre> |
| * |
| * Should produce: |
| * |
| * <pre><code> |
| * PortB --> out --\ |
| * RunnerSource --> out --> RunnerParDo |
| * \--> PortA |
| * PortA --> out --\ |
| * CustomSource --> out --> SdkParDo |
| * \--> PortB |
| * </code></pre> |
| */ |
| public class CreateRegisterFnOperationFunction |
| implements Function<MutableNetwork<Node, Edge>, MutableNetwork<Node, Edge>> { |
| private final IdGenerator idGenerator; |
| private final Supplier<Node> portSupplier; |
| private final Function<MutableNetwork<Node, Edge>, Node> registerFnOperationFunction; |
| private final boolean useExecutableStageBundleExecution; |
| |
| /** |
| * Constructs a function which is able to break up the instruction graph into SDK and Runner |
| * harness portions utilizing the passed in {@code portSupplier} to create communication ports |
| * between the SDK and Runner harness and the {@code registerFnOperationFunction} to notify the |
| * SDK of functions it is required to execute. |
| * |
| * @param portSupplier A {@link Supplier} which produces nodes to communicate data between the SDK |
| * and Runner harness. |
| * @param registerFnOperationFunction A {@link Function} which takes a {@link MutableNetwork} and |
| * produces a {@link Node} that is able to register the SDK functions within the SDK harness. |
| */ |
| public CreateRegisterFnOperationFunction( |
| IdGenerator idGenerator, |
| Supplier<Node> portSupplier, |
| Function<MutableNetwork<Node, Edge>, Node> registerFnOperationFunction, |
| boolean useExecutableStageBundleExecution) { |
| this.idGenerator = idGenerator; |
| this.portSupplier = portSupplier; |
| this.registerFnOperationFunction = registerFnOperationFunction; |
| this.useExecutableStageBundleExecution = useExecutableStageBundleExecution; |
| } |
| |
| @Override |
| public MutableNetwork<Node, Edge> apply(MutableNetwork<Node, Edge> network) { |
| |
| // Record all SDK nodes, and all root nodes. |
| Set<Node> runnerRootNodes = new HashSet<>(); |
| Set<Node> sdkNodes = new HashSet<>(); |
| Set<Node> sdkRootNodes = new HashSet<>(); |
| for (ParallelInstructionNode node : |
| Iterables.filter(network.nodes(), ParallelInstructionNode.class)) { |
| if (executesInSdkHarness(node)) { |
| sdkNodes.add(node); |
| if (network.inDegree(node) == 0) { |
| sdkRootNodes.add(node); |
| } |
| } else if (network.inDegree(node) == 0) { |
| runnerRootNodes.add(node); |
| } |
| } |
| |
| // If nothing executes within the SDK harness, return the original network. |
| if (sdkNodes.isEmpty()) { |
| return network; |
| } |
| |
| // Represents the set of nodes which represent gRPC boundaries from the Runner to the SDK. |
| Set<Node> runnerToSdkBoundaries = new HashSet<>(); |
| // Represents the set of nodes which represent gRPC boundaries from the SDK to the Runner. |
| Set<Node> sdkToRunnerBoundaries = new HashSet<>(); |
| |
| ImmutableNetwork<Node, Edge> originalNetwork = ImmutableNetwork.copyOf(network); |
| |
| // Update the network with outputs which are meant to bridge the instructions |
| // that execute in different harnesses. One output per direction of information |
| // flow from runner to SDK and SDK to runner per original output node. |
| for (InstructionOutputNode outputNode : |
| Iterables.filter(originalNetwork.nodes(), InstructionOutputNode.class)) { |
| |
| // Categorize all predecessor instructions |
| Set<Node> predecessorRunnerInstructions = new HashSet<>(); |
| Set<Node> predecessorSdkInstructions = new HashSet<>(); |
| for (Node predecessorInstruction : originalNetwork.predecessors(outputNode)) { |
| if (sdkNodes.contains(predecessorInstruction)) { |
| predecessorSdkInstructions.add(predecessorInstruction); |
| } else { |
| predecessorRunnerInstructions.add(predecessorInstruction); |
| } |
| } |
| |
| // Categorize all successor instructions |
| Set<Node> successorRunnerInstructions = new HashSet<>(); |
| Set<Node> successorSdkInstructions = new HashSet<>(); |
| for (Node successorInstruction : originalNetwork.successors(outputNode)) { |
| if (sdkNodes.contains(successorInstruction)) { |
| successorSdkInstructions.add(successorInstruction); |
| } else { |
| successorRunnerInstructions.add(successorInstruction); |
| } |
| } |
| |
| // If there is data that will be flowing from the Runner to the SDK, rewire network to have |
| // nodes connected across a gRPC node. Also add the gRPC node as an SDK root. |
| if (!predecessorRunnerInstructions.isEmpty() && !successorSdkInstructions.isEmpty()) { |
| runnerToSdkBoundaries.add( |
| rewireAcrossSdkRunnerPortNode( |
| network, outputNode, predecessorRunnerInstructions, successorSdkInstructions)); |
| } |
| |
| // If there is data that will be flowing from the SDK to the Runner, rewire network to have |
| // nodes connected across a gRPC node. |
| if (!predecessorSdkInstructions.isEmpty() && !successorRunnerInstructions.isEmpty()) { |
| sdkToRunnerBoundaries.add( |
| rewireAcrossSdkRunnerPortNode( |
| network, outputNode, predecessorSdkInstructions, successorRunnerInstructions)); |
| } |
| |
| // Remove original output node if it was rewired because it will have become disconnected |
| // through the new output node. |
| if (network.inDegree(outputNode) == 0) { |
| network.removeNode(outputNode); |
| } |
| } |
| |
| // Create the subnetworks that represent potentially multiple fused SDK portions and a single |
| // fused Runner portion replacing the SDK portion that is embedded within the Runner portion |
| // with a RegisterFnOperation, adding edges to maintain proper happens before relationships. |
| Set<Node> allRunnerNodes = |
| Networks.reachableNodes( |
| network, Sets.union(runnerRootNodes, sdkToRunnerBoundaries), runnerToSdkBoundaries); |
| if (this.useExecutableStageBundleExecution) { |
| // When using shared library, there is no grpc node in runner graph. |
| allRunnerNodes = |
| Sets.difference(allRunnerNodes, Sets.union(runnerToSdkBoundaries, sdkToRunnerBoundaries)); |
| } |
| MutableNetwork<Node, Edge> runnerNetwork = Graphs.inducedSubgraph(network, allRunnerNodes); |
| |
| // TODO: Reduce the amount of 'copying' of SDK nodes by breaking potential cycles |
| // between the SDK networks and the Runner network. Cycles can occur because entire |
| // SDK subnetworks are replaced by a singular node within the Runner network. |
| // khines@ suggested to look at go/priority-based-fusion for an algorithm based upon |
| // using poison paths. |
| for (Node sdkRoot : Sets.union(sdkRootNodes, runnerToSdkBoundaries)) { |
| Set<Node> sdkSubnetworkNodes = |
| Networks.reachableNodes(network, ImmutableSet.of(sdkRoot), sdkToRunnerBoundaries); |
| MutableNetwork<Node, Edge> sdkNetwork = Graphs.inducedSubgraph(network, sdkSubnetworkNodes); |
| Node registerFnNode = registerFnOperationFunction.apply(sdkNetwork); |
| |
| runnerNetwork.addNode(registerFnNode); |
| // Create happens before relationships between all Runner and SDK nodes which are in the |
| // SDK subnetwork; direction dependent on whether its a predecessor of the SDK subnetwork or |
| // a successor. |
| if (this.useExecutableStageBundleExecution) { |
| // When using shared library, there is no gprc node in runner graph. Then the registerFnNode |
| // should be linked directly to 2 OutputInstruction nodes. |
| for (Node predecessor : Sets.intersection(sdkSubnetworkNodes, runnerToSdkBoundaries)) { |
| predecessor = network.predecessors(predecessor).iterator().next(); |
| runnerNetwork.addEdge(predecessor, registerFnNode, HappensBeforeEdge.create()); |
| } |
| for (Node successor : Sets.intersection(sdkSubnetworkNodes, sdkToRunnerBoundaries)) { |
| successor = network.successors(successor).iterator().next(); |
| runnerNetwork.addEdge(registerFnNode, successor, HappensBeforeEdge.create()); |
| } |
| } else { |
| for (Node predecessor : Sets.intersection(sdkSubnetworkNodes, runnerToSdkBoundaries)) { |
| runnerNetwork.addEdge(predecessor, registerFnNode, HappensBeforeEdge.create()); |
| } |
| for (Node successor : Sets.intersection(sdkSubnetworkNodes, sdkToRunnerBoundaries)) { |
| runnerNetwork.addEdge(registerFnNode, successor, HappensBeforeEdge.create()); |
| } |
| } |
| } |
| |
| return runnerNetwork; |
| } |
| |
| /** |
| * Rewires the given set of predecessors and successors across a gRPC port surrounded by output |
| * nodes. Edges to the remaining successors are copied over to the new output node that is placed |
| * before the port node. For example: |
| * |
| * <pre><code> |
| * predecessors --> outputNode --> successors |
| * \--> existingSuccessors |
| * </pre></code> becomes: |
| * |
| * <pre><code> |
| * |
| * outputNode -------------------------------\ |
| * \ \ |
| * |-> existingSuccessors \ |
| * / \ |
| * predecessors --> newPredecessorOutputNode --> portNode --> portOutputNode --> successors}. |
| * </code></pre> |
| */ |
| private Node rewireAcrossSdkRunnerPortNode( |
| MutableNetwork<Node, Edge> network, |
| InstructionOutputNode outputNode, |
| Set<Node> predecessors, |
| Set<Node> successors) { |
| |
| InstructionOutputNode newPredecessorOutputNode = |
| InstructionOutputNode.create( |
| outputNode.getInstructionOutput(), outputNode.getPcollectionId()); |
| InstructionOutputNode portOutputNode = |
| InstructionOutputNode.create( |
| outputNode.getInstructionOutput(), outputNode.getPcollectionId()); |
| Node portNode = portSupplier.get(); |
| network.addNode(newPredecessorOutputNode); |
| network.addNode(portNode); |
| for (Node predecessor : predecessors) { |
| for (Edge edge : ImmutableList.copyOf(network.edgesConnecting(predecessor, outputNode))) { |
| network.removeEdge(edge); |
| network.addEdge(predecessor, newPredecessorOutputNode, edge); |
| } |
| } |
| |
| // Maintain edges for existing successors. |
| List<Node> existingSuccessors = |
| ImmutableList.copyOf(Sets.difference(network.successors(outputNode), successors)); |
| for (Node existingSuccessor : existingSuccessors) { |
| List<Edge> existingSuccessorEdges = |
| ImmutableList.copyOf(network.edgesConnecting(outputNode, existingSuccessor)); |
| for (Edge existingSuccessorEdge : existingSuccessorEdges) { |
| network.addEdge(newPredecessorOutputNode, existingSuccessor, existingSuccessorEdge.clone()); |
| } |
| } |
| // Rewire the requested successors over the port node. |
| network.addEdge( |
| newPredecessorOutputNode, |
| portNode, |
| MultiOutputInfoEdge.create(new MultiOutputInfo().setTag(idGenerator.getId()))); |
| network.addEdge( |
| portNode, |
| portOutputNode, |
| MultiOutputInfoEdge.create(new MultiOutputInfo().setTag(idGenerator.getId()))); |
| for (Node successor : successors) { |
| for (Edge edge : ImmutableList.copyOf(network.edgesConnecting(outputNode, successor))) { |
| network.addEdge(portOutputNode, successor, edge.clone()); |
| } |
| } |
| return portNode; |
| } |
| |
| /** |
| * Returns true iff the {@link ParallelInstruction} represents an instruction that executes within |
| * the SDK harness as part of the fused sub-graph. For details on how node locations are deduced |
| * refer to {@link DeduceNodeLocationsFunction#executesInSdkHarness}. |
| */ |
| private static boolean executesInSdkHarness(ParallelInstructionNode node) |
| throws IllegalStateException { |
| Nodes.ExecutionLocation location = node.getExecutionLocation(); |
| if (location == Nodes.ExecutionLocation.UNKNOWN) { |
| throw new IllegalStateException("Node must not have unknown execution location: " + node); |
| } |
| return location == Nodes.ExecutionLocation.SDK_HARNESS; |
| } |
| } |
