Examples with RemoteInputDestination org.apache.beam.runners.fnexecution.data.RemoteInputDestination used on opensource projects

Example 1 with RemoteInputDestination

use of org.apache.beam.runners.fnexecution.data.RemoteInputDestination in project beam by apache.

the class SdkHarnessClientTest method verifyCacheTokensAreUsedInNewBundleRequest.

@Test
public void verifyCacheTokensAreUsedInNewBundleRequest() throws InterruptedException {
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(CompletableFuture.<InstructionResponse>completedFuture(InstructionResponse.newBuilder().build()));
    ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
    List<RemoteInputDestination> remoteInputs = Collections.singletonList(RemoteInputDestination.of(FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM));
    BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs);
    when(dataService.send(any(), any())).thenReturn(mock(CloseableFnDataReceiver.class));
    StateRequestHandler stateRequestHandler = Mockito.mock(StateRequestHandler.class);
    List<BeamFnApi.ProcessBundleRequest.CacheToken> cacheTokens = Collections.singletonList(BeamFnApi.ProcessBundleRequest.CacheToken.newBuilder().getDefaultInstanceForType());
    when(stateRequestHandler.getCacheTokens()).thenReturn(cacheTokens);
    processor1.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mock(RemoteOutputReceiver.class)), stateRequestHandler, BundleProgressHandler.ignored());
    // Retrieve the requests made to the FnApiControlClient
    ArgumentCaptor<BeamFnApi.InstructionRequest> reqCaptor = ArgumentCaptor.forClass(BeamFnApi.InstructionRequest.class);
    Mockito.verify(fnApiControlClient, Mockito.times(1)).handle(reqCaptor.capture());
    List<BeamFnApi.InstructionRequest> requests = reqCaptor.getAllValues();
    // Verify that the cache tokens are included in the ProcessBundleRequest
    assertThat(requests.get(0).getRequestCase(), is(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE));
    assertThat(requests.get(0).getProcessBundle().getCacheTokensList(), is(cacheTokens));
}

Example 2 with RemoteInputDestination

use of org.apache.beam.runners.fnexecution.data.RemoteInputDestination in project beam by apache.

the class SdkHarnessClientTest method testRegister.

@Test
public void testRegister() throws Exception {
    ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
    List<RemoteInputDestination> remoteInputs = Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM));
    sdkHarnessClient.getProcessor(descriptor1, remoteInputs);
    verify(fnApiControlClient).registerProcessBundleDescriptor(descriptor1);
}

Example 3 with RemoteInputDestination

use of org.apache.beam.runners.fnexecution.data.RemoteInputDestination in project beam by apache.

the class SdkHarnessClientTest method testRegisterWithStateRequiresStateDelegator.

@Test
public void testRegisterWithStateRequiresStateDelegator() throws Exception {
    ProcessBundleDescriptor descriptor = ProcessBundleDescriptor.newBuilder().setId("test").setStateApiServiceDescriptor(ApiServiceDescriptor.newBuilder().setUrl("foo")).build();
    List<RemoteInputDestination> remoteInputs = Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM));
    thrown.expect(IllegalStateException.class);
    thrown.expectMessage("containing a state");
    sdkHarnessClient.getProcessor(descriptor, remoteInputs);
}

Example 4 with RemoteInputDestination

use of org.apache.beam.runners.fnexecution.data.RemoteInputDestination in project beam by apache.

the class ProcessBundleDescriptors method fromExecutableStageInternal.

private static ExecutableProcessBundleDescriptor fromExecutableStageInternal(String id, ExecutableStage stage, ApiServiceDescriptor dataEndpoint, @Nullable ApiServiceDescriptor stateEndpoint) throws IOException {
    // Create with all of the processing transforms, and all of the components.
    // TODO: Remove the unreachable subcomponents if the size of the descriptor matters.
    Map<String, PTransform> stageTransforms = stage.getTransforms().stream().collect(Collectors.toMap(PTransformNode::getId, PTransformNode::getTransform));
    Components.Builder components = stage.getComponents().toBuilder().clearTransforms().putAllTransforms(stageTransforms);
    ImmutableList.Builder<RemoteInputDestination> inputDestinationsBuilder = ImmutableList.builder();
    ImmutableMap.Builder<String, Coder> remoteOutputCodersBuilder = ImmutableMap.builder();
    WireCoderSetting wireCoderSetting = stage.getWireCoderSettings().stream().filter(ws -> ws.getInputOrOutputId().equals(stage.getInputPCollection().getId())).findAny().orElse(WireCoderSetting.getDefaultInstance());
    // The order of these does not matter.
    inputDestinationsBuilder.add(addStageInput(dataEndpoint, stage.getInputPCollection(), components, wireCoderSetting));
    remoteOutputCodersBuilder.putAll(addStageOutputs(dataEndpoint, stage.getOutputPCollections(), components, stage.getWireCoderSettings()));
    Map<String, Map<String, SideInputSpec>> sideInputSpecs = addSideInputs(stage, components);
    Map<String, Map<String, BagUserStateSpec>> bagUserStateSpecs = forBagUserStates(stage, components.build());
    Map<String, Map<String, TimerSpec>> timerSpecs = forTimerSpecs(stage, components);
    lengthPrefixAnyInputCoder(stage.getInputPCollection().getId(), components);
    // Copy data from components to ProcessBundleDescriptor.
    ProcessBundleDescriptor.Builder bundleDescriptorBuilder = ProcessBundleDescriptor.newBuilder().setId(id);
    if (stateEndpoint != null) {
        bundleDescriptorBuilder.setStateApiServiceDescriptor(stateEndpoint);
    }
    if (timerSpecs.size() > 0) {
        // By default use the data endpoint for timers, in the future considering enabling specifying
        // a different ApiServiceDescriptor for timers.
        bundleDescriptorBuilder.setTimerApiServiceDescriptor(dataEndpoint);
    }
    bundleDescriptorBuilder.putAllCoders(components.getCodersMap()).putAllEnvironments(components.getEnvironmentsMap()).putAllPcollections(components.getPcollectionsMap()).putAllWindowingStrategies(components.getWindowingStrategiesMap()).putAllTransforms(components.getTransformsMap());
    return ExecutableProcessBundleDescriptor.of(bundleDescriptorBuilder.build(), inputDestinationsBuilder.build(), remoteOutputCodersBuilder.build(), sideInputSpecs, bagUserStateSpecs, timerSpecs);
}

Example 5 with RemoteInputDestination

use of org.apache.beam.runners.fnexecution.data.RemoteInputDestination in project beam by apache.

the class SdkHarnessClientTest method testRegisterCachesBundleProcessors.

@Test
public void testRegisterCachesBundleProcessors() throws Exception {
    ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
    ProcessBundleDescriptor descriptor2 = ProcessBundleDescriptor.newBuilder().setId("descriptor2").build();
    List<RemoteInputDestination> remoteInputs = Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM));
    BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs);
    BundleProcessor processor2 = sdkHarnessClient.getProcessor(descriptor2, remoteInputs);
    assertNotSame(processor1, processor2);
    // Ensure that caching works.
    assertSame(processor1, sdkHarnessClient.getProcessor(descriptor1, remoteInputs));
}