Examples with RemoteGrpcPortWriteOperation org.apache.beam.runners.dataflow.worker.fn.data.RemoteGrpcPortWriteOperation used on opensource projects

Example 1 with RemoteGrpcPortWriteOperation

use of org.apache.beam.runners.dataflow.worker.fn.data.RemoteGrpcPortWriteOperation in project beam by apache.

the class BeamFnMapTaskExecutor method createProgressTracker.

private ProgressTracker createProgressTracker() {
    ReadOperation readOperation;
    RemoteGrpcPortWriteOperation grpcWriteOperation;
    RegisterAndProcessBundleOperation bundleProcessOperation;
    try {
        readOperation = getReadOperation();
    } catch (Exception exn) {
        readOperation = null;
        LOG.info("Unable to get read operation.", exn);
        return new NullProgressTracker();
    }
    // RegisterAndProcessBundleOperation we know they have the right topology.
    try {
        grpcWriteOperation = Iterables.getOnlyElement(Iterables.filter(operations, RemoteGrpcPortWriteOperation.class));
        bundleProcessOperation = Iterables.getOnlyElement(Iterables.filter(operations, RegisterAndProcessBundleOperation.class));
    } catch (IllegalArgumentException | NoSuchElementException exn) {
        // TODO: Handle more than one sdk worker processing a single bundle.
        grpcWriteOperation = null;
        bundleProcessOperation = null;
        LOG.debug("Does not have exactly one grpcWRite and bundleProcess operation.", exn);
    }
    if (grpcWriteOperation != null && bundleProcessOperation != null) {
        return new SingularProcessBundleProgressTracker(readOperation, grpcWriteOperation, bundleProcessOperation);
    } else {
        return new ReadOperationProgressTracker(readOperation);
    }
}

Example 2 with RemoteGrpcPortWriteOperation

use of org.apache.beam.runners.dataflow.worker.fn.data.RemoteGrpcPortWriteOperation in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createOperationTransformForGrpcPortNodes.

private Function<Node, Node> createOperationTransformForGrpcPortNodes(final Network<Node, Edge> network, final FnDataService beamFnDataService, final OperationContext context) {
    return new TypeSafeNodeFunction<RemoteGrpcPortNode>(RemoteGrpcPortNode.class) {

        @Override
        public Node typedApply(RemoteGrpcPortNode input) {
            RegisterAndProcessBundleOperation registerFnOperation = (RegisterAndProcessBundleOperation) Iterables.getOnlyElement(Iterables.filter(network.adjacentNodes(input), OperationNode.class)).getOperation();
            // The coder comes from the one and only adjacent output node
            Coder<?> coder = Iterables.getOnlyElement(Iterables.filter(network.adjacentNodes(input), OutputReceiverNode.class)).getCoder();
            // We figure out whether we are outputting some where if the output node is a
            // successor.
            Iterable<OutputReceiverNode> outputReceiverNodes = Iterables.filter(network.successors(input), OutputReceiverNode.class);
            Operation operation;
            if (outputReceiverNodes.iterator().hasNext()) {
                OutputReceiver[] outputReceivers = new OutputReceiver[] { Iterables.getOnlyElement(outputReceiverNodes).getOutputReceiver() };
                operation = new RemoteGrpcPortReadOperation<>(beamFnDataService, input.getPrimitiveTransformId(), registerFnOperation::getProcessBundleInstructionId, (Coder) coder, outputReceivers, context);
            } else {
                operation = new RemoteGrpcPortWriteOperation<>(beamFnDataService, input.getPrimitiveTransformId(), registerFnOperation::getProcessBundleInstructionId, (Coder) coder, context);
            }
            return OperationNode.create(operation);
        }
    };
}

Example 3 with RemoteGrpcPortWriteOperation

use of org.apache.beam.runners.dataflow.worker.fn.data.RemoteGrpcPortWriteOperation in project beam by apache.

the class SingularProcessBundleProgressTrackerTest method testProgressInterpolation.

@Test
public void testProgressInterpolation() throws Exception {
    ReadOperation read = Mockito.mock(ReadOperation.class);
    RemoteGrpcPortWriteOperation grpcWrite = Mockito.mock(RemoteGrpcPortWriteOperation.class);
    RegisterAndProcessBundleOperation process = Mockito.mock(RegisterAndProcessBundleOperation.class);
    when(grpcWrite.processedElementsConsumer()).thenReturn(elementsConsumed -> {
    });
    SingularProcessBundleProgressTracker tracker = new SingularProcessBundleProgressTracker(read, grpcWrite, process);
    when(read.getProgress()).thenReturn(new TestProgress("A"), new TestProgress("B"), new TestProgress("C"));
    when(grpcWrite.getElementsSent()).thenReturn(1, 10, 20, 30);
    // This test ignores them, directly working on mocked getInputElementsConsumed
    when(process.getProcessBundleProgress()).thenReturn(CompletableFuture.completedFuture(BeamFnApi.ProcessBundleProgressResponse.getDefaultInstance()));
    when(process.getInputElementsConsumed(any(Iterable.class))).thenReturn(1L, 4L, 10L).thenThrow(new RuntimeException());
    // Initially no progress is known.
    assertEquals(null, tracker.getWorkerProgress());
    // After reading, and writing, and processing one element, the progress is aligned at A.
    tracker.updateProgress();
    assertEquals(new TestProgress("A"), tracker.getWorkerProgress());
    // We've read up to B (10 elements) but only consumed 4.  Progress remains at A.
    tracker.updateProgress();
    assertEquals(new TestProgress("A"), tracker.getWorkerProgress());
    // Once 10 elements have been consumed, advance to B.
    tracker.updateProgress();
    assertEquals(new TestProgress("B"), tracker.getWorkerProgress());
    // An exception is thrown, default to latest read progress.
    tracker.updateProgress();
    assertEquals(new TestProgress("C"), tracker.getWorkerProgress());
}