use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor 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));
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor 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);
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor 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);
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor in project beam by apache.
the class ProcessBundleDescriptorsTest method testLengthPrefixingOfKeyCoderInStatefulExecutableStage.
/**
* Tests that a stateful stage will wrap the key coder of a stateful transform in a
* LengthPrefixCoder.
*/
@Test
public void testLengthPrefixingOfKeyCoderInStatefulExecutableStage() throws Exception {
// Add another stateful stage with a non-standard key coder
Pipeline p = Pipeline.create();
Coder<Void> keycoder = VoidCoder.of();
assertThat(ModelCoderRegistrar.isKnownCoder(keycoder), is(false));
p.apply("impulse", Impulse.create()).apply("create", ParDo.of(new DoFn<byte[], KV<Void, String>>() {
@ProcessElement
public void process(ProcessContext ctxt) {
}
})).setCoder(KvCoder.of(keycoder, StringUtf8Coder.of())).apply("userState", ParDo.of(new DoFn<KV<Void, String>, KV<Void, String>>() {
@StateId("stateId")
private final StateSpec<BagState<String>> bufferState = StateSpecs.bag(StringUtf8Coder.of());
@TimerId("timerId")
private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@ProcessElement
public void processElement(@Element KV<Void, String> element, @StateId("stateId") BagState<String> state, @TimerId("timerId") Timer timer, OutputReceiver<KV<Void, String>> r) {
}
@OnTimer("timerId")
public void onTimer() {
}
})).apply("gbk", GroupByKey.create());
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto);
Optional<ExecutableStage> optionalStage = Iterables.tryFind(fused.getFusedStages(), (ExecutableStage stage) -> stage.getUserStates().stream().anyMatch(spec -> spec.localName().equals("stateId")));
checkState(optionalStage.isPresent(), "Expected a stage with user state.");
ExecutableStage stage = optionalStage.get();
PipelineNode.PCollectionNode inputPCollection = stage.getInputPCollection();
// Ensure original key coder is not a LengthPrefixCoder
Map<String, RunnerApi.Coder> stageCoderMap = stage.getComponents().getCodersMap();
RunnerApi.Coder originalMainInputCoder = stageCoderMap.get(inputPCollection.getPCollection().getCoderId());
String originalKeyCoderId = ModelCoders.getKvCoderComponents(originalMainInputCoder).keyCoderId();
RunnerApi.Coder originalKeyCoder = stageCoderMap.get(originalKeyCoderId);
assertThat(originalKeyCoder.getSpec().getUrn(), is(CoderTranslation.JAVA_SERIALIZED_CODER_URN));
// Now create ProcessBundleDescriptor and check for the LengthPrefixCoder around the key coder
BeamFnApi.ProcessBundleDescriptor pbd = ProcessBundleDescriptors.fromExecutableStage("test_stage", stage, Endpoints.ApiServiceDescriptor.getDefaultInstance()).getProcessBundleDescriptor();
Map<String, RunnerApi.Coder> pbsCoderMap = pbd.getCodersMap();
RunnerApi.Coder pbsMainInputCoder = pbsCoderMap.get(pbd.getPcollectionsOrThrow(inputPCollection.getId()).getCoderId());
String keyCoderId = ModelCoders.getKvCoderComponents(pbsMainInputCoder).keyCoderId();
RunnerApi.Coder keyCoder = pbsCoderMap.get(keyCoderId);
ensureLengthPrefixed(keyCoder, originalKeyCoder, pbsCoderMap);
TimerReference timerRef = Iterables.getOnlyElement(stage.getTimers());
String timerTransformId = timerRef.transform().getId();
RunnerApi.ParDoPayload parDoPayload = RunnerApi.ParDoPayload.parseFrom(pbd.getTransformsOrThrow(timerTransformId).getSpec().getPayload());
RunnerApi.TimerFamilySpec timerSpec = parDoPayload.getTimerFamilySpecsOrThrow(timerRef.localName());
RunnerApi.Coder timerCoder = pbsCoderMap.get(timerSpec.getTimerFamilyCoderId());
String timerKeyCoderId = timerCoder.getComponentCoderIds(0);
RunnerApi.Coder timerKeyCoder = pbsCoderMap.get(timerKeyCoderId);
ensureLengthPrefixed(timerKeyCoder, originalKeyCoder, pbsCoderMap);
}
use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor 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);
}
Aggregations