use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.
the class BeamFnMapTaskExecutorFactory method create.
/**
* Creates a new {@link DataflowMapTaskExecutor} from the given {@link MapTask} definition using
* the provided {@link ReaderFactory}.
*/
@Override
public DataflowMapTaskExecutor create(InstructionRequestHandler instructionRequestHandler, GrpcFnServer<GrpcDataService> grpcDataFnServer, Endpoints.ApiServiceDescriptor dataApiServiceDescriptor, GrpcFnServer<GrpcStateService> grpcStateFnServer, MutableNetwork<Node, Edge> network, PipelineOptions options, String stageName, ReaderFactory readerFactory, SinkFactory sinkFactory, DataflowExecutionContext<?> executionContext, CounterSet counterSet, IdGenerator idGenerator) {
// TODO: remove this once we trust the code paths
checkArgument(DataflowRunner.hasExperiment(options.as(DataflowPipelineDebugOptions.class), "beam_fn_api"), "%s should only be used when beam_fn_api is enabled", getClass().getSimpleName());
// Swap out all the InstructionOutput nodes with OutputReceiver nodes
Networks.replaceDirectedNetworkNodes(network, createOutputReceiversTransform(stageName, counterSet));
if (DataflowRunner.hasExperiment(options.as(DataflowPipelineDebugOptions.class), "use_executable_stage_bundle_execution")) {
LOG.debug("Using SingleEnvironmentInstanceJobBundleFactory");
JobBundleFactory jobBundleFactory = SingleEnvironmentInstanceJobBundleFactory.create(StaticRemoteEnvironmentFactory.forService(instructionRequestHandler), grpcDataFnServer, grpcStateFnServer, idGenerator);
// If the use_executable_stage_bundle_execution is enabled, use ExecutableStage instead.
Networks.replaceDirectedNetworkNodes(network, createOperationTransformForExecutableStageNode(network, stageName, executionContext, jobBundleFactory));
} else {
// Swap out all the RegisterFnRequest nodes with Operation nodes
Networks.replaceDirectedNetworkNodes(network, createOperationTransformForRegisterFnNodes(idGenerator, instructionRequestHandler, grpcStateFnServer.getService(), stageName, executionContext));
// Swap out all the RemoteGrpcPort nodes with Operation nodes, note that it is expected
// that the RegisterFnRequest nodes have already been replaced.
Networks.replaceDirectedNetworkNodes(network, createOperationTransformForGrpcPortNodes(network, grpcDataFnServer.getService(), // TODO: Set NameContext properly for these operations.
executionContext.createOperationContext(NameContext.create(stageName, stageName, stageName, stageName))));
}
// Swap out all the FetchAndFilterStreamingSideInput nodes with operation nodes
Networks.replaceDirectedNetworkNodes(network, createOperationTransformForFetchAndFilterStreamingSideInputNodes(network, idGenerator, instructionRequestHandler, grpcDataFnServer.getService(), dataApiServiceDescriptor, executionContext, stageName));
// Swap out all the ParallelInstruction nodes with Operation nodes
Networks.replaceDirectedNetworkNodes(network, createOperationTransformForParallelInstructionNodes(stageName, network, options, readerFactory, sinkFactory, executionContext));
// Collect all the operations within the network and attach all the operations as receivers
// to preceding output receivers.
List<Operation> topoSortedOperations = new ArrayList<>();
for (OperationNode node : Iterables.filter(Networks.topologicalOrder(network), OperationNode.class)) {
topoSortedOperations.add(node.getOperation());
for (Node predecessor : Iterables.filter(network.predecessors(node), OutputReceiverNode.class)) {
((OutputReceiverNode) predecessor).getOutputReceiver().addOutput((Receiver) node.getOperation());
}
}
if (LOG.isDebugEnabled()) {
LOG.info("Map task network: {}", Networks.toDot(network));
}
return BeamFnMapTaskExecutor.withSharedCounterSet(topoSortedOperations, counterSet, executionContext.getExecutionStateTracker());
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.
the class BeamFnMapTaskExecutorFactory method createParDoOperation.
private OperationNode createParDoOperation(Network<Node, Edge> network, ParallelInstructionNode node, PipelineOptions options, DataflowExecutionContext<?> executionContext, DataflowOperationContext operationContext) throws Exception {
ParallelInstruction instruction = node.getParallelInstruction();
ParDoInstruction parDo = instruction.getParDo();
TupleTag<?> mainOutputTag = tupleTag(parDo.getMultiOutputInfos().get(0));
ImmutableMap.Builder<TupleTag<?>, Integer> outputTagsToReceiverIndicesBuilder = ImmutableMap.builder();
int successorOffset = 0;
for (Node successor : network.successors(node)) {
for (Edge edge : network.edgesConnecting(node, successor)) {
outputTagsToReceiverIndicesBuilder.put(tupleTag(((MultiOutputInfoEdge) edge).getMultiOutputInfo()), successorOffset);
}
successorOffset += 1;
}
ParDoFn fn = parDoFnFactory.create(options, CloudObject.fromSpec(parDo.getUserFn()), parDo.getSideInputs(), mainOutputTag, outputTagsToReceiverIndicesBuilder.build(), executionContext, operationContext);
OutputReceiver[] receivers = getOutputReceivers(network, node);
return OperationNode.create(new ParDoOperation(fn, receivers, operationContext));
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.
the class BatchDataflowWorker method doWork.
/**
* Performs the given work; returns true if successful.
*
* @throws IOException Only if the WorkUnitClient fails.
*/
@VisibleForTesting
boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) throws IOException {
LOG.debug("Executing: {}", workItem);
DataflowWorkExecutor worker = null;
SdkWorkerHarness sdkWorkerHarness = sdkHarnessRegistry.getAvailableWorkerAndAssignWork();
try {
// Populate PipelineOptions with data from work unit.
options.setProject(workItem.getProjectId());
final String stageName;
if (workItem.getMapTask() != null) {
stageName = workItem.getMapTask().getStageName();
} else if (workItem.getSourceOperationTask() != null) {
stageName = workItem.getSourceOperationTask().getStageName();
} else {
throw new RuntimeException("Unknown kind of work item: " + workItem.toString());
}
CounterSet counterSet = new CounterSet();
BatchModeExecutionContext executionContext = BatchModeExecutionContext.create(counterSet, sideInputDataCache, sideInputWeakReferenceCache, readerRegistry, options, stageName, String.valueOf(workItem.getId()));
if (workItem.getMapTask() != null) {
MutableNetwork<Node, Edge> network = mapTaskToNetwork.apply(workItem.getMapTask());
if (LOG.isDebugEnabled()) {
LOG.debug("Network as Graphviz .dot: {}", Networks.toDot(network));
}
worker = mapTaskExecutorFactory.create(sdkWorkerHarness.getControlClientHandler(), sdkWorkerHarness.getGrpcDataFnServer(), sdkHarnessRegistry.beamFnDataApiServiceDescriptor(), sdkWorkerHarness.getGrpcStateFnServer(), network, options, stageName, readerRegistry, sinkRegistry, executionContext, counterSet, idGenerator);
} else if (workItem.getSourceOperationTask() != null) {
worker = SourceOperationExecutorFactory.create(options, workItem.getSourceOperationTask(), counterSet, executionContext, stageName);
} else {
throw new IllegalStateException("Work Item was neither a MapTask nor a SourceOperation");
}
workItemStatusClient.setWorker(worker, executionContext);
DataflowWorkProgressUpdater progressUpdater = new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options);
executeWork(worker, progressUpdater);
workItemStatusClient.reportSuccess();
return true;
} catch (Throwable e) {
workItemStatusClient.reportError(e);
return false;
} finally {
if (worker != null) {
try {
worker.close();
} catch (Exception exn) {
LOG.warn("Uncaught exception while closing worker. All work has already committed or " + "been marked for retry.", exn);
}
}
if (sdkWorkerHarness != null) {
sdkHarnessRegistry.completeWork(sdkWorkerHarness);
}
}
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.
the class RemoveFlattenInstructionsFunctionTest method testFlattenMultiplePCollectionsHavingMultipleConsumers.
@Test
public void testFlattenMultiplePCollectionsHavingMultipleConsumers() {
Node a = ParallelInstructionNode.create(new ParallelInstruction().setName("A"), Nodes.ExecutionLocation.UNKNOWN);
Node aPCollection = InstructionOutputNode.create(new InstructionOutput().setName("A.out"), PCOLLECTION_ID);
Edge aOutput = DefaultEdge.create();
Node b = ParallelInstructionNode.create(new ParallelInstruction().setName("B"), Nodes.ExecutionLocation.UNKNOWN);
Edge bOutput = DefaultEdge.create();
Node bPCollection = InstructionOutputNode.create(new InstructionOutput().setName("B.out"), PCOLLECTION_ID);
Node flatten = ParallelInstructionNode.create(new ParallelInstruction().setName("Flatten").setFlatten(new FlattenInstruction()), Nodes.ExecutionLocation.UNKNOWN);
Node flattenPCollection = InstructionOutputNode.create(new InstructionOutput().setName("Flatten.out"), PCOLLECTION_ID);
Node c = ParallelInstructionNode.create(new ParallelInstruction().setName("C"), Nodes.ExecutionLocation.UNKNOWN);
Edge cOutput = DefaultEdge.create();
Node cPCollection = InstructionOutputNode.create(new InstructionOutput().setName("C.out"), PCOLLECTION_ID);
Node d = ParallelInstructionNode.create(new ParallelInstruction().setName("D"), Nodes.ExecutionLocation.UNKNOWN);
Edge dOutput = DefaultEdge.create();
Node dPCollection = InstructionOutputNode.create(new InstructionOutput().setName("D.out"), PCOLLECTION_ID);
// A --\
// -> Flatten --> C
// B --/-------------> D
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(a);
network.addNode(aPCollection);
network.addNode(b);
network.addNode(bPCollection);
network.addNode(flatten);
network.addNode(flattenPCollection);
network.addNode(c);
network.addNode(cPCollection);
network.addEdge(a, aPCollection, aOutput);
network.addEdge(aPCollection, flatten, DefaultEdge.create());
network.addEdge(b, bPCollection, bOutput);
network.addEdge(bPCollection, flatten, DefaultEdge.create());
network.addEdge(bPCollection, d, DefaultEdge.create());
network.addEdge(flatten, flattenPCollection, DefaultEdge.create());
network.addEdge(flattenPCollection, c, DefaultEdge.create());
network.addEdge(c, cPCollection, cOutput);
network.addEdge(d, dPCollection, dOutput);
// A --\
// -> C
// B --/-> D
assertThatFlattenIsProperlyRemoved(network);
}
use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.
the class RemoveFlattenInstructionsFunctionTest method testRemoveFlattenOnMultiOutputInstruction.
@Test
public void testRemoveFlattenOnMultiOutputInstruction() {
Node a = ParallelInstructionNode.create(new ParallelInstruction().setName("A"), Nodes.ExecutionLocation.UNKNOWN);
Node aOut1PCollection = InstructionOutputNode.create(new InstructionOutput().setName("A.out1"), PCOLLECTION_ID);
Node aOut2PCollection = InstructionOutputNode.create(new InstructionOutput().setName("A.out2"), PCOLLECTION_ID);
Node aOut3PCollection = InstructionOutputNode.create(new InstructionOutput().setName("A.out3"), PCOLLECTION_ID);
Edge aOut1 = MultiOutputInfoEdge.create(new MultiOutputInfo().setTag("out1"));
Edge aOut2 = MultiOutputInfoEdge.create(new MultiOutputInfo().setTag("out2"));
Edge aOut3 = MultiOutputInfoEdge.create(new MultiOutputInfo().setTag("out3"));
Edge aOut1PCollectionEdge = DefaultEdge.create();
Node b = ParallelInstructionNode.create(new ParallelInstruction().setName("B"), Nodes.ExecutionLocation.UNKNOWN);
Node bOut1PCollection = InstructionOutputNode.create(new InstructionOutput().setName("B.out1"), PCOLLECTION_ID);
Node bOut2PCollection = InstructionOutputNode.create(new InstructionOutput().setName("B.out1"), PCOLLECTION_ID);
Edge bOut1 = MultiOutputInfoEdge.create(new MultiOutputInfo().setTag("out1"));
Edge bOut2 = MultiOutputInfoEdge.create(new MultiOutputInfo().setTag("out2"));
Edge bOut1PCollectionEdge = DefaultEdge.create();
Node flatten = ParallelInstructionNode.create(new ParallelInstruction().setName("Flatten").setFlatten(new FlattenInstruction()), Nodes.ExecutionLocation.UNKNOWN);
Node flattenPCollection = InstructionOutputNode.create(new InstructionOutput().setName("Flatten.out"), PCOLLECTION_ID);
Node c = ParallelInstructionNode.create(new ParallelInstruction().setName("C"), Nodes.ExecutionLocation.UNKNOWN);
Edge cOutput = DefaultEdge.create();
Node cPCollection = InstructionOutputNode.create(new InstructionOutput().setName("C.out"), PCOLLECTION_ID);
Node d = ParallelInstructionNode.create(new ParallelInstruction().setName("D"), Nodes.ExecutionLocation.UNKNOWN);
Edge dOutput = DefaultEdge.create();
Node dPCollection = InstructionOutputNode.create(new InstructionOutput().setName("D.out"), PCOLLECTION_ID);
Node e = ParallelInstructionNode.create(new ParallelInstruction().setName("E"), Nodes.ExecutionLocation.UNKNOWN);
Edge eOutput = DefaultEdge.create();
Node ePCollection = InstructionOutputNode.create(new InstructionOutput().setName("E.out"), PCOLLECTION_ID);
// /-out1-> C
// A -out2-\
// \-out3--> Flatten --> D
// B -out2-/
// \-out1-> E
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(a);
network.addNode(aOut1PCollection);
network.addNode(aOut2PCollection);
network.addNode(aOut3PCollection);
network.addNode(b);
network.addNode(bOut1PCollection);
network.addNode(bOut2PCollection);
network.addNode(flatten);
network.addNode(flattenPCollection);
network.addNode(c);
network.addNode(cPCollection);
network.addNode(d);
network.addNode(dPCollection);
network.addNode(e);
network.addNode(ePCollection);
network.addEdge(a, aOut1PCollection, aOut1);
network.addEdge(a, aOut2PCollection, aOut2);
network.addEdge(a, aOut3PCollection, aOut3);
network.addEdge(aOut1PCollection, c, aOut1PCollectionEdge);
network.addEdge(aOut2PCollection, flatten, DefaultEdge.create());
network.addEdge(aOut3PCollection, flatten, DefaultEdge.create());
network.addEdge(b, bOut1PCollection, bOut1);
network.addEdge(b, bOut2PCollection, bOut2);
network.addEdge(bOut1PCollection, e, bOut1PCollectionEdge);
network.addEdge(bOut2PCollection, flatten, DefaultEdge.create());
network.addEdge(flatten, flattenPCollection, DefaultEdge.create());
network.addEdge(flattenPCollection, d, DefaultEdge.create());
network.addEdge(c, cPCollection, cOutput);
network.addEdge(d, dPCollection, dOutput);
network.addEdge(e, ePCollection, eOutput);
// /-out1-> C
// A -out2-\
// \-out3--> D
// B -out2-/
// \-out1-> E
assertThatFlattenIsProperlyRemoved(network);
}
Aggregations