Examples with ProcessBundleResponse org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse used on opensource projects

Example 1 with ProcessBundleResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse in project beam by apache.

the class SdkHarnessClientTest method handleCleanupWhenAwaitingOnClosingOutputReceivers.

@Test
public void handleCleanupWhenAwaitingOnClosingOutputReceivers() throws Exception {
    Exception testException = new Exception();
    InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
    CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
    FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
    BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)));
    when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
    when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
    doThrow(testException).when(mockOutputReceiver).awaitCompletion();
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
    try {
        try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockProgressHandler)) {
            // Correlating the ProcessBundleRequest and ProcessBundleResponse is owned by the underlying
            // FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
            // the response.
            // 
            // Currently there are no fields so there's nothing to check. This test is formulated
            // to match the pattern it should have if/when the response is meaningful.
            BeamFnApi.ProcessBundleResponse response = BeamFnApi.ProcessBundleResponse.getDefaultInstance();
            processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build());
        }
        fail("Exception expected");
    } catch (Exception e) {
        assertEquals(testException, e);
    }
}

Example 2 with ProcessBundleResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse in project beam by apache.

the class SdkHarnessClientTest method testBundleCheckpointCallback.

@Test
public void testBundleCheckpointCallback() throws Exception {
    InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
    CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
    FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
    BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)));
    when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
    when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
    BundleSplitHandler mockSplitHandler = mock(BundleSplitHandler.class);
    BundleCheckpointHandler mockCheckpointHandler = mock(BundleCheckpointHandler.class);
    BundleFinalizationHandler mockFinalizationHandler = mock(BundleFinalizationHandler.class);
    ProcessBundleResponse response = ProcessBundleResponse.newBuilder().addResidualRoots(DelayedBundleApplication.getDefaultInstance()).build();
    try (ActiveBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), Collections.emptyMap(), (request) -> {
        throw new UnsupportedOperationException();
    }, mockProgressHandler, mockSplitHandler, mockCheckpointHandler, mockFinalizationHandler)) {
        processBundleResponseFuture.complete(InstructionResponse.newBuilder().setProcessBundle(response).build());
    }
    verify(mockProgressHandler).onCompleted(response);
    verify(mockCheckpointHandler).onCheckpoint(response);
    verifyZeroInteractions(mockFinalizationHandler, mockSplitHandler);
}

Example 3 with ProcessBundleResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse in project beam by apache.

the class ProcessBundleHandler method embedOutboundElementsIfApplicable.

private void embedOutboundElementsIfApplicable(ProcessBundleResponse.Builder response, BundleProcessor bundleProcessor) {
    if (bundleProcessor.getOutboundAggregators().isEmpty()) {
        return;
    }
    List<Elements> collectedElements = new ArrayList<>(bundleProcessor.getOutboundAggregators().size());
    boolean hasFlushedAggregator = false;
    for (BeamFnDataOutboundAggregator aggregator : bundleProcessor.getOutboundAggregators().values()) {
        Elements elements = aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
        if (elements == null) {
            hasFlushedAggregator = true;
        }
        collectedElements.add(elements);
    }
    if (!hasFlushedAggregator) {
        Elements.Builder elementsToEmbed = Elements.newBuilder();
        for (Elements collectedElement : collectedElements) {
            elementsToEmbed.mergeFrom(collectedElement);
        }
        response.setElements(elementsToEmbed.build());
    } else {
        // Since there was at least one flushed aggregator, we have to use the aggregators that were
        // able to successfully collect their elements to emit them and can not send them as part of
        // the ProcessBundleResponse.
        int i = 0;
        for (BeamFnDataOutboundAggregator aggregator : bundleProcessor.getOutboundAggregators().values()) {
            Elements elements = collectedElements.get(i++);
            if (elements != null) {
                aggregator.sendElements(elements);
            }
        }
    }
}

Example 4 with ProcessBundleResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse in project beam by apache.

the class ExecutableStageDoFnOperator method open.

@Override
public void open() throws Exception {
    executableStage = ExecutableStage.fromPayload(payload);
    hasSdfProcessFn = hasSDF(executableStage);
    initializeUserState(executableStage, getKeyedStateBackend(), pipelineOptions);
    // TODO: Wire this into the distributed cache and make it pluggable.
    // TODO: Do we really want this layer of indirection when accessing the stage bundle factory?
    // It's a little strange because this operator is responsible for the lifetime of the stage
    // bundle "factory" (manager?) but not the job or Flink bundle factories. How do we make
    // ownership of the higher level "factories" explicit? Do we care?
    stageContext = contextFactory.get(jobInfo);
    stageBundleFactory = stageContext.getStageBundleFactory(executableStage);
    stateRequestHandler = getStateRequestHandler(executableStage);
    progressHandler = new BundleProgressHandler() {

        @Override
        public void onProgress(ProcessBundleProgressResponse progress) {
            if (flinkMetricContainer != null) {
                flinkMetricContainer.updateMetrics(stepName, progress.getMonitoringInfosList());
            }
        }

        @Override
        public void onCompleted(ProcessBundleResponse response) {
            if (flinkMetricContainer != null) {
                flinkMetricContainer.updateMetrics(stepName, response.getMonitoringInfosList());
            }
        }
    };
    finalizationHandler = BundleFinalizationHandlers.inMemoryFinalizer(stageBundleFactory.getInstructionRequestHandler());
    checkpointHandler = getBundleCheckpointHandler(hasSdfProcessFn);
    minEventTimeTimerTimestampInCurrentBundle = Long.MAX_VALUE;
    minEventTimeTimerTimestampInLastBundle = Long.MAX_VALUE;
    super.setPreBundleCallback(this::preBundleStartCallback);
    super.setBundleFinishedCallback(this::finishBundleCallback);
    // This will call {@code createWrappingDoFnRunner} which needs the above dependencies.
    super.open();
}

Example 5 with ProcessBundleResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse in project beam by apache.

the class RegisterAndProcessBundleOperation method start.

@Override
public void start() throws Exception {
    try (Closeable scope = context.enterStart()) {
        super.start();
        // Only register once by using the presence of the future as a signal.
        if (registerFuture == null) {
            InstructionRequest request = InstructionRequest.newBuilder().setInstructionId(idGenerator.getId()).setRegister(registerRequest).build();
            registerFuture = instructionRequestHandler.handle(request);
        }
        checkState(registerRequest.getProcessBundleDescriptorCount() == 1, "Only one bundle registration at a time currently supported.");
        InstructionRequest processBundleRequest = InstructionRequest.newBuilder().setInstructionId(getProcessBundleInstructionId()).setProcessBundle(ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId(registerRequest.getProcessBundleDescriptor(0).getId())).build();
        deregisterStateHandler = beamFnStateDelegator.registerForProcessBundleInstructionId(getProcessBundleInstructionId(), this::delegateByStateKeyType);
        processBundleResponse = getRegisterResponse(registerFuture).thenCompose(registerResponse -> instructionRequestHandler.handle(processBundleRequest));
    }
}